Interpret what the logit lens reveals when you unembed intermediate layers

MCQ·Hard·4.0 · 0·~1 min·Asked atBanana Dev Canva Observe Ai

Attempt it

The **logit lens** is a small but powerful interpretability trick. You take the residual stream after layer k of a trained transformer, apply the model's own final norm, then multiply by the unembedding matrix W_U. The result is a vocabulary-sized distribution: what the model would predict if the network ended at layer k. What you typically see is that the prediction sharpens with depth. Early layers might rank generic tokens at the top (the, of, a). Middle layers start to look on-topic. By two thirds of depth, the correct next token is often already at rank 1, and the remaining layers mostly increase its confidence. This convergence is what motivates **early-exit** techniques: if the answer is locked in at layer 20, why run layers 21-40? The lens is informative because pre-norm transformers maintain one shared residual stream that every layer reads from and writes to. The output head reads the final state of this stream. Earlier states are not exactly head-compatible, but they are close enough that the unembedding gives a meaningful signal. The three distractor options each miss the point in a different way.

**Logit lens** takes the residual stream `x_k` after layer k, applies the model's own final RMSNorm or LayerNorm, then multiplies by the unembedding W_U to produce vocabulary logits: $$\text{lens}(x_k) = \text{softmax}\!\big(\text{Norm}(x_k) \cdot W_U\big)$$ The result reads as the model's current best guess for the next token if the network ended at layer k. Plot the rank of the eventual top-1 prediction across depth and you typically see it falling from random to rank 1 well before the final layer. **Why this works.** Pre-norm transformers maintain a single shared residual stream of width d_model. Every attention and FFN sublayer reads from this stream (via its input norm) and writes back additively. The output head reads the final stream state through one last norm and W_U. Mathematically, the stream after layer k is the partial sum of all writes so far; the head's computation is the same regardless of which layer's residual you feed it. The signal is meaningful because the head was trained to read this exact additive structure. **Typical findings.** Predictions sharpen monotonically through depth in the middle of the network. Factual associations often resolve in the middle layers (a fact circuit's payload appears at layer 15-25 in a 32-layer model). Style and surface choices keep refining in the last few layers. **Failure modes.** Post-norm models have intermediate streams in a different distribution than the head expects, so raw lens noise is much higher.

TL;DR

Logit lens applies the model's own final norm and unembedding head to intermediate residual streams, revealing how the next-token guess evolves and often converges before the top layer.

Memory aid

Easy to grasp

Imagine a long relay of artists drawing the same picture, each one adding strokes on top of the last. If you stop the relay after the third artist and ask the gallery to frame what is on the canvas right now, you get a picture: rougher, fewer details, but recognizable. Keep stopping at later artists and the picture gets sharper. The gallery's framing process is the same every time. The logit lens does this for a language model. After each layer, you ask the model's own output head what it would say if the network ended here. Early layers give blurry guesses. Mid-to-late layers often have the answer locked in already, with the remaining layers just refining confidence.

Key concepts

Transformer Architecture

Concept explanation~2 min read

Everything you need to truly understand this topic: intuition, mechanics, step by step explanation, code, formulas, and worked example. Click to expand.

The logit lens is one of the cheapest and most useful interpretability tools in the transformer toolkit. The operation is trivial to describe (take an intermediate residual, push it through the model's own output head) but the implications are surprisingly deep: it reveals that modern language models often decide on their next token several layers before the network ends, and that the residual stream is a meaningfully readable workspace, not just an opaque pipeline.

The lens works because pre-norm transformers maintain a single shared residual stream that every layer writes into additively. The output head reads the final state of this stream. Nothing in the math forbids you from running the head on an earlier state; you just get a less refined version of the prediction. That less-refined version is itself a window into how the model assembles its answer.

This walkthrough explains exactly what the lens does, why it works structurally, what patterns you see in practice, and the known limits including tuned lens as the cleaner modern version.

The operation in detail

A pre-norm transformer's forward pass through L layers produces a sequence of residual stream states x_0, x_1, ..., x_L. The final output is computed by applying a final norm (RMSNorm or LayerNorm) and then the unembedding matrix W_U:

p(\text{next token}) = \text{softmax}\!\big(\text{Norm}(x_L) \cdot W_U\big)

The logit lens applies the same two operations to an intermediate state:

\text{lens}(x_k) = \text{softmax}\!\big(\text{Norm}(x_k) \cdot W_U\big)

That is the entire technique. No training. No probes. No additional parameters. You reuse the model's own final norm and unembedding head, swap in an earlier residual, and read off a vocabulary distribution.

The interpretation: this is the prediction the model would make if the network ended at layer k. It is meaningful because the head was trained to read residual states in the additive shared-stream sense, and the intermediate state is just a partial sum of the same writes.

A single line of code in PyTorch:

code

with torch.no_grad():
    hidden = model.transformer.h[k](prev_hidden)
    lens_logits = model.lm_head(model.transformer.ln_f(hidden))
    lens_top1 = lens_logits.argmax(dim=-1)

That is it. The interpretability community has spent fifteen years writing about it because the patterns it reveals are nontrivial, not because the operation itself is hard.

Why it works: the shared residual stream

Typical patterns and what they tell you

Limits, distractors, and practical use

Putting numbers to it: 2026 frontier-model context

Situations where this technique stops working.

2–4 min · Everything important, quickly.

Real products, models, and research that use this idea.

nostalgebraist's original logit lens post on GPT-2 (2020) is the canonical reference for the technique and the early-convergence observation.
Belrose et al. 2023 'Tuned Lens' paper adds per-layer linear probes that clean up the noise and gives the version most interpretability researchers now cite.

What an interviewer would ask next. Try answering before peeking at the approach.

QHow does tuned lens improve on logit lens, and what assumptions does the tuned version make?

Trains a small affine probe per layer to map intermediate residuals into the head's expected input distribution; assumes a low-rank linear bridge suffices, which holds empirically for most layers in pre-norm decoders.

2 more follow-ups an interviewer would ask next. Sign in to reveal them.

Red flags & common mistakes

The phrases that signal junior thinking. Click to expand.

Most common mistake

Thinking the lens recovers the input tokens at every layer, or that it works only on post-norm models. The residual stream is a write-target for predictions, not a verbatim copy of the input, and pre-norm is where the lens works best.

60 second bullets to scan on the way to the call.

The mathematical operation: residual after layer k, apply final norm, multiply by W_U, softmax
Why pre-norm's shared residual stream makes the operation meaningful

Primary sources. Browse if you want the original framing.