Unembedding and the Logit Lens

The unembedding matrix U projects a model's final residual stream vector into logits over the vocabulary, feeding softmax to produce the next-token distribution. It is normally applied exactly once, after the last transformer block and its final normalisation. In 2020 the pseudonymous researcher nostalgebraist noticed something that should not obviously work: take the residual stream at an intermediate layer, say layer 15 of 32, and apply that same final unembedding matrix directly to it, skipping every remaining layer. The result is often a coherent, if less confident and sometimes incorrect, guess at the next token, one that visibly sharpens and corrects itself layer by layer as you move deeper. This diagnostic is the logit lens.

Why an early read is even meaningful

The trick works because of the residual stream's structure (see the-residual-stream): every layer only adds to the stream, it never resets or changes basis. The coordinate system at layer 15 is the same coordinate system the final unembedding was trained to read at layer 32, just with fewer terms summed in. Reading it early is a legitimate partial sum, not a different space entirely. The technique is cleanest when the model uses weight tying, since then the very matrix used to look up input tokens is also the one doing this readout, reinforcing that early and late layers share one consistent geometry.

What it is used for

Because it is cheap, a single matrix multiply per layer, the logit lens became a standard first diagnostic for asking "at what depth does this model settle on its answer". It has been used to study induction heads and copying behaviour, to trace where factual recall becomes visible in the residual stream, and to watch when refusal or safety-relevant decisions become expressed in the output distribution, all without any training or fine-tuning of a separate probe.

Later work refined the raw technique. The tuned lens (Belrose et al., 2023) learns a small per-layer affine correction before applying the unembedding, correcting for the fact that intermediate layers are not perfectly calibrated to the final layer's expected scale and basis, which produces materially cleaner early-layer readings than the raw logit lens.

When it falls down

• Early layers often look uninformative even when relevant computation is happening there. A flat or wrong distribution at layer 3 is not evidence the model has done nothing useful yet, only that whatever it has done is not yet expressed in the final unembedding's basis.
• It assumes intermediate layers are close enough to the final basis to be meaningfully read this way. That assumption gets weaker in models without weight tying, or in layers where the residual stream's role has shifted substantially from what the final layer expects.
• The raw version ignores per-layer scale mismatches. Layernorm statistics differ across depth, so early logits from the untuned lens can be systematically miscalibrated, which the tuned lens exists specifically to correct.
• A plausible mid-network guess is observational, not causal. Seeing a sensible next-token distribution appear early suggests, but does not prove, that this is the representation the model actually relies on; establishing that requires a causal intervention such as activation patching, not just reading the residual stream.