A Short Course for Modern Large Language ModelsGlossary

Lesson 21 of 23 · Transformer block

Learn an update, then add it to the state.

Your win: calculate both residual additions in a pre-norm block and explain their shape requirement.

12 minutesNeeds: vector additionOutcome: read residual equations
Mission link: attention and the feed-forward network do not replace the hidden state directly; each produces an update added to its input.

The identity route and the learned route meet

\[ \text{Attention residual:}\qquad Y=X^{(\ell)}+O_{\text{attn}}, \] \[ \text{FFN residual:}\qquad X^{(\ell+1)}=Y+\operatorname{FFN}(\operatorname{Norm}(Y)). \]

A residual sublayer has the pattern \(X+F(X)\), or \(X+F(\operatorname{Norm}(X))\) in a pre-norm block. The sublayer learns a correction \(F\) while the identity route carries \(X\) directly to the addition.

Numbers substituted: one token

Suppose a token's state and attention update are

\[ \mathbf x=\begin{bmatrix}1&-2&0.5\end{bmatrix}, \qquad \mathbf o_{\text{attn}}=\begin{bmatrix}0.2&0.3&-0.1\end{bmatrix}. \] \[ \mathbf y=\mathbf x+\mathbf o_{\text{attn}} =\begin{bmatrix}1.2&-1.7&0.4\end{bmatrix}. \]

If the feed-forward update is \(\mathbf f=(-0.1,0.5,0.2)\), then

\[ \mathbf x^{(\ell+1)}=\mathbf y+\mathbf f =\begin{bmatrix}1.1&-1.2&0.6\end{bmatrix}. \]

Scale the learned update while the identity path stays fixed

update scale = 1.0
Identity \(X\)
+
Update \(F\)
=
New state \(Y\)

Addition forces the output width

\(X:B\times T\times d_{\text{model}}\)+\(F:B\times T\times d_{\text{model}}\)=\(Y:B\times T\times d_{\text{model}}\)

This is why multi-head attention uses \(W_O\) to return to \(d_{\text{model}}\), and why the feed-forward network projects its expanded hidden width back down before addition.

Precise intuition: the identity path provides a direct route for the current state and for gradients. It does not guarantee that every original fact survives all later computations unchanged.

Code checkpoint · the complete pre-norm block skeleton

The identity path is visible in the repeated x + ... form. Each sublayer receives a normalized view, but its update is added to the unnormalized running state.

Show both residual updates
class TransformerBlock(nn.Module):
    def forward(self, x):
        x = x + self.attention(self.attn_norm(x))
        x = x + self.feed_forward(self.ffn_norm(x))
        return x

Trace it: Point to the two identity routes and the two learned routes in these three lines.

Retrieval check

What shape must an attention update have before adding it to \(X\in\mathbb R^{B\times T\times d_{\text{model}}}\)?

Practice before moving on

  1. Add \((2,-1,4)\) and \((-0.5,2,1)\).
  2. If the learned update is the zero vector, what does a residual sublayer output?
  3. For \(X\) of shape \(2\times5\times8\), what shape must \(O_{\text{attn}}\) have?
  4. Why is a raw concatenated head tensor sometimes reshaped before residual addition?
  5. Write the two residual equations of a pre-norm decoder block.
  6. Explain the phrase “learn a correction” using \(X+F(X)\).
Check solutions
  1. \((1.5,1,5)\).
  2. \(X\), because \(X+0=X\).
  3. \(2\times5\times8\).
  4. The head and feature axes must be concatenated to recover model width, followed by \(W_O\) in standard attention.
  5. \(Y=X+\operatorname{Attn}(\operatorname{Norm}(X))\); \(X^{(\ell+1)}=Y+\operatorname{FFN}(\operatorname{Norm}(Y))\).
  6. \(F(X)\) specifies the change to add rather than an entirely new replacement state.

Primary sources: He et al., Deep Residual Learning; Xiong et al., pre-norm analysis.