Evolution of ML Architectures

Three axes that define an architecture

Almost every architecture argument in deep learning is a disagreement about one question: how should information at position i reach position j? Everything else — depth, width, normalization, activation function — is tuning. The mixing operator is the architecture.

There have really only been five answers to that question in sixty-seven years. Don't mix at all and let a dense matrix memorize everything (the MLP). Mix locally with a shared kernel and stack until the receptive field is big enough (the CNN). Mix through a single hidden vector carried step by step (the RNN). Mix everything with everything, weighted by content similarity (attention). Or mix through a linear recurrence whose coefficients are themselves computed from the input (selective state spaces). Each answer buys something and mortgages something else, and the mortgage always comes due in one of three places.

Axis 1 — path length. How many operations separate two positions that need to interact? Attention says one. A stacked CNN says O(N/k) layers, or O(log N) if you dilate. An RNN says O(N), which is why it forgets. Short paths mean gradients survive; long paths mean information is repeatedly re-encoded and degraded.

Axis 2 — training parallelism. Can position t be computed without waiting for t−1? Convolution and attention: yes, the whole sequence is one matmul. Vanilla recurrence: no, and that single fact — not accuracy — is what killed the LSTM for large-scale pretraining. Modern state space models re-earn parallelism by restricting the recurrence to be linear, which makes it an associative scan.

Axis 3 — inference state. What must you carry to generate token N+1? A Transformer carries every key and value it has ever seen: state grows linearly with context and decoding becomes memory-bandwidth-bound long before it becomes compute-bound. An RNN or SSM carries a fixed-size vector: constant memory, constant time per token, but a hard information ceiling. This is the axis that current research is fighting over.

Sixty-seven years in one picture

The field did not progress smoothly. It moved in bursts, each one triggered by a constraint being lifted: a learning rule (1986), a dataset and a GPU (2012), a way to build depth without vanishing gradients (2015), a way to remove sequential dependency from training (2017), and a way to make a recurrence both selective and hardware-efficient (2023). The animation below pans across that history and draws the lineage arcs as it goes.

Two things stand out when you lay it out this way. First, most "new" architectures are old ideas that became affordable: convolution predates backpropagation-trained CNNs by a decade, attention is a 2014 machine-translation patch that took three years to become a standalone architecture, and state space models are control theory from the 1960s wearing a GPU kernel. Second, the gaps get shorter. There were eleven years between LeNet and AlexNet, five between AlexNet and the Transformer, and roughly eighteen months between Mamba and a production hybrid stack.

The dense layer: perceptrons and MLPs

The perceptron is a dot product and a threshold. Rosenblatt's 1958 version learned by a local update rule, could only separate linearly separable classes, and was famously shown by Minsky and Papert to be unable to compute XOR. The fix — stacking layers with a differentiable nonlinearity between them, and assigning credit through the chain rule — arrived properly in 1986 with backpropagation, and it is still, in every architecture below, the only learning algorithm anyone uses.

Walk the figure through a forward pass, a loss, and a backward pass. Note in stage 2 that a neuron is doing nothing more mysterious than a projection onto a learned direction followed by a squashing function; the "intelligence" is entirely in which directions get learned.

Why the MLP is both universal and useless on its own

A two-layer MLP with enough hidden units can approximate any continuous function on a compact set. That theorem is true and almost entirely unhelpful, because it says nothing about how many units, how much data, or how the function generalizes. A dense layer treats input coordinate 7 and coordinate 8 as no more related than coordinate 7 and coordinate 700. For a 224×224×3 image, the first dense layer alone would need 150,528 × d weights — about 154M parameters at d = 1024 — and it would have to relearn every edge detector separately at every pixel offset, because nothing ties the weights at one location to the weights at another.

That is the whole motivation for every architecture that follows. Structure in the data — locality, translation, sequence order, permutation symmetry — is either baked into the operator as a prior, or it must be paid for in parameters and data.

Convolution: the first good prior

Convolution makes two assumptions and gets an enormous amount in return. First, locality: whatever matters at a pixel is decided by its neighborhood, so the kernel is small. Second, stationarity: an edge detector useful in the top-left is useful in the bottom-right, so the same weights slide everywhere. Together these replace a dense N×N matrix with k shared weights, and turn "translation" from a nuisance the model must memorize into a symmetry the operator respects for free.

The figure below slides a 3×3 kernel across an image and builds the feature map cell by cell, then shows what stride, padding, dilation and multiple channels each do to the arithmetic. Step it with Next if you want to read the numbers inside each patch.

Equivariance, invariance, and the thing pooling actually does

Convolution is translation equivariant: shift the input, and the feature map shifts identically. It is not translation invariant — that only arrives when you pool, stride, or global-average-pool at the end, which deliberately throws away position. The distinction matters: for classification you want invariance at the top; for segmentation and detection you want equivariance preserved all the way down, which is exactly why U-Net keeps skip connections from encoder to decoder instead of relying on the bottleneck alone.

The depth problem and the residual fix

Stacking convolutions grows the receptive field linearly, so reaching global context on a 224-pixel image needs depth. But plain deep stacks got worse — not from overfitting, from optimization: the signal had to survive a long multiplicative chain. ResNet's answer was to make the identity the default path, y = x + F(x), so the Jacobian of a block is I + ∂F/∂x and gradients have an unobstructed route to the input. That one line is arguably the most reused idea in the field; every architecture after 2015, including every Transformer and every Mamba block, is a stack of residual branches.

Recurrence: memory as a bottleneck

Sequences need a different prior. Elman's recurrent network keeps a hidden state and applies the same transition at every step, which is the temporal analogue of weight sharing: h_t = σ(W_h h_{t−1} + W_x x_t + b). It handles arbitrary lengths, it has constant memory, and at inference it is genuinely elegant — one matmul per token, forever, no matter how long the context.

Training is where it falls apart, in two separate ways. Run the figure: stages 1–5 unroll the forward pass, stage 6 sends the gradient back through time, and stages 7–8 show why it arrives as noise.

