The Transformer Block

Two rounds of attention collapse into one

Everything asked for at the start is built. The sentence arrives as with its positions folded in, each head turns it into , and , the mask settles which positions a query may read, and folds the heads back into one vector per word. chased leaves holding 0.60 of mouse, 0.18 of cat, and a little of the rest.

The masking chapter set the model a task: read The cat chased the and give the word that follows. Answering mouse means using two facts that now sit in the same vector — the subject is a cat, the verb is chased. Attention put both there by averaging, and averaging is not deriving a third fact from two. Gathering information and computing with it are different operations, and the model so far does only the first.

Why not simply run attention again? Its output has the shape of its input, so it can be fed straight back in. Follow a single head and write the weight table as , so one round is

where mixes the rows and reshapes each row on its own. Freeze the weights for a moment and put two rounds one after the other:

is one table and is one projection, so two rounds have the form of one. In a real stack is computed from what the first round produced, so the two are not literally identical — but that dependence lives in the proportions of the blend, never in the vectors being blended. Every value crosses the mechanism without passing through a nonlinear function. The softmax decides how much of each word to take, never what any word is.

Geometrically: every row of is positive and sums to one, so a head's answer is an average of the value vectors, and a head can move a word anywhere inside the region those vectors span and nowhere else.

Give every word its own network

Attention has already done all the sideways looking, and it is the only part of the model that needs to. So whatever comes next should read one word's vector and nothing else. That step is two matrix multiplications with a function in between:

is and is , so the vector widens and comes back. The usual choice is , which is 36 for the nine-wide example here. is applied to each number on its own, most often ReLU, which replaces negatives with zero.

That makes each column of a small test. It compresses the word's vector to one number, and ReLU zeroes it unless a specific combination of features is present — something carrying both a cat and a chase direction, say. takes whichever tests fired and writes their conclusions back into the vector. Widening to is room to run thousands of tests at once; narrowing back is packaging the result so the rest of the model can read it.

Adding instead of replacing

The natural move is to repeat the pair: attention, feed-forward, attention again, as many times as the budget allows. Stacking that way has a cost. Each sublayer hands back a different vector than it was given, so a dozen blocks in a row means the word has been overwritten a dozen times, and anything about it still needed has to be rebuilt at every step.

The fix is to add rather than replace:

or, over the whole sentence, and then . All three matrices are , so both additions line up row by row. This is a residual connection. Nothing is thrown away: the vector that arrived is still in the sum, and a sublayer with nothing useful to add can return values near zero and let the vector through unchanged.

The second advantage shows up only during training. A correction travels backward from the loss to every earlier block, and in a plain stack it is reshaped by every sublayer on the way. With the addition, the derivative of is . That is an express lane straight past the sublayer's arithmetic.

Keeping the numbers in range

The stream solves one problem and creates another. Crossing blocks, the vector accumulates contributions, so its numbers tend to grow with depth. Queries and keys are built from that stream, and , so if doubles then both double and the score quadruples. Stretched gaps are exactly what took mouse from 0.60 to 0.96 with everything else rounding away, and collapses with it. Dividing by does not help: it answers growth from the dimension, not growth from depth.

So put the numbers back in range before a sublayer reads them. Take one word's vector, compute the mean of its components and their spread, subtract the mean and divide by the spread:

guards against a division by zero when every component is identical. and are learned vectors of width , applied element-wise, so the model is not pinned to zero mean and unit variance — it can learn its own scale and shift, independent of how much depth has accumulated.

One detail carries the whole idea: every sum runs over the features of a single word, to . Not over the words in the sentence, not over the batch. chased is normalised from its own nine numbers and would come out identical if the rest of the batch were thrown away.

That matters because the common alternative, batch normalisation, takes its statistics across the examples in the batch, which breaks here three times over. Sentences are padded, so statistics at a late position would be computed partly over [PAD] slots. Sentences differ in length, so how many real words reach a position depends on what else is in the batch. And at generation time there is a single sequence and no batch at all.

Where the norm goes is a second decision. The original Transformer normalised the sum, , which is post-norm. Almost everything since normalises only the copy the sublayer reads, , which is pre-norm. The two differ by one box moved, and that box decides how the gradient flows: in post-norm the norm sits on the stream and rescales every correction travelling back; in pre-norm the stream is untouched. The price is one extra at the very end, since the stream now reaches the output unnormalised.

The block

Attention to move information between words, a feed-forward network to compute with it, an addition so neither overwrites what came before, a normalisation so the numbers stay in range. A transformer block is those four things:

Two sublayers, two norms, two additions. The norms are separate, each with its own and , which is why they carry different subscripts. All of , and are , so a block hands back exactly what it was given: five words in, five words out, nine numbers each.

def layer_norm(X, gamma, beta, eps=1e-5):
    """Statistics from the last axis alone, so every position is on its own."""
    mu = X.mean(dim=-1, keepdim=True)
    var = X.var(dim=-1, unbiased=False, keepdim=True)
    return (X - mu) / torch.sqrt(var + eps) * gamma + beta

def block(X, attn, ffn, ln_1, ln_2, mask=None):
    """One pre-norm block. X is (n, d_model), and so is the result."""
    Z = X + multi_head_attention(layer_norm(X, *ln_1), *attn, mask)
    return Z + feed_forward(layer_norm(Z, *ln_2), *ffn)

Each claim above is a line of that function. The sublayer reads layer_norm(X, ...) and never the stream, so the norm sits on the branch. What it returns is added rather than substituted. And the result is shaped like X, which is what lets the last line of one block be the first argument of the next.

Counting parameters for the nine-wide example: attention holds , , and , which across all heads come to weights, or 324. The feed-forward holds and , which at come to , or 648. Two thirds of a block is the part that does not look sideways at all.

Stacking these is the next chapter, along with the two pieces still missing between a stack of blocks and a word.