Reading Attention and the Logit Lens

Heads as little routing programs

Multi-head attention gives each layer several heads, and a recurring finding is that individual heads specialise. Interpreting a head (注意力 head 解釋) means working out which tokens it moves information from and to, and what it copies. Some heads track syntax (a verb attending to its subject), some copy rare tokens forward, some attend to the previous token, some to the start-of-sequence position. The catalogue of these attention head roles is one of the most legible parts of a transformer — partly because attention patterns are easy to visualise.

Induction heads: how models learn in context

The celebrated example is the induction head (歸納 head). Given a sequence that contains '...[A][B]...[A]', an induction head notices the earlier [A][B] pair and predicts [B] again after the second [A]. In plain terms: it completes patterns it has seen earlier in the same prompt. A pair of heads working together — a previous-token head feeding an induction head — implements this copy-the-continuation algorithm, and its emergence in training lines up with a jump in the model's in-context abilities.

Induction heads matter beyond a cute trick: the same machinery, studied across models in the broader induction heads literature, is a leading mechanistic candidate for in-context learning itself — the ability to pick up a format or task from a few examples in the prompt. It is the clearest case so far of a named, reusable circuit explaining a high-level capability.

Seeing — and doubting — attention patterns

Heads are easy to look at: render the attention weights as a grid of which token attends to which, and patterns leap out — diagonals for previous-token heads, vertical stripes for heads parked on the start position. This attention pattern visualization (注意力模式視覺化) is the fastest way to form a hypothesis about a head's job. But it is only a hypothesis. A vivid pattern shows where a head looks, not whether that lookup changes the output; the warning against attention-as-explanation applies in full force here.

The logit lens: peeking at every layer

Because of the additive residual stream, the hidden state at any layer lives in the same space the final layer reads from. The logit lens (logit 透鏡) exploits this: apply the model's final unembedding to an intermediate hidden state and read off which tokens it currently favours. Run this at every layer and you watch a prediction form — early layers guessing vaguely, middle layers committing to a topic, late layers settling on the exact word. It is a one-line diagnostic that turns the depth of the network into a timeline.

The logit lens reads any layer's residual state in the final vocabulary basis: layer-normalise, multiply by the unembedding, then softmax.

# Logit lens: decode each layer's residual as if it were the last
for L in range(num_layers):
    h = residual[L][:, -1, :]          # hidden state at the final position
    logits = unembed(final_norm(h))     # reuse the model's output head
    print(L, top_token(logits))         # how the guess evolves with depth