Logit Lens

The core mechanism involves taking the residual stream activation at a specific layer l and multiplying it by the model's unembedding matrix (W_U). This operation projects the high-dimensional hidden state directly into the vocabulary space, yielding a probability distribution over all possible tokens. This bypasses the final layer norm and the remaining transformer blocks, providing a snapshot of the model's current 'belief' about the next token at that exact computational stage.

By applying the logit lens to every layer, researchers can observe the gradual refinement of predictions. Early layers typically produce low-confidence, semantically diffuse outputs. As information flows through the network, predictions sharpen. A common finding is that the model's final answer often solidifies surprisingly early, with later layers primarily acting to boost confidence or refine formatting rather than changing the core semantic decision.

A key limitation of the standard logit lens is that it assumes the unembedding matrix is a valid decoder for all layers, despite it being trained only on the final layer's output. The Tuned Lens addresses this by learning a separate, affine transformation for each layer to decode its hidden state. This compensates for the distributional shift across layers, providing a more faithful representation of the model's intermediate predictions, especially in the early stages of computation.

The technique is crucial for pinpointing the exact layer where a model commits to a factual answer. For example, when queried 'The capital of France is', the logit lens might reveal that the token 'Paris' becomes the top-ranked prediction as early as layer 15 out of 32. This demonstrates that the remaining layers are not performing the core retrieval but are engaged in post-hoc confidence calibration and output formatting.

The logit lens can be used as a diagnostic tool for hallucination detection. By analyzing the intermediate token probability distributions, researchers can observe if a model initially retrieves the correct fact but later 'overrules' it with a confabulation, or if the error is present from the earliest layers. This distinguishes between a retrieval failure and a flawed reasoning override, providing insight into the mechanistic origin of factual errors.

The technique is highly effective when combined with activation patching or corrupted inputs. By comparing the logit lens output for a clean prompt versus one with a corrupted subject, researchers can trace the causal pathway of factual recall. This reveals how information from specific input tokens is progressively copied and transformed through the residual stream to influence the final prediction, isolating the function of specific attention heads.

The theory that neural networks represent more independent features than they have dimensions in a given layer by encoding them in overlapping, nearly orthogonal directions in activation space.

• Why it matters: Explains why individual neurons rarely correspond to single concepts. Models compress representations by exploiting high-dimensional geometry.
• Evidence: Models can represent up to 5-10x more features than available dimensions before interference becomes catastrophic.
• Logit Lens connection: When Logit Lens shows coherent predictions from intermediate layers, it's reading from these superposed representations. The unembedding matrix partially disentangles them, but full decomposition requires tools like SAEs.

The central information highway in transformer architectures where each layer reads from and writes to a shared representation space via additive updates.

• Mechanism: Each attention and MLP layer adds its output to the residual stream rather than replacing it. This creates a linear, cumulative representation.
• Why Logit Lens works: The residual stream at any layer is a sum of all previous layer contributions. Applying the unembedding matrix directly to this accumulated state reveals the model's evolving prediction.
• Key property: The residual stream's dimensionality remains constant across layers, making it the ideal location for diagnostic probing with the Logit Lens technique.