The Transformer Block and the Residual Stream

Overview & Motivation

The attention topic read attention as a kernel smoother: each token’s output is a weighted average of the values of the tokens it attends to. Averaging has a consequence when we do it twice. Average a set of points, average the averages, and keep going, and every point drifts toward one consensus. Take ten tokens in the plane and send them through twelve layers of attention with nothing else in the block. At a moderate temperature their cloud, whose widest pair starts $2.574$ apart, shrinks until every token sits at the same point to within $10^{-10}$ . Add one term to each layer, the token itself, and the same twelve layers leave every token distinct, with the cloud $2862$ times wider than it began.

That added term is the residual connection, and this topic is about what it builds. With it, a transformer block stops being a map that replaces its input and becomes a write into a shared vector, the residual stream, which every later layer can read. We develop the consequences in order:

Attention alone collapses. A row-stochastic matrix shrinks the cloud of tokens by at least a computable factor, so wide-kernel attention stacks lose every distinction between tokens. With the skip, and value maps of norm at most $1$ , each layer keeps at least one minus that factor of the cloud’s diameter.
LayerNorm is a sphere in a hyperplane. It projects onto the hyperplane orthogonal to the all-ones vector, then radially onto a sphere, and its Jacobian forgets the radial direction.
The identity path. The Jacobian of a residual stack multiplies out into a sum over every subset of layers, and the empty subset is the identity. Every layer receives the top gradient unchanged, and where LayerNorm sits decides what happens to the rest.
The stream as a channel. An attention-only model’s logits are exactly a sum over chains of heads, and two heads talk only through the subspaces where one writes and the other reads.
Induction heads. We train a two-layer model from scratch and watch two heads compose through the stream, a previous-token head and a head whose keys read it, into a circuit that completes [A][B] … [A] with [B]. One layer cannot do this, and we show why.

Then the finance case. On filings about companies that never occur in training, a model trained on the filings completes none of their multi-word names, and the induction model, which never saw a filing, completes $0.768$ of them.

The block, assembled

We write a sequence of $n$ tokens as a matrix $X \in \mathbb{R}^{n \times d}$ , one row $x_i$ per position, and let matrices act on the right, $x \mapsto xW$ . Write $\mathbf{1}$ for the all-ones vector.

Definition 1 (The transformer block).

A head $h$ has maps $W_Q^h, W_K^h, W_V^h \in \mathbb{R}^{d \times d_h}$ and $W_O^h \in \mathbb{R}^{d_h \times d}$ and sends $X$ to

\operatorname{head}_h(X) = \operatorname{softmax}\!\Bigl(\frac{XW_Q^h (XW_K^h)^\top}{\sqrt{d_h}} + B\Bigr)\, X W_V^h W_O^h,

with the softmax taken row by row and $B$ either $0$ or the causal mask ( $-\infty$ above the diagonal). Rotary positions (Definition 4 of the positional-encodings topic) may rotate the queries and keys. Multi-head attention is $\operatorname{MHA}(X) = \sum_h \operatorname{head}_h(X)$ , and an MLP is $\operatorname{MLP}(x) = \sigma(xW_1 + b_1)W_2 + b_2$ on each row, with $\sigma = \operatorname{ReLU}$ here. The pre-LN block is

X \leftarrow X + \operatorname{MHA}(\operatorname{LN} X), \qquad X \leftarrow X + \operatorname{MLP}(\operatorname{LN} X),

and the post-LN block is $X \leftarrow \operatorname{LN}(X + \operatorname{MHA} X)$ , then $X \leftarrow \operatorname{LN}(X + \operatorname{MLP} X)$ , with $\operatorname{LN}$ the LayerNorm of Definition 3. The residual stream is the sequence of matrices $X_0, X_1, \dots$ that the blocks pass along: each sublayer reads it and writes back by addition.

The post-LN block is Vaswani and colleagues’ (2017) original; the pre-LN block is what most large models use now, for a reason Theorem 2 makes precise. The skip itself is He and colleagues’ (2016).

Proposition 1 (Heads are smoothers that write; the MLP is a memory).

(i) Without rotary positions or query and key biases, $\operatorname{MHA}(X) = \sum_h \operatorname{Attn}(XW_Q^h, XW_K^h, XW_V^h)\, W_O^h$ , where $\operatorname{Attn}$ is the attention topic’s scaled dot-product attention. Each head is a Nadaraya–Watson smoother, and the block adds $H$ smoothed copies of the stream, each mapped into it by its own $W_O^h$ .

(ii) $\operatorname{MLP}(x) = b_2 + \sum_{i=1}^{w} \sigma(\langle x, k_i \rangle + b_{1,i})\, v_i$ , where the key $k_i$ is the $i$ -th column of $W_1$ and the value $v_i$ is the $i$ -th row of $W_2$ .

Proof.

For (i), the head’s formula is the attention topic’s Definition 2 applied to $XW_Q^h$ , $XW_K^h$ , and $XW_V^h$ , followed by $W_O^h$ ; summing over heads is block matrix multiplication of the concatenated head outputs by the stacked output maps $W_O^h$ , and that attention is a Nadaraya–Watson smoother is that topic’s Proposition 1. For (ii), the entries of $xW_1 + b_1$ are $\langle x, k_i\rangle + b_{1,i}$ , and multiplying the vector $\sigma(xW_1 + b_1)$ by $W_2$ sums its entries times the rows $v_i$ .

∎

