04 — Neural Layers & Architectures

The building blocks of a neural network — dense, convolutional, recurrent and attention layers — form a chain: each one addresses a specific failure in the one before it, so the most durable way to understand any of them is to derive it from that failure.

Each layer is covered in three parts: the assumption it makes about the data, what it costs, and what breaks when the assumption is false. The goal is to be able to look at a dataset, name the layer its structure justifies and say why, and estimate a transformer’s parameter count without a calculator.

Vocabulary. Nine terms, used throughout without re-explanation.

• Layer. One step of the computation: an array of numbers goes in, an array of numbers comes out.
• Parameters (equivalently weights). The adjustable numbers that control what a layer does. Training changes these and nothing else.
• Activations. The array a layer produces. “The activations of layer 3” means the numbers coming out of layer 3.
• Loss. One number saying how wrong the network’s output is on the examples in front of it. Lower is better.
• Training. Repeatedly measuring the loss and nudging every parameter in the direction that lowers it.
• Gradient. The collection of partial derivatives of the loss with respect to each parameter — one number per parameter, saying which way to nudge it and how hard.
• Backpropagation. The algorithm that computes the gradient. It starts at the loss and walks backwards through the stack.
• Jacobian. The matrix of partial derivatives of one layer’s outputs with respect to its inputs. Backpropagation multiplies by one of these per layer on the way back.
• Chain rule. The rule that says you get the whole gradient by chaining those per-layer multiplications together.

Read dy/dx aloud as “how much the output y changes per unit change in the input x”.

One consequence of the chain rule causes many of the failures in this chapter. The gradient that reaches layer 1 of an L-layer stack is a product of L matrices. Multiply many numbers smaller than 1 and the result collapses toward zero; multiply many numbers larger than 1 and it blows up. Sections 5, 6 and 10 are all restatements of that one fact.

A layer is an assumption about structure. Convolution assumes the same pattern can appear anywhere; recurrence assumes the past compresses into a fixed state; attention assumes the right thing to read depends on what you are looking for. Choose the wrong assumption and the model needs orders of magnitude more data to learn what the right one would have supplied for free.

If you want the linear algebra underneath in more depth, Matrix calculus derived not memorized derives the chain rule and the Jacobian properly, and Vectors norms and what cosine similarity actually measures derives the dot product as a geometric object, which Attention derived as content based lookup here uses. Neither is required reading; both facts are restated where this chapter needs them.

0. What goes in and what comes out

Every layer below is a transformation from one named shape to another, so the shapes need pinning down first.

A tensor is the only data type in this chapter: a rectangular block of numbers with a shape, written as a tuple of sizes. A single number is a shape-() tensor, a list of 512 numbers is shape (512), and a stack of 32 such lists is shape (32, 512).

Two concrete end-to-end examples follow. The rest of the chapter alternates between them.

An image classifier.

• In: one photograph as a (224, 224, 3) tensor — 224 pixel rows, 224 pixel columns, and 3 channels. A channel is one measurement per pixel position; here it is the red, green and blue (RGB) brightness of that pixel.
• Out: a (1000) tensor — one score per candidate class, so the largest entry names the predicted object.
• In between: each layer turns one (height, width, channels) tensor into another, generally shrinking height and width while growing the channel count.

A language model.

• In: a sequence of n integers, one per token. A token is a chunk of text (roughly a word or word-piece) that the model treats as an indivisible symbol, and the integer is that token’s row number in a fixed vocabulary of V distinct symbols.
• Out: at each position, a (V) tensor of scores over the whole vocabulary, from which the next token is chosen.
• In between: the sequence is carried as an (n, d) tensor — one vector of d numbers per token, where d is the model width (also called hidden size).

In both cases the whole thing runs on a batch — several independent examples stacked so one pass of arithmetic serves all of them — so real shapes carry a leading batch dimension B: (B, 224, 224, 3) and (B, n, d).

The table below maps the chapter. Each row gets its own section; use the table as a summary once you have read them.

Read a row left to right as a sentence: this layer turns this shape into this shape, on the assumption that this is true about the data, and when it is not, this is what goes wrong.

1. Why depth beats width

You can grow a network by widening one layer or by stacking more of them. Both add capacity. Depth adds it exponentially faster, and the clearest demonstration is one-dimensional.

How to count expressive power

The argument needs one unit of measurement first.

ReLU is rectified linear unit, the function max(0, z). It passes positive numbers through unchanged and flattens everything negative to zero.

A network built out of ReLU units is a piecewise-linear function: it is made of straight segments joined at corners. Each corner is a kink, and each straight stretch between kinks is a piece (the formal name is linear region). More pieces means the network can express a more intricate shape, so counting pieces is a way of counting expressive power.

Start with the shallow case. A single hidden layer of w ReLU units, fed a single number, produces a function with at most w + 1 pieces — each unit contributes one kink, and n kinks cut a line into n + 1 segments. Width buys pieces linearly: double the units, roughly double the pieces.

The fold: why composing beats adding

Now stack layers instead of widening one. The whole argument rests on a single map.

A width-2 ReLU layer can compute x -> |2x - 1| — read that as “x maps to the absolute value of two x minus one”. Two units do it: ReLU(2x - 1) + ReLU(1 - 2x) equals |2x - 1|, because whichever of the two arguments is positive is the absolute value and the other is clamped to zero.

Call it the fold. Here is what it does to the interval [0, 1]:

x         0.00   0.25   0.50   0.75   1.00
|2x - 1|  1.00   0.50   0.00   0.50   1.00

Read the row as two halves. As x sweeps 0 -> 0.5 the output sweeps 1 -> 0. As x sweeps 0.5 -> 1 the output sweeps back 0 -> 1. Every value in [0,1] comes out exactly twice, once from each half of the input. The fold copies its input range onto the output range twice.

That is why composition doubles. Whatever the next layer computes on [0,1] — one kink, ten kinks, a whole staircase — the fold makes the input trace it twice, so the piece count doubles.

Apply the fold to itself, |2·|2x - 1| - 1|:

x                       0.000  0.125  0.250  0.375  0.500  0.625  0.750  0.875  1.000
inner  |2x - 1|         1.000  0.750  0.500  0.250  0.000  0.250  0.500  0.750  1.000
outer  |2·inner - 1|    1.000  0.500  0.000  0.500  1.000  0.500  0.000  0.500  1.000

The inner row has one kink (at x = 0.5) and two pieces. The outer row goes down-up-down-up, with kinks at 0.25, 0.5 and 0.75 — three kinks, four pieces. Two layers, 2^2 = 4.

Width cannot do this. Two units in the same layer are added, and adding kinks is arithmetic — one unit, one more kink. Stacking layers composes, and composing a 2-to-1 map copies the piece count. So L folded layers give 2^L pieces, read as “two to the power of the number of layers”.

Price both routes to the same 1,024 pieces:

                       pieces   parameters
1 hidden layer,  w=1023   1024   1·1023 + 1023 + 1023·1 + 1 = 3,070
10 hidden layers, w=2     1024   (1·2+2) + 9·(2·2+2) + (2·1+1) = 61

same number of linear regions, 50x fewer parameters
push to 20 layers:  2^20 = 1,048,576 pieces from 121 parameters;
                    the shallow net would need ~1M units, ~3.1M parameters

Both parameter counts are the same sum done twice: for each layer, the number of weights (inputs × outputs) plus one bias per output.

• Shallow. One input feeds 1,023 units (1·1023 weights, 1023 biases), and those 1,023 units feed one output (1023·1 weights, 1 bias). Total 3,070.
• Deep. The first layer is 1 -> 2 (2 weights + 2 biases = 4). Nine more hidden layers are 2 -> 2 (4 weights + 2 biases = 6 each, so 54). The output layer is 2 -> 1 (2 weights + 1 bias = 3). Total 61.

3,070 / 61 ≈ 50. Same expressive power, fifty times fewer knobs.

The general bound

The toy example is one-dimensional. The general statement (Montufar et al., 2014) covers any input size.

First, one symbol. Write d_in for the input dimension — how many numbers the network is fed. This is a different quantity from the model width d of What goes in and what comes out, which is why it gets its own name.

The theorem: a ReLU network with input dimension d_in and L layers of width w can carve the input space into

Omega( (w/d_in)^(d_in·(L-1)) · w^d_in )   linear regions

Read Omega(...) as “grows at least as fast as”. It is the standard notation for a lower bound with constant factors ignored — the theorem says such a network can achieve at least this many regions, not that it is capped there.

That direction matters. A lower bound has to be witnessed by an actual construction, and the construction is the fold.

The shape of the expression is the point: L sits in the exponent, w sits in the base. Exponential in depth, polynomial in width. Composition reuses every feature the previous layer built; width just adds more features at the same level of abstraction.

Substitute the toy example into the theorem and the two collapse into one statement. The deep net is d_in = 1, w = 2, L = 10:

(w/d_in)^(d_in(L-1)) · w^d_in  =  (2/1)^(1·9) · 2^1  =  2^9 · 2  =  1,024      <- the deep net
                               =  (1023/1)^(1·0) · 1023^1  =  1 · 1023 = 1,023 <- w=1023, L=1

The first line is exactly the 1,024 pieces the fold gave. The second line is the shallow net, where L = 1 makes the exponent d_in·(L-1) = 0, the first factor collapses to 1, and all that survives is w^1 = 1,023 — the shallow net’s w + 1, up to the constant Omega throws away.

The theorem and the toy example are the same fact. The exponent d_in(L-1) is where depth lives — it puts L upstairs — and the trailing w^d_in is everything width gets.

Expressivity is not learnability

This distinction sets up much of the chapter. Being able to represent a function is not the same as gradient descent being able to find it.

A 50-layer plain network can represent more than a 10-layer one and trains worse. The reason is the chain rule again: the gradient has to survive 50 Jacobian multiplications, chain-ruled into a single product, and a product of 50 matrices that each shrink a vector slightly arrives at layer 1 as nothing.

Two devices make depth trainable rather than merely expressive, and both get their own section below.

• A residual connection adds a layer’s input back to its output, so there is always a route the gradient can take that skips the layer entirely (Residual connections the gradient highway).
• Normalization rescales a layer’s activations to a consistent mean and spread before the next layer sees them, so values cannot drift into ranges where gradients die (Normalization and why transformers use layer norm).

The decision tree this chapter expands

So the question that picks a layer is not “how much capacity do I need” but what structure does the input have.

The diagram below is that question as a decision tree, and it is the outline of the rest of the chapter. Each leaf names a layer family. Its bet: line is the structural assumption the layer makes; its fails: line is the failure mode that forces the next section to exist.

flowchart TD
    IN{"What structure does<br/>the input have?"}
    IN -->|"none - each feature<br/>is meaningful on its own"| DENSE["Dense<br/>bet: nothing<br/>fails: params blow up,<br/>no prior, needs data"]
    IN -->|"local patterns that can<br/>appear at any position"| CONV["Convolution<br/>bet: weight sharing<br/>fails: fixed local window,<br/>translation only"]
    IN -->|"a sequence"| R{"how long are the<br/>dependencies?"}
    R -->|"short, streaming,<br/>fixed memory budget"| RNN["Recurrent / LSTM<br/>bet: history compresses<br/>fails: sequential in T,<br/>gradient decays with span"]
    R -->|"long, and you can<br/>train in parallel"| ATT["Attention<br/>bet: read by content<br/>fails: O(n^2) cost,<br/>no position by itself"]

    style DENSE fill:#1d3557,color:#fff
    style CONV fill:#2d6a4f,color:#fff
    style RNN fill:#bc6c25,color:#fff
    style ATT fill:#40916c,color:#fff

