Attention is one part of a Transformer. This chapter builds the rest (the residual stream, normalization, the feed-forward network, and position) and then counts every parameter and every FLOP of the block that every large model is built from.

Section 4.7 ended with a warning. A stack of pure attention layers loses rank doubly exponentially fast with depth, until every token carries nearly the same vector (Dong et al., 2021). Something has to carry each token's own information up through dozens of layers. You met the fix in §0.8, where it let ResNets train at 152 layers at a time when a plain 56-layer network had a higher training error than a 20-layer one (He et al., 2016). Instead of asking each layer to compute a new representation, ask it to compute a correction and add it to its input, $\mathbf{x} + f(\mathbf{x})$ instead of $f(\mathbf{x})$.

Every Transformer since 2017 is built on the same trick. In a Transformer it deserves a better name than "skip connection", because it changes how you should picture the whole model. The vector that gets corrected, layer after layer, is called the residual stream, and the rest of this chapter hangs off it.

Section 0.8 showed why the identity path keeps gradients from vanishing. This section goes further. We will see the stream as a running sum that every sublayer reads from and writes to, split a model's output exactly into per-layer contributions, and find that the same identity path lets gradients explode unless each branch shrinks with depth.

A block computes a correction, not a replacement

Write the $T$ token vectors of a sequence as the rows of $\mathbf{X} \in \R^{T \times d}$, as in Chapter 4, where $d$ is the model width. A residual layer computes

where $\mathbf{X}^{(\ell)}$ is the stream before sublayer $\ell$ and $f_\ell$ is the sublayer (an attention layer or a feed-forward network, each with its own parameters). The superscript $(\ell)$ always means "after $\ell$ sublayers". The stream starts as $\mathbf{X}^{(0)}$, the token embeddings (plus position information, which we add in §5.4).

In plain words, each sublayer looks at the current state of every token and proposes a change. The change is added, nothing is overwritten.

Two things follow immediately. First, if $f_\ell$ outputs zeros, the layer is exactly the identity. A deep residual network at initialization (where each $f_\ell$ is small) behaves like a shallow one, and depth is something the network grows into rather than something it must fight. That is the cure for the 56-layer puzzle. Second, the output of a block is the same shape as its input, $T \times d$, always. That is why you can stack 12, 96, or 128 identical blocks without any glue.

A Transformer block contains two residual sublayers, attention followed by a feed-forward network (§5.3), so a model with $L$ blocks has $2L$ sublayers writing into the stream. In this section we count sublayers and write $S$ for their number, so $S = 2L$ for a Transformer. Attention is the only place where tokens exchange information. The feed-forward network works on each token separately.

Unroll it, and the stream becomes a sum

Apply the update rule $S$ times and expand,

In plain words, the final representation of every token is its embedding plus the sum of everything every sublayer ever wrote.

Here is a tiny example you can check by hand. Take $d = 3$ and one token with embedding $(1, 0, 0)$. The first sublayer writes $(0, 0.5, 0)$, the second writes $(0, 0, 0.2)$, and the third writes $(-0.3, 0, 0)$. The stream goes $(1, 0, 0) \to (1, 0.5, 0) \to (1, 0.5, 0.2) \to (0.7, 0.5, 0.2)$. Nothing is lost. You could read off the third layer's write from the difference of two consecutive states, and an earlier write is still sitting in the final vector unless a later layer deliberately subtracts it.

The widget below makes this concrete for a toy model with six sublayers and an 8-dimensional stream. Each dimension's bar is stacked from the contributions of every component that has written so far, colored by who wrote it.

Switching a layer off removes its color and, in this toy, also changes what later layers read, because each later write depends on the stream it sees. The logit change is therefore its direct share plus an indirect effect through later layers. Keep that distinction in mind. It is the difference between direct and total effects, and interpretability researchers fight over it constantly.

Let's build the same thing in NumPy. Each sublayer here is a stand-in (a random matrix and a tanh) rather than real attention, but the bookkeeping is identical.

import numpy as np

