Logit Lens › Limitations of the Logit Lens

Understanding what the logit lens does not tell you is as important as knowing what it does. Several fundamental limitations constrain its interpretation.

The core assumption of the logit lens is that the unembedding matrix is a sensible decoder for intermediate hidden states. This assumption holds most strongly for the final layer (where the matrix was explicitly trained) and weakens for earlier layers, where representations may encode information in geometric arrangements the unembedding matrix was not designed to decode.

Early-layer logit lens distributions are often noisy precisely because the hidden states have not yet been processed into a form that the unembedding matrix understands. The distribution at layer 2 may show meaningless or misleading predictions not because the model lacks information at that layer, but because the information is encoded in a way that projects poorly through the output unembedding. The tuned lens addresses this by learning per-layer decoders, but it introduces its own assumptions (that a linear transformation is sufficient to bridge the domain gap) and requires training data and computation.

This mismatch is not evenly distributed. For models with strong residual connections and pre-norm architecture, the early-layer hidden states are more similar to the final-layer hidden state (because each layer only adds a small delta), and the mismatch is less severe. For models where individual layers perform large transformations of the residual, the mismatch is more pronounced.

Seeing a token appear in the logit lens at a given layer does not tell you that the computation at that layer is responsible for placing it there. A token may appear as the top prediction at layer 7 because layers 1-6 have already done most of the work, and layer 7 is simply passing the information through with no modification. The logit lens cannot distinguish between a layer that actively computes a fact and a layer that passively carries it forward.

This is a fundamental limitation of any representation-reading technique. If you want causal claims about which computations are responsible for a prediction, you need techniques like activation patching (covered in the Activation Patching chapter). Activation patching replaces activations at specific layers with corrupted versions and measures the effect on the final output, establishing causal responsibility rather than mere correlation. The logit lens provides a correlation-based view: "at this layer, the representation is consistent with predicting X." Patching provides a causal view: "this layer's computation is causally necessary for predicting X." Both are valuable, but they answer different questions.

The logit lens works most cleanly for models that use a residual stream architecture where the final hidden state is the sum of contributions from all layers. GPT-2 and most modern transformers have this property, but not all architectures do. The interpretation of logit lens results also depends on whether the model uses pre-norm or post-norm layer normalization, as this affects what LayerNorm parameters to use when decoding intermediate states.

The choice of layer normalization placement matters. Pre-norm models (LayerNorm before each sublayer) are somewhat more cleanly analyzed by the logit lens than post-norm models (LayerNorm after each sublayer), because in pre-norm models the residual stream maintains more consistent scale across layers.

The logit lens operates at the token level, showing what the model predicts for the next token at each position. For tasks that require understanding at the sentence or document level, such as text classification, question answering over long documents, or summarization, the logit lens provides limited insight. The relevant computations for these tasks are distributed across many positions and do not reduce cleanly to next-token predictions at any single position.

Even for next-token prediction tasks, the logit lens is most informative when the correct prediction is a single token. For multi-token answers, the logit lens shows the model predicting the first token of the answer but cannot reveal the reasoning process that led to the full multi-token completion.

Applying the unembedding matrix collapses a $d$-dimensional representation into a $V$-dimensional probability distribution. This projection necessarily discards information. The hidden state at any layer may encode rich, multi-dimensional contextual representations that do not correspond cleanly to any single token. A logit lens prediction of "Paris" means the hidden state is most similar to the "Paris" token embedding, but the hidden state may simultaneously encode information about France, the Eiffel Tower, tourism, architecture, plus dozens of other related concepts that do not surface in the top-k predictions.

This bottleneck means the logit lens gives you a "headline" summary of what the model is thinking but misses details of the internal representation. Two hidden states that produce the same logit lens output can be very different vectors that happen to project similarly onto the unembedding matrix.

The logit lens assigns probabilities at the token level, and many meaningful words or phrases correspond to multiple tokens in byte-pair encoding schemes. When the correct completion is a multi-token word or phrase, the logit lens shows only the first token's probability, which may not capture the full probability of the phrase. This can make the logit lens misleading for evaluating the model's confidence in specific completions that span multiple tokens.