Logit Lens

The core mechanism involves taking the residual stream activation at a specific layer and token position and multiplying it by the model's unembedding matrix (the linear map from the final hidden state to the vocabulary logits). This bypasses the final layer norm and all subsequent transformer blocks, producing a raw probability distribution over the vocabulary that represents the model's 'best guess' at that intermediate stage. The technique reveals how predictions crystallize layer by layer, often showing that the final token identity is determined surprisingly early in deep networks.

By applying the Logit Lens at every transformer layer, researchers can trace the evolution of a prediction from input embedding to final output. This creates a trajectory showing how the model's confidence in the correct token increases and how competing hypotheses are pruned. Key observations include:

• Early layers often represent syntactic possibilities
• Middle layers resolve semantic ambiguity
• Final layers perform calibration and sharpening This trajectory provides a window into the model's progressive refinement process without requiring any auxiliary classifiers.

Unlike probing classifiers which require training a separate diagnostic model on extracted representations, the Logit Lens requires no additional training. It leverages the model's own pre-trained unembedding matrix, making it a lightweight and faithful tool. This eliminates the risk that a trained probe is memorizing task-specific patterns rather than reflecting the model's true internal state. The technique is immediately applicable to any pre-trained transformer with a language modeling head, including GPT, LLaMA, and Mistral architectures.

A striking finding from Logit Lens analysis is that for many factual and straightforward prompts, the model's final prediction is locked in at 30-50% of the network's depth. The remaining layers often serve to increase confidence or refine the output distribution rather than change the top prediction. This has profound implications for early exiting and speculative decoding—if the prediction stabilizes early, computation in later layers could theoretically be skipped, reducing inference latency without sacrificing accuracy.

The standard Logit Lens suffers from a distributional mismatch: intermediate residual stream states are not optimized to be directly unembedded by the final matrix. The Tuned Lens addresses this by learning a separate affine transformation (a learned linear map and bias) for each layer to translate its residual stream into a form the unembedding matrix can interpret. This improves prediction accuracy at early layers but sacrifices the Logit Lens's key advantage of requiring no training. The choice between them depends on whether fidelity or zero-shot applicability is prioritized.

The Logit Lens is a critical tool in mechanistic interpretability workflows for validating hypothesized circuits. After identifying a circuit component (e.g., an attention head or MLP neuron) through causal tracing, researchers can ablate that component and use the Logit Lens to observe how the prediction trajectory changes across all layers. This reveals whether the component's contribution is direct (immediate effect on the logit lens output) or indirect (mediated through later layers). It helps distinguish necessary computational steps from redundant or backup pathways.

The primary information highway that the Logit Lens reads from. Each transformer layer performs a read-modify-write operation on this accumulating state. The Logit Lens applies the unembedding matrix directly to this stream at any layer to project it into vocabulary space, revealing the model's intermediate predictions before the final layer norm is applied. This bypasses the final normalization step to show raw, evolving beliefs.

The linear transformation that maps the model's hidden dimension back to the vocabulary size. In a Logit Lens, this matrix is applied directly to intermediate residual stream states. It is the transpose of the embedding matrix in models with weight tying. This operation converts a high-dimensional vector into a probability distribution over tokens, showing what the model would predict if it stopped at that layer.

The process of identifying the minimal subgraph of a network responsible for a behavior. The Logit Lens is a key diagnostic tool in this workflow. By tracking how predictions evolve layer-by-layer, researchers can isolate which attention heads and MLP layers contribute to a specific output. For example, a sudden jump in prediction accuracy at a specific layer indicates the circuit component that computes that feature has just written to the stream.

A critical technical detail of the Logit Lens technique. The standard forward pass applies a final layer normalization before the unembedding matrix. The Logit Lens applies the unembedding directly to the pre-norm residual stream state. This is necessary because the final norm's scale and shift parameters are optimized for the last layer's distribution, and applying them to intermediate states would distort the signal, making early-layer logits uninterpretable.

A related architectural concept where a model is trained with multiple output heads attached to intermediate layers. The Logit Lens provides a post-hoc analysis of this capability without any training. By examining the logits at each layer, one can determine the earliest point at which the model converges on the correct answer, revealing the computational depth required for different types of queries—from simple syntax to complex reasoning.

A causal intervention technique often paired with Logit Lens analysis. After identifying a layer where a specific prediction crystallizes, researchers can replace that layer's activation with a cached version from a corrupted prompt. If the Logit Lens shows the prediction reverting, it confirms the causal role of that layer. This combination of probing and patching validates that the observed representations are actually used by the model.