def run_stream(X0, blocks):
    """X0: (T, d); blocks: list of f(X) -> (T, d). Returns states (S+1, T, d), writes (S, T, d)."""
    raise NotImplementedError

Reading and writing, the stream as a shared channel

Look at what a sublayer actually does with the stream. It never sees "the previous layer's output" as a separate object. It sees the sum. An attention layer reads the stream through its query, key, and value projections $\mathbf{W}_Q, \mathbf{W}_K, \mathbf{W}_V \in \R^{d \times d}$ and writes back through its output projection $\mathbf{W}_O \in \R^{d \times d}$. The feed-forward network reads through its first matrix and writes through its second. Every read is a linear projection of the stream (after a normalization, §5.2), and every write is a vector added to it.

Elhage et al. (2021) turned this observation into a way of thinking that has shaped interpretability research ever since. The residual stream is a communication channel. Sublayers do not talk to each other directly. They leave messages in the stream, in some subspace, and later sublayers pick them up by projecting onto that subspace. Three consequences are worth having in your head.

Composition is a product of matrices. Suppose sublayer $a$ writes $\mathbf{h}\,\mathbf{W}_{\text{out}}^{(a)}$ into the stream (with $\mathbf{h}$ its internal activation) and a later sublayer $b$ reads $\mathbf{x}\,\mathbf{W}_{\text{in}}^{(b)}$. The part of $b$'s input that came from $a$ is $\mathbf{h}\,\mathbf{W}_{\text{out}}^{(a)}\mathbf{W}_{\text{in}}^{(b)}$. The product $\mathbf{W}_{\text{out}}^{(a)}\mathbf{W}_{\text{in}}^{(b)}$ is a "virtual weight" that connects the two sublayers directly, even though they may be twenty layers apart. This is how attention heads in different layers cooperate, for example in the induction heads previewed in §4.7.

The direct path is a model too. Set every $f_\ell$ to zero. What remains is embedding followed by unembedding, logits $= \mathbf{x}^{(0)}\mathbf{W}_U$ with $\mathbf{W}_U \in \R^{d \times V}$ the unembedding matrix. That zero-layer model can only use the current token to predict the next, so the best it can learn is bigram statistics. A trained Transformer keeps a version of this bigram table in the product $\mathbf{E}\mathbf{W}_U$ of its embedding and unembedding matrices, and every layer adds refinements on top of it.

Bandwidth is limited. All $S$ sublayers share the same $d$ dimensions. GPT-2 small has $d = 768$ and 24 sublayers, many of which compute far more than 768 internal features. The stream cannot give each feature its own axis, so models store more features than dimensions by packing them into nearly orthogonal directions, the high-dimensional trick from §1.3. This is the setting for superposition, where a model represents more concepts than it has neurons. Elhage et al. also observed sublayers whose writes appear to erase information that is no longer needed, a kind of memory management.

Decomposing the output, exactly

Because the stream is a sum, the final logits are almost a sum too. The "almost" is the final normalization. GPT-style models compute logits as $\operatorname{LN}(\mathbf{x}^{(S)})\,\mathbf{W}_U$, where for one token $\mathbf{x} \in \R^{d}$

Here $\boldsymbol{\gamma}, \boldsymbol{\beta} \in \R^d$ are a learned gain and shift, $\mathbf{1}$ is the all-ones vector, and $\epsilon$ is a small constant (§5.2 derives all of this properly). The mean subtraction is linear in $\mathbf{x}$. The division by $s$ is not, but $s$ is a single number per token. If we treat $s$ as fixed for the token we are explaining, the logit of vocabulary item $v$ splits exactly into one term per component $\mathbf{c}$ of the stream (the embedding and every write),

In plain words, center each write, scale it by the shared $1/s$ and the gain, and project it onto the unembedding column of the token you care about. That number is the write's direct contribution to the logit. This is called direct logit attribution, and it is the first tool people reach for when asking "which layer made the model say Paris?".

