How Does a Transformer Block Update a Token’s Representation?

English · 中文 · Français

First Principles of LLMs & RL · Article 07 · Previous article

Writing note: The ideas and initial material are mine. ChatGPT assisted with drafting, editing, translation, numerical checks, and the original figures and animation sources.

At the end of the previous article, we had two vectors for position 2. One was the representation it brought into the layer; the other was what it read from the positions it could see:

Suppose the next line of code has to pass a vector to the next layer. Which one should it pass?

We could pass the second one. The shapes agree, so the code would run. But that line would already make a decision: replace what the position had with what it just read. We could also let the read update what was already there. The difference between those two lines of code is a good place to start building a Transformer block.

If a position can read itself, why add its input back?

Recall where the second vector came from. Position 2 could read position 1 and itself, but not the future position 3. In our constructed example, it put half the weight on each visible position:

Some of the position’s own information is already in that sum. Follow the route it took, though. The input first went through a Value projection. Its Value then received the weight assigned to position 2 and was mixed with a contribution from position 1. The Value at position 2 happened to equal the original input in this example. Another projection could change it. Even with the projection held fixed, the head would bring back less of that Value if it gave more weight to position 1.

There is a practical reason to offer another route: a position should be able to devote its attention to something elsewhere without having to reserve some of that attention for carrying its own input forward.

Give the input a direct path to the sum. For the moment, use one head whose output width matches the input width, take the output projection to be the identity, and leave out normalization. Write u₂ for the position after the update:

The direct path around the read branch is the residual connection. Why “residual”? If the desired output differs from the input by a certain amount, that difference is all the branch needs to supply. The first coordinate started at zero and gained one. The second started at one; the branch only had to supply another half. With different parameters an update can be negative, too.¹

Here is a test that requires no guess about what a coordinate means: set the branch output to zero. Replacing the representation with the read sends zero to the next layer. Adding the read sends the original input on unchanged. At the other extreme, a branch that returns the negative of the input can cancel it completely. What the residual connection preserves is a direct path for the input. The sum need not let us recover each contribution separately.

Both paths start with the same input and the same branch output. Set the branch output to zero and follow what reaches the next layer.

What does one set of attention weights leave out?

Look again at the factor of one-half in front of the first Value. It applies to every coordinate of that vector. What if we want more of the first coordinate from position 1 and more of the second coordinate from position 2? Can one head make those two choices independently?

Call the weight on position 1 a. The weight on position 2 is then 1 − a. With the two Values fixed, every possible read has the form

Turn up a a little. The first coordinate rises, but the second falls. Draw the two Values as points in the plane: the output can only move along the line segment joining them. Getting 3/2 in the first coordinate requires a to be 3/4. Getting 3/4 in the second coordinate requires that same a to be 1/4. More precision in choosing the weight wouldn’t help. We need a second weight that we can choose independently.

Switch to two heads, each returning one number. To keep their roles visible, give A a Value projection that picks out the first coordinate from the preceding example, and B one that picks out the second. They compute their own matching scores. For position 2, choose these weights by hand:

Now both numbers are available at once. A and B need not agree: each head has its own scores and its own softmax. We have changed the reading rule; the entries in the table were not measured from the model in Article 06. Picking out a coordinate is a teaching convenience, too. A learned Value projection can combine coordinates. With our fixed pair of Values, the extra freedom is to read different features from the positions in different proportions. A single head with a different Value projection might well produce those same two numbers on this one input. The example isolates what changes about the read.

After the heads finish, we still have to decide what to write. Concatenation puts their outputs side by side. To see what the output projection can do next, keep the convention from Article 06: an individual position vector is a column; its transpose is a row in the sequence matrix. Call the concatenated vector s₂ and choose a small output matrix:

The first column writes the difference between the two heads; the second writes B’s contribution. The attention update and the resulting representation are therefore

The heads determine how much to take from each position. The output projection mixes what they bring back. The residual addition writes that mixture into the current representation. Four heads returning two numbers each, for instance, give eight numbers after concatenation. The output projection maps them to the width of the input so the two paths can be added coordinate by coordinate. This concatenation followed by a projection is the multi-head construction in the original Transformer.²

The weights are constructed for this example. Follow the output of each head through concatenation and projection, then look for the separate residual addition.

Once the information is here, what can the position do with it?

Position 2 now holds 3/4 and 7/4. Give it a small task: if the second coordinate exceeds the first, add the excess to the first coordinate. Otherwise leave the representation alone.

Does the model have to consult the context again? Both numbers it needs are already here. Call the coordinates p and q. One hidden unit suffices to subtract, clip a negative result to zero, and send the result to the first output coordinate:

In our example the difference is exactly one. After adding the update, we have