Read (ii) beside attention. Attention takes a weighted sum of values whose keys come from the context, with weights normalized by a softmax. The MLP takes a sum of values whose keys are parameters, gated by an unnormalized $\sigma$ . Geva and colleagues (2021) read trained MLPs this way, as memories keyed on patterns in the input. The notebook checks (i) head by head against the attention topic’s attention function and (ii) row by row, both to $10^{-12}$ .

Attention alone collapses

Every row of an attention matrix is a probability vector, so the matrix is row-stochastic, a Markov transition matrix on positions. formalML’s random-walks topic proves that a mixing chain forgets where it started (its Theorem 3), by a spectral argument about one fixed chain. A stack applies a different matrix at every layer, so we need a bound that survives a product of different matrices, and Dobrushin’s contraction is that tool. Write $\operatorname{diam}(X) = \max_{i,j} \lVert x_i - x_j \rVert$ for the diameter of the cloud of rows of $X$ , in any norm on $\mathbb{R}^d$ , and $\lVert W \rVert$ for the operator norm of $x \mapsto xW$ in the same norm.

Definition 2 (The Dobrushin coefficient; the smoothing stack).

For a row-stochastic $A \in \mathbb{R}^{n \times n}$ , the Dobrushin coefficient is

\tau(A) = \tfrac12 \max_{i,j} \lVert A_{i\cdot} - A_{j\cdot} \rVert_1 = 1 - \min_{i,j} \sum_k \min(A_{ik}, A_{jk}) \in [0, 1].

The pure attention stack is $X_{l+1} = A_l X_l W_l$ , and the stack with the skip is $X_{l+1} = X_l + A_l X_l W_l$ .

The two forms of $\tau$ agree because $\lvert a_k - b_k \rvert = a_k + b_k - 2\min(a_k, b_k)$ for any numbers, and two probability vectors each sum to $1$ . So $\tau(A) = 0$ exactly when all rows are equal, and $\tau(A) = 1$ when two rows have disjoint supports. Dobrushin (1956) introduced it for nonstationary Markov chains; Seneta’s book (1981) calls it a coefficient of ergodicity.

Theorem 1 (Attention contracts toward consensus).

Let $A$ be row-stochastic.

(i) $\operatorname{diam}(AX) \le \tau(A)\operatorname{diam}(X)$ , in every norm.

(ii) If each row of a score matrix $S$ has range $\max_k S_{ik} - \min_k S_{ik} \le r$ , then $\tau(\operatorname{softmax} S) \le 1 - e^{-r}$ .

(iii) The pure attention stack satisfies $\operatorname{diam}(X_L) \le \bigl(\prod_{l \lt L} \tau(A_l)\, \lVert W_l \rVert\bigr) \operatorname{diam}(X_0)$ .

Proof.

For (i), fix rows $i$ and $j$ and let $c = A_{i\cdot} - A_{j\cdot}$ . Then $\sum_k c_k = 0$ and $\lVert c \rVert_1 = 2\tau_{ij}$ , with $\tau_{ij} \le \tau(A)$ . If $\tau_{ij} = 0$ the two rows are equal. Otherwise the positive and negative parts of $c$ , which sum to zero together, have equal mass $\tau_{ij}$ , so $c = \tau_{ij}(p - q)$ for two probability vectors $p$ and $q$ . Then

(AX)_i - (AX)_j = \sum_k c_k x_k = \tau_{ij}\Bigl(\sum_k p_k x_k - \sum_m q_m x_m\Bigr) = \tau_{ij} \sum_{k,m} p_k q_m (x_k - x_m),

whose norm is at most $\tau_{ij} \sum_{k,m} p_k q_m \lVert x_k - x_m \rVert \le \tau(A) \operatorname{diam}(X)$ . Maximize over $i$ and $j$ .

For (ii), every entry of row $i$ satisfies $A_{ik} = e^{S_{ik}} / \sum_m e^{S_{im}} \ge e^{\min_m S_{im}} / (n\, e^{\max_m S_{im}}) \ge e^{-r}/n$ . So $\sum_k \min(A_{ik}, A_{jk}) \ge e^{-r}$ for every pair of rows, and Definition 2 gives $\tau \le 1 - e^{-r}$ .

For (iii), $\lVert x_i W - x_j W \rVert \le \lVert W \rVert\, \lVert x_i - x_j \rVert$ , so $\operatorname{diam}(XW) \le \lVert W \rVert \operatorname{diam}(X)$ . By (i), $\operatorname{diam}(A_l X_l W_l) \le \tau(A_l)\operatorname{diam}(X_l W_l) \le \tau(A_l)\lVert W_l \rVert \operatorname{diam}(X_l)$ , and we induct on $l$ .

∎

Part (ii) is the attention topic’s temperature talking. Scores at inverse temperature $\beta$ have range proportional to $\beta$ , so a wide kernel (small $\beta$ , high temperature) keeps $\tau$ well below $1$ and the collapse in (iii) is geometric. Dong, Cordonnier, and Loukas (2021) prove more: pure self-attention with value maps loses rank doubly exponentially in depth. We prove the geometric rate and cite theirs.