The same picture gives the logit lens (nostalgebraist, 2020). Apply the final normalization and the unembedding to an intermediate state $\mathbf{x}^{(\ell)}$ instead of $\mathbf{x}^{(S)}$, and you get the model's current best guess partway through. In trained models those guesses sharpen layer by layer, and the logit lens has become one of the basic tools for reading a model from the inside.

The two numbers agree to every printed digit. The decomposition is not an approximation once $s$ is frozen. What it does not capture is that a write also changes what later layers read, which is the indirect effect you saw in the widget.

import numpy as np

def logit_attribution(parts, gamma, beta, W_U, v, eps=1e-5):
    """parts: (C, d). Returns contrib (C,), bias_term (float)."""
    raise NotImplementedError

Gradients ride the identity

Now the reason residual networks train at all, redone for a stack of $S$ sublayers so that we can see how depth enters. Take one token's stream $\mathbf{x}^{(\ell)} \in \R^d$ and the Jacobian of one sublayer, $\mathbf{J}_\ell = \partial f_\ell / \partial \mathbf{x}^{(\ell)} \in \R^{d \times d}$ (we use the column convention of §0.2, entry $(i,j)$ is $\partial f_i/\partial x_j$). Differentiate the update rule,

Multiply out the product and the first term is $\mathbf{I}$. Then come the sums of single Jacobians, then pairs, and so on,

Backpropagation sends the row vector $\mathbf{g}^{(S)} = \partial \mathcal{L}/\partial \mathbf{x}^{(S)}$ down by multiplying on the right, $\mathbf{g}^{(\ell)} = \mathbf{g}^{(\ell+1)}(\mathbf{I} + \mathbf{J}_\ell)$. So

In plain words, every layer receives the loss gradient at the top, untouched, plus whatever the sublayers above add to it. In a plain network $\mathbf{x}^{(\ell+1)} = f_\ell(\mathbf{x}^{(\ell)})$, the gradient is only the last term, a product of $S - \ell$ Jacobians, and products of many matrices shrink or blow up exponentially.

How fast? Here is a calculation you can do in an interview. Model each Jacobian as a random matrix with independent entries $\mathcal{N}(0, \alpha^2/d)$, so it scales a vector by about $\alpha$ (formally $\mathbb{E}\lVert\mathbf{g}\mathbf{J}\rVert^2 = \alpha^2\lVert\mathbf{g}\rVert^2$). For one residual step,

because the cross term has mean zero. After $S$ independent steps the gradient's root-mean-square size is $(1+\alpha^2)^{S/2}$ for the residual network and $\alpha^S$ for the plain one.

The plain network vanishes for any $\alpha < 1$ and explodes for any $\alpha > 1$. Only the knife edge $\alpha = 1$ survives, and real nonlinearities (a ReLU zeroes half its inputs) knock you off it. The residual network can never vanish, since $(1 + \alpha^2) \ge 1$. But look carefully, it can explode, $(1 + \alpha^2)^{S/2} \approx e^{\alpha^2 S/2}$. The way out is to keep each branch small in proportion to depth. With $\alpha^2 = 1/S$ the growth is $(1 + 1/S)^{S/2} \to e^{1/2} \approx 1.65$ no matter how deep the model is. Remember that $1/\sqrt{S}$, because it comes back as the GPT-2 initialization rule in §5.7, where $S = 2L$ gives a factor of $1/\sqrt{2L}$.

At 64 layers, the plain network's gradient at $\alpha = 0.5$ is about $10^{-20}$, numerically dead. The residual network's never drops below 1, but at $\alpha = 1$ it is about $10^{9}$, which is why nobody trains a deep residual network with full-size branches. At $\alpha = 1/\sqrt{S}$ it sits near 1.6, matching the theory.

import numpy as np

def backprop_stack(g_top, Js, residual=True):
    """g_top: (d,), Js: (S, d, d). Returns G: (S+1, d), G[S] = g_top."""
    raise NotImplementedError

The forward signal grows too

