Logit Lens

A pre-norm transformer’s residual stream is a running sum: every block writes an additive contribution and nothing overwrites. The unembedding is therefore meaningful at any depth, not only the last, and applying it early gives a legible trajectory of the developing prediction.

The technique is unreliable in exactly one predictable way: it assumes intermediate states live in the same basis as the final one. Where they do not — commonly in the early layers, and in models without a pre-norm residual path — the lens reports noise. A learned affine probe per layer (the tuned lens) fixes the basis mismatch and little else.

A negative result from the logit lens is not evidence of absence. It is evidence that the information, if present, is not yet written in the unembedding’s basis.

import torch
from torch import Tensor


@torch.no_grad()
def logit_lens(model, hidden_states: list[Tensor], top: int = 5):
    """hidden_states: per-layer residual stream [B, N, D]."""
    ln_f, unembed = model.ln_f, model.lm_head
    trajectory = []
    for layer, h in enumerate(hidden_states):
        logits = unembed(ln_f(h[:, -1]))        # last position only
        probs = logits.softmax(-1)
        vals, idx = probs.topk(top, dim=-1)
        trajectory.append((layer, idx[0].tolist(), vals[0].tolist()))
    return trajectory

The final layer norm must be applied even to early states — skipping it produces logits whose scale drifts by orders of magnitude across depth and makes the probabilities meaningless.