Failure one: the gradient is a matrix power

Backpropagation through time multiplies the same Jacobian once per step. Over k steps the gradient carries a factor ∏ W_hᵀ diag(σ′), whose magnitude behaves like λ^k where λ is roughly the largest singular value of that product. If λ < 1 the signal vanishes geometrically — at λ = 0.9, a dependency 100 steps back arrives at 0.0027 of its strength. If λ > 1 it explodes, which is at least easy to see and easy to patch with gradient clipping. Vanishing is the dangerous one because training looks fine; the model simply never learns anything long-range.

Failure two: the sequential wall

This is the one that ended the era. Computing h_t requires h_{t−1}, so a length-2048 sequence is 2048 dependent kernel launches, each a small matmul with terrible arithmetic intensity. An A100 doing a 1024×1024 matvec is running at a percent or two of peak. Attention replaced recurrence not because it was more accurate per parameter — early comparisons were close — but because it turned N sequential steps into one big parallel matmul, which meant you could spend 100× the compute in the same wall-clock time. Architecture progress and hardware utilization are the same story after 2017.

Gates: LSTM, GRU and the constant error carousel

The LSTM's insight is that the vanishing problem is caused by repeatedly multiplying the state by a learned matrix, so build a path where the state is added to instead. The cell state c_t is updated as c_t = f_t ⊙ c_{t−1} + i_t ⊙ g_t: a per-channel forget gate scales it, and new content is added. When f_t ≈ 1 the derivative ∂c_t/∂c_{t−1} ≈ 1 and gradient flows without decay. Hochreiter and Schmidhuber called it the constant error carousel; today you would recognize it as a residual connection along the time axis, six years before ResNet used the same trick along depth.

The GRU merges the forget and input gates into one (z_t decides "keep old" vs "take new") and drops the separate cell state, giving three gate matrices instead of four — roughly 25% fewer parameters, usually within noise on accuracy, occasionally better on small data.

Gating fixed axis 1 (path length becomes effectively short along the carousel) but did nothing for axis 2. That is the whole reason the field moved on. And note what survived: the gate is a data-dependent, elementwise multiplicative control on what enters a running state. Hold onto that sentence — it is exactly what Mamba reinvents in §11.

Attention: content-addressed routing

Attention began as a repair to sequence-to-sequence translation. Encoding a whole sentence into one vector and asking the decoder to reconstruct it was clearly a bottleneck; Bahdanau's fix was to let the decoder, at each output step, compute a similarity score against every encoder state and read a weighted average. The reading pattern is not fixed by the architecture — it is computed from the content, per query, per layer, per head. That is the entire idea, and it is qualitatively different from everything before it.

A convolution says "look 3 to the left, always." An RNN says "look at the compressed summary of everything, always." Attention says "score all candidates and decide right now." It is soft, differentiable random access into memory.

Why the √d_k, and why softmax

Two vectors of dimension d_k with unit-variance entries have a dot product with variance d_k. Feed that into a softmax at d_k = 128 and you get logits with standard deviation ≈ 11, which saturates the softmax into a near-one-hot distribution and kills the gradient. Dividing by √d_k restores unit variance. The softmax itself is doing two jobs: normalizing so the output is a convex combination of values (keeping activations in range regardless of N), and creating competition between positions, so that attending strongly to one token necessarily means attending less to others. Both matter — several "linear attention" variants that dropped the softmax lost measurable quality on retrieval tasks precisely because they dropped the competition.

Multi-head is not just capacity

One N×N attention map can only express one routing pattern per layer. Splitting d into H heads of size d/H costs the same FLOPs but produces H independent routing patterns that are then concatenated and mixed by W_O. Empirically, heads specialize hard: some become positional (attend to t−1), some syntactic (verb → subject), some become induction heads that implement "find the previous occurrence of this token and copy what followed it" — the mechanism behind most in-context learning. That behaviour is not expressible by a single head, and it is not expressible at all by a fixed kernel.

The Transformer block, piece by piece

The 2017 paper's real claim was in its title: attention is all you need — drop the recurrence entirely, and what remains is a residual stack alternating two operations. A token-mixing operation (attention, the only place positions talk) and a channel-mixing operation (a position-wise MLP, which is where most of the parameters live). Norm, residual, repeat.

Which term dominates depends entirely on your ratio N/d. At d = 4096 and N = 2048, the FFN is doing more FLOPs than attention; at N = 128k, attention is 8× the FFN and utterly dominates. Every efficiency paper you read is implicitly targeting one regime or the other, so always ask which one before believing a speedup number.

Inference economics: the KV cache

Training complexity gets all the attention in papers; serving cost is what decides deployments. Generation has two phases with completely different characteristics, and conflating them is the most common analysis error I see.

Prefill processes the prompt in one parallel pass: O(N²d) FLOPs, compute-bound, high utilization, and it is the phase that determines time-to-first-token. Decode then produces one token at a time. Each step is a single query attending over all cached keys and values, so per-token compute is O(N·d) — but crucially, every one of those cached bytes must be read from HBM to compute a handful of FLOPs. At batch size 1 the arithmetic intensity is around 1 operation per byte, against an H100's roughly 300 ops/byte break-even. Decoding is not a compute problem, it is a memory-bandwidth problem, and the size of the state you carry is the whole ballgame.

That table is why every serious 2023–2026 model uses grouped-query attention or a latent-compressed variant: MQA shares one KV head across all query heads (maximum saving, some quality loss), GQA shares one KV head per group of 8 (nearly free), and DeepSeek's multi-head latent attention caches a low-rank joint projection instead of full K and V. It is also why paged allocation (vLLM) mattered so much: naive contiguous per-request cache allocation wasted 60–80% of the memory to fragmentation and over-reservation.

The efficiency era: sparse, linear, fused, sparse again

Between 2019 and 2023 the community produced dozens of sub-quadratic attention variants. Almost none of them are in the models you use today, and the reason is instructive: most traded asymptotic complexity for constant factors and quality, on a workload where N was not yet large enough for the asymptotics to pay. The ones that survived changed the memory hierarchy rather than the math.

