What Attention Computes

Dive into Deep Learning · §10.6

What attention computes
the residual stream · QK and OV circuits · induction heads · in-context learning in miniature

The residual stream

Attention patterns and value transformations

Each position carries a d-dimensional vector from embedding to logits. Layers never overwrite it; heads only add:

\mathbf{h}_i^{(\ell)} = \mathbf{h}_i^{(\ell-1)} + \sum_{h=1}^{H} \sum_{j \leq i} \alpha_{ij}^{\ell h}\, \mathbf{W}_{\mathrm{OV}}^{\ell h}\, \mathbf{h}_j^{(\ell-1)}

Attention-only, no FFN, no LayerNorm (Elhage et al., 2021): with the patterns fixed, the map from embeddings to logits is linear and can be expanded exactly. TinyCharLM uses this restricted architecture.

Where and what: the QK and OV circuits

Each head is described by two matrices that can be read from the checkpoint:

s_{ij} = \frac{\mathbf{h}_i^\top \mathbf{W}_{\mathrm{QK}} \mathbf{h}_j}{\sqrt{d_h}}, \qquad \mathbf{W}_{\mathrm{QK}} = \mathbf{W}_q^\top \mathbf{W}_k, \qquad \mathbf{W}_{\mathrm{OV}} = \mathbf{W}_o \mathbf{W}_v

• QK circuit: determines the attention scores.
• OV circuit: transforms the attended residual stream.
• Both d \times d, rank \leq d_h: a head moves a 32-dimensional slice.

W_QK: shape (128, 128), rank 32
W_OV: shape (128, 128), rank 32

What depth can express

With attention patterns as if–then rules, the one-layer logits expand into paths:

\mathbf{z}_i = \underbrace{\mathbf{E}\,\mathbf{e}_{x_i}}_{\textrm{low-rank bigram}} + \sum_{h}\sum_{j \leq i} \alpha_{ij}^{h}\; \underbrace{\mathbf{E}\,\mathbf{W}_{\mathrm{OV}}^{h}\,\mathbf{e}_{x_j}}_{\textrm{skip-trigram}}

• Zero layers: bigram statistics — constrained to the tied head’s low-rank Gram form \mathbf{e}_b^\top\mathbf{e}_a. One layer: skip-trigrams (“[\mathrm{A}] before [\mathrm{B}] → boost [\mathrm{C}]”), incl. copying.
• Not expressible in one layer: attend to the token after the previous [\mathrm{A}] because the key at j contains no information about position j-1.

The induction circuit needs two layers

Layer 1 writes each token’s predecessor into its stream; layer 2 matches “I am [\mathrm{A}]” against “follows [\mathrm{A}]” and copies. Whether gradient descent finds this circuit is an empirical question.

A repeated-sequence task

Each batch contains newly sampled tiled patterns. Because average bigram statistics are uniform, predicting later copies requires retrieval from the context:

pattern lengths: [6, 6]
tensor([[4, 0, 3, 4, 0, 0, 4, 0, 3, 4, 0, 0, 4, 0, 3, 4],
        [1, 0, 3, 5, 0, 7, 1, 0, 3, 5, 0, 7, 1, 0, 3, 5]])

A fixed-period result

With a fixed pattern length of 32, a one-block model obtains nearly perfect accuracy. This experiment uses bias=True:

first copy 4.42 nats, second copy 0.09 nats, second-copy accuracy 1.00

The learned fixed-offset rule

We probe the same model with patterns of length 24.

patterns of length 24: second-copy accuracy 0.02
attention mass at offset 31: 0.96
attention mass at offset 23: 0.00

• The model learned a RoPE head that attends exactly 31 positions back, without content-based matching.
• The qkv bias makes that head cheap: a constant, content-free query. With biases disabled, as in every model analyzed below, the shortcut is much less accurate.

A synthetic benchmark may admit a simpler circuit than the intended mechanism. Sampling the period removes this fixed-offset solution.

Two blocks learn to look things up

Pattern lengths are uniform from 16 to 32, so no fixed offset solves every example.

2 blocks: first copy 4.50, later copies 0.45, accuracy 0.87
1 block: first copy 4.34, later copies 3.65, accuracy 0.11

• Two blocks: later copies below half a nat, ~9 of 10 tokens right.
• One block: later-copy loss remains above three nats; copying heads provide only limited improvement.

An abrupt transition during training

• The curve remains close to the one-block curve for hundreds of steps, then falls by several nats within one or two hundred steps.
• The displayed fixed-seed run establishes neither the frequency nor the timing of this transition. Olsson et al. (2022) report a related induction transition in larger language models.

Learned attention patterns

Block 1 has a sharp line one step below the diagonal, consistent with a previous-token head. Block 2 has a stripe at offset L-1 beginning with the second copy, consistent with an induction head.

Scoring every head

Average mass on the previous token vs. on the induction target j = i + 1 - L:

block 1: previous token 0.00 0.95 0.02 0.01 | induction target 0.00 0.00 0.01 0.00
block 2: previous token 0.04 0.05 0.04 0.04 | induction target 0.67 0.71 0.66 0.65

• In this run, previous-token attention is concentrated in block 1 and induction-target attention in block 2.

Pattern completion is in-context learning

Every evaluation pattern is newly sampled, so the result is consistent with a content-based copying rule rather than memorized training sequences:

One exposure suffices for the displayed examples; restart locations are sampled independently of the token values.

Verifying copying in the weights

If block 2 copies, attending to token a should boost logit a: C_{ab} = \mathbf{e}_b^\top \mathbf{W}_{\mathrm{OV}} \mathbf{e}_a should be diagonal-dominant. This property can be checked without a forward pass:

rows whose largest entry is the diagonal: 0.94 (chance 1/64)

What attention weights do not tell you

• Weights specify where a head attends, but its effect also depends on the OV circuit. A sharp head can have little effect, while a diffuse head can implement a precise computation.
• Different weights, same predictions (Jain & Wallace, 2019); the rebuttal (Wiegreffe & Pinter, 2019) sharpened, not settled, the debate.
• Our identification used behavior (the loss split), causal probes (period shift, ablation), and weight-level checks (the OV diagonal). The attention stripes identify candidate locations but not their effects.

Feed-forward layers, normalization, and superposed features make the same analysis harder in full transformers. Only a limited set of circuits, including induction heads, has comparable supporting evidence.

Summary

• The residual stream carries representations from embeddings to logits, and attention heads add their outputs to it.
• Every head factors into a QK circuit (where) and an OV circuit (what), each with rank at most d_h.
• Increasing depth expands the representable interactions from bigrams to skip-trigrams and then to composed induction circuits.
• Trained on repetition, TinyCharLM learns the induction circuit after an abrupt loss reduction; a fixed-period task is solved by a positional shortcut instead.
• Pattern completion over never-seen tokens = in-context learning in miniature and uses a mechanism associated with in-context learning in larger language models.