The Logit Lens: Reading Predictions From Middle Layers

The logit lens takes the residual stream at layer 12 of a 24-layer model, applies the final layer norm and the unembedding matrix, and reads off a distribution over tokens — what the model would say if it had to answer from there. It is about thirty lines and it is the cheapest window into a model’s internals that exists.

The idea

A transformer’s residual stream runs the length of the network. Each layer reads it, computes something, and adds the result back. At the end, one layer norm and one matrix multiplication turn the final residual state into logits over the vocabulary.

The logit lens applies that same final operation early. Nothing about the architecture says you are entitled to — the intermediate state is not the final state — but in many models the result is legible: the top tokens at layer 8 are vaguely related, at layer 16 they are on topic, and by layer 22 they are the answer. The technique was published by the pseudonymous researcher nostalgebraist in 2020 and named there.

Why this is allowed at all

Two properties of the architecture make it more than a trick.

First, the residual stream is additive. Layer n’s output is h_n = h_(n-1) + f_n(h_(n-1)), so the state at every depth lives in the same space and is a running sum of contributions. There is no basis change between layers to undo.

Second, the unembedding is a fixed linear map from that space to vocabulary logits. Applying it to a partial sum is a legitimate projection of that partial sum — it tells you which tokens the contributions written so far point towards. The interpretive leap is treating that as “what the model currently believes”, which assumes the intermediate state is in roughly the same coordinate system the unembedding expects. That assumption is what fails in the models discussed below.

The script

Complete and runnable against gpt2 on a CPU. The two model-specific names are the final norm and the unembedding; for GPT-2 they are transformer.ln_f and lm_head, and for a Llama-style model they are model.norm and lm_head. If you are unsure, print model and read the last two modules.

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

name = "gpt2"
tok = AutoTokenizer.from_pretrained(name)
model = AutoModelForCausalLM.from_pretrained(name).eval()

final_norm = model.transformer.ln_f     # Llama-style: model.model.norm
unembed    = model.lm_head

@torch.no_grad()
def logit_lens(text, k=5, position=-1):
    ids = tok(text, return_tensors="pt")
    out = model(**ids, output_hidden_states=True)

    print(f"prompt: {text!r}")
    print(f"{'layer':>6}  {'top tokens at this depth':<52}  entropy")
    # hidden_states[0] is the embedding output; [i+1] is the output of layer i
    for i, h in enumerate(out.hidden_states):
        logits = unembed(final_norm(h[0, position]))
        probs  = logits.softmax(-1)
        top    = probs.topk(k)
        parts  = [f"{tok.decode([t]).strip()!r} {p:.2f}"
                  for t, p in zip(top.indices, top.values)]
        ent = -(probs * probs.clamp_min(1e-12).log()).sum().item()
        label = "embed" if i == 0 else f"L{i-1}"
        print(f"{label:>6}  {'  '.join(parts):<52}  {ent:5.2f}")

logit_lens("The capital of France is")

The entropy column is worth as much as the tokens. It falls as the model commits, and where it falls sharply is where the decision was made. Two prompts of the same shape with the commitment happening at different depths is a genuine observation about the model, not a visualisation.

Reading the output

On a factual completion in GPT-2 you will typically see three regimes. Early layers give high-entropy noise: frequent tokens, punctuation, nothing to do with the prompt. Middle layers give something topically adjacent — on a France prompt, other place names. Late layers converge on the answer, usually several layers before the end, after which the distribution mostly sharpens rather than changes.

Two things are worth doing with that. Find the layer where the correct answer first enters the top five, and vary the prompt to see how it moves — harder prompts commit later. And watch for reversals: a token that is top-one at layer 14 and gone by layer 20 means a later component actively suppressed it, which is a lead worth chasing with activation patching. Suppression is real; circuit work has repeatedly found heads whose job is to push down a candidate the rest of the model proposed.

When it produces noise

The logit lens is not architecture-independent, and this is the part most tutorials leave out. Belrose and colleagues reported in 2023 that on several models the plain logit lens gives essentially uninterpretable output at intermediate layers — not a weaker version of the GPT-2 picture, but noise. The technique works well on some model families and fails on others.

The explanation is the assumption in the section above. Nothing constrains intermediate representations to be expressed in the basis the unembedding reads. A model can perfectly well hold its intermediate state in a rotated or rescaled form and only bring it into the output basis in the last few layers. When that happens the projection is meaningless, and — this is the dangerous part — it does not look meaningless. It looks like a model that thinks about punctuation for twenty layers.

Practical rule: before drawing any conclusion from a logit lens on a model you have not used it on before, check it on prompts where you know the answer and know roughly where it should appear. If the correct token never enters the top ten until the final layer, the lens is not working on that model, whatever the middle layers appear to say.

The tuned lens

The fix, from Belrose and colleagues in 2023: instead of applying the unembedding directly, learn a small affine map per layer that translates that layer’s residual state into the final layer’s basis, then unembed. The map is trained to match the model’s own final output distribution, so it is calibrated rather than assumed.

What it is actually good for

• Locating where a decision is made. A cheap first sweep before an expensive patching experiment: if the answer appears at layer 18, patch layers 14 to 20 rather than all of them.
• Comparing prompts. Two phrasings that produce the same answer at different depths tell you something about how much work each one saved the model.
• Spotting suppression. Candidates that rise and then fall are the clearest lens-visible sign of a component doing something deliberate.
• Building intuition quickly. Half an hour with this script does more for understanding what depth means in a transformer than any diagram.

What it is not good for: any causal claim. The lens reads a state; it does not establish that the state caused anything. Pair it with the instrumentation in the full-instrumentation walkthrough if you want the rest of the picture, and with patching if you want a claim that survives review.