Sparsity: restrict who talks to whom

Sliding-window attention limits each query to w neighbours, giving O(N·w) and a receptive field that grows w per layer — a convolution with content-based weights. Longformer and BigBird add a few global tokens so information can still cross the sequence in two hops, and BigBird proves that random + window + global retains universal approximation. Mistral shipped pure sliding-window at w = 4096, and current frontier models interleave: a few full-attention layers among many local ones, which keeps exact long-range recall while paying the quadratic cost only a handful of times.

Linearization: drop the softmax, keep an associative form

If you replace exp(qᵀk) with φ(q)ᵀφ(k) for some feature map, associativity lets you compute φ(Q)(φ(K)ᵀV) instead of (φ(Q)φ(K)ᵀ)V. The N×N matrix never exists; cost becomes O(N d²) and inference becomes a recurrence over a d×d state matrix. This is the crucial bridge: a linear-attention Transformer is an RNN, and every modern SSM is a variation on the same identity. The catch is quality — without softmax competition, associative recall degrades, which the Zoology line of work quantified as a direct trade between recurrent state size and recall accuracy.

Fusion: FlashAttention changed everything by changing nothing

FlashAttention computes exactly standard attention, same numbers, same gradients. It just never writes the N×N matrix to HBM: it tiles the computation, keeps blocks in SRAM, and uses the online-softmax trick to accumulate a running max and sum. Memory drops from O(N²) to O(N), and because attention was memory-bound, wall-clock drops 2–4× too. It is the clearest lesson of the era: on modern accelerators, the relevant complexity measure is often bytes moved, not FLOPs.

Conditional compute: mixture of experts

MoE attacks a different axis — decoupling parameters from FLOPs. Replace the FFN with E experts and a router that sends each token to the top-k (usually 1 or 2). Total parameters scale with E; compute per token does not. Mixtral 8×7B has 47B parameters and the forward cost of about 13B. The price is paid in memory (all experts must be resident), in load-balancing losses to stop the router collapsing onto a favourite expert, and in all-to-all communication when experts are sharded across devices.

State spaces: HiPPO, S4, and Mamba

The state space line comes at sequences from control theory. A linear time-invariant system is h′(t) = A h(t) + B x(t), y(t) = C h(t). Discretize it with step size Δ and you get a linear recurrence h_t = Ā h_{t−1} + B̄ x_t — an RNN with no nonlinearity in the loop. That restriction is what buys everything: a linear recurrence can be unrolled into a convolution, and it can be evaluated by a parallel associative scan in O(log N) depth.

The remaining question is what A should be, since a random linear recurrence still forgets exponentially. HiPPO answers it with mathematics rather than learning: derive the A whose state holds the optimal coefficients of a polynomial approximation to the entire input history under a chosen measure. S4 made that matrix trainable and numerically stable (a diagonal-plus-low-rank parameterization with an FFT-based convolution kernel), and set records on the Long Range Arena — 16k-length benchmarks where Transformers were failing outright.

The selectivity problem, and Mamba's answer

S4 is time-invariant: the same Ā, B̄, C at every step. That is exactly what makes the convolutional view possible, and it is also a hard ceiling. A time-invariant system cannot decide to ignore a token, cannot reset its state at a boundary, and cannot perform selective copying — because its behaviour does not depend on the content it is processing. This is the same limitation a CNN has, in a different costume.

Mamba makes Δ, B and C functions of the current input: Δ_t = softplus(Linear(x_t)), and similarly for B_t, C_t. Now a large Δ_t means "this token matters, write it into the state and reset the decay", while a small Δ_t means "filler, let the state coast". Look at the discretization Ā_t = exp(Δ_t A): it is a per-channel, input-dependent forget gate. Twenty-six years after the LSTM, the gate is back — but now the recurrence around it stays linear, so the scan is still parallel.

The cost of selectivity is that Ā_t now varies with t, so there is no fixed convolution kernel and no FFT trick. Mamba pays for that with a hardware-aware implementation: expand into the high-dimensional state inside SRAM, run a fused selective scan, and never materialize the (B, N, d, s) state tensor in HBM. Same lesson as FlashAttention — the algorithm is chosen to fit the memory hierarchy, and it recomputes in the backward pass rather than storing.

Mamba-2 then closed the circle: it showed that selective SSMs and a family of linear attentions are the same computation viewed through a structured-matrix decomposition (semiseparable matrices), which let the scan be rewritten as block matmuls that actually use tensor cores — 2–8× faster than Mamba-1 while allowing a much larger state dimension.

Vision in the state-space era: scan order is the whole problem

Everything in the previous section assumed a sequence: tokens arrive in an order, and that order is real. An image has no such order. It is a grid, it is not causal, and the two pixels immediately above and below each other are genuinely adjacent in a way that no flattening preserves. So when people ask whether Mamba works for vision, the interesting part of the answer is not about state sizes or selectivity — it is about what you do to a 2D array to make it look like a 1D stream, and what that costs you.

Flatten a H×W patch grid in raster order and two vertically adjacent patches land W steps apart. For a 224px image at patch 16 that is 14 steps; at 1024px with patch 16 it is 64. A causal scan is worse still: half the image is simply unavailable to any given patch, which is fine for language and absurd for a picture. Every vision SSM is, at bottom, a proposal for how to fix those two problems.

The two families

Vision Mamba (Vim) takes the minimal step: keep the raster flattening, drop causality, and run the selective scan bidirectionally. Each block scans the token sequence forward and backward and sums the two output streams, so every patch can see every other patch, and a class token placed in the middle of the sequence rather than at the front sits within a short scan distance of everything. It is architecturally a ViT with the attention block swapped out — same isotropic stack, same patch embedding, same fixed resolution throughout.