A prior, in the fails: lines above, is a structural assumption built into the layer’s wiring rather than learned from data. A strong prior is worth data: whatever the layer assumes for free is something the model does not have to see examples of.

Three of those edges deserve a note before the sections that own them.

Edge 1: none - each feature is meaningful on its own. That is tabular data. Picture a row where column 7 is “days since signup” and column 8 is “country code”. There is no locality between them, no order to run along, and nothing that would still mean the same thing if you shifted it.

There is nothing to bet on, so the dense layer’s total absence of a prior is the correct choice. But gradient-boosted trees are usually a better one still, because the fails: line under Dense is real — params blow up, and with no prior the layer has to learn from data what a prior would have given it for free.

A gradient-boosted tree model is a sum of many small decision trees. Each tree is a chain of threshold tests — “is column 7 above 30?” — that routes a row down to a predicted value. The trees are fitted one at a time, and each new tree is trained to predict the error left over by the sum of all the trees before it.

That model assumes nothing about the columns either, which is what makes it the right comparison. It wins on tables for two reasons: axis-aligned thresholds on raw columns are exactly what heterogeneous tabular features want, and no gradient has to travel through a deep stack to fit them.

Edge 2: long, and you can train in parallel. This edge is about training throughput, not accuracy. Attention scores every position against every other in a single matrix multiplication, so a whole sequence is processed in one pass. Recurrence must walk it one step at a time. Lstm and gru replacing the product with a gate prices the difference out.

Edge 3: fixed memory budget. This is the one axis on which recurrence still wins outright. A recurrent network’s state is a fixed-size vector however much history has gone by; attention must keep every past token around. Parameter arithmetic you should be able to do in your head prices that at 512 KiB per token (KiB is kibibyte, 1,024 bytes). Every other branch of this chapter goes against recurrence; that one does not.

2. Dense — the layer with no prior

Start at the decision tree’s first leaf: the layer that assumes nothing, and pays the full cost of that.

A dense layer — also called fully connected or linear — computes y = Wx + b. Read that as “the output vector y is the weight matrix W times the input vector x, plus the bias vector b”. A bias is one constant per output, added to shift that output up or down.

Every input is connected to every output, so a layer with in inputs and out outputs holds

in·out weights (one per input-output pair)  +  out biases  =  in·out + out parameters

It is the only layer in this chapter that assumes nothing about the input, which is both why it is universal and why it is expensive.

The assumption it does make is easy to miss. A dense layer has a separate weight for every input coordinate, so it assumes each coordinate means the same thing every time — that position 7 of the input vector is always the same quantity.

For a table of features that is exactly true: column 7 is always “days since signup”. For an image it is false, and the failure is total.

The number that killed dense-only vision. Flatten a 224×224 RGB image into one long vector and you get

224 · 224 · 3                  = 150,528 input values
150,528 · 1,000 + 1,000        = 150,529,000 parameters in one 1,000-unit dense layer

That is 150M weights in the first layer, before any learning has happened.

The second failure is worse than the cost. Because there is a separate weight per pixel position, a cat shifted three pixels right lands on a completely different set of weights. The layer must learn “cat” independently at every location.

Both problems have the same root — no weight sharing — and one fix.

3. Convolution — weight sharing, and what it buys

The fix is weight sharing, and the layer built around it is convolution.

A convolution slides a small kernel (equivalently a filter) — a little grid of weights, typically 3×3 — over the input. It reuses the same weights at every position and produces one output number per position.

The block below has two parts: the formula for one output value, and the parameter count that follows from it.

y[i] = sum_k  w[k] · x[i + k]        (cross-correlation; the ML "convolution")

3x3 kernel, C_in = 3, C_out = 64:  3·3·3·64 + 64 = 1,792 parameters
same job as the dense layer above:               150,529,000 parameters
                                                 ~84,000x fewer

Read the formula as: output position i is the sum, over every kernel tap k (one weight in the grid), of that tap’s weight times the input value k steps along from i.

Two symbols in the count. C_in is the number of input channels — 3 for red, green and blue. C_out is the number of different kernels the layer learns, each producing its own output channel.

The parameter count is then kernel height × kernel width × input channels × output channels, plus one bias per output channel:

3 · 3 · 3 · 64  = 1,728 weights
              + 64 biases
              = 1,792 parameters        vs 150,529,000 for the dense layer -> 84,000x fewer

The output of a convolution layer is called a feature map: one grid per output channel, where each entry says how strongly that kernel’s pattern matched at that position.

The assumption. Convolution assumes two things:

1. The patterns worth detecting are local, so a 3×3 window is enough to see one.
2. A pattern means the same thing wherever it appears, so one set of weights can serve every position.

Images satisfy both. An edge is an edge whether it is in the corner or the centre, and pixels far apart are related only through the pixels between them.

Translation equivariance falls straight out of the sharing

Equivariance means “shift the input, and the output shifts the same way”. The three-step chain below proves that convolution has it, using nothing but the fact that the weights do not depend on position.

Let T_s mean “shift by s positions”, and let * mean convolution:

(w * T_s x)[i] = sum_k w[k] · x[i + k - s] = (w * x)[i - s] = T_s(w * x)[i]

Read that chain left to right.

• (w * T_s x)[i] is “convolve the input after shifting it, then look at position i”.
• sum_k w[k] · x[i + k - s] is what that expands to — the same weights, reading the original input s steps earlier.
• (w * x)[i - s] says that sum is just the unshifted convolution read at position i - s.
• T_s(w * x)[i] says that is the same as shifting the unshifted convolution.

So convolving the shifted input gives exactly the shifted convolution of the original. A dense layer cannot do this, because the middle equality needs w[k] to be independent of i — and “the weights do not depend on position” is weight sharing.

Equivariance is not a bonus property of convolution. It is the same fact as weight sharing, written differently.

Note it is equi-variance, not in-variance. Invariance would mean the output does not change at all when the input shifts. Equivariance means it changes in step: the feature map moves. Invariance — the class score not changing when the cat moves — comes from pooling or global averaging on top, which is Pooling.

Two more knobs: stride and padding

Both appear in every formula for the rest of this section, so pin them down now.

Stride is how far the kernel hops between evaluations. Stride 1 evaluates at every position. Stride 2 evaluates at every other one, halving the output’s height and width.

Padding is the ring of extra values (usually zeros) added around the border, so the kernel has something to read when it is centred on an edge pixel. same padding is whatever amount makes the output the same size as the input.

Receptive field arithmetic

How much of the input can one output unit see? That number decides whether your network can see the object at all, and it comes from tracking two numbers about the current feature map:

• r — the receptive field: how many input pixels feed one unit of this map.
• j — the jump: how many input pixels lie between two adjacent units of this map. It happens to equal the product of all strides so far, but that is how you compute it, not what it means.

r_l = r_{l-1} + (k_l - 1) · j_{l-1}          j_l = j_{l-1} · s_l

Read the subscripts as layer indices: r_l is the receptive field after layer l, and r_{l-1} the one before it. k_l is that layer’s kernel size and s_l its stride.

Both halves fall straight out of what j means, so derive them rather than memorizing them.

Why r grows by (k - 1) · j. A k-tap kernel reads k adjacent units of the previous map. Beyond the one it is centred on, it reaches k - 1 neighbours in each direction — and each neighbour is j_{l-1} input pixels away, because that is what j means. So the reach grows by (k - 1) · j_{l-1}.

Why j multiplies by s. A stride of s keeps only every s-th output. Two units that survive are therefore s times further apart, in input pixels, than two adjacent units were before.

The ordering trap

This is the common mistake. r is updated using the jump entering the layer. Only then does j update.

The awkward case is a strided layer, because it has to use its own pre-stride jump. Take P1 in the diagram below — a 2×2 pool at stride 2:

arrives with          r = 5,  j = 1
update r first:       r = 5 + (2 - 1) · 1 = 6      <- uses j = 1, the jump on the way IN
then update j:        j = 1 · 2 = 2

Advance j first and you would compute r = 5 + 1·2 = 7, which is wrong there and stays wrong at every layer after it.

The non-strided case is easier to check against the picture. L3 is a 3×3 stride-1 conv following P1:

r_L3 = r_P1 + (3 - 1) · j_P1 = 6 + 2·2 = 10

which is the value printed in the L3 node.

Running the recurrence down a real stack

The diagram below applies those two lines to eleven layers in a row. Every node shows the running r and j after that layer. Note how much each conv adds to r, and how that amount changes each time a pool doubles j.

flowchart LR
    L0["input<br/>r=1, j=1"] --> L1["conv 3x3 s1<br/>r=3, j=1"]
    L1 --> L2["conv 3x3 s1<br/>r=5, j=1"]
    L2 --> P1["pool 2x2 s2<br/>r=6, j=2"]
    P1 --> L3["conv 3x3 s1<br/>r=10, j=2"]
    L3 --> L4["conv 3x3 s1<br/>r=14, j=2"]
    L4 --> P2["pool 2x2 s2<br/>r=16, j=4"]
    P2 --> L5["conv 3x3 s1<br/>r=24, j=4"]
    L5 --> L6["conv 3x3 s1<br/>r=32, j=4"]
    L6 --> P3["pool 2x2 s2<br/>r=36, j=8"]
    P3 --> L7["conv 3x3 s1<br/>r=52, j=8"]
    L7 --> L8["conv 3x3 s1<br/>r=68, j=8"]

    style P1 fill:#bc6c25,color:#fff
    style P2 fill:#bc6c25,color:#fff
    style P3 fill:#bc6c25,color:#fff
    style L8 fill:#2d6a4f,color:#fff

Each node reads conv 3x3 s1 or pool 2x2 s2 — a 3×3 kernel at stride 1, or a 2×2 pooling window at stride 2. The orange nodes are the stride-2 pools, the only layers that change j; the green node is the final layer whose receptive field the section is chasing.

Read the increments a 3×3 conv makes to r:

before any pool   j = 1   ->  each conv adds (3-1)·1 =  2      1 -> 3 -> 5
after one pool    j = 2   ->  each conv adds (3-1)·2 =  4      6 -> 10 -> 14
after two pools   j = 4   ->  each conv adds (3-1)·4 =  8     16 -> 24 -> 32
after three pools j = 8   ->  each conv adds (3-1)·8 = 16     36 -> 52 -> 68

Now look at what the pools themselves contribute. A stride-2 2×2 pool adds only (k-1)·j = 1·j = j pixels directly: 1 at P1, 2 at P2, 4 at P3 — exactly the 5→6, 14→16 and 32→36 steps in the diagram. Tiny.

But it doubles j, which multiplies the contribution of every layer downstream. That is why downsampling, not depth, is what buys global context cheaply.

Put a number on “cheaply”. Eight convs plus three pools reach 68 pixels. Reaching 224 with stride-1 3×3 convs alone — where every conv adds a flat 2 — would take ceil((224 - 1) / 2) = 112 layers. (ceil means round up to the next whole layer; you cannot build a fraction of one.)

Padding never enters the recurrence. same padding changes the output size — it stops the map shrinking by k-1 per layer — but a unit still reads k neighbours at spacing j, so r and j are untouched. Receptive field is a statement about which input positions a unit depends on, not about how many units there are.