The laboratory’s cloud makes this concrete. Ten tokens in the plane pass through twelve layers of $A_l = \operatorname{softmax}(\beta X_l X_l^\top)$ with $W = I$ , so queries, keys, and values are the tokens themselves. The first layer’s $\tau$ is $0.188$ , $0.369$ , $0.667$ , $0.940$ , and $0.999$ at $\beta = 0.25$ , $0.5$ , $1$ , $2$ , and $4$ . For $\beta \le 2$ , the pure stack’s diameter falls below $10^{-10}$ by layer twelve, under the bound of (iii) at every layer. At $\beta = 4$ , $\tau$ stays above $0.97$ through layer twelve, so the bound falls slowly: it allows $2.333$ at layer twelve, and the diameter is $1.997$ . Attention here is not close to the identity. The scores are dot products, so long tokens win: six of the ten tokens put less than half their weight on themselves, and every row’s largest weight in the first layer falls on one of four long tokens, which act as leaders. By layer twelve the cloud has gathered into two tight clusters, of eight tokens and two, with a relative spread of $0.614$ . It looks frozen, and it is not. Run on, and the clusters merge: the diameter is $0.208$ at layer sixteen and $1.07 \cdot 10^{-8}$ at layer eighteen, still under the bound. Theorem 1 holds throughout; over twelve layers its rate is too slow to force the collapse. Geshkovski and colleagues (2023) observe two such timescales numerically in a related dynamics, a continuous-time flow of tokens on a sphere: clusters form fast and merge slowly, toward the single point they prove is the long-time limit.

Proposition 2 (The skip keeps the diameter).

For $X' = X + AXW$ , $\operatorname{diam}(X') \ge (1 - \tau(A)\lVert W\rVert)\operatorname{diam}(X)$ , and every pair satisfies $\lVert x'_i - x'_j \rVert \ge \lVert x_i - x_j \rVert - \tau(A)\lVert W\rVert \operatorname{diam}(X)$ . Over a stack, if every $\tau(A_l)\lVert W_l \rVert \le 1$ , then

\operatorname{diam}(X_L) \ge \prod_{l \lt L} \bigl(1 - \tau(A_l)\lVert W_l\rVert\bigr) \operatorname{diam}(X_0) \ge \Bigl(1 - \sum_{l \lt L} \tau(A_l)\lVert W_l \rVert\Bigr)\operatorname{diam}(X_0).

Proof.

For any pair, $x'_i - x'_j = (x_i - x_j) + \bigl((AXW)_i - (AXW)_j\bigr)$ , and by Theorem 1 (i) applied to $XW$ , with $\operatorname{diam}(XW) \le \lVert W \rVert \operatorname{diam}(X)$ , the second term has norm at most $\tau(A)\lVert W \rVert \operatorname{diam}(X)$ . The reverse triangle inequality gives the pairwise claim, and a pair with $\lVert x_i - x_j \rVert = \operatorname{diam}(X)$ gives the first claim. Iterating gives the product, and $\prod_l (1 - a_l) \ge 1 - \sum_l a_l$ for $a_l \in [0, 1]$ by induction.

∎

Pure attention multiplies the diameter by at most $\tau\lVert W\rVert$ per layer. The skip adds each token back, so every old difference survives inside the new one, beside a correction of norm at most $\tau\lVert W\rVert \operatorname{diam}(X)$ , and a stack of small updates, with $\sum_l \tau(A_l)\lVert W_l \rVert \lt 1$ , cannot erase the diameter. A smaller gap can go. Take $X = (0, 0.1, 10)^\top$ , $W = I$ , and $A$ with rows $(0.9, 0, 0.1)$ , $(0.91, 0, 0.09)$ , and $(0.9, 0, 0.1)$ , so $\tau(A) = 0.01$ . Then $X + AX = (1, 1, 11)^\top$ : the first two tokens merge, and the diameter stays $10$ . The proposition protects the widest gap, and any gap wider than the correction, not every pair. In all five of the laboratory’s runs with the skip, the diameter never shrinks.

What the skip keeps is the diameter, though, not its share of the stream. If every row of $A$ is the same, $AX$ is one common vector, and $X + AX$ adds that vector to every token: the differences stay where they were while the common part grows. A wide kernel heads there, and not by staying wide. At $\beta = 0.25$ the first $A$ is close to uniform, but the stream grows, the scores grow with it, and $\tau$ runs $0.188$ , $0.246$ , $0.354$ , $0.559$ , $0.713$ , $0.162$ , $0.001$ , $0$ over the first eight layers. It falls at the end because the rows agree, not because they are flat: within seven layers every row sits on one token, the same for all of them. From there the cloud only translates. Its diameter stops at $3.58$ times its start, and its relative spread $\lVert X - \mathbf{1}\bar{x}^\top \rVert_F / \lVert X \rVert_F$ halves each layer, to $0.0016$ at layer twelve. Proposition 2 is a statement about the first quantity, not the second.

LayerNorm is a sphere in a hyperplane

Every sublayer of a pre-LN block reads the stream through LayerNorm, so its geometry is the geometry of every read.

Definition 3 (LayerNorm and RMSNorm).

For $x \in \mathbb{R}^d$ , $d \ge 2$ , not a multiple of $\mathbf{1}$ , with mean $\mu$ and variance $\sigma^2 = \frac1d \sum_k (x_k - \mu)^2 \gt 0$ , LayerNorm is $\operatorname{LN}(x) = (x - \mu\mathbf{1})/\sigma$ , with its learned gain $1$ and bias $0$ , as at initialization. For $x \ne 0$ , RMSNorm is $\operatorname{RMS}(x) = \sqrt d\, x/\lVert x \rVert$ .

