Logit Lens › Summary

The logit lens is one of the most accessible tools in the transformer interpretability toolkit. Its elegance lies in using the model's own output machinery as a probe into its intermediate computations, requiring no training, no additional parameters, and no architectural changes.

The key ideas are:

• Transformers produce output by applying LayerNorm followed by an unembedding matrix to the final hidden state. This fixed structure can be applied to any intermediate hidden state.
• By applying this decoder at every layer, we read out an implicit "prediction" at each stage of processing, revealing how the model's answer evolves from generic and uncertain to specific and confident.
• The residual stream view explains why this works: each layer adds a delta to the same shared hidden state, so intermediate states are partial sums of the total computation.
• Common patterns include gradual refinement (slow convergence for hard predictions), sudden jumps (rapid shifts caused by specific computations), and the late layer effect (correct answer visible before the final layer).
• Entropy analysis quantifies uncertainty across layers, identifying which layers perform the most decisive reductions in prediction uncertainty.
• The tuned lens improves on the original by learning per-layer affine transformations that better calibrate intermediate predictions to the final output distribution.
• Key limitations include the representation mismatch problem (early layers may not be well-decoded by the final unembedding), the correlation-not-causation issue (visible information is not necessarily computed at that layer), and the vocabulary projection bottleneck (rich representations are collapsed to token predictions).

The logit lens is most powerful when used alongside other interpretability tools. Combine it with probing classifiers to ask "what is encoded?" independently of the model's prediction. Combine it with activation patching to establish which layers are causally responsible for which predictions. Use it as a first-pass diagnostic to identify interesting layers worth deeper investigation, then apply mechanistic interpretability techniques to understand what computations those layers perform. Together, these tools build toward a detailed, mechanistic understanding of how transformer language models process language and knowledge.

The logit lens cannot tell you why a particular layer produces a particular prediction, nor can it identify the causal mechanisms behind it. For that, we need tools like activation patching and the mechanistic interpretability methods covered in the upcoming chapters. But as a first diagnostic tool, the logit lens is fast, requires no training, and often provides immediate insight into where and how a model is solving a given problem.