Dilation is the alternative when you cannot afford to lose resolution. A dilated kernel skips d-1 positions between its taps, so a k-tap kernel covers d·(k-1) + 1 pixels while still holding only k weights.

Stack four 3×3 convs at dilations 1, 2, 4, 8, all stride 1 so j stays at 1:

d=1   effective size 1·2+1 = 3    r = 1 +  2 =  3
d=2   effective size 2·2+1 = 5    r = 3 +  4 =  7
d=4   effective size 4·2+1 = 9    r = 7 +  8 = 15
d=8   effective size 8·2+1 = 17   r = 15 + 16 = 31

31 pixels of receptive field in four layers, at full resolution, with no pooling anywhere.

The code below is the same two lines, run over the whole stack. The only thing to look at is the order of the two assignments inside the loop.

def receptive_field(layers):
    """layers: list of (kernel, stride). Returns (r, j) after each layer."""
    r, j, out = 1, 1, []
    for k, s in layers:
        r = r + (k - 1) * j
        j = j * s
        out.append((r, j))
    return out

vgg_block = [(3, 1), (3, 1), (2, 2), (3, 1), (3, 1), (2, 2),
             (3, 1), (3, 1), (2, 2), (3, 1), (3, 1)]
assert receptive_field(vgg_block)[-1] == (68, 8)

vgg_block is the stack drawn in the diagram, written in the layer pattern of VGG — an early convolutional image network built entirely from 3×3 convolutions and 2×2 pools.

The order of the two assignments is the whole point. r = r + (k - 1) * j reads the old j, and only the next line advances it. Swap the two lines and vgg_block comes out at 75 instead of 68. The assert — a statement that raises an error if its condition is false — exists to catch exactly that swap.

Worked example: a ResNet stem

Now run the same recurrence on an architecture that is not in the diagram: the opening layers of a ResNet, the residual network of Residual connections the gradient highway. Its stem — the first few layers before the main body — is 7×7 s2, then 3×3 s1, then 3×3 s2.

Each row below applies r = r + (k-1)·j with the incoming j, then updates j:

layer     r                        j
(7, 2)     1 + (7-1)·1 =  7        1·2 = 2
(3, 1)     7 + (3-1)·2 = 11        2·1 = 2
(3, 2)    11 + (3-1)·2 = 15        2·2 = 4

Three layers reach 15 pixels, where three stride-1 3×3 convs would reach only 7. The reason: the second 3×3 is worth 4 pixels rather than 2 because a stride came before it and doubled j.

What breaks when the assumption breaks

The window is fixed and local, so a dependency wider than the receptive field is invisible. The classic bug is a segmentation model — one that labels every pixel with the object it belongs to — that cannot use context outside 68 pixels and confidently mislabels a large uniform region, because from inside the window a patch of sky and a patch of sea look identical.

Equivariance is to translation only. Rotate or rescale the input and nothing is shared, because the same weights no longer line up with the same pattern. That is why augmentation exists: you train on deliberately rotated, scaled and flipped copies of each image, so the model learns from data the invariances the layer does not give it for free.

4. Pooling

The receptive-field arithmetic just showed that downsampling is what buys global context cheaply. The cheapest way to downsample destroys information to do it.

Pooling downsamples a feature map by taking the maximum or the mean over a small window. A 2×2 window at stride 2 replaces every block of four numbers with one, halving height and width. It has no parameters at all.

It has three effects, in order of importance:

1. It multiplies j, so every later layer’s receptive field grows faster — the Convolution weight sharing and what it buys arithmetic.
2. It cuts activation memory 4× per 2×2 stage, because four numbers become one.
3. It makes the representation locally invariant to small shifts. The max over a window does not change if the peak moves inside that window.

The assumption is one sentence. Pooling bets that which position inside the window held the peak does not matter, only that the peak was there. Where that is true it is free downsampling. Where it is false it is destruction.

Three ways to halve a feature map, compared on what they keep and what they cost:

The third column is the parameterized alternative. Instead of a fixed max or mean, run an ordinary convolution at stride 2 and let the network learn what to keep, at a cost of k^2·C_in·C_out weights.

Failure mode: pooling discards where. For classification that is the point — you want “there is a cat”, not “there is a cat at pixel (91, 140)”. For detection or segmentation it is a loss you spend the rest of the architecture recovering.

Two standard recoveries:

• Skip connections, as in U-Net — an architecture that wires each downsampling stage’s full-resolution map directly across to the matching upsampling stage, so the fine detail is still available when the output is reassembled.
• Feature pyramids, which make predictions from several resolutions at once rather than only the coarsest.

Modern networks often replace pooling with stride-2 convolutions, so the downsampling is learned rather than fixed. They also replace the dense classifier head with global average pooling — collapsing each channel to one number by averaging its whole map.

That last swap is worth the arithmetic. On a 7×7×512 feature map:

flatten then dense:   7·7·512 = 25,088  ->  25,088 · 1,000 = 25,088,000 params
global avg pool:      7·7·512 ->    512  ->     512 · 1,000 =    512,000 params

A 25.1M-parameter head becomes a 512K one, and it no longer cares what size the input image was.

5. Recurrence, and the vanishing gradient derived

Now move from images to sequences. The classic sequence layer cannot learn long-range dependencies, and the reason is an exponential you can derive rather than an empirical quirk.

A recurrent neural network (RNN) reads a sequence one step at a time and compresses everything it has seen into a fixed-size vector called the hidden state:

h_t = tanh(W · h_{t-1} + U · x_t + b)

Read that as: the state at step t is a squashing function applied to (the old state h_{t-1} times a weight matrix W) plus (the new input x_t times a second weight matrix U) plus a bias.

tanh is the hyperbolic tangent, an S-shaped function that maps any number into the range (-1, 1). It is what keeps the state from growing without bound.

The same two matrices W and U are reused at every step. That is weight sharing along time, exactly as convolution shares weights along space.

The assumption: everything the future needs to know about the past fits in one fixed-size vector, and the sequence must be read in order. That is the right assumption for streaming and the wrong one for long dependencies. The rest of the section derives why.

The gradient is a product, so it is an exponential

Backpropagation through time (BPTT) is ordinary backpropagation applied to the unrolled sequence, treating each time step as a layer. So it multiplies one Jacobian per step.

Three lines, and the third is the one that matters:

dh_t / dh_{t-1} = D_t · W        where D_t = diag(1 - tanh^2(a_t)),  entries in (0, 1]

dL_T / dh_k = (dL_T / dh_T) · prod_{t=k+1..T} (D_t · W)