LayerNorm is Ba, Kiros, and Hinton’s (2016), who add a small constant to $\sigma^2$ that we take to be $0$ ; RMSNorm is Zhang and Sennrich’s (2019). Picture $d = 3$ : the plane $x + y + z = 0$ through the origin, and on it the circle of radius $\sqrt 3$ . LayerNorm drops each point onto the plane, then moves it along its ray to the circle, outward or inward.

Proposition 3 (LayerNorm is a sphere in a hyperplane).

Let $d \ge 2$ , $P = I - \mathbf{1}\mathbf{1}^\top/d$ , and $Px \ne 0$ . Then $\operatorname{LN}(x) = \sqrt d\, Px / \lVert Px \rVert$ : the orthogonal projection onto the hyperplane $\mathbf{1}^\perp$ , followed by the radial projection onto the sphere of radius $\sqrt d$ inside it. Consequently $\operatorname{LN}(cx + s\mathbf{1}) = \operatorname{LN}(x)$ for $c \gt 0$ and any $s$ , $\operatorname{LN} \circ \operatorname{LN} = \operatorname{LN}$ , and $\operatorname{LN}(x) = \operatorname{RMS}(Px)$ .

Proof.

$Px = x - \mu\mathbf{1}$ , and $\sigma^2 = \lVert Px \rVert^2/d$ , so $(x - \mu\mathbf{1})/\sigma = Px / (\lVert Px \rVert/\sqrt d)$ . Next, $P(cx + s\mathbf{1}) = cPx$ , and normalizing removes $c \gt 0$ . A point $y$ already on the sphere in $\mathbf{1}^\perp$ has $Py = y$ and $\lVert y \rVert = \sqrt d$ , so $\operatorname{LN}(y) = y$ . The last identity is Definition 3 applied to $Px$ .

∎

Brody, Alon, and Yahav (2023) give this reading and argue that each of its two parts matters to the attention that follows. The first part is a projection of the kind formalStatistics’ linear-regression topic builds least squares from (its Theorem 1): centering is the projection off the intercept direction.