That is a computation a position-wise feed-forward network (FFN) can perform. The first linear map extracts “second coordinate minus first.” ReLU lets the feature contribute only when the difference is positive. The second linear map chooses where to write it. Without ReLU the two linear maps could be collapsed into one, and this one-sided response would be lost. Nothing requires two coordinates in a trained model to admit this interpretation. We chose the parameters so that we could check what they do.

Compare two positions in the sequence. Attention can make the result at position 2 depend on content at position 1. The FFN, given the updated vector at position 2, transforms that vector alone. Apply the same FFN at position 1 and the parameters are the same, but the result can differ because the input differs. A position-wise operation need not be processing an isolated word: context has already arrived through attention.²

The original Transformer uses an FFN with ReLU. In my implementation, I use SwiGLU. At each position, one projection supplies candidate features; another goes through an activation and modulates them before the result is projected back to the model width. This changes what happens within one position. It still doesn’t read the next one.³

Which path goes through normalization?

So far we have left out normalization to make the additions easy to follow. If we stack these layers, each one writes an update. The vectors arriving at later branches need not have the same scale as those at the entrance.

That matters even for a single read. Consider bias-free dot-product attention. Double every input vector and hold everything else fixed. The projected Queries and Keys both double, so their dot products grow fourfold. Softmax can turn a slight initial preference into a much more concentrated read. The division by √dₖ in Article 06 addressed the scale associated with the dot-product dimension. Here the dimension hasn’t changed at all; the change comes from the magnitude of the inputs.

One design lets the residual stream carry on while each branch works with a normalized version of its input. Here we will use RMSNorm. For a vector z of width d at one position, let γ denote the learned per-coordinate gains and ε a small positive number that keeps the denominator away from zero:

The denominator is computed across the coordinates of this position, not across tokens. When ε is negligible relative to the mean square, multiplying the whole vector by a positive factor scales numerator and denominator by approximately the same factor. A longer input vector alone need no longer produce more extreme scores. Unlike LayerNorm, RMSNorm does not subtract the mean.⁴

We can now write the block. Let X contain the input sequence: one row for each of n positions and one column for each of the model’s d_model coordinates. U is the sequence after attention writes its update; Y is the block output. N₁ and N₂ are RMSNorm operations with separate parameters, applied to each row. MHA includes causal multi-head attention and its output projection. Omitting the batch dimension and dropout:

Look at the additions. The first adds to X itself. The second adds to U, which already contains the attention update. Each branch normalizes its input; the direct path bypasses that normalization. This arrangement is called Pre-Norm. The hand calculations above isolated particular operations. Feeding normalized inputs into the full block can change the numerical results. If we leave the cache logic out of the implementation, the two writes read like this:

u = x + attention(attention_norm(x))
y = u + feed_forward(feed_forward_norm(u))

The overview follows the normalized branches; the earlier hand calculations omit normalization. The FFN transforms each position separately. The final representation supplies the input to the next layer’s projections.⁵

Try the opening test again. If both branches return zero on their respective inputs, the first line gives back the input and so does the second. Omit the first x + and the output becomes zero, even though every tensor still has the right shape. This is a test of this particular block, before any separate final normalization of the model. In the original Transformer’s Post-LN arrangement, normalization follows the residual addition. Even if a sublayer returns zero, the input still passes through LayerNorm; we cannot carry the identity conclusion over unchanged.²⁶

This also connects to backpropagation in Article 03. Flatten the whole input of one residual sublayer into a column vector x. Call its branch F and its output y. If g is the column gradient arriving from the loss, the chain rule gives, wherever the branch is differentiable,

Here J_F is the Jacobian of the branch. The first copy of g on the right doesn’t have to pass through that branch; it comes from the direct addition in the forward pass. The other term can reinforce or cancel it. The expression tells us which path the gradient can take, not how large it must remain.¹

Where does the next layer’s Query come from?

Return to I like and They like. In the setup of the previous article, like has the same token ID in the same position. At the first layer it can therefore begin with the same representation and issue the same request. The words before it differ, so that request can still bring back different content.

After one block, the ID of like has not changed. Its row has. The representation has received a write from attention, and the FFN has computed another update from the features now at that position. This updated row is what the next layer uses, with its own normalization and projections, to form its Query. A difference in the prefix can now affect both what came back on the first read and the cues used to look for information on the second. The next projection may ignore that difference; the architecture makes the dependency possible without prescribing a grammatical interpretation.

That’s what stacking two layers adds to the picture. The second read can begin with information that the first read brought home.

At the very last layer, one thing still hasn’t happened. The final position gives us a representation vector, and we’re waiting for a word. How does that vector yield a distribution over the vocabulary? Once a token is chosen and appended, which calculations have to be done again, and which can we keep? That’s where generation begins.