|| dL_T / dh_k ||  <=  ( max|tanh'| · ||W|| )^(T-k)   :=  s^(T-k)

Line 1 — one step’s Jacobian. It is the weight matrix W scaled by D_t, a diagonal matrix (zero everywhere except the diagonal) whose entries are the derivative of tanh at that step. Each entry is between 0 and 1.

Line 2 — chaining the steps. The gradient of the loss at the final step T, with respect to the state at an earlier step k, is a product (prod) of all the Jacobians in between. Go back 100 steps and you multiply 100 matrices.

Line 3 — bounding that product. ||·|| is a norm, one number measuring how much a matrix can stretch a vector. := means “define this to be”, naming the per-step factor s. The whole product is bounded by s raised to the number of steps.

So the gradient over a span of T-k steps is an exponential in that span. If s < 1 it vanishes; if s > 1 it explodes. There is no setting of s that stays put for all T.

Putting numbers on the vanishing side

Work out a realistic s.

The derivative of tanh is tanh'(z) = 1 - tanh(z)^2. It is 1 at z = 0, but only 1 - tanh(1)^2 = 0.42 at z = 1 — and healthy pre-activations (the values z entering the nonlinearity) sit around 1.

So take max|tanh'| ≈ 0.42 and a generously large ||W|| = 2. Then s ≈ 0.42 · 2 = 0.84.

Now raise 0.84 to a few span lengths:

span     0.84^span      what it means
  10      1.7e-1        learnable
  25      1.3e-2        marginal
  50      1.6e-4        swamped by nearer-term gradients
 100      2.7e-8        numerically zero in fp32 alongside gradients of order 1e-2

1.7e-1 is scientific notation for 0.17. fp32 is 32-bit floating point, the standard number format for training.

Read the last row carefully. A gradient of 2.7e-8 is not small — it is nothing, sitting next to nearby-in-time gradients around 1e-2. Six orders of magnitude down, it contributes nothing to the update.

A dependency the network cannot get gradient for is a dependency it cannot learn, no matter how long you train. This is not slow convergence. It is a signal that never arrives.

The exploding side, which is the easy one

Flip s above 1 and the same exponential runs the other way: s = 1.2 over 100 steps gives 1.2^100 = 8.3e7. One step at that magnitude blows the weights up and the loss goes to NaN.

The fix is gradient clipping: whenever the gradient’s norm exceeds a threshold (conventionally around 1.0), rescale the whole gradient down to that threshold.

Note the asymmetry. Clipping can rescale a too-large gradient back into range. Nothing can rescale a gradient that has already underflowed to zero — there is no direction information left in it to recover.

6. LSTM and GRU — replacing the product with a gate

The vanishing gradient was a product of Jacobians, so the fix is to change what the gradient is a product of: give the state a path through time that is additive rather than a repeated matrix multiplication.

The long short-term memory (LSTM) cell does it with a second state vector and three gates. A gate is a vector of numbers between 0 and 1 that the network computes itself and then multiplies into something else, deciding position by position how much to let through.

The six lines below are the entire cell. The top four compute the gates and the candidate content; the bottom two are where the state actually updates.

f_t = sigmoid(W_f · [h_{t-1}, x_t] + b_f)      forget: what to keep from the old cell
i_t = sigmoid(W_i · [h_{t-1}, x_t] + b_i)      input:  how much new content to write
g_t = tanh   (W_g · [h_{t-1}, x_t] + b_g)      candidate content
o_t = sigmoid(W_o · [h_{t-1}, x_t] + b_o)      output: what to expose as h_t

c_t = f_t (*) c_{t-1}  +  i_t (*) g_t          (*) = elementwise
h_t = o_t (*) tanh(c_t)

Four things in that block need a gloss.

sigmoid is the S-shaped function that maps any number into (0, 1), so its output reads directly as “what fraction to keep”.

c_t is the cell state — the long-term memory running along the top of the cell. h_t is the hidden state, which is what the rest of the network actually sees. The LSTM has two state vectors where the plain RNN had one.

Elementwise, written (*) here and ⊙ in prose, means multiply matching positions of two equal-length vectors, one pair at a time. No mixing across positions.

[h_{t-1}, x_t] is the two vectors stacked end to end into one vector of length d_h + d_in, where d_h is the state width and d_in the input width. So each W is d_h × (d_h + d_in): one matrix reading the state and the new input together, rather than two matrices added.

That last point is where the parameter count comes from:

one gate's W:   d_h · (d_h + d_in)  =  d_h^2 + d_in·d_h  weights
one gate's b:                          d_h               biases
four gates:     4 · (d_in·d_h + d_h^2 + d_h)

All four lines are the same computation with different weights. Only the nonlinearity differs, and it differs for a reason: the three sigmoids produce numbers in (0,1) because they are fractions to keep, while g_t’s tanh produces numbers in (-1,1) because it is a value to write.

Why the cell state is called a highway

The diagram below prices the same 100 steps two ways. Top row: the plain RNN. Bottom row: the LSTM cell state. Look at what sits on each arrow, and at the two end nodes.

flowchart LR
    subgraph V["Vanilla RNN - gradient path is a matrix product"]
        A1["h_t-2"] -->|"D·W"| A2["h_t-1"] -->|"D·W"| A3["h_t"]
        A3 --> AN["prod of Jacobians<br/>0.84^100 = 2.7e-8"]
    end
    subgraph L["LSTM - cell state is an additive highway"]
        B1["c_t-2"] -->|"x f_t-1"| B2["c_t-1"] -->|"x f_t"| B3["c_t"]
        B3 --> BN["prod of forget gates<br/>0.99^100 = 0.37"]
    end

    style AN fill:#9d0208,color:#fff
    style BN fill:#2d6a4f,color:#fff

Top row: every edge is a multiplication by D·W, the Recurrence and the vanishing gradient derived Jacobian, so the chain terminates in 0.84^100 = 2.7e-8. Bottom row: every edge is a multiplication by a forget gate, so the chain terminates in 0.99^100 = 0.37.

But the diagram calls the bottom row an “additive highway” while its own edges are multiplications. Reconciling those two is the whole point of this section.

“Additive” describes how new content enters the cell. Look again at c_t = f_t ⊙ c_{t-1} + i_t ⊙ g_t. New content is added on, so writing to the cell does not push the old contents through a transformation the way the vanilla RNN’s D·W does.

The gradient path is still a product in both rows. The difference is what it is a product of:

• Vanilla RNN: a product of D·W. One matrix, shared across all steps, whose scale the network cannot pick freely — it is the same W that has to do the actual computing.
• LSTM: a product of forget gates. Numbers between 0 and 1 that the network computes per step and per channel, independently of the content path.

A vanilla RNN cannot make D·W equal 1. An LSTM can hold f_t near 0.99 on exactly the channels that need to remember. That is the entire difference between 2.7e-8 and 0.37 over the same 100 steps.

The derivation is one line. Differentiate c_t = f_t ⊙ c_{t-1} + i_t ⊙ g_t with respect to c_{t-1} and the second term drops out (it does not contain c_{t-1}), leaving

dc_t / dc_{t-1} = diag(f_t)

No weight matrix. No tanh'. The gradient across T steps is prod_t f_t, a product of numbers the network chooses:

f = 0.99   ->  0.99^100 = 0.366     signal survives 100 steps
f = 0.95   ->  0.95^100 = 0.0059    marginal
f = 0.50   ->  0.50^100 = 7.9e-31   dead
vanilla                  = 2.7e-8   dead, and not adjustable

Why you initialize the forget-gate bias to 1.0

That table is also the reason for a standard trick: initialize b_f, the forget-gate bias, to 1.0 or 2.0 rather than 0.

Follow the arithmetic. At initialization the weights are small random numbers, so W_f · [h, x] is near zero and the gate is essentially sigmoid(b_f):

b_f = 0.0   ->  sigmoid(0) = 0.50    the 0.50 row above: memory gone in ~10 steps
b_f = 1.0   ->  sigmoid(1) = 0.73
b_f = 2.0   ->  sigmoid(2) = 0.88

With b_f = 0 the untrained cell sits on the 0.50 row and destroys its own memory long before any gradient could teach it not to — the signal that would carry that lesson has already vanished.

You are not setting the model’s memory. You are setting where gradient descent starts looking for it — and the search only works if it starts from the survivable end of that table.

GRU: the same idea with one less gate

The gated recurrent unit (GRU) does the LSTM’s job with one state vector instead of two and two gates instead of three.

The update gate z_t interpolates directly:

h_t = (1 - z_t) ⊙ h_{t-1} + z_t ⊙ g_t

One number now does the job of both f_t and i_t. Note the cost: how much you forget and how much you write are forced to trade off (1 - z against z) rather than being independently controllable. That is the GRU’s one real loss of expressiveness.

The reset gate r_t multiplies h_{t-1} before it enters the candidate:

g_t = tanh(W_g · [r_t ⊙ h_{t-1}, x_t] + b_g)

This lets the cell compute new content while ignoring its history, without also erasing that history — the two are separate decisions.

Merging h and c then costs nothing structurally. The LSTM’s output gate o_t existed only to decide how much of the hidden cell to expose as h_t; with a single state there is nothing to hide.

Three weight matrices instead of four is where the 25% saving in the table comes from:

The last row substituted, so you can check it:

d_in·d_h + d_h^2 + d_h  =  262,144 + 262,144 + 512  =  524,800
LSTM   4 · 524,800 = 2,099,200
GRU    3 · 524,800 = 1,574,400        3/4 of the LSTM, so 25% fewer

What actually ended recurrence

It was not gradients. It was parallelism.

h_t requires h_{t-1}, which requires h_{t-2}, and so on. Training over a length-1,000 sequence is 1,000 sequential steps, no matter how many graphics processing units (GPUs) you own — you cannot start step 500 before step 499 finishes.

Attention computes all positions in one matrix multiplication, so the whole sequence trains at once. That branch of Why depth beats width’s decision tree — long, and you can train in parallel — not accuracy, is why transformers took over.

Recurrence still wins the fixed memory budget branch.

(h, c) is a fixed-size pair of vectors, so serving cost per token is O(1) in memory — constant, independent of how long the stream has been running.

A transformer’s key-value (KV) cache — the stored keys and values of every token seen so far, kept so that each decoding step does not recompute them — is O(n), growing linearly with the number of tokens. At d = 4,096 that is 512 KiB per token (Parameter arithmetic you should be able to do in your head), so 4 GiB by 8k tokens, and unbounded after that.

Recurrence pays for the whole history once, at compression time. Attention pays for it again at every step. That is why recurrent and state-space architectures keep reappearing for streaming workloads even though the training argument above is settled.

7. Attention, derived as content-based lookup

Attention can be built from a dictionary lookup by removing, one at a time, every property that blocks gradient descent. Derived this way, each piece of the final formula is forced rather than chosen.

Attention lets every position in a sequence read from every other position, choosing what to read by content rather than by distance.

Start from an ordinary dictionary lookup. Four steps below, and the last line is the finished formula:

step 0   hard lookup:   out = V[j*]  where j* = argmax_j 1[k_j == q]
                        -> not differentiable, and needs exact equality

step 1   soft match:    score_j = q · k_j            similarity, not equality
step 2   soft select:   a = softmax(score)           normalized, differentiable
step 3   soft read:     out = sum_j a_j · v_j        a convex combination of values

attention(Q, K, V) = softmax( Q·K^T / sqrt(d_k) ) · V

Step 0 reads: return the value V at whichever index j* has a key exactly equal to the query. 1[...] is an indicator — 1 when the condition holds, 0 otherwise — and argmax_j means “the j that maximizes”.

Step 0 is a Python dict lookup, and it is useless to gradient descent. It is a step function: nudge a key slightly and nothing changes, until suddenly everything does. There is no gradient to learn from.

Step 3’s “convex combination” means a weighted average whose weights are non-negative and sum to 1. The consequence is that the output always lies somewhere inside the set of values being averaged, never outside it.

In the final line, Q, K and V are the queries, keys and values for every position, stacked into matrices. K^T is K transposed (rows and columns swapped), so Q·K^T compares every query against every key in one multiplication. d_k is the length of a single query or key vector.

Two of those steps need unpacking before the rest of the section works.

Why a dot product means “relevance”

The dot product q · k multiplies two vectors position by position and sums the results.

It also equals ||q|| · ||k|| · cos(theta), where ||q|| is the length of q and theta is the angle between the two vectors. So the score is large exactly when query and key point in the same direction, and negative when they oppose.

Exact equality has been replaced by alignment: relevance is an angle.

The two magnitudes ||q|| and ||k|| ride along as a multiplier, and that multiplier is exactly what sqrt(d_k) below has to keep under control. The geometry is derived in Vectors norms and what cosine similarity actually measures.

Why softmax

Softmax is softmax(s)_j = e^{s_j} / sum_m e^{s_m}. Exponentiate every score so everything is positive, then divide by the total so the weights sum to 1.

It is the smooth stand-in for argmax. The largest score still takes the most weight, but every score takes some, and — the part that matters — every score has a nonzero derivative. That is precisely the property step 0 lacked.

Every step is forced, which is why this is a derivation and not a design. Differentiability requires soft selection. Soft selection requires normalized weights. Normalizing unbounded scores requires a softmax.

Where q, k and v come from

Each token has exactly one vector x_i. Three learned matrices turn it into three:

q_i = W_q x_i        k_i = W_k x_i        v_i = W_v x_i

The same three matrices are used at every position. That is why one (n, d) × (d, d) matrix multiplication produces all n queries at once, and why attention is a handful of matrix multiplications rather than a loop over positions.

Q, K and V are not three arbitrary projections. They are the three roles a lookup has:

• q — what I am asking for
• k — what I advertise
• v — what I hand over

Q and K have to be separate matrices so that a token can ask for something different from what it offers. Intuitively: a verb hunting for its subject wants to advertise “I am a verb” while asking for “a noun before me”. Tie W_q = W_k and one vector would have to do both jobs.

The assumption, and it is unusually weak

Attention assumes only that relevance between two positions can be read off a comparison of their content. Nothing about distance. Nothing about order. Nothing about locality.

That weakness is its strength on language and its cost everywhere else. Three consequences, all of which show up later in this section:

1. With no prior, it has to learn structure from data rather than getting it for free.
2. It pays for comparing all n positions against all n — the O(n^2) cost.
3. It cannot tell word order apart at all unless position is supplied to it explicitly.

A lookup, worked

To turn every symbol above into a number, take the four tokens the cat sat on with d_k = 4 and compute what sat reads. The numbers are illustrative, but they have the shape of what a trained head holds.

The block below has three parts: the query, then a key and a value for each of the four tokens, then the four-column computation — dot product, scale, exponentiate, normalize.

q_sat = (1, 2, 0, 1)          what "sat" is asking for

        key                    value
the   (0, 0, 1, 0)      (0.1, 0.0, 0.2, 0.0)
cat   (1, 2, 0, 1)      (0.9, 0.2, 0.1, 0.4)
sat   (0, 1, 0, 1)      (0.2, 0.8, 0.0, 0.3)
on    (1, 0, 0, 0)      (0.0, 0.1, 0.7, 0.1)

        q_sat · k_j        / sqrt(4) = / 2       exp        softmax
the     0+0+0+0 =  0             0.0            1.000        0.037
cat     1+4+0+1 =  6             3.0           20.086        0.738
sat     0+2+0+1 =  3             1.5            4.482        0.165
on      1+0+0+0 =  1             0.5            1.649        0.061
                                                ------       -----
                                                27.216       1.000

out = 0.037·v_the + 0.738·v_cat + 0.165·v_sat + 0.061·v_on
    = (0.701, 0.285, 0.124, 0.351)

Follow one row across. For cat, the dot product of q_sat = (1, 2, 0, 1) with k_cat = (1, 2, 0, 1) is 1 + 4 + 0 + 1 = 6; dividing by sqrt(4) = 2 gives 3.0; exponentiating gives 20.086; dividing by the column total 27.216 gives the weight 0.738.

Now read what happened.

k_cat is the key most aligned with q_sat — it is q_sat, so its dot product is the largest of the four — and 74% of the output weight went there. The result (0.701, 0.285, 0.124, 0.351) sits most of the way toward v_cat = (0.9, 0.2, 0.1, 0.4), dragged slightly toward the other three.

sat read mostly from cat. That is what “the verb attends to its subject” is, numerically.

Two details in that row are worth noticing.

sat gave 0.165 of its weight to itself. This is self-attention — attention where queries, keys and values all come from the same sequence — and it includes the query’s own position. A token routinely reads part of its own value back.

Nothing here is a hard choice. the still contributed 3.7%, even though it matched badly. That nonzero contribution is what gives the’s score a gradient, which is the whole reason for step 2.

The same row with a causal mask

The row above is unmasked, which is what an encoder would compute. Now add the causal mask of Encoder decoder encoder decoder — the rule that a position may read only from itself and earlier positions, never later ones.

on is at position 3, which is in sat’s future, so its score is set to -inf (negative infinity) before the softmax. Then e^(-inf) = 0, and it leaves the denominator entirely:

denominator  27.216 - 1.649 = 25.567
the    1.000 / 25.567 = 0.039
cat   20.086 / 25.567 = 0.786
sat    4.482 / 25.567 = 0.175
on                      0        (masked out before the softmax)
                        -----
                        1.000

The three survivors renormalize and still sum to 1.

That is why the mask goes in before the softmax rather than zeroing weights after it. Zero out on’s 0.061 afterwards and the row sums to 0.939 — the position is reading only 94% of a value vector, and every downstream layer sees a systematically shrunken input.

The same row without the / 2

Now put the mask aside and delete the scaling instead. The same four scores (0, 6, 3, 1) go straight into the softmax:

exp:      1.000   403.429   20.086   2.718     sum = 427.23
softmax:  0.002     0.944    0.047   0.006

That is already nearly one-hot — nearly all the weight on a single entry and almost none anywhere else — and this is only d_k = 4, where the scale factor is a mere 2. The next subsection is what the same effect looks like at d_k = 128.

The diagram below is the finished mechanism as a pipeline. Every node is a line of the derivation above; the green node is the one the next subsection exists to justify.

flowchart LR
    X["token i"] --> Q["Q_i = W_q x_i<br/>what am I looking for"]
    XS["all tokens j"] --> K["K_j = W_k x_j<br/>what do I advertise"]
    XS --> V["V_j = W_v x_j<br/>what do I hand over"]
    Q --> S["scores = Q·K^T<br/>variance = d_k"]
    K --> S
    S --> SC["divide by sqrt(d_k)<br/>restores unit variance"]
    SC --> SM["softmax<br/>weights sum to 1"]
    SM --> O["out = sum_j a_j · V_j"]
    V --> O

    style SC fill:#2d6a4f,color:#fff
    style O fill:#1d3557,color:#fff

Walk it once, node by node.

• The three projection nodes — W_q x_i, W_k x_j, W_v x_j — are one token vector sent through three different learned matrices to play three different roles.
• scores = Q·K^T is the all-against-all comparison. Its second line, variance = d_k, is the problem the next node exists to solve.
• divide by sqrt(d_k) shrinks the scores back to a spread the softmax can work with.
• softmax turns them into weights that sum to 1.
• out = sum_j a_j · V_j reads the values out using those weights.

Why sqrt(d_k) exists — the derivation

The scaling factor falls out of the statistics of a dot product, and removing it breaks the layer in the forward direction and the backward direction at once.

Step 1: the spread of a raw score

Two words first. Variance measures how spread out a set of numbers is; its square root is the standard deviation (sd), which is in the same units as the numbers themselves. E[...] is the expectation, the long-run average value.

Assume the components of q and k are independent, with mean 0 and variance 1 — which is roughly what initialization gives you. Then:

score = q · k = sum_{i=1..d_k} q_i · k_i
E[score]   = 0
Var(score) = sum_i Var(q_i k_i) = d_k · 1 · 1 = d_k
sd(score)  = sqrt(d_k)

Line by line: the score averages zero (each term is a product of two independent zero-mean numbers). Its variance is the sum of d_k independent terms, each of variance 1, so the variance is d_k. Its typical spread is therefore sqrt(d_k), read “root d_k”.

The score’s spread grows with the head dimension, purely from summing more terms. Nothing about the content changed; you just added more coordinates to the sum.

At d_k = 128 that means sd = sqrt(128) = 11.3. Two scores drawn from that distribution routinely sit a full 11.3 apart, and a softmax does not survive an 11.3 gap.

One clarification, because it trips people up. Under the stated model there is no “matching key” anywhere — just scores scattered around zero with spread sqrt(d_k). A key that genuinely matched, k = q, would score ||q||^2, which averages d_k = 128. That is an order of magnitude further out again, so the estimate below is on the mild side.

Step 2: what an 11.3 gap does to the softmax

Take the smallest case that shows the problem: two competing keys, separated by one standard deviation of a single score.

With only two competitors the softmax collapses to the logistic function, p_max = 1 / (1 + e^-s), where s is the separation between the two scores. A logit is a raw score before the softmax turns it into a probability, so s is the difference between two logits.

                     score separation   p_max      dp/ds = p(1-p)
unscaled, d_k = 128              11.3   0.99999            1.2e-5
scaled by sqrt(128)               1.0   0.7311             0.1966

gradient through the attention weights: ~16,000x smaller unscaled

Both rows use the same convention: the separation is set to the one-score standard deviation, which is sqrt(d_k) = 11.3 unscaled and exactly 1.0 after scaling. So the ratio between the rows is precisely what the sqrt(d_k) division removes.

dp/ds is how much the winner’s weight moves per unit change in the separation. For a two-way softmax that derivative is p(1-p), and that derivative is the gradient signal reaching the score. At 1.2e-5, the score has essentially stopped being trainable.

(If you prefer the spread of the difference of two scores rather than of one score, it is wider: two independent scores give Var(gap) = 2·d_k, so sd(gap) = sqrt(256) = 16.0 at d_k = 128. That makes the unscaled row worse, not better. The table is the conservative reading.)

Step 3: both directions break at once

Forward — the layer stops blending. An almost-one-hot distribution is a hard lookup. A token can read from exactly one position and cannot combine evidence from several, which is the one thing the whole derivation was built to allow.

Backward — the pattern freezes. p(1-p) ≈ 1e-5 means the score gradients are effectively zero. The attention pattern stays in whatever configuration random initialization produced and never trains out of it.

The symptom in practice: a model whose loss drops for a few hundred steps — that is the feed-forward network learning — and then flatlines.

Why sqrt(d_k) and not d_k

Because the target is unit variance, and sqrt(d_k) is what hits it:

Var( score / sqrt(d_k) ) = d_k / d_k = 1        <- correct
Var( score / d_k )       = d_k / d_k^2 = 1/d_k  -> sd = 1/sqrt(128) = 0.088

Over-divide and you break the layer the other way. With sd = 0.088 every score is nearly identical, the softmax comes out nearly uniform, and every token attends equally to everything. The layer computes a global average and carries no content signal at all.

Both failures are softmax saturation, in opposite directions, and sqrt(d_k) is the unique scale that avoids both.

The code below is the finished layer. Three of its lines are the three steps just derived, in order.

import math
import torch
import torch.nn.functional as F

def attention(q, k, v, causal=True):
    """q,k,v: (batch, heads, seq, d_head)"""
    scores = q @ k.transpose(-2, -1) / math.sqrt(q.size(-1))   # the derivation above
    if causal:
        n = q.size(-2)
        mask = torch.ones(n, n, dtype=torch.bool).triu(1)      # block the future
        scores = scores.masked_fill(mask, float("-inf"))
    return F.softmax(scores, dim=-1) @ v

def saturation_demo(d_k=128):
    q, k = torch.randn(1000, d_k), torch.randn(1000, d_k)
    raw = (q * k).sum(-1)
    return raw.std().item(), (raw / math.sqrt(d_k)).std().item()   # ~11.3, ~1.0

The three load-bearing lines of attention, in order:

• q @ k.transpose(-2, -1) — @ is matrix multiplication — scores every query against every key in one operation. It produces a (seq, seq) matrix, and that matrix is the entire O(n^2) cost of attention.
• / math.sqrt(q.size(-1)) is the sqrt(d_k) scale. Note it is read off the tensor’s own last dimension, so it can never silently disagree with a head width someone configured elsewhere.
• masked_fill(mask, -inf) with triu(1) — the strictly upper triangle, meaning row i from column i+1 onward — sets every future position’s score to -inf. Then e^(-inf) = 0 and the softmax gives it exactly zero weight.

Masking before the softmax rather than zeroing after is what keeps the surviving weights summing to 1. Zero them afterwards and every row is short by however much mass the future was holding.

saturation_demo is the variance claim run as an experiment rather than assumed. It draws 1,000 independent (q, k) pairs from N(0,1) — the standard normal distribution, mean 0 and variance 1 — takes their dot products with (q * k).sum(-1), and returns the standard deviation before and after the scale.

The quantities it is estimating are sqrt(128) = 11.314 and exactly 1.

It is not seeded, and 1,000 samples is not many, so the printed numbers wander. Across twenty runs the first lands anywhere in roughly 10.7 to 11.7, and the second in roughly 0.95 to 1.04. Read the closed forms as the claim and the output as evidence for them, not the other way round. If you want a number you can quote, call torch.manual_seed first or raise the sample count.

Multi-head, and the two things attention cannot do alone

One softmax produces exactly one convex combination. That is the limitation.

If a token needs both the previous word and the subject noun, a single head has to split its weight between them, and the output is a blend of two different things rather than either one of them.

Multi-head attention is the fix. Run h heads — h independent attention computations, each of width d/h — in parallel, concatenate their outputs, and let a fourth matrix W_O mix the concatenation back into one vector of width d.

(That per-head width d/h is the same quantity as d_k above. Parameter arithmetic you should be able to do in your head and the cheat sheet call it d_head. One symbol, three names, all meaning “the width of one head”.)

Why splitting into heads is free

Total parameters are 4·d^2 regardless of h, and the reason h drops out is worth seeing.

W_q has to produce a query for every head. So it maps d inputs to h · (d/h) = d outputs — the h cancels, and the matrix is d × d whatever h you pick. Same for W_k, W_v and W_o.

d = 4,096, h = 32, d_head = 128
4 · 4,096^2 = 4 · 16,777,216 = 67,108,864 parameters

You would get that identical count from one head of width 4,096, or from 4,096 heads of width 1. Heads partition a fixed projection rather than adding to it.

What you actually pay is representational, not parametric. A width-128 head computes dot products in a 128-dimensional space, so it can draw fewer distinctions than one 4,096-wide head could. You are trading each head’s resolution for the ability to attend to several places at once, and empirically that trade is worth making.

Thing 1 attention cannot do: know where anything is

Attention is permutation-equivariant. Permutation means reordering the positions; equivariant means the output reorders identically.

Trace it: permute the input rows, and Q, K and V permute identically. The score matrix Q·K^T then permutes in both axes. The softmax is row-wise, so it survives the permutation. The output permutes the same way as the input.

So attention alone cannot tell “the dog bit the man” from “the man bit the dog”. Nothing in the mechanism looks at an index.

Position has to be injected, and there are two standard ways:

• Positional encodings — sinusoidal or learned vectors that depend only on the index — added to the embeddings, the per-token vectors x_i that entered this section, each looked up from a table by token id (Embeddings derives that table).
• Rotary position embedding (RoPE), the modern default, which rotates Q and K by an angle proportional to position, so that the score between two positions depends only on the offset between them and not on their absolute indices.

Thing 2: it cannot be cheap

Cost is O(n^2·d) in both time and memory, because of the (n, n) score matrix.

The KV cache, sized in Parameter arithmetic you should be able to do in your head, grows linearly with context on top of that.

That is the same machinery seen from the serving side in Attention and why context costs what it does, which derives the pricing consequences of the mechanism derived here.

8. The transformer block

Attention alone is not a model. Wrapped together with a small dense network, normalization and residual adds, it becomes the unit that is stacked dozens of times to make one — and where the normalization sits inside that unit decides whether the stack trains at all.

A transformer is a stack of identical blocks. Each block holds two sublayers — an attention sublayer and a small position-wise network — and each sublayer is wrapped in a normalization step and a residual add.

The diagram below is one block. Read it top to bottom, and notice the two arrows that bypass everything: x -> A1 and A1 -> A2. Those are the residual paths, and they are what makes the stack trainable.

flowchart TD
    X(["x"]) --> N1["LayerNorm / RMSNorm"]
    N1 --> MHA["Multi-head attention<br/>4·d^2 params"]
    MHA --> A1(("+"))
    X --> A1
    A1 --> N2["LayerNorm / RMSNorm"]
    N2 --> FF["FFN: d -> 4d -> d<br/>8·d^2 params<br/>GELU or SwiGLU"]
    FF --> A2(("+"))
    A1 --> A2
    A2 --> OUT(["out"])

    style A1 fill:#2d6a4f,color:#fff
    style A2 fill:#2d6a4f,color:#fff
    style MHA fill:#1d3557,color:#fff
    style FF fill:#1d3557,color:#fff

In the diagram, blue marks the two sublayers and green the residual adds. The sublayers are the two things wrapped, and they differ in what they compute, not in how they are wired.

Now the nodes.

Multi-head attention is Attention derived as content based lookup’s mechanism, with its four d × d projections counted as 4·d^2 parameters.

FFN: d -> 4d -> d is the feed-forward network: two dense layers applied independently to each position. The first widens the vector from d to 4d, the second brings it back to d, and a nonlinearity sits in between — GELU or SwiGLU, which Activations and the dying relu mechanism argues about. Two matrices of d × 4d is 8·d^2 parameters.

LayerNorm / RMSNorm are the two normalization choices of Normalization and why transformers use layer norm.

The two + circles are the residual adds. The output is the input plus what the sublayer computed, not the sublayer’s output alone. The arrows reaching them directly from x and from A1, bypassing everything, are Residual connections the gradient highway’s identity paths, and they are why this block can be stacked 80 deep without the gradient dying.

Note where the norms sit: the diagram is drawn pre-LN. Each norm is inside a branch, before its sublayer. Nothing normalizes the trunk running down the left-hand side. That choice gets its own subsection below.

So a block is two sublayers, two residual adds and two norms. The division of labour between the sublayers is worth stating plainly:

Attention moves information between positions. The FFN transforms information within a position. Neither can do the other’s job, which is why removing either one collapses the model.

The FFN’s 4× expansion is not arbitrary. It is where most of the parameters live — 8·d^2 against attention’s 4·d^2 — and widening it is the cheapest way to add capacity, because it adds no cross-position cost.

Pre-LN versus post-LN, and why everyone moved

These two names say where the normalization sits relative to the residual add.

post-LN (original):   x = LN( x + Sublayer(x) )      norm is ON the residual path
pre-LN  (modern):     x = x + Sublayer( LN(x) )      norm is inside the branch

Post-LN puts a normalization on the residual path. So the identity route from the loss back to the embeddings gets rescaled once per block — it is no longer an identity. At initialization the deeper blocks dominate the output, and post-LN transformers need learning-rate warmup (starting with a very small step size and raising it over the first few thousand steps) plus careful initialization, or they diverge within a few hundred steps.

Pre-LN leaves the residual stream untouched, so Residual connections the gradient highway’s identity path is genuinely an identity. That is what removed the dependence on careful initialization and made very deep stacks routine.

It is not a licence to drop warmup. Learning rate schedules traces a 12-layer pre-LN transformer collapsing to uniform at learning rate 3e-4 with no warmup, and the reason is outside this section: the early instability is in Adam’s second-moment estimate, not in the block. Pre-LN widens the margin. Warmup is still mandatory.

The price of pre-LN: the residual stream grows

The residual stream is the running vector carried down the trunk from block to block. Under pre-LN nothing rescales it, so its variance grows with depth.

Count the additions. Each block adds to the stream twice — once per sublayer, which is exactly the two + circles in the diagram.

If those additions are roughly independent and each has unit variance, variances add:

variance entering block 1     = 1
each block adds 2 sublayers   -> +2 variance per block
after L blocks                = 1 + 2L

L = 32   ->  1 + 64 = 65        sd = sqrt(65) = 8.06

(Simulated at 8192×1024, that comes out at variance 65.0 and sd 8.06. Counting one add per block instead — a common slip, and wrong for this diagram — would give 33 and 5.75.)

Two consequences follow from that single number.

Each block contributes proportionally less as depth grows. A unit-variance addition barely moves a stream whose standard deviation is already 8.06. That is a soft form of the “deep blocks matter less” effect.

A final norm after the last block is mandatory, not stylistic. The activations arriving at the unembedding — the (d, V) matrix that turns the final hidden vector into one logit per vocabulary token, the output-side counterpart of Embeddings’s embedding table — are about 8× the scale that layer was initialized for.

The other half of the fix lives at initialization: shrink each branch by 1/sqrt(2L) so the total stays bounded in depth. Same 2, same reason, derived in Initialization.

9. Normalization — and why transformers use layer norm

The block just drawn contains two normalization layers, which raises the question of which kind — and the candidates are all one formula, differing only in which numbers they average over.

Normalization rescales a layer’s activations so they have a consistent mean and spread before the next layer sees them. That keeps values from drifting into ranges where gradients die.

Every normalization layer in this section computes the same formula:

(x - mu) / sqrt(var + eps) · gain + bias

Subtract the mean mu, divide by the standard deviation, then apply a learned gain (a multiplier) and bias (an offset) so the layer can undo the normalization if that turns out to be useful. eps is a tiny constant added to stop division by zero.

They differ only in which axis the statistics are taken over, and every practical consequence follows from that one choice.

Activations here are shaped (B, T, D):

• B — how many examples are in the batch
• T — how many tokens are in each example
• D — how many channels each token’s vector has

The same 24 numbers, read two ways

Take B = 2 sequences of T = 3 tokens with D = 4 channels. The block below is those 24 numbers with LayerNorm’s statistics on the right and BatchNorm’s underneath. Compare how many statistics each one produces, and where each one gets its numbers from.

                    d0   d1   d2   d3        LayerNorm reduces along ->
example 0   t0  [   4    1   -1    4  ]      mu = 2.0   sd = 2.12
            t1  [   1    0    3    0  ]      mu = 1.0   sd = 1.22
            t2  [   3    0    6    3  ]      mu = 3.0   sd = 2.12
example 1   t0  [   2    2    4    0  ]      mu = 2.0   sd = 1.41
            t1  [   0    3    0    1  ]      mu = 1.0   sd = 1.22
            t2  [   2    0    6    4  ]      mu = 3.0   sd = 2.24

BatchNorm    mu   2.0  1.0  3.0  2.0
reduces      sd  1.29 1.15 2.71 1.73
downwards |
          v

That is six statistics against four.

LayerNorm computes one (mu, sd) per token — six pairs here, each taken from the four numbers in its own row and nothing else. Check the first one: row [4, 1, -1, 4] has mean (4+1-1+4)/4 = 2.0, and spread sqrt(((4-2)^2 + (1-2)^2 + (-1-2)^2 + (4-2)^2)/4) = sqrt(18/4) = 2.12.

BatchNorm computes one per channel — four pairs, each pooled from the six numbers in its own column, which means from both examples. Check d0: the column is 4, 1, 3, 2, 0, 2, mean 12/6 = 2.0, sd sqrt(10/6) = 1.29.

“Channel” is the word Convolution weight sharing and what it buys used for an image’s RGB planes. For a (B, T, D) activation it means an index of D, one coordinate of the hidden vector. Same idea — a feature index — different axis.

The experiment that separates them

Now change example 0 only. Replace its t2, d2 entry, 6, with 12 — the kind of change a different batch composition would produce.

The block below shows what happens to a token in example 1, which was not touched at all:

                             LN out, example 1 t0        BN out, example 1 t0
before   x[0, 2, 2] =  6     ( 0.00, 0.00, 1.41, -1.41)  ( 0.00, 0.87, 0.37, -1.15)
after    x[0, 2, 2] = 12     ( 0.00, 0.00, 1.41, -1.41)  ( 0.00, 0.87, 0.00, -1.15)

Not one number in example 1 changed, and its BatchNorm output changed anyway.

Here is the mechanism. Channel 2’s column statistics moved:

column d2 before:  -1, 3,  6, 4, 0, 6   ->  mu = 3.0,  sd = 2.71
column d2 after:   -1, 3, 12, 4, 0, 6   ->  mu = 4.0,  sd = 4.28

example 1, t0, channel 2 holds the value 4, unchanged:
  before   (4 - 3.0) / 2.71 = 0.37
  after    (4 - 4.0) / 4.28 = 0.00

Its LayerNorm output is bit-identical, because that row’s statistics came only from that row.

Every row of the table below, and all four reasons after it, follow from that one difference: where each layer draws its statistics.

The “running stats” entry names BatchNorm’s other complication: at inference there may be no batch to compute statistics from, so it must instead reuse a running average of the means and variances it saw during training.

Four independent reasons transformers cannot use batch norm, and any one of them is disqualifying:

1. Variable length. Reducing over T mixes real tokens with padding — filler positions added to make every sequence in a batch the same length — and the statistic shifts with whatever length distribution the batch happened to have. Masking helps but position t still has a different number of contributing samples than position t'.
2. Batch size 1 at inference. Autoregressive decoding — generating one token at a time, each conditioned on all the ones before it — processes one token for one sequence. There is no batch to compute statistics from, so BatchNorm must fall back to running averages collected under a different length and content distribution — a train/inference mismatch on every activation.
3. Small effective batch. Language-model training runs long sequences and few sequences per device (say 4 × 8192). Tokens within a sequence are highly correlated, so the effective sample count behind a per-channel estimate is far smaller than B·T suggests.
4. Cross-example coupling. BatchNorm makes each example’s output depend on its batch-mates — the 0.37 -> 0.00 above — which breaks an independence assumption that three separate systems rely on:
   • Gradient accumulation runs several small micro-batches and sums their gradients to simulate one large batch. That is only equivalent if each micro-batch produces the gradient it would have produced alone, which BatchNorm’s pooled statistics destroy.
   • Pipeline parallelism puts different layers on different devices with several micro-batches in flight at once, so a BatchNorm layer would be computing statistics over whichever fragment happened to be resident.
   • Reproducible generation requires that a prompt produce the same output regardless of what else was batched with it. BatchNorm cannot promise that.

LayerNorm has none of these problems because its statistic comes from the single token being normalized. The function is literally identical at training and inference, at batch size 1 or batch size 1,024.

RMSNorm: the same thing, one pass cheaper

Root mean square normalization (RMSNorm) drops the mean subtraction and the bias:

LayerNorm:  (x - mu) / sqrt(var + eps) · gain + bias
RMSNorm:     x       / sqrt(mean(x^2) + eps) · gain

It divides by the root of the mean squared value rather than by a standard deviation measured about the mean. Empirically the centering contributes almost nothing.

Dropping it removes one reduction pass over the vector and one subtraction, and it saves D parameters per norm layer (the bias).

The pass matters more than it sounds, because norm layers are memory-bandwidth bound: the time they take is set by moving data to and from memory, not by arithmetic. They read and write the whole activation tensor to do only O(D) arithmetic per element. Removing a pass is a real wall-clock win, which is why every recent open model uses RMSNorm.

What interviewers probe: “Why not batch norm in a transformer?” Do not answer “it doesn’t work well.” Answer with the axis: BN’s statistic is over the batch and time axes, so it depends on padding, on batch composition, and on there being a batch at all — none of which hold during autoregressive decoding. LN’s statistic is per token, so the function is the same at train and inference.

10. Residual connections — the gradient highway

The other wrapper in the block is the residual add — the device Why depth beats width promised would make depth trainable rather than merely expressive.

A residual connection (also skip connection) replaces y = F(x) with y = x + F(x). The block computes a correction to its input rather than a replacement for it.

That one change rewrites the Jacobian, and the rewrite is the whole section:

dy/dx = I + dF/dx

over L blocks:  prod_{l=1..L} (I + J_l)
              = I + sum_l J_l + sum_{l<m} J_m·J_l + ...
                ^ a path of length zero: reaches the input undiminished

I is the identity matrix — the matrix that leaves any vector unchanged. J_l is block l’s Jacobian.

Expand the product (I + J_1)(I + J_2)···(I + J_L) the way you would expand (1+a)(1+b)(1+c), and you get one term per subset of blocks. The very first term, from picking I at every factor, is a bare I: the route from the loss to the input that passes through no block at all.

There is always an identity term, so the gradient at layer 1 cannot decay exponentially in L.

Comparing the two floors

Compare with a plain stack, where the gradient is just the bare product prod J_l:

||J_l|| = 0.8, L = 50    plain:     0.8^50 = 1.4e-5     layer 1 gets nothing
                         residual:  >= 1.0 from the identity term alone

Those two numbers are different kinds of object, and the difference is the point.

1.4e-5 is the plain net’s gradient exactly. It is the whole product; there is nothing else in it.

1.0 is a floor for the residual net. The length-zero path contributes I whatever the blocks do, and every other term in the expansion adds on top of it.

So the fair reading: the residual net’s worst case already beats the plain net’s actual value by 1 / 1.4e-5 = 70,000×, and the true gap is larger. Push L higher and the plain net degrades exponentially while the residual floor does not move at all.

The second thing residuals buy: capacity that cannot go backwards

This one is about representation, not gradients.

A residual block can learn F = 0, which makes it the identity — it passes its input through unchanged. So a 50-block residual net contains a 25-block one as a special case, with the other 25 blocks set to zero.

The consequence: adding depth cannot make the achievable training loss worse.

That is precisely the degradation the ResNet paper observed and fixed. Plain 56-layer nets had higher training error than 20-layer ones. That is not overfitting — memorizing the training set at the expense of new data — because the training error itself was worse. It was an optimization failure.

The unraveled view

Expanding prod (I + J_l) gives one term per subset of blocks, and each term is a path through the network: skip some blocks, pass through others.

So L residual blocks form an implicit ensemble — a collection of models whose predictions combine — of 2^L paths of different lengths. At L = 50 that is 2^50 = 1.1e15 paths, ranging from the length-zero identity to the one that passes through all 50, and most of the gradient flows through the short ones.

Depth therefore buys an ensemble, not just a longer function. That is also why dropping a random residual block at inference barely hurts accuracy: you removed some paths, and 2^49 remain.

The assumption: that the useful function is reachable as a series of small corrections to the input. That is a mild assumption and it is why residual connections are near-universal, but it is not vacuous — it is what makes the identity a sensible default for a block that has learned nothing yet.

11. Dropout

Dropout was standard in one era of deep learning and is switched off in another.

Dropout is a regularizer: a deliberate handicap applied during training to stop the model from relying too heavily on any one part of itself. At training time it zeroes each unit independently with probability p and scales the survivors by 1/(1-p); at test time it does nothing at all:

train:  y = (x (*) mask) / (1 - p),   mask ~ Bernoulli(1 - p)
E[y] = x · (1-p)/(1-p) = x            -> train and test agree in expectation

Read mask ~ Bernoulli(1 - p) as “the mask is drawn from a coin flip that comes up 1 with probability 1 - p”. Read E[y] as the average value of y over those coin flips.

Substitute numbers. Take p = 0.1 and a unit whose clean output is x = 2.0:

survives (prob 0.9):   2.0 / (1 - 0.1) = 2.0 / 0.9 = 2.222
zeroed   (prob 0.1):   0

average:  0.9 · 2.222 + 0.1 · 0 = 2.0     <- exactly the clean value

2.0 is precisely what test time computes, because test time does nothing at all.

The 1/(1-p) factor is not a heuristic. It is what makes “do nothing at test time” the correct thing to do.

You could equally scale by 1-p at test time and leave training alone. That is the original formulation. The train-time version (“inverted dropout”) won only because it keeps inference a plain matrix multiplication, with no dropout-aware code in it at all.

Why zeroing units helps. Dropout samples a different subnetwork every step — 2^n of them for n units. A single width-512 layer has 2^512 ≈ 1.3e154 possible subnetworks, so no two training steps ever draw the same one.

All those subnetworks share weights. So no unit can rely on a specific co-adaptation partner — another unit whose presence it has learned to depend on — actually being there.

Failure mode 1: dropout with batch norm. Dropout changes the variance of activations between train and test, but BatchNorm’s running statistics were estimated under train-time dropout. The mismatch shows up as a model that scores well during training-mode evaluation and worse in .eval(), the call that switches a framework’s layers into inference behaviour. Put dropout after BatchNorm, or use one and not both.

Failure mode 2 (the modern one): using it at all during large-scale pretraining.

Dropout fights memorization, and memorization requires seeing the same example repeatedly. An epoch is one full pass over the training set.

A single-epoch run over trillions of tokens never revisits an example, so there is nothing to memorize. The model is underfit, not overfit. Dropout there only adds gradient noise and slows convergence.

Dropout is an epoch-count decision, not a model-size decision. It comes straight back for fine-tuning on a small dataset with many epochs.

12. Embeddings

Every layer so far consumed vectors, but a language model’s raw input is a list of integers. The bridge between the two is a lookup table — a surprisingly large one, and the place where the rarest symbols misbehave.

An embedding table is how a model turns a symbol into numbers. It is a lookup table of shape (V, d) — one learned d-dimensional vector per vocabulary entry — so token number 4,021 is simply row 4,021.

Mathematically it is a matrix multiplication against a one-hot vector (all zeros except a single 1 at the token’s index). In practice it is implemented as a gather, which copies the row and skips the arithmetic entirely.

The assumption is that tokens are unstructured symbols. Nothing about the integer 4,021 says anything about its meaning. Token 4,022 is not “close to” it. So the geometry that puts related tokens near each other has to be learned entirely from data.

That table is expensive. At V = 128,000 and d = 4,096:

128,000 · 4,096 = 524,288,000 parameters  ->  524.3M
524.3M / 6,967M = 7.5% of a 7B model

7.5% of the weights are a dictionary.

Tying the input embedding to the output unembedding — using one matrix for both directions instead of two — saves the same amount again. Small models nearly always tie. Very large ones often do not, because the output head benefits from having its own weights once the block total 12·L·d^2 dwarfs V·d.

Failure mode: frequency imbalance. A token seen 50 times in training gets 50 gradient updates, so its vector barely moves from its random initialization. A common token gets billions.

Rare-token embeddings are therefore close to noise. That is one mechanism behind the “glitch token” behaviours people find in production models, and it connects directly to the tokenizer discussion in Tokens: a rare identifier fragments into rare pieces, each carrying an undertrained vector.

13. Activations, and the dying-ReLU mechanism

One component has been taken on trust since the fold in Why depth beats width: the function between the linear layers. The most popular choice can permanently kill part of a network, and the mechanism takes three lines to derive.

An activation function is the elementwise nonlinearity applied between two linear layers.

Without one, depth buys nothing. Stack two linear layers and you get W_2(W_1 x + b_1) + b_2 = (W_2 W_1)x + (W_2 b_1 + b_2), which is one linear layer with different weights. The activation is the component that makes a deep network actually deep.

Six of them, compared. The column that matters is the third one — what the gradient does on the negative side — because that column is the whole of the next subsection.

The dying-ReLU mechanism

This is the one to be able to derive on a whiteboard.

Unit j computes a pre-activation z = w·x + b, then outputs max(0, z).

Suppose z < 0 for every example in the data distribution. Then the ReLU’s derivative is 0 everywhere the unit is ever evaluated, and backpropagation multiplies by that 0:

dL/dw = dL/dz · x = 0        and        dL/db = dL/dz = 0

The gradient that would revive the unit is gated by the unit being alive. Both the weights and the bias get exactly zero update, forever.

That is an absorbing state — a condition that, once entered, cannot be left. No amount of further training escapes it.

Here is how a unit gets there:

pre-activation  z ~ N(+0.5, 1)   ->  P(alive) = 0.69
one bad step with lr = 0.5 and dL/db = +5  ->  b moves by -2.5
                z ~ N(-2.0, 1)   ->  P(alive) = 0.023
a couple more steps            ->  z < 0 on every batch  ->  dead forever

Read z ~ N(+0.5, 1) as “the pre-activation is normally distributed with mean 0.5 and variance 1”. Under that distribution, 69% of examples land above zero and keep the unit alive.

lr is the learning rate, the multiplier on each gradient step. At lr = 0.5 with a gradient of 5, the bias moves by 0.5 · 5 = 2.5 in the shrinking direction. The mean of z drops from +0.5 to -2.0, where only 2.3% of examples still activate the unit — and those 2.3% are not enough to pull it back before the next couple of steps finish the job.

With too high a learning rate this happens to many units at once, because they all take the same oversized step on the same batch. What you lose is not one unit but a substantial fraction of the layer.

How large a fraction depends on your learning rate, your data and your initialization. It is not a constant worth quoting — measure it.

The symptom is a training loss that plateaus above where it should. The measurement is direct: log the fraction of examples for which each unit outputs exactly zero. An activation histogram that is a single spike at exactly zero names the dead units for you.

Fixes, in this order and for this reason — first remove the cause, then remove the possibility, then reduce the exposure:

1. Lower the learning rate, or add warmup. The cause is one oversized step; a smaller step never creates the state.
2. Switch to GELU / SiLU / LeakyReLU. These have nonzero gradient for z < 0, so the state stops being absorbing — a unit that drifts negative can still climb back.
3. He initialization (Var(W) = 2/fan_in, where fan_in is the number of inputs feeding the unit). Starts pre-activations centered so fewer units begin near the cliff.
4. Add normalization. Re-centering z each step means a unit cannot drift permanently negative in the first place.

SwiGLU, and where odd numbers like 11,008 come from

SwiGLU is the current transformer default. It is a gated variant — one branch multiplies another elementwise, exactly the gating idea of Lstm and gru replacing the product with a gate:

FFN(x) = ( SiLU(x·W_gate) (*) (x·W_up) ) · W_down

That is three matrices instead of the plain FFN’s two. To keep the parameter count unchanged, d_ff — the FFN’s inner width — is set to (8/3)·d rather than 4·d:

plain FFN (2 matrices at 4d):   2 · d · 4d       = 8·d^2
SwiGLU    (3 matrices at 8/3d): 3 · d · (8/3)d   = 8·d^2      same

Then one more step, because the hardware has an opinion:

(8/3) · 4,096            = 10,922.67
ceil(10,922.67 / 256)    = 43           round up to a multiple of 256 so the
43 · 256                 = 11,008       matmul tiles cleanly on the hardware

Tiling is how a matrix multiplication is chopped into fixed-size blocks for the hardware. A width that is not a multiple of the block size leaves a partly-empty block and wastes that block’s work.

That is where real models get odd d_ff values like 11,008 instead of 4d = 16,384. The 8/3 sets the target, and the tiling constraint rounds it.

14. Encoder, decoder, encoder-decoder

A transformer stack can be wired three ways, and the wiring — not the layer inventory — decides whether the result classifies, generates, or translates. The diagram below has one column per wiring; read each column top to bottom: what the attention mask allows, what the training objective is, and what comes out.

flowchart TD
    subgraph E["Encoder-only - BERT"]
        E1["bidirectional attention<br/>every token sees every token"] --> E2["masked-LM objective<br/>~15% of positions supervised"]
        E2 --> E3["output: one vector per token<br/>classify, retrieve, rerank"]
    end
    subgraph D["Decoder-only - GPT, Claude"]
        D1["causal attention<br/>token i sees 0..i"] --> D2["next-token objective<br/>100% of positions supervised"]
        D2 --> D3["output: generation<br/>KV cache reusable"]
    end
    subgraph ED["Encoder-decoder - T5, Whisper"]
        F1["bidirectional encoder<br/>over the source"] --> F2["causal decoder<br/>+ cross-attention"]
        F2 --> F3["output: seq-to-seq<br/>encode source once"]
    end

    style E3 fill:#1d3557,color:#fff
    style D3 fill:#2d6a4f,color:#fff
    style F3 fill:#40916c,color:#fff

The three column headings name the three families. Encoder-only - BERT: BERT stands for Bidirectional Encoder Representations from Transformers, and bidirectional means every token may read every other token in both directions. Decoder-only - GPT, Claude: GPT stands for Generative Pre-trained Transformer, and these models read only leftwards. Encoder-decoder - T5, Whisper: T5 is the Text-to-Text Transfer Transformer and Whisper a speech recognition model; both pair a bidirectional reader of the source with a leftward-only writer of the output.

Four labels inside those boxes carry the comparison and none of them are self-explanatory.

Label 1: causal attention, or token i sees 0..i. Position i’s scores against every j > i are set to -inf before the softmax, so those weights come out exactly 0 and the future contributes nothing. That is the triu(1) mask in Attention derived as content based lookup’s code.

Because the future is blocked, every position can be asked to predict its own successor without having already seen it. That is why a decoder-only model gets 100% of positions supervised from one forward pass.

An encoder-only model like BERT gets only ~15%, because its masked language modelling (masked-LM) objective hides a random 15% of tokens and trains the model to guess them from both sides. The other 85% of positions produce no training signal.

Label 2: cross-attention. This is Attention derived as content based lookup’s mechanism with the three roles drawn from two different places:

self-attention:   Q, K, V  all from the same stream
cross-attention:  Q from the decoder's own positions
                  K, V from the encoder's outputs

Cross-attention takes the question from the decoder and the answer from the encoder. That is the only structural difference between the two.

Label 3: encode source once. That split is exactly what this buys. The bidirectional encoder runs a single time over the source and produces a fixed set of states. Every decoding step then cross-attends against those same states rather than re-reading the source.

So an n_src-token source is encoded once for an n_tgt-token output, not once per generated token.

Label 4: output: one vector per token versus output: generation. This is the shape difference that decides what each architecture is for.

An encoder emits one vector per token — a contextual representation of that input position — and then stops. You pool those vectors to retrieve with, classify them, or rerank a query-document pair, meaning score how well a candidate document answers a query so a shortlist can be reordered. There is no next step to take.

A decoder emits a distribution over the vocabulary at the final position, then feeds its own sample back in as input. That loop is what makes its KV cache reusable across steps.

Representations versus continuations.

ASR in the last cell is automatic speech recognition, turning audio into text — a fixed source producing a new sequence, which is exactly the encoder-decoder shape.

Why decoder-only won

The main argument is supervision density.

masked LM:   ~15% of positions produce a training target
causal LM:   100% of positions produce a training target
                                    100 / 15 ≈ 6.7x the learning signal per token processed

A masked language model throws away the compute it spent on the other 85%. A causal one gets a prediction target at every position from the same forward pass.

Three more things stack on top of that.

• One stack rather than two. Fewer moving parts, and all the parameters serve the same objective.
• Prefix caching works, because attention is causal. A token’s representation cannot depend on anything after it, so a cached prefix stays valid whatever follows (Prompt caching derived).
• In-context learning absorbs the encoder’s job. In-context learning is the ability to pick up a task from examples placed in the prompt, with no weight update at all. “Encode the source, then decode conditioned on it” becomes “put the source in the prompt”, and one stack does both halves.

The case is closed for general-purpose models.

Encoder-only is not obsolete. It is the right architecture whenever you need a representation rather than a continuation.

The cross-encoder reranker in Embeddings and why dense search misses err_4021 is exactly that: query and document fed through a bidirectional stack together. A causal model structurally cannot do this as well, because the query sits before the document and therefore cannot see it.

15. Parameter arithmetic you should be able to do in your head

A whole transformer reduces to one formula — params ≈ 12·L·d^2 — and four separate engineering numbers fall out of it. Take one concrete configuration and hold it for the whole section:

d       = 4,096      model width
n_heads = 32         so d_head = 4,096 / 32 = 128
d_ff    = 4d = 16,384
L       = 32         number of blocks
V       = 128,000    vocabulary size

Now count. The block below adds up one block first, then the whole model.

per block
  attention  W_q, W_k, W_v, W_o     4 · d^2  = 4 · 16,777,216 =  67,108,864
  FFN        W_1 (d->4d), W_2 (4d->d)  8 · d^2                = 134,217,728
  2 norms                              2 · 2d = 16,384        (negligible)
  ----------------------------------------------------------------------
  total per block                12 · d^2 + the two norms     = 201,342,976

whole model
  32 blocks       32 · 201,342,976 = 6,442,975,232
  embeddings         128,000·4,096 =   524,288,000
  ----------------------------------------------------
  tied embeddings                     6,967,263,232  ->  "a 7B model"

Where the 12 comes from:

attention   4 matrices of d × d      ->  4·d^2
FFN         2 matrices of d × 4d     ->  8·d^2
                                         ------
                                         12·d^2 per block

The norms add only two gains and two biases of length d each — 2 · 2d = 16,384 against 201 million, which is why they round away.

The code below is that same count, with tie_embeddings exposed as a flag so you can see what untying costs.

def transformer_params(d, layers, vocab, ffn_mult=4, tie_embeddings=True):
    attn = 4 * d * d                      # W_q, W_k, W_v, W_o
    ffn = 2 * d * (ffn_mult * d)          # up and down projections
    norms = 2 * 2 * d                     # two norms, gain + bias each
    per_block = attn + ffn + norms
    emb = vocab * d * (1 if tie_embeddings else 2)
    return per_block * layers + emb

assert round(transformer_params(4096, 32, 128_000) / 1e9, 2) == 6.97

Four rules fall out of params ≈ 12·L·d^2, and each one gets its own subsection below.

Rule 1: compute per token

FLOPs are floating-point operations, the standard unit of compute.

For a model of N parameters, a forward pass costs about 2·N FLOPs per token — one multiply and one add per weight. Training costs about 6·N: the forward pass’s 2N, plus 4N for the backward pass, which has to compute gradients with respect to both the activations and the weights.

7B model, generating:   2 · 7e9 = 1.4e10 = ~14 GFLOPs per token
7B model, training:     6 · 7e9 = 4.2e10 = ~42 GFLOPs per token

Rule 2: when attention actually dominates

“Attention is quadratic” is true and usually not what you are paying for. Here is where the crossover is.

Per token per layer, there are two costs:

parameter matmuls:   2 · 12·d^2 = 24·d^2      grows with model size
attention itself:    about 4·n·d              grows with context length n

Set them equal and solve:

24·d^2 = 4·n·d   ->   n = 6d   ->   at d = 4,096, n = 24,576 tokens

Below about 24k tokens of context, the FFN is what you are paying for, not the attention.

Rule 3: KV cache size

The cache holds one key and one value per token per layer:

bytes per token = 2 (K and V) · L · d_kv · bytes_per_number

d_kv = n_kv_heads · d_head is the total width of the keys — or of the values — stored per layer.

Full multi-head attention keeps one KV head per query head, so d_kv = 32 · 128 = 4,096 = d. At fp16 (16-bit floating point, 2 bytes per number):

2 · 32 · 4,096 · 2 = 524,288 bytes = 512 KiB per token
8,192 tokens · 512 KiB = 4 GiB     for a single sequence

Grouped-query attention (GQA) keeps all 32 query heads but shares each KV head across a group of them. With 8 KV heads:

d_kv = 8 · 128 = 1,024
2 · 32 · 1,024 · 2 = 131,072 bytes = 128 KiB per token   ->  1 GiB at 8k

A 4× cut, and that is the entire reason GQA exists.

Multi-query attention (MQA) takes the idea to its limit: one KV head for all 32 queries, d_kv = 128, a 32× cut. The quality cost is large enough that 8 groups is the usual compromise.

The serving consequences are in The kv cache the most important mechanism in this chapter.

Rule 4: sanity-checking a claimed parameter count

Someone says “40 layers, d = 5,120”. Estimate it in your head:

12 · 40 · 5,120^2 = 12 · 40 · 26,214,400 = 12,582,912,000  ->  ~12.6B, plus embeddings

If their stated number is far off that, one of two things is true: d_ff is not 4d, or the embeddings are untied.

16. Cheat sheet

Every derivation above leaves a fingerprint you can observe in a training log or an evaluation run. Read the left column as the fingerprint, the middle as the mechanism behind it, and the right as the fix that follows.

Three abbreviations appear only here. NaN is “not a number”, the floating-point value produced by an undefined operation, which then contaminates everything it touches. bf16 is bfloat16, a 16-bit format with the same exponent range as fp32 and fewer precision bits, so it overflows far less readily than fp16. And a quantized KV cache stores the cached keys and values in fewer bits per number — 8 instead of 16, say — trading a little precision for half the memory.

The one-line version: depth buys exponentially more structure than width but makes gradients travel further, so residuals give the gradient an identity path and normalization keeps activations in range; convolution and recurrence are structural bets that trade generality for sample efficiency; and attention is a differentiable dictionary lookup whose sqrt(d_k) exists solely to keep the softmax off its saturated tails.

Next: 05 — Training & Optimization.