Deep Learning Diagrams

Despite being mathematical systems, deep learning models are not systematically approached. Instead, they are understood in an ad-hoc manner -- every particular feature or insight is derived independently of others. Proper analysis requires various abstractions. Mathematically, we are interested in how functions feed into each other, how they are parallelized, and how functions relate to training and inference forms. Practically, we are interested in the resource costs given some mathematical goal, allowing us to find the optimal execution strategy on parallelized GPU hardware, and even to engage in systematized software-hardware co-design. Typical approaches struggle to consider these different lenses.

Category theory's tools for studying composition and abstraction are well-suited to the challenge of formally understanding deep learning models. It allows us to work with an abstract, mathematical model, and layer on meaning as required. This meaning could be the inference (KV-cache) form of a model, the training (backpropagated) form, or assignment to logical and physical hardware levels. Furthermore, category theory allows us to develop a rigorous diagrammatic language which reflects these abstractions. In the era of agentic mathematics, this is especially critical. Proofs are cheap, but nonetheless require an appropriate language in which to operate, which category theory gladly offers.

Our methods have been successfully applied to systematically derive hardware-aware FlashAttention in FlashAttention on a Napkin. These methods have been applied to quickly generate a kernel for an attention variant. The underlying mathematics is formalized in our paper Weaves, Wires, and Morphisms, which is reflected in two open-source packages PyNCD and TSNCD.

Interactive Diagrams of Complete Architectures

Each diagram below is an expression written with PyNCD and drawn by TSNCD. The diagrams below include foundational, tutorial models such as standard attention variants. This builds up to classic models: Attention is All You Need, Mixtral-8x7B, and DeepSeek-V3. Finally, there are the modern models: GLM-5.3, and the absolute artwork which is DeepSeek-V4.1-Flash, with all its mixed streams, batched indexers, and engram hashings.

There are three forms of diagrams: arrows and boxes, arrows and broadcasted, and all broadcasted. These go from how models are typically expressed without broadcasting information, through to the full NCD expressions which are akin to "seeing like a GPU". Different "views" of the same model are possible, and each is procedurally derived from the base form by the tools in PyNCD. A quantized abstract representation is sufficient to derive the unquantized form, the training algorithm, and the cached algorithm. You will also find a link to the notebooks which generates the models. So, go ahead and start exploring.

Attention

Scaled dot-product attention and its training step

Explore the interactive diagram

Drag to pan, Ctrl + scroll to zoom, and hover over a block to inspect it.Drag to pan, pinch to zoom, and tap a block to inspect it.

Scaled dot-product attention, written as a morphism. Its training step is the forward pass, which saves the values the backward pass reads, followed by the backward pass, which returns the gradients of the queries, the keys and the values.