VMamba attacks the geometry instead. Its 2D-selective-scan block (SS2D, built on what the paper calls the Cross-Scan Module) unrolls the grid four ways — rows left-to-right, rows right-to-left, columns top-to-bottom, columns bottom-to-top — runs a selective SSM along each, and merges the four resulting maps back onto the grid. The point is not more compute; it is that any two patches are close in scan distance under at least one of the four orders, so the arbitrary penalty that raster order imposes on vertical neighbours gets cancelled by the column scans. VMamba also drops the isotropic ViT shape for a hierarchical, four-stage pyramid with downsampling, which is what actually makes it competitive on detection and segmentation.

Everything since is a variation on those two moves. LocalMamba scans inside windows before scanning globally, on the observation that four full-image scans still separate local neighbours more than a 3×3 convolution would. Mamba-ND generalises the idea to n dimensions by alternating which axis is unrolled at each layer instead of paying for all orders at every layer. PlainMamba keeps the isotropic shape but uses a continuous zigzag so that consecutive tokens are always spatially adjacent. VideoMamba and Vivim extend the scan across time; U-Mamba, SegMamba and Swin-UMamba put the blocks into encoder–decoder segmentation backbones, which is where the linear scaling matters most.

The honest scorecard

Reported ImageNet numbers put VMamba-T around 82–83% top-1 at roughly 30M parameters, which is comparable to Swin-T and ConvNeXt-T rather than decisively better, and the gap that does appear tends to shrink under matched training recipes. Several careful ablations since have argued that most of the advantage comes from the hierarchical design, the large effective receptive field, and modern training rather than from selectivity itself — replacing the scan with a cheap global-mixing operator recovers much of the gain. Meanwhile the four-way scan multiplies constant factors by roughly four, and the fused scan kernel is far less optimised than FlashAttention, so the asymptotic advantage only converts into wall-clock advantage somewhere in the low thousands of tokens.

There is also a subtler cost. Attention is permutation-equivariant, and position enters only through an explicit encoding you control. A scan hard-codes an order, and the model's sensitivity to two patches depends on their distance along that order. Adding scans is how you buy the bias back down, which means vision SSMs spend parameters and FLOPs undoing a problem attention never had. That is the trade in one sentence: linear cost, purchased with an ordering prior that images do not actually possess.

Hybrids and what is actually winning

The honest 2026 summary is that neither pure attention nor pure recurrence wins, and the interesting models are mixtures. The mechanism is easy to state: a fixed-size state cannot store N tokens losslessly, so pure SSMs will always lose on exact retrieval; but most tokens do not need exact retrieval, so paying quadratic cost on every layer is wasteful. Interleave them and you get almost all of both.

The recipes converge on roughly one attention layer per six to eight recurrent or local layers. Jamba interleaves Mamba, attention and MoE at a 1:7 ratio and fits a 256k context on one 80GB GPU with a KV cache an order of magnitude smaller than a comparable Transformer. Samba alternates Mamba with sliding-window attention and extrapolates far past its training length. Griffin pairs gated linear recurrences with local attention and matches Llama-2 quality on far fewer tokens. Even conventional models hybridize within attention: interleaving local and global layers is the same idea with a different local operator.

How big is a model, really? Parameters, bytes and FLOPs

Every architectural choice above has a price you can compute exactly, on paper, before you write any code. Three numbers matter and they are different numbers: how many parameters the design implies, how many bytes those parameters occupy at each stage of the pipeline, and how many FLOPs a forward and backward pass costs. People routinely conflate the first two and then wonder why an 8B model will not train on an 80GB card. It is worth doing the arithmetic once, slowly.

Counting a Transformer decoder

A decoder layer holds two weight groups. Attention has four projections; with grouped-query attention the key and value projections are narrower than the query and output ones, because n_kv heads are shared across n_h query heads:

The feed-forward block is the larger half. A classic GELU MLP has two matrices of width d_ff = 4d; a SwiGLU MLP has three — gate, up, down — and the convention is to set d_ff ≈ (8/3)d so that three matrices cost the same as two at 4d:

So the rule of thumb is N ≈ 12·L·d² plus embeddings — and the "plus embeddings" is not a rounding error at small scale. GPT-2 small is 124M parameters of which 39M are the token embedding: nearly a third of the model is a lookup table. At 70B the same table is under 1%. This is why tying input and output embeddings was standard practice for years and then quietly stopped mattering.

A worked example that lands on a real number

Take the published Llama-3-8B configuration: L = 32, d = 4096, d_ff = 14336, n_h = 32, d_h = 128, n_kv = 8, V = 128256, untied head.

That is the shipped parameter count, from six numbers and no measurement. The same exercise for a Mamba block gives roughly 6d² — the input projection to 2Ed and the output projection back cost 3Ed², which is 6d² at the standard expansion E = 2, with the convolution, Δ and B, C projections contributing a few percent. That is exactly half a Transformer layer, which is why Mamba stacks two blocks where a Transformer has one: it is a deliberate choice to make parameter-matched comparisons meaningful. A convolution, for the record, is k²·C_in·C_out per layer and is almost always the smallest term in anything modern.

Every other family, the same arithmetic

Nothing about the counting is Transformer-specific. Every layer is a set of matrices, and the only question is which ones and how many times they are reused. Writing them all out with a common width d makes the families directly comparable for the first time:

Read down that column and a few things fall out immediately. An LSTM layer costs two thirds of a Transformer layer at the same width — the four gates are expensive, and the old intuition that recurrence is "cheap" was never about parameters. A 3×3 convolution is nine dense layers stacked in one, which is why channel counts in CNNs stay small: VGG's widest block is 512 channels where a language model is at 4096. And a Mamba block is exactly half a Transformer layer, which is a deliberate design choice so that parameter-matched comparisons mean something.

The right-hand columns are where the families actually separate, and it is not the parameter count. Every recurrent family carries a state that does not depend on how much text it has seen; the Transformer alone carries a state proportional to the conversation. A CNN carries nothing between images but pays enormous activation memory during a forward pass, because a 224×224 feature map at 64 channels is 3.2M activations per layer before you have learned anything. Those are three different scarcities, and no single "how big is it" number captures them.