The same arithmetic applies on the way up. If each write were a random linear function of the stream, $\lVert\mathbf{x}^{(S)}\rVert^2$ would grow like $(1 + \alpha^2)^S$, exponentially. Real Transformers avoid that by normalizing what each sublayer reads (the pre-LN arrangement of §5.2), so each write has a roughly fixed size no matter how large the stream has become. Then the writes add like independent random vectors and the squared norm grows linearly, $\lVert\mathbf{x}^{(S)}\rVert^2 \approx \lVert\mathbf{x}^{(0)}\rVert^2 + S\,c$ for a per-write squared size $c$. Linear growth is survivable, but it has consequences. The later a sublayer sits, the smaller its write is relative to the stream, and the final normalization must divide by a number that grows like $\sqrt{S}$. Section 5.2 shows how this shapes the gradients, and §5.7 how initialization keeps it in check.

A residual network is an ensemble of paths

Go back to the expanded product, $\mathbf{I} + \sum \mathbf{J}_k + \sum \mathbf{J}_{k'}\mathbf{J}_k + \cdots$. Each term corresponds to a path from layer $\ell$ to the top that passes through some subset of the sublayers and skips the rest. With $S$ sublayers there are $2^S$ paths, and $\binom{S}{k}$ of them pass through exactly $k$ sublayers.

Veit, Wilber & Belongie (2016) took this "unraveled" view seriously. If each sublayer on a path multiplies the gradient by a typical factor $\alpha$, then the total contribution of length-$k$ paths is proportional to $\binom{S}{k}\alpha^k$. Normalizing by $\sum_k \binom{S}{k}\alpha^k = (1 + \alpha)^S$ (the binomial theorem) gives a distribution over path lengths with mean

In plain words, when sublayers shrink the gradient ($\alpha < 1$), most of the gradient travels along paths that use only a fraction $\alpha/(1+\alpha)$ of the layers. Veit et al. measured this in a 110-layer ResNet and found that most of the gradient came from paths only 10 to 34 layers deep. They also found that deleting a single block from a trained ResNet at test time barely hurts accuracy, while deleting a layer from a plain VGG network destroys it. The network behaves less like one very deep function and more like a large ensemble of shallower ones that share weights.

Large language models show a version of the same robustness. Gromov et al. (2024) removed large fractions of the deepest layers of open models with modest loss on question-answering benchmarks after a small amount of fine-tuning. This is an empirical finding, and it is double-edged. It suggests that deep layers in today's pre-LN models are underused, which is part of the motivation for the depth-aware normalization and scaling tricks in §5.2 and §5.7.

import numpy as np

def path_length_weights(S, alpha):
    """Returns w (S+1,), normalized, and the mean path length."""
    raise NotImplementedError

In PyTorch

In PyTorch the residual add is one line, and forward hooks let you record every write without touching the model's code. This is how interpretability tooling captures the stream.

Plain, highway, residual

Residual connections were not the only proposal for training deep networks. Highway networks (Srivastava, Greff & Schmidhuber, 2015) came first, with a learned gate that mixes the input and the transformation.

The residual form won because the identity path is exact and free. A highway gate can close, and concatenation makes the width grow with every layer, which a 96-layer Transformer cannot afford.

Read more

• He, Zhang, Ren & Sun (2016). Deep Residual Learning for Image Recognition. The degradation problem and the residual fix.
• He, Zhang, Ren & Sun (2016). Identity Mappings in Deep Residual Networks. Why a clean identity path matters, derived from the backward pass.
• Elhage et al. (2021). A Mathematical Framework for Transformer Circuits. The residual stream as a communication channel, virtual weights, and composition.
• Veit, Wilber & Belongie (2016). Residual Networks Behave Like Ensembles of Relatively Shallow Networks. The unraveled view and the lesion experiments.
• nostalgebraist (2020). Interpreting GPT, the logit lens. Reading intermediate residual states with the unembedding.
• Gromov et al. (2024). The Unreasonable Ineffectiveness of the Deeper Layers. Pruning deep layers of LLMs and what it suggests about depth.
• Srivastava, Greff & Schmidhuber (2015). Highway Networks. The gated predecessor of residual connections.

Watch next

▶

▶