Logit Lens

The technique operates by taking the hidden state from a specific layer's output and multiplying it by the unembedding matrix (the final linear projection). This bypasses all subsequent layers, producing a probability distribution over the vocabulary as if the model stopped at that point. The core insight is that the residual stream is a linear space where the unembedding matrix can be meaningfully applied at any depth, not just the final layer.

Applying the Logit Lens at successive layers reveals a model's progressive refinement of its final prediction. Early layers often show nonsensical or low-confidence tokens, while middle layers begin to surface semantically related candidates. The final layers typically perform confidence calibration, sharpening the probability distribution. This provides a window into the model's step-by-step construction of an answer before it is formally output.

Unlike linear probing, the Logit Lens requires no auxiliary classifier training. It reuses the model's native unembedding matrix, making it a lightweight, immediate diagnostic. This eliminates the risk of the probe itself learning a task the model hasn't, providing a more faithful readout of the model's current belief state. It is a purely inference-time intervention that leaves all model weights frozen.

By examining the top predicted tokens at each layer, researchers can identify the exact point where a model 'locks in' on a fact or a grammatical structure. For example, in a factual recall task, the correct answer token often first appears in the top predictions of middle layers, long before the final output. This helps localize where specific computations, such as entity binding or relation extraction, occur in the network.

The Logit Lens assumes the residual stream's norm and geometry are consistent with the final layer's expectations, which is not guaranteed. Early layer representations often have a different L2 norm scale, leading to underconfident probability distributions. A common mitigation is to apply a learned temperature scalar or a norm correction factor to calibrate the logits before the softmax, improving the interpretability of early-layer predictions.

The standard Logit Lens uses the same unembedding matrix for all layers, while the Tuned Lens variant addresses the norm mismatch by learning a separate, layer-specific affine transformation. This learned correction is trained via distillation to match the final layer's output distribution, providing a more calibrated view of intermediate states. The choice between them involves a trade-off between the Logit Lens's simplicity and the Tuned Lens's accuracy.