One more comparison worth holding on to. LeNet-5 is roughly 60k parameters; AlexNet 62M, of which 58M are in three fully-connected layers at the end and only 4M in all five convolutions; VGG-16 is 138M with 124M in the classifier head. The entire history of CNN design after 2014 is the story of deleting those dense layers — global average pooling in place of a flatten, then residual stacks, then depthwise separable convolutions. ResNet-50 gets better accuracy than VGG-16 with 25M parameters. Almost none of that progress shows up in a FLOP count; it shows up in where the parameters were allowed to sit.

From parameters to bytes

Now the part people get wrong. Weights at bf16 are 2 bytes each, so 8B parameters is 16GB — that is the inference number, and it fits on one card. Training with Adam in mixed precision needs, per parameter: bf16 weights (2), bf16 gradients (2), an fp32 master copy (4), and the two fp32 optimizer moments (4 + 4). Sixteen bytes per parameter, 128GB for the same 8B model, before a single activation is stored. This is the entire reason ZeRO and FSDP exist: that 16 bytes is what gets sharded across ranks, and it divides cleanly by world size while the activations do not.

For Llama-3-8B that KV cache is 128KB per token: 1GB at 8k context, 16GB at 128k — the cache alone matches the weights. That single ratio is why grouped-query attention went from a curiosity to universal in about eighteen months, and it is the same pressure that makes a constant-state recurrence attractive at long context.

From parameters to FLOPs

A forward pass costs about 2N FLOPs per token — one multiply and one add per parameter — and the backward pass costs roughly twice the forward, giving the familiar C ≈ 6ND for training on D tokens. Attention's sequence-dependent term sits on top:

At d = 4096 and a 4k context that ratio is about 17%: attention is a real but minority cost, and the matmuls in the MLP dominate. Push the context to 24k and the two are equal; past that, the quadratic term owns the budget and everything in the efficiency section starts to pay for itself. The crossover is a property of N_ctx / d, not of the model size — which is why long-context work and width scaling pull in opposite directions.

Play with it

The playground below implements the arithmetic above with nothing hidden. Move the width, depth, vocabulary and context sliders, switch the architecture between multi-head attention, grouped-query attention, a Mamba stack and a sparse mixture of experts, and step through the five views: where the parameters live, what the weights weigh, what training costs, how the cache grows, and how long the run takes. The read-outs are the same formulas printed above, so anything surprising is a fact about the design, not about the widget.

A few things are worth finding by hand. Doubling d quadruples the parameters but only doubles the KV cache. Doubling L doubles both. Raising the vocabulary from 32k to 256k adds 1.7B parameters at d = 4096 — more than a third of a small model — while changing FLOPs per token barely at all, because the output head is one matmul against a very wide matrix. Switching to MoE multiplies total parameters by the expert count while leaving active parameters and therefore FLOPs almost unchanged, which is precisely the trade it exists to make: memory is cheap and bandwidth-bound serving is not, so you buy capacity in the dimension that does not cost you compute.

The backward pass, derived

Everything above described what an architecture computes. None of it explains how the weights get there, and the answer is not incidental to the design: several of the architectural choices in this article exist only because of what they do to the gradient. The residual connection, the gates in an LSTM, the 1/√d_k in attention, the choice of RMSNorm over LayerNorm — each is a decision about the backward pass wearing the costume of a forward-pass detail. So it is worth deriving the whole thing carefully, once, for every block in the article.

Reverse mode in one page

A network is a composition L = f_n ∘ f_{n−1} ∘ … ∘ f_1 (x) ending in a scalar loss. The chain rule says the derivative is a product of Jacobians, and the only question is what order to multiply them in:

Matrix multiplication is associative, so both give the same answer, but the cost is wildly different. With one scalar output and 10⁹ parameters, reverse mode costs one sweep and forward mode costs 10⁹ of them. That asymmetry is the entire reason deep learning trains at all, and it is why the loss must be a scalar.

The crucial practical point: reverse mode never forms a Jacobian. Each node implements a vector–Jacobian product — given the gradient of the loss with respect to its output, return the gradient with respect to its input. Write X̄ ≡ ∂L/∂X for the adjoint of any tensor X; it always has exactly the same shape as X. Then a layer's backward pass is a function Ȳ ↦ X̄, and it typically costs about the same as the forward pass, never the d² that materialising ∂Y/∂X would.

Primitive 1: the linear layer

Everything else is decoration on this one. Let Y = XW + b with X ∈ ℝ^{n×k}, W ∈ ℝ^{k×m}. In index notation Y_ab = Σ_c X_ac W_cb + b_b, so:

Two consequences worth internalising. First, the backward pass of a matmul is two matmuls of the same size as the forward one, which is where the familiar "backward costs twice the forward" comes from, and hence 6N FLOPs per token per parameter for training against 2N for inference. Second, W̄ needs X, the layer's input. That is rule 3 in its most expensive form: every activation on the forward path has to survive until its matmul's backward runs.

Primitive 2: elementwise nonlinearities

If y = f(z) elementwise, the Jacobian is diagonal and the VJP is a Hadamard product — no matrix ever appears:

The ¼ on the sigmoid is not a footnote; it is the arithmetic reason the pre-2010 deep networks did not train. Stack six sigmoids and the best case gradient multiplier is 4⁻⁶ ≈ 2×10⁻⁴ before any weight matrix is involved. ReLU's derivative being exactly one on the active half is the whole of its contribution, and every modern activation (GELU, SiLU) is a smoothed version of that same property — near-unit slope on the positive side, a small non-zero slope on the negative side so that dead units can recover.

Primitive 3: softmax and cross-entropy

These are derived together because the composition collapses into something far simpler than either piece. With p_i = e^{z_i} / Σ_j e^{z_j}:

Now compose with cross-entropy against a one-hot target, L = −log p_y, whose adjoint is p̄_i = −y_i / p_i. Then pᵀp̄ = Σ_i p_i(−y_i/p_i) = −Σ_i y_i = −1, and:

