Vol. XXVIII · Interpretability Plate 1 of 9logit lens, residual stream, intermediate predictions, tuned lens, layers, KL divergenceAE

The Logit Lens.

Read the residual stream as if the model had stopped early.

Every layer of a Transformer adds to a residual stream that the output head reads at the end (plate V.2). The logit lens (nostalgebraist, 2020) applies that final norm and head to the stream after each intermediate layer and asks what the model would predict if it stopped there. Because the stream is a running sum, its early states are already in roughly the coordinates the head understands, and the lens shows a prediction being refined layer by layer. The tuned lens (Belrose et al., 2023) trains a small affine map per layer to correct for the drift in those coordinates.

The course model micro (four layers, d = 128; plate I.9) is opened up in the page. The workshop applies the lens after the embedding and after each of micro’s four layers, over every position of 705 tokens of held-out course text, English and Chinese, and measures how often the intermediate top prediction matches the final one and the actual next token, and the KL divergence from the final distribution.

Below: agreement with the final prediction, accuracy on the next token, and KL divergence from the final distribution, at each depth.

What you will learn

Describe the logit lens.
Measure how predictions form across layers.
Read KL to the final distribution.
Name the tuned lens.

Workshop

micro on 705 tokens of held-out course text, English and Chinese, read in windows of its 128-token context. Depth 0 is the embedding; depth 4 is the final output. Agreement: the lens’s top token equals the final top token. Accuracy: the lens’s top token equals the actual next token. KL in nats.

A prediction assembled in steps

Over 705 positions, the lens after micro’s first layer agrees with the final prediction 7% of the time and is right about the next token 4% of the time; after the third layer, 26% and 11%, against 25% for the full model. The KL divergence from the final distribution falls from 4.40 nats after one layer to 2.11 after three: each layer moves the stream measurably towards the final answer, and none of them holds it alone.

The lens is the cheapest interpretability tool and a useful first look: it shows where in depth a prediction becomes available, which matters for early-exit decoding (plate XXIV.5) and for locating where a fact is retrieved. Its limits are that intermediate layers need not use the head’s coordinates — more so in large models, which is what the tuned lens corrects — and that “what the head would read” is not the same as “what the layer computes”. The next plates ask more directly what is represented and what causes what.

micro ties its embedding and output head, so the lens at depth 0 just maps each token back to itself; the interesting depths are the middle ones.

Terms

Logit lens: The final norm and head applied to an intermediate residual stream.
Tuned lens: A per-layer learned map that corrects the logit lens.
Residual stream: The running sum each layer reads from and writes to.
KL divergence: How far one distribution is from another, in nats.

The logit lens reads intermediate layers as predictions; in micro the answer forms gradually across depth.

Sources

nostalgebraist, “interpreting GPT: the logit lens”, LessWrong, 2020.
Belrose et al., “Eliciting Latent Predictions from Transformers with the Tuned Lens”, arXiv:2303.08112, 2023.
Elhage et al., “A Mathematical Framework for Transformer Circuits”, Anthropic, 2021.