Corollary 1 (LayerNorm's Jacobian).

Let $d \ge 2$ , $Px \ne 0$ , and $u = Px / \lVert Px \rVert$ . Then

\frac{\partial \operatorname{LN}}{\partial x} = \frac{\sqrt d}{\lVert Px \rVert}\bigl(P - uu^\top\bigr),

a symmetric matrix that sends both $\mathbf{1}$ and $x$ itself to $0$ . Its singular values are $\sqrt d / \lVert Px \rVert$ , with multiplicity $d - 2$ , and $0$ , twice.

Proof.

With $c = Px$ , the differential of $\sqrt d\, c/\lVert c \rVert$ is $(\sqrt d / \lVert c \rVert)(I - uu^\top)\, dc$ , and $dc = P\, dx$ . Since $u \perp \mathbf{1}$ , $uu^\top P = uu^\top$ , which gives the formula. $P\mathbf{1} = 0$ and $u^\top \mathbf{1} = 0$ , so $\mathbf{1}$ is sent to $0$ . $Px = c$ and $u^\top x = u^\top Px = \lVert c \rVert$ , so $(P - uu^\top)x = c - \lVert c \rVert u = 0$ . Finally $P - uu^\top$ is the orthogonal projection onto the $(d-2)$ -dimensional complement of $\operatorname{span}\{\mathbf{1}, u\}$ .

∎

LayerNorm does not care how long its input is, and its gradient says so: it loses the radial component and is scaled by $1/\lVert Px \rVert$ . That scale is where depth enters. For a standard Gaussian vector, $x \sim \mathcal{N}(0, I_d)$ , $\lVert Px \rVert$ concentrates near $\sqrt d$ , the thin shell of formalRAG’s high-dimensional-geometry topic (its Theorem 1), so at large $d$ LayerNorm’s radial step barely moves such a vector: the notebook’s median $\lVert \operatorname{LN}(x) - Px \rVert / \lVert Px \rVert$ is below $0.025$ at $d = 512$ and above $0.1$ at $d = 8$ . Measured from $x$ itself, so that the centering counts too, the median at $d = 512$ is $0.044$ . The radius belongs to unit variance: at $\mathcal{N}(0, 9I_d)$ the same radial change is about $2/3$ . LayerNorm acts when the stream’s norm drifts away from $\sqrt d$ , and in a pre-LN stack it does.

The identity path

Theorem 2 (The identity path).

(i) For a residual stack $x_{k+1} = x_k + f_k(x_k)$ , with $J_k$ the Jacobian of $f_k$ at the point the forward pass visits,

\frac{\partial x_L}{\partial x_l} = (I + J_{L-1}) \cdots (I + J_l) = \sum_{\{k_1 \lt \dots \lt k_m\} \subseteq \{l, \dots, L-1\}} J_{k_m} \cdots J_{k_1},

a sum of $2^{L-l}$ terms in which the empty subset contributes $I$ . So the gradient of a loss at $x_l$ is the gradient at $x_L$ , unchanged, plus a term for every nonempty set of sublayers.

(ii) In a pre-LN stack with $d \ge 3$ , $f_k = F_k \circ \operatorname{LN}$ , so $J_k = J_{F_k}\, J_{\operatorname{LN}}(x_k)$ with $\lVert J_{\operatorname{LN}}(x_k) \rVert = \sqrt d / \lVert P x_k \rVert$ . As the stream grows, that LayerNorm factor of $J_k$ shrinks, and the identity term has no such factor. This bounds one factor of each branch, not a whole gradient.

(iii) In a post-LN stack, $x_{k+1} = \operatorname{LN}(x_k + F_k(x_k))$ and $\partial x_{k+1}/\partial x_k = J_{\operatorname{LN}}(x_k + F_k(x_k))\,(I + J_{F_k})$ . Every path, the skip-only path included, passes through $L - l$ LayerNorm Jacobians.

Proof.

For (i), the multivariate chain rule of formalCalculus’ Jacobian topic (its Theorem 2) gives $\partial x_{k+1}/\partial x_k = I + J_k$ and composes these by multiplication. Multiplying out the product of $L - l$ binomials chooses, at each layer, either $I$ or $J_k$ , which gives one term per subset with its factors in layer order. Parts (ii) and (iii) are the chain rule applied to the two block forms, with Corollary 1 for the norm.

∎

Veit, Wilber, and Belongie (2016) read (i) as an ensemble of paths of different lengths. It is also the attention topic’s one-hop property (its Proposition 5) one level up: there, the gradient from a later position reached an earlier one in a single hop through an attention weight; here, the gradient from the top of the stack reaches every layer through a single identity.

Parts (ii) and (iii) are where LayerNorm’s placement matters. Xiong and colleagues (2020, their Theorem 1) show by mean-field arguments that at initialization the gradient of the last block’s MLP weights is $O(d\sqrt{\ln d})$ in a post-LN stack and $O(d\sqrt{\ln d / L})$ in a pre-LN one, and they trace post-LN’s need for learning-rate warmup to the first bound. We do not prove their theorem. We measure the scaling it predicts, on random stacks with $d = 64$ , four heads, and a ReLU MLP of width $256$ , read out by a random linear map through one final LayerNorm, as root-mean-squares over eight draws:

depth $L$	pre-LN, last $\lVert\partial\mathcal{L}/\partial W_2\rVert_F$	post-LN	pre-LN $\lVert x_L \rVert$
$6$	$10.169$	$22.019$	$24.659$
$12$	$7.673$	$23.416$	$33.641$
$24$	$5.218$	$22.718$	$49.281$

From $L = 6$ to $24$ the pre-LN gradient falls by a factor of $1.95$ , against $\sqrt 4 = 2$ ; the post-LN gradient does not move, and at $L = 24$ it is $4.35$ times the pre-LN one. The last column is the mechanism, and it is the readout’s. In a pre-LN stack each block adds to the stream, its squared norm grows linearly in depth, and the gradient reaches the last block through the final LayerNorm, whose Jacobian scales as $1/\lVert Px_L \rVert$ , so like $1/\sqrt L$ : the first and last columns multiply to $251$ , $258$ , and $257$ . It is a statement about the last block, not every branch. The first block’s gradient is $23.403$ , $24.208$ , and $23.904$ at the same depths, and does not shrink. A post-LN stack pins the stream at $\lVert x \rVert = \sqrt d = 8$ . The shape of the gradient across layers depends on the initialization, and ours differs from the figure Xiong and colleagues report; we claim the scaling, not the shape.

The stream as a channel

To see the stream at work we remove everything that is not attention.

Definition 4 (The attention-only model; circuits; composition).

The attention-only model has an embedding $E \in \mathbb{R}^{V \times d}$ and an unembedding $U \in \mathbb{R}^{d \times V}$ ; it reads tokens $t_1, \dots, t_n$ as $x_0 = E[t]$ (one row per token), updates $x_{l+1} = x_l + \operatorname{MHA}_l(x_l)$ with causal heads, rotary positions, and query and key biases, and outputs logits $x_L U$ . It has no MLP and no LayerNorm. Head $(l, h)$ writes through its OV circuit $W_{OV}^{lh} = W_V^{lh} W_O^{lh} \in \mathbb{R}^{d \times d}$ , of rank at most $d_h$ . A path is a chain of heads $(l_1, h_1), \dots, (l_m, h_m)$ with $l_1 \lt \dots \lt l_m$ . The composition score of heads $(0, h)$ and $(1, h')$ is $\lVert W_{OV}^{0h} M \rVert_F / (\lVert W_{OV}^{0h}\rVert_F \lVert M \rVert_F)$ , with $M = W_K^{1h'}$ for K-composition, $W_Q^{1h'}$ for Q-composition, and $W_{OV}^{1h'}$ for V-composition.

This is the setting of Elhage and colleagues’ (2021) framework, which names the residual stream and the two circuits. The OV circuit’s rank bound is the singular-value picture of formalML’s SVD topic (its Theorem 1), and it is the same kind of bound as the attention topic’s one-head rank bound (its Proposition 7), which limits $\operatorname{rank} \log A$ , seen from the write side. The composition score measures how much of what one head writes the other reads.

Theorem 3 (The path expansion).

Let $A^{lh}$ be the attention matrix of head $(l, h)$ in the forward pass. Then

x_L U = \sum_{\text{paths}} A^{l_m h_m} \cdots A^{l_1 h_1}\; x_0\; W_{OV}^{l_1 h_1} \cdots W_{OV}^{l_m h_m}\; U,

where the empty path contributes the direct path $x_0 U = E[t]U$ . With $L$ layers of $H$ heads there are $(1 + H)^L$ paths.

Proof.

Each layer is $x_{l+1} = x_l\, I + \sum_h A^{lh} x_l W_{OV}^{lh}$ . An attention matrix acts on the token index, from the left, and $W_{OV}$ on the feature index, from the right, so $A(XW) = (AX)W$ and the two kinds of factor can be collected separately. Substituting $x_l$ into $x_{l+1}$ for each $l$ and distributing chooses, at each layer, either the skip or one head, and the chosen heads form a path. Collecting the left factors and the right factors gives the formula; at each layer there are $1 + H$ choices.

∎

The direct path is a function of the current token alone: a bigram model, built into the embedding and unembedding. The one-head paths copy what a head attends to; the two-head paths compose. The expansion is exact, and the notebook checks it against the forward pass of the trained model of the next section, whose two layers of two heads give $9$ paths, to about $2 \cdot 10^{-14}$ . Dong, Cordonnier, and Loukas use a decomposition of this kind to prove their collapse rate. One caution keeps it honest: the $A^{lh}$ depend on the stream, so the expansion is linear in the values, not in the attention. Q- and K-composition live inside the patterns, and only V-composition appears in the products of OV circuits.

Proposition 4 (Orthogonal subspaces do not talk).

If $W_{OV}^{0h} W_K^{1h'} = 0$ for every layer-0 head $h$ , then the keys of head $(1, h')$ computed from $x_1$ equal those computed from $x_0$ , and its K-composition scores are $0$ . If also $W_{OV}^{0h} W_Q^{1h'} = 0$ for every $h$ , the same holds for its queries, and its attention pattern is the same with or without layer 0’s output.

Proof.

$x_1 W_K^{1h'} = x_0 W_K^{1h'} + \sum_h A^{0h} x_0\, W_{OV}^{0h} W_K^{1h'} = x_0 W_K^{1h'}$ , and the key bias and the rotation are applied after this product. The same computation with $W_Q^{1h'}$ gives the queries, and the pattern depends on the stream only through queries and keys.

∎

The stream is one $d$ -dimensional space shared by every head, and heads communicate only through the subspaces where one writes and another reads. The notebook builds a model whose layer-0 heads write into the first $16$ of $32$ coordinates while the layer-1 keys read only the last $16$ , and computing those keys from $x_0$ leaves the pattern unchanged to $10^{-12}$ . The keys are only half of a score, though. The queries still read $x_1$ , and removing layer 0’s output moves each row of the pattern by a mean total variation of $0.104$ . Restrict the queries to the last $16$ coordinates as well, and the pattern with and without layer 0’s output agrees to $10^{-12}$ .

Induction heads

Here is a task that needs two heads to talk. Draw a random block of tokens from a vocabulary of $V = 256$ , of period $p$ between $6$ and $24$ , and repeat it to $48$ tokens. From position $p$ on, the next token is predictable, by a rule Olsson and colleagues (2022) call induction: find where the current token occurred before, and copy what came after it, [A][B] … [A] → [B]. We train attention-only models with $d = 32$ and two heads of width $16$ per layer, rotary positions, and hand-written gradients, for $1000$ Adam steps on fresh sequences, and test on $1000$ more, scoring the positions from $p$ on:

model	loss (nats)	accuracy
two layers	$0.157$	$0.974$
one layer	$4.280$	$0.108$
two layers, layer 1’s keys read from $x_0$	$6.135$	$0.075$

Guessing uniformly costs $\ln 256 = 5.545$ nats. The two-layer model solves the task; the one-layer model does not come close. The last row cuts one connection: layer 1’s keys are computed from the embeddings instead of the stream, so they cannot read anything layer 0 wrote, and the solved model falls apart, below even the one-layer model.

The two-layer model’s heads say how it works. Head $(0, 0)$ puts $0.887$ of its weight on the previous token, on average. Its content-free score, the part that comes from its query and key biases alone, is $b_q^\top R(-\Delta)\, b_k / \sqrt{d_h}$ for a key $\Delta$ positions back: the positional-encodings topic’s kernel over offsets, and it peaks at $\Delta = 1$ , with $7.783$ against $6.709$ at $\Delta = 0$ . So head $(0,0)$ is a positional head: through position $37$ its weight on the previous token is about $0.98$ , and it writes a copy of each token’s predecessor into the stream at that token. The mean hides a drop. From position $39$ the mean weight is under $0.5$ , shared about evenly with the key $39$ positions back, where the kernel has its next peak, $6.687$ . The model’s accuracy does not fall there: it stays between $0.968$ and $0.983$ from position $39$ on. Both layer-1 heads put $0.593$ and $0.594$ of their weight on the position $i - p + 1$ , just after the current token’s occurrence one period back, the induction target, and $0.943$ and $0.940$ on the positions after any earlier occurrence. Nearly all of those hold a correct copy of the next token. A block can repeat a token by chance, and then, at $4.85\%$ of the scored positions, one of them holds another token; the heads put about $0.02$ of their weight on such positions. Every K-composition score is between $0.245$ and $0.272$ , against $0.046$ to $0.075$ for Q-composition and $0.036$ to $0.076$ for V-composition, and above the largest K-score the same model had at its initialization, $0.188$ . Those scores do not separate the two layer-0 heads, so we cut them one at a time. Give layer 1’s keys $x_0$ plus head $(0, 0)$ ‘s write alone, and the model stays solved, at $0.228$ nats and accuracy $0.969$ . Give them head $(0, 1)$ ‘s write alone, and it fails, at $5.391$ nats and $0.084$ . Their keys read what head $(0, 0)$ wrote.

Proposition 5 (What one layer's keys see).

In a one-layer model, the score $s_h(i, j)$ of head $h$ is a function of $(t_i, i, t_j, j)$ alone. In particular it does not change when the token at position $j - 1$ changes. In two layers, the key at $j$ reads $x_1[j] = x_0[j] + \sum_h \sum_k A^{0h}_{jk}\, x_0[k]\, W_{OV}^{0h}$ , which contains $x_0[j-1]$ when some layer-0 head attends to the previous token.

Proof.

In one layer, the query at $i$ is the rotation by $i$ of $E[t_i]W_Q + b_q$ and the key at $j$ is the rotation by $j$ of $E[t_j]W_K + b_k$ . The two-layer formula is Definition 4’s update at $l = 0$ .

∎

Induction must attend to the position whose predecessor is the current token, and a one-layer score cannot see predecessors, so it cannot pick that position out by its content. Two layers can, through K-composition. Proposition 5 is about information, not a bound on loss: a one-layer model can still raise the probability of every token it has seen, and how close it gets to that in training is a measurement. The notebook changes the token before position $j$ and recomputes the last layer’s score at $(i, j)$ : in the one-layer model it moves by less than $10^{-12}$ , and in the two-layer model by $4.45$ nats on average.

Theorem 3 shows where the answer lives. At the scored positions, the paths that end in a layer-1 head carry more than $90\%$ of the correct token’s centered logit (the median over positions), and the direct path and layer-0 paths alone rank the correct token first less than $5\%$ of the time. The logit lens (nostalgebraist 2020) reads the stream through the unembedding between layers, $x_l U$ . This model has no final LayerNorm, so the lens is its own readout, and the probability of the correct token after the embedding, after layer $0$ , and after layer $1$ is $0.004$ , $0.004$ , and $0.958$ , with chance $1/256 = 0.004$ . The answer exists only after the second layer reads the first. On models with a final LayerNorm the lens reads $U$ applied to $\operatorname{LN}(x_l)$ , a heuristic that Belrose and colleagues (2023) refine by training a probe for each layer.

The stream, interactively

The laboratory has two figures. Figure A is the geometry: the ten tokens of the collapse experiment, layer by layer, with and without the skip at five temperatures, beside their diameters and Theorem 1’s bound, and the LayerNorm picture at $d = 3$ , where you can center a cloud, normalize it, and scale it without moving the result. Figure B is the trained model: the previous-token head’s and the induction head’s attention on a test sequence of period $11$ , the correct token’s logit split into its $9$ paths at any position you pick, the logit lens, the ablation, and the filing from the case study below.

A · Smoothing without a self — collapse, the skip, LayerNorm, and the identity path

layerl = 12

pure attentionwith the skipTheorem 1 bound, ∏τ · diam(X₀)

diameter

τ of layer 12

0.000

relative spread ‖X − 1x̄ᵀ‖/‖X‖

0.0000

diam(X₁₂) / diam(X₀)

A wide kernel (small β) keeps τ below 1, and the pure stack shrinks the cloud to a point, under Theorem 1’s bound. A narrow kernel (β = 4) lets a few long tokens win the rows: τ stays near 1, the bound falls slowly, and at layer 12 the cloud holds two clusters, which merge by layer 18, past this slider. With the skip the diameter never shrinks in these runs (Proposition 2 is the floor beneath them), but at β = 0.25 the relative spread still falls: the skip keeps the diameter, not its share.

B · Two heads that compose — the trained two-layer model on a sequence of period 11

position ii = 47

token at i, then the next

84 → 65

largest weight in row i

0.55 on j = 26

induction target i − p + 1

The previous-token head’s weight sits just left of the diagonal, at j = i − 1, until position 39, where part of it moves to j = i − 39. The induction head attends to the earlier positions whose predecessor is the current token: i − p + 1 (outlined), and one or more periods further back on a long repetition. Unless the block repeats a token by chance, each holds a copy of the next token. Its keys can find them because the previous-token head wrote predecessors into the stream (Proposition 5).

Finance case study: completing a company’s name

§ Finance case study: completing a company's name

A filing names a company in full, “Tamsin Retail Group”, and names it again a paragraph later. Completing the second mention from the first is induction on real document structure, and it asks whether a model’s knowledge of a name lives in its weights or in the text in front of it.

The shared synthetic corpus lets us separate the two. Its roster of public companies splits $75/25$ by sector, and we build three sets of filings: $300$ about the training companies to train on, $100$ new ones about the same companies, and $100$ about the held-out companies, whose names never occur in training. Words are hashed to $256$ symbols, by the hash the positional-encodings topic reduced to $8$ . A query is a word of a multi-word company name, other than its last, whose name the same filing already gave in full within the $48$ words before it; the answer is the name’s next word. There are $269$ queries on the seen companies and $336$ on the unseen ones. We ask five readers: the bigram table counted on the training filings; a two-layer model of the previous section’s shape trained on the training filings, and that model’s direct path alone; and the one- and two-layer induction models of the previous section, trained only on repeated random sequences.

reader	seen companies	unseen companies
bigram table	$0.576$	$0.000$
filing-trained model, direct path	$0.717$	$0.000$
filing-trained model	$0.814$	$0.000$
induction model, one layer	$0.056$	$0.071$
induction model, two layers	$0.710$	$0.768$

Everything the filing-trained model knows about names is in its weights. It completes $81\%$ of the names it was trained on, more than the bigram table, and none of the names it was not, exactly like the bigram table. Its direct path alone gets $72\%$ of the seen names, so most of that knowledge sits in the bigram part of the expansion. The induction model never saw a word of finance. It completes $71\%$ of the seen names and $77\%$ of the unseen ones, because it reads each name off the filing’s own earlier text. This is the attention topic’s kernel cache made into a circuit.

The filing-trained model has the same two layers and could have learned to copy. It did not. A natural guess at the reason is that a name was never new in training, so memorizing names was the cheaper solution, but we have not tested that. The laboratory’s filing is a held-out filing that names Tamsin Retail Group in full eight times; the two-layer induction model completes $11$ of its $12$ queries, and the bigram table and the filing-trained model complete none.

Where this goes next

Contextual embeddings. The same block without the causal mask, trained to fill in masked tokens, gives contextual word vectors; the contextual-embeddings topic builds it on this one.
Scaling laws. Depth and width are the axes along which models scale, and the growing pre-LN stream of Theorem 2 (ii) is one reason depth is not free; the scaling-laws topic takes the empirical side.
Euler steps. A residual block $x + f(x)$ is one step of Euler’s method for the flow $x' = f(x)$ , whose global error formalCalculus’ numerical-ODEs topic bounds (its Theorem 1); a deep stack of small updates discretizes a trajectory of the tokens, and Geshkovski and colleagues study the clusters of a related flow.
Superposition. A $32$ -dimensional stream carried $256$ embeddings, drawn at random and then trained, that the induction head still copied. Many nearly orthogonal directions fit in few dimensions: the trained table’s root-mean-square cosine is $0.201$ . That is the concentration idea behind formalRAG’s Johnson–Lindenstrauss topic, though its Theorem 3 needs more than $133$ dimensions for $256$ points, so it does not cover this case.

Honest caveats

⚠ Rigor flag

The collapse rate here is geometric. Theorem 1 (iii) proves it; Dong, Cordonnier, and Loukas’s doubly exponential rate is cited, not reproved. At narrow kernels, where $\tau$ is near $1$ , the bound falls slowly: at $\beta = 4$ it still allows $2.333$ at layer twelve, where the cloud holds two clusters, and they merge by layer eighteen.
The skip keeps the diameter, not every pair and not proportions. Proposition 2’s example merges two tokens outright. At $\beta = 0.25$ the laboratory’s cloud grows $3.58$ times while its relative spread falls to $0.0016$ .
Xiong and colleagues’ bounds are cited and measured, not proved. They are mean-field results at initialization under a concentration assumption. We measure the predicted $L^{-1/2}$ scaling on one configuration, in the last block, where it comes from the readout’s LayerNorm; the first block’s gradient does not shrink, and the per-layer profile’s shape depends on the initialization and is not claimed.
The path expansion is exact given the patterns. It is linear in the values and not in the attention, and with LayerNorm it is exact only with each token’s scale frozen. The trained models here have no LayerNorm for that reason.
Readings, not theorems. “Channel,” “key–value memory,” and “logit lens” name ways of looking. The lens is exact for the models here, which have no final LayerNorm, and a heuristic elsewhere.
Proposition 5 limits sight, not loss. The one-layer model’s $4.280$ nats is a measurement after $1000$ steps, not a lower bound.
Toy models. The induction heads are shown in attention-only models with $d = 32$ . Their role in the in-context learning of large models is Olsson and colleagues’ empirical claim, with causal evidence for small attention-only models and correlational evidence for large ones.
Synthetic filings. The case study uses the shared synthetic corpus with hashed words. It shows a mechanism, not a name tagger, and its unseen-company misses ( $23\%$ of queries) are not analyzed here.

Implementation

The companion notebook (notebookPath) imports transformer_block_residual_stream.py, which owns every number on this page. That module imports the positional-encodings topic’s module for its rotations, and through it the attention topic’s attention function and rounding conventions, the neural-language-models topic’s gradient checker and Adam constants, and the shared finance corpus. It trains three attention-only models from scratch in NumPy, two on repeated random sequences and one on filings, with hand-written gradients, and runs a harness of $22$ assertion-based tests in about twenty-five seconds on a laptop CPU. The tests cover:

Proposition 1 against the attention topic’s attention function, and the MLP as a key–value sum.
Proposition 3 and Corollary 1 at $d = 2$ , $3$ , $8$ , and $64$ , against central differences, and the thin shell.
Theorem 1 on $300$ random instances in three norms, and Theorem 1 and Proposition 2 on the laboratory’s cloud and on $100$ random stacks with random value maps; the narrow-kernel run through layer eighteen, the wide-kernel run landing on one token, and the merged pair.
Theorem 2 against the chain rule and the subset sum; the pre-LN and post-LN gradient scaling.
Theorem 3 on the trained model and on a random three-layer model; Proposition 4 on a constructed model, keys alone and keys with queries; Proposition 5 on both trained models.
The induction experiments, the heads by position, the composition scores, the per-head cut, the paths, and the logit lens; the finance table.
Every hand-written gradient against central differences: the attention-only model at random coordinates, and both stacks at every coordinate.

The laboratory reads its data from a JSON file that one of those tests requires to equal the module’s viz_payload() exactly, so the figure cannot drift from the code.

Overview & Motivation

The block, assembled

Attention alone collapses

LayerNorm is a sphere in a hyperplane

The identity path

The stream as a channel

Induction heads

The stream, interactively

Finance case study: completing a company’s name

Where this goes next

Honest caveats

Implementation

References & Further Reading