The prediction minus the target, and nothing else. Note that the softmax Jacobian appears in the answer only through the two-term structure p ⊙ (p̄ − pᵀp̄) — a scaled subtraction of a row-mean. That exact shape reappears verbatim inside attention's backward pass, and inside normalisation's, because all three are computing a gradient through a quantity that was normalised to sum to one.

This also explains the fused cross-entropy kernels every large-vocabulary model now uses. Materialising logits for a 128k vocabulary at 8k tokens is 1B floats per layer of nothing; but since z̄ = p − y depends only on p and the target, you can compute the loss and its gradient in tiles and never store the full matrix.

Primitive 4: normalisation, the one people get wrong

LayerNorm over a feature dimension of size H: μ = mean(x), σ = √(var(x) + ε), x̂ = (x − μ)/σ, y = γ ⊙ x̂ + β. The subtlety is that μ and σ both depend on every component of x, so changing one input coordinate moves every output coordinate. Writing g = ȳ ⊙ γ for the adjoint of x̂, and differentiating through all three paths (direct, through μ, through σ):

Read the bracket geometrically: the gradient is g with its component along 1 removed and its component along x̂ removed. LayerNorm projects the incoming gradient onto the subspace orthogonal to both — it is literally impossible for a gradient to ask the layer below to shift the mean or scale the variance, because the normaliser will undo it. This is the mechanism behind its stabilising effect and behind the fact that the effective learning rate of a normalised layer is scale-invariant: x̄ carries a 1/σ, so if the incoming activations grow, the gradient shrinks in proportion, automatically.

RMSNorm drops the mean subtraction, and its backward drops exactly one of the three terms:

Which is most of the reason RMSNorm won: one fewer reduction in the forward pass, one fewer in the backward, no measurable quality cost.

Primitive 5: residuals, and what they do to gradients

The identity term is the whole point. Composing n residual blocks gives ∏(I + J_k), which expands into I plus a sum of products — there is always a path from the loss to any layer that passes through no Jacobians at all. Depth stops being a multiplicative hazard and becomes an additive one. The gradient at layer ℓ is at least ȳ, no matter how many blocks sit above it, which is precisely what the residual-stack panel in the depth figure is showing.

This also settles the pre-LN versus post-LN question, which is purely a backward-pass argument. Post-LN puts the normaliser on the residual path (y = LN(x + F(x))), so every skip connection is followed by a projection that removes gradient components — the identity path is no longer clean, and deep post-LN Transformers need learning-rate warmup to train at all. Pre-LN (y = x + F(LN(x))) keeps the residual stream unnormalised end to end, so the gradient highway is exact. Every large model since roughly 2020 is pre-LN for this reason alone.

The MLP, end to end

With the primitives in hand, a two-layer MLP with a softmax head is a five-line backward pass. Forward: a¹ = xW¹ + b¹, h = f(a¹), a² = hW² + b², p = softmax(a²).

Every deep network in this article is this pattern with more structure in the middle. Notice the alternation: a matmul backward is two matmuls, an elementwise backward is a multiply, and the adjoint marches leftward at the same width the activations marched rightward.

Convolution: weight sharing means gradient summing

A convolution (strictly, a cross-correlation) is Y[o,i,j] = Σ_{c,u,v} K[o,c,u,v] · X[c, i+u, j+v]. Because the same kernel entry is used at every spatial position, rule 2 applies with a vengeance — its adjoint is a sum over all of them:

Three things fall out. The gradient with respect to the input is itself a convolution, which is why the backward pass runs on the same hardware kernels as the forward one. The kernel gradient is a sum over H·W positions, so a 3×3 kernel with 9 parameters receives gradient signal from every pixel in the image — the reason CNNs are so sample-efficient is not just that they have fewer parameters, it is that each parameter gets H·W times more gradient per image. And the flip in X̄ is where the name "transposed convolution" comes from: it is the transpose of the forward operator's matrix, nothing more.

Backpropagation through time

An RNN unrolled is a very deep MLP with tied weights, so BPTT is the MLP rule plus rule 2 applied across time. With a_t = W_h h_{t−1} + W_x x_t + b and h_t = tanh(a_t), each hidden state feeds both its own output and the next step, so its adjoint has two contributions:

Unrolling the recursion gives the product that decides everything:

A geometric series in the number of steps. If σ_max(W_h) < 1 the gradient vanishes exponentially; if it exceeds 1 it explodes exponentially; there is no stable middle because the same matrix is applied every step. Both failure modes have the same cause and different fixes — exploding gradients are cured by clipping the global norm, which is a hack that works because the direction is usually right even when the magnitude is absurd; vanishing gradients cannot be clipped away, because the information is gone rather than merely large. Truncated BPTT bounds the cost by cutting the sum at k steps, which is honest about the fact that the terms beyond k were numerically zero anyway.

Why gates fix it, exactly

Now the same derivation for an LSTM. Forward: c_t = f_t ⊙ c_{t−1} + i_t ⊙ g_t, h_t = o_t ⊙ tanh(c_t). The cell-state adjoint picks up two terms, and the second one is the important one:

Compare that with the RNN's ∏ diag(1−h²) W_h. The LSTM's cell path has no weight matrix in it at all — the multiplier is the forget gate, a number in (0,1) that the network chooses per timestep and per channel. Set f ≈ 1 and the product is ≈ 1 over hundreds of steps: that is the constant error carousel, and it is a statement about the backward pass, not the forward one. The gates are a learned gradient-routing mechanism. This is also why the forget-gate bias is conventionally initialised to +1: it starts f near 1, which starts the network with an open gradient highway rather than asking it to discover one.

GRU works the same way with one gate fewer, ∂h_t/∂h_{t−1} carrying a (1 − z_t) factor along the interpolation path.

Attention, and the trick that makes FlashAttention possible

Forward, for one head: S = QKᵀ/√d_k, P = softmax(S) row-wise, O = PV. Apply the matmul rule twice and the softmax rule once:

Now the 1/√d_k earns its keep. It appears in Q̄ and K̄ directly, but its real job is upstream: without it, S has variance d_k, the softmax saturates, P approaches one-hot, and P ⊙ (P̄ − rowsum(P̄ ⊙ P)) goes to zero because p_i(1 − p_i) → 0. The scaling factor is not there to make the forward pass prettier; it is there to keep the softmax Jacobian from collapsing and killing the gradient. Same reasoning, same shape, as the sigmoid's ¼.

The identity that makes the memory-efficient version work is worth its own line. Define D_i = rowsum(P̄ ⊙ P)_i, the scalar subtracted in each row. Then:

The row-correction depends only on the output and its adjoint — both N×d — so it can be computed in O(Nd) without ever touching the N×N matrix. That is what lets FlashAttention's backward pass recompute S and P tile by tile in SRAM (using the saved row-wise logsumexp statistics) instead of reading back a stored attention matrix. The forward pass stores O and one scalar per row; the backward pass rebuilds everything else. Memory drops from O(N²) to O(N), FLOPs rise by roughly 30%, and the whole thing gets faster because attention was memory-bound, not compute-bound.

For multi-head attention the heads are independent up to the concatenation, so W̄_O = concat(O_h)ᵀ Ō_final and each head receives its slice. Under GQA the shared K and V projections are used by several query heads, so — rule 2 again — their adjoints are the sum over the group.

Selective state spaces: the backward pass is another scan

Strip an SSM layer to its recurrence, h_t = ā_t h_{t−1} + b̄_t, y_t = c_t h_t. The adjoint obeys the same kind of recurrence, running in the opposite direction:

This is the structural payoff of linearity. The RNN's backward pass is sequential because its per-step Jacobian depends on h through tanh′; the SSM's backward pass is a linear recurrence with known coefficients, which is associative, which means it is a parallel scan in O(log N) depth just like the forward pass. Linearity buys parallelism in both directions, which is the entire reason a linear recurrence is trainable at scale and a tanh recurrence is not.

The selection mechanism then chains through the discretisation. With Ā = exp(ΔA) and Δ = softplus(W_Δ x + b_Δ):

Two things follow. The exponential is well-behaved for A < 0, which is why A is parameterised as −exp(A_log) — it forces the multiplier into (0,1) and makes the state a decay rather than a divergence, structurally guaranteeing what an RNN could only hope for. And because Δ_t is data-dependent, the gradient through the state is data-dependent too: the model can learn to open the gate (Δ large, ā → 0, state overwritten) or close it (Δ small, ā → 1, state and gradient preserved). It is the LSTM's forget gate again, in exponential coordinates and with an associative-scan-friendly algebra.

One implementation note that matters: the fused selective-scan kernel does not store the intermediate states h_t. They are recomputed during the backward pass from the saved inputs, exactly as with FlashAttention. Activation memory is O(N·d) rather than O(N·d·s), and that factor of s is the difference between a state size of 16 being free and being unaffordable.

Mixtures of experts, and a genuinely non-differentiable step

With y = Σ_{k ∈ top-K} g_k · E_k(x) and g = softmax(x W_r) restricted to the chosen experts:

The top-K selection itself has zero gradient almost everywhere — it is a discrete argmax, and no amount of chain rule will tell the router about an expert it did not choose. The router only learns to re-weight among experts it already picked. Every practical fix is a way around this hole: an auxiliary load-balancing loss that penalises uneven routing directly, noise added to the router logits so that near-ties get explored, or the expert-choice formulation that inverts the assignment so experts pick tokens instead. This is one of the few places in a modern network where the gradient genuinely does not exist and the training recipe has to compensate.

What the backward pass costs, and how to pay less

Collecting the accounting from the linear-layer rule: forward is one matmul, backward is two, so a training step is roughly 3× the forward FLOPs — the 6N per token from the sizing section. Memory is the harder constraint, because rule 3 means every activation on the forward path must be live until its backward runs. The standard mitigations are all trades against that one fact:

Two numerical footnotes complete the picture. In mixed precision the loss is multiplied by a scale factor S before backward() and the gradients divided by S afterwards, because bf16/fp16 gradients underflow to zero long before they become meaningless — the scaling shifts them into representable range and is exactly invertible. And gradient clipping by global norm, ḡ ← ḡ · min(1, c/‖ḡ‖), rescales the whole gradient vector rather than clipping elementwise, which preserves the direction; that distinction is why global-norm clipping is standard and per-element clipping is not.

Lab: the same image through three architectures

Descriptions of inductive bias stay abstract until you watch the same pixels being consumed three different ways. The lab below takes one 16×16 image and pushes it through a dense layer, a convolution stack, and a vision Transformer. Switch architecture with the buttons, step with Previous/Next, and use Shift the object to translate the input by a few pixels — then watch which representations move with it and which ones fall apart.

The takeaways are visible rather than argued. The MLP's first layer needs 32,768 weights to the same hidden width that the convolution reaches with 9 shared ones, and when you shift the object every one of those weights is now looking at the wrong thing. The CNN's feature map simply translates with the input — the representation is equivariant, so the classifier above it barely notices. The ViT keeps position only through an added embedding, so a shift changes the tokens' positional codes but not their content; it recovers translation robustness from data and augmentation instead of from the operator, which is precisely why ViTs need far more images (or distillation, or heavy augmentation) to match a CNN at small scale, and why they overtake it once data is abundant.

Lab: the same sentence through four architectures

Now the sequence case, which is where the three axes really bite. The sentence has a long-range agreement dependency: the verb must agree with keys, eleven tokens earlier, across four distracting nouns. Choose an architecture, then step through the tokens and watch three things: what the model is holding at each step, what it can see when it reaches the verb, and what that costs.

The RNN compresses everything into one vector, and by the time it reaches the verb the trace of keys has decayed through eleven multiplications — the bar showing its surviving contribution is the direct visual of λ^k. The LSTM's forget gate holds the subject on the cell highway, so it survives — this is genuinely what gating buys, and why LSTMs handled agreement tasks well long before Transformers existed. The Transformer does not compress at all: every token is still sitting in the cache, so the verb reaches back in exactly one hop with a sharp attention weight — but notice the cache bar growing linearly and the per-token cost growing with it. Mamba keeps the constant state of the RNN, but the selection gate opens on keys and nearly closes on the filler tokens, so the subject is written strongly into the state and the distractors decay — long-range memory without a growing cache, at the price of a state that is still finite, and still lossy if you ask it to reproduce an exact 40-token string it saw once.

Lab: the complexity explorer

Asymptotics decide nothing on their own; crossover points do. Drag the sliders and read where the lines actually cross for your model dimension and your context length. The three tabs correspond to the three questions worth asking separately: what does a training step cost, what does generating one more token cost, and how much state has to sit in memory while you do it.

Two crossovers are worth memorizing. Attention's N²d term overtakes the FFN's Nd² term at N ≈ d — for a 4096-dim model, quadratic cost stops being a rounding error somewhere around 4k tokens and dominates completely by 32k. And the memory picture flips much earlier than the compute picture: KV cache passes model weights in serving footprint at surprisingly modest batch × context products, which is the real reason long-context serving is expensive.

Master comparison table

Per layer, sequence length N, model width d, kernel k, SSM state size s, batch 1. "Depth" is sequential dependency during training — the number of steps that cannot be parallelized.

Green marks the axis a design wins on, red the one it pays for. r = attention-layer ratio in a hybrid (typically 6–8); w = window; E = experts; s = SSM state dimension (16–128).

Design-choice nuances that matter more than the block

Two models with the same "architecture" can differ by more than a Transformer differs from a Mamba. These are the choices that actually move numbers, and they are mostly invisible in the block diagram.

Where you inject position

Attention is permutation-equivariant, so position is an additive design decision with real consequences. Sinusoidal and learned absolute embeddings add a vector at the input, which makes them structurally unable to extrapolate past the trained length. RoPE rotates Q and K by an angle proportional to position, so the dot product depends on relative offset — it extrapolates better, composes with FlashAttention, and can be stretched at inference (NTK-aware and YaRN scaling) to extend context after training. ALiBi adds a linear distance penalty to the scores instead, giving a soft recency prior and clean extrapolation at some cost to long-range precision. Recurrent and convolutional models need none of this: order is built into the operator.

Normalization placement and type

Pre-norm for trainable depth, post-norm for slightly better final loss, and modern large runs often use RMSNorm — which drops the mean-centering and the bias, keeping only the scale — because it is measurably cheaper on a memory-bound op and works just as well. Some 2024–2026 models also add QK-norm inside attention to stop logit magnitudes drifting during long training runs, a cheap fix for a failure mode that otherwise appears only after weeks of compute.

Activation and the FFN's shape

ReLU → GELU → SwiGLU is not fashion. Gated linear units split the up-projection into a gate and a value branch and multiply them, which adds a third weight matrix; to keep parameters constant you shrink d_ff from 4d to about (8/3)d. At equal parameters and equal FLOPs, gated variants consistently reach lower loss. The mechanism is the same one that shows up in LSTMs and Mamba: multiplicative, data-dependent interaction is more expressive per parameter than an elementwise nonlinearity.

The head-count / cache trade

Query heads are free at inference; key-value heads are not, because each one multiplies the cache. GQA with 8 KV heads is the current default because it is the point where cache size stops being the binding constraint and quality loss is still negligible. Latent attention pushes further by caching a shared low-rank projection. If you remember one serving heuristic: the number of KV heads is a memory decision, not a modelling decision.

Depth versus width

At fixed parameter count, deeper models generally reason better and wider models memorize and parallelize better. Width is friendlier to tensor parallelism (fewer sequential all-reduces) and to arithmetic intensity; depth increases pipeline stages and sequential latency at decode. Most production models sit in a narrow, empirically-found band of aspect ratios rather than at either extreme.

What you actually optimize

Scaling laws changed architecture practice more than most architectures did. Chinchilla's correction — roughly 20 tokens per parameter for compute-optimal training — reallocated budget from parameters to data, and inference-aware practice has pushed far past even that, since a smaller model trained on more tokens is cheaper to serve forever. An architecture that is 10% better per FLOP but 3× slower to decode loses in production, every time.

How to choose

Compressed decision rules, stated as bluntly as I can make them.

Fixed-size, tabular, or already-featurized input: an MLP with residuals and normalization. Gradient-boosted trees still beat it on small tabular data; check before you build.

Images, small to medium dataset, latency-sensitive: a modern CNN. ConvNeXt-style designs match ViTs at ImageNet scale with better throughput and no patch-size headaches, and the locality prior is worth several multiples of data.

Images or video at large scale, or anything multimodal: a ViT, almost always. The reason is not accuracy alone — it is that the output is a sequence of tokens that a language model can consume directly, which no CNN feature pyramid gives you as cleanly.

Text generation, general purpose: a Transformer with GQA, RoPE, RMSNorm, pre-norm and SwiGLU. This is the boring answer and it remains the right one; the ecosystem advantage — kernels, quantization, serving, fine-tuning tooling — is worth more than a few percent of architectural elegance.

Very long context, high throughput, streaming: a hybrid. Mamba or gated linear recurrence for most layers, a few full-attention layers for exact recall, sliding windows in between. This is where the frontier is moving and where the cost curves are most favourable.

Signals: audio, sensors, DNA, long time series: a selective SSM. Continuous-time formulation, principled long-memory initialization, linear scaling — this is the regime SSMs were designed for and they dominate it.

Small data and a real symmetry in your problem: encode the symmetry in the operator. A convolution, a graph network, or an equivariant layer will beat a general-purpose block that has to learn the symmetry from examples you do not have.

References

Grouped by topic, every entry linked to the paper. Where a work exists in several versions, the link goes to the canonical one.