Logit Lens: Reading Transformer Hidden States

When a transformer processes text, it runs each token representation through dozens of layers before producing the final output. For most of machine learning's history, those intermediate layers were treated as black boxes, their representations inscrutable. The logit lens changes that. It asks a simple but powerful question: if we stopped the model mid-computation and asked it "what word do you predict right now?", what would it say?

The answer turns out to be deeply informative. By projecting intermediate hidden states directly into vocabulary space, the logit lens reveals how a model's prediction evolves layer by layer, how early layers are confused or wrong, and how later layers converge on the final answer. This makes it one of the most direct tools for understanding what transformers do as they process language.

As we saw in the Attention Visualization and Probing Classifiers chapters, understanding what information is encoded at each layer is central to transformer interpretability. The logit lens offers a complementary perspective: rather than asking what a linear probe can classify, it directly reads out the model's implicit prediction at each layer. You install no extra machinery, train nothing, and make no architectural modifications. The logit lens is purely an analytical technique that repurposes the model's existing output machinery to peer inside its processing.

The logit lens was introduced by nostalgebraist in a 2020 blog post as a tool for examining GPT-2's internal computations. What motivated it was a basic puzzle: transformer language models produce strikingly capable outputs, but their internal processing is opaque. We can measure what the model outputs at the final layer, but we have no direct window into what it is "thinking" at intermediate layers.

Before the logit lens, the main tools for understanding intermediate representations were probing classifiers and dimensionality reduction techniques like PCA and t-SNE. These techniques require training additional models (probes) or involve lossy projections that discard most of the representational structure. The logit lens sidesteps these problems by using the model's own vocabulary-projection machinery as the decoder.

The key insight is that transformers have a specific, well-understood structure for converting hidden states to vocabulary predictions, and that structure consists of a layer normalization followed by a matrix multiplication. Its fixed parameters require no additional training. The question it answers, "what token would the model predict if it had to answer right now?", is interpretable and directly meaningful.

This technique connects two research traditions in mechanistic interpretability. The first asks about representations: what information is stored in the model's activations at each layer? The second asks about computations: what operations does the model perform to transform inputs to outputs? The logit lens leans toward the first tradition but illuminates the second by revealing when information becomes accessible in a form the model can act on.

The core idea is elegantly simple: transformers end with a specific structure that converts hidden states to vocabulary probabilities, and we can apply that same structure to any intermediate hidden state.

Before diving into the logit lens, it helps to understand exactly how the final layer of a transformer produces a probability distribution over the vocabulary.

A standard transformer language model maintains a hidden state $\mathbf{h} \in \mathbb{R}^d$ for each token position, where $d$ is the hidden dimension, also called the model dimension or $d_{\text{model}}$. After all transformer layers have processed the sequence, the final hidden state passes through two components to produce a prediction.

The first component is layer normalization. The final hidden state $\mathbf{h}^{(L)}$ is normalized to have zero mean and unit variance across its $d$ dimensions, then rescaled by learned parameters:

where:

• $\mu = \frac{1}{d} \sum_{i=1}^{d} h_i^{(L)}$ is the mean of the hidden state components
• $\sigma = \sqrt{\frac{1}{d} \sum_{i=1}^{d} (h_i^{(L)} - \mu)^2 + \epsilon}$ is the standard deviation with a small stability constant $\epsilon$
• $\boldsymbol{\gamma} \in \mathbb{R}^d$ and $\boldsymbol{\beta} \in \mathbb{R}^d$ are learned scale and shift parameters

The second component is the unembedding matrix multiplication:

where $W_U \in \mathbb{R}^{V \times d}$ is the unembedding matrix and $V$ is the vocabulary size.

The resulting logits vector $\mathbf{logits} \in \mathbb{R}^V$ assigns a raw score to each vocabulary token. Applying softmax converts these to probabilities:

where:

• $\mathbf{logits}_v$ is the score for vocabulary token $v$
• $V$ is the total vocabulary size
• The denominator normalizes the scores to form a valid probability distribution over all tokens

This pipeline, LayerNorm followed by unembedding followed by softmax, is the entire machinery that converts a hidden state to a next-token prediction. The logit lens borrows this machinery and applies it to every layer, including intermediate ones.

The logit lens applies this same final-layer decoder, consisting of LayerNorm followed by the unembedding matrix, to the hidden state at each intermediate layer. For a model with $L$ layers, the forward pass produces a sequence of hidden states at each token position:

where $\mathbf{h}^{(0)}$ is the embedding layer output (token embedding plus positional encoding, before any transformer layers) and $\mathbf{h}^{(L)}$ is the final hidden state that produces the model's output.

The logit lens at layer $\ell$ is defined as:

where:

• $\mathbf{h}^{(\ell)} \in \mathbb{R}^d$ is the hidden state after the $\ell$-th transformer layer
• $\text{LayerNorm}$ uses the same learned parameters ($\boldsymbol{\gamma}$, $\boldsymbol{\beta}$) as the model's actual final layer normalization
• $W_U \in \mathbb{R}^{V \times d}$ is the fixed unembedding matrix, shared across all layers
• The result is a vector of logits $\in \mathbb{R}^V$ over the vocabulary, which can be converted to probabilities via softmax

This gives us a probability distribution over the vocabulary at every single layer, which we can inspect and visualize across layers. We are not re-running the model from each intermediate state; we are simply applying the output machinery at different points in the processing pipeline.

The logit lens connects naturally to the residual stream view of transformers. In a transformer with pre-norm architecture (like GPT-2 and most modern language models), each layer adds a delta to the residual stream:

where $\Delta_{\text{attn}}^{(\ell)}$ is the contribution from the multi-head attention sublayer and $\Delta_{\text{ffn}}^{(\ell)}$ is the contribution from the feed-forward network sublayer at layer $\ell$. Because each layer adds to the same residual stream rather than replacing it, the final hidden state $\mathbf{h}^{(L)}$ is the sum of contributions from all layers plus the initial embedding.

The logit lens at any intermediate layer $\ell$ shows what prediction is implied by the partial sum of contributions accumulated up to that layer. In this sense, the logit lens acts as a running tally of how the cumulative residual stream converges toward the final prediction. This makes the technique especially natural for residual-stream architectures: intermediate states live in the same geometric space as the final state, and the unembedding matrix was trained to decode from this space.

In many transformer language models, including GPT-2, GPT-Neo, and others, the token embedding matrix $W_E \in \mathbb{R}^{V \times d}$ and the unembedding matrix $W_U \in \mathbb{R}^{V \times d}$ are the same matrix. This is called weight tying, and it was originally introduced to reduce parameter count and improve sample efficiency. A convenient side effect is that it gives the logit lens a clean geometric interpretation.

When $W_U = W_E$, the logit for token $v$ at layer $\ell$ becomes:

where:

• $\mathbf{e}_v \in \mathbb{R}^d$ is the embedding vector for vocabulary token $v$
• The dot product measures geometric alignment between the normalized hidden state and each token's embedding vector

A high logit for token $v$ means the hidden state at that layer is geometrically close to $v$'s embedding in the $d$-dimensional representation space. This is an intuitive way to think about what "the model is thinking" at each layer: the model implicitly compares its current partial computation to every token it knows and reports which ones are most similar. Early layers produce hidden states that are close to many tokens' embeddings (high uncertainty). Later layers produce hidden states that concentrate geometric similarity onto the correct answer token (low uncertainty).

For models without weight tying, such as some larger-scale transformers where weight tying is removed for training stability, this geometric interpretation does not apply directly, but the operational definition of the logit lens (applying LayerNorm then $W_U$) remains valid.

The logit lens shows you an implicit prediction at each layer, but it does not show you the hidden state itself, the attention patterns at that layer, or what computation that layer performed. It is a read-out mechanism, not a causal analysis. The token that ranks highest at layer 7 was not necessarily computed at layer 7; it may have been present in the residual stream since layer 3, with layer 7 simply leaving it undisturbed.

This distinction matters for interpretation. When you see "Paris" emerge as a top prediction in a late layer, you should read that as "by this point, the accumulated computation is pointing toward Paris." You should not read it as "this layer is the one that retrieved the Paris fact." Establishing causal responsibility requires different tools, which we cover in the Activation Patching chapter.

Let's trace through a concrete example to build intuition before diving into code. Consider a language model processing the prompt "The Eiffel Tower is located in the city of". At the last token position, after "of", the model should eventually rank "Paris" among the most likely continuations.

If we apply the logit lens at each layer of GPT-2, we observe a characteristic progression. At layer 0 (the raw embeddings), the hidden state for the last "of" token is essentially its token embedding with positional encoding mixed in. Decoding that state produces token fragments and generic continuations rather than a context-sensitive answer. There is no evidence of geographic knowledge at this stage, because the embedding has not yet been processed by any attention mechanism.

By layer 3, the model has run three rounds of attention and feed-forward processing. The representation has started to absorb context from the surrounding tokens: "Eiffel", "Tower", and "located" are all visible to the attention heads at the last position. The highest-ranked continuations are still generic, so the residual stream is not yet exposing the factual answer through the raw lens.

In layers 9 through 11, geographically specific candidates rise sharply. Paris and London briefly trade the lead among the tracked city names, while Berlin receives a smaller peak. This is evidence of progressive contextual integration, but not of smooth or monotonic convergence.

By the final layer (layer 12 in GPT-2 Small), "Paris" becomes the top prediction. Its probability remains modest because several city names are plausible continuations, which is a useful reminder that top rank and high confidence are different claims.

This layer-by-layer narrative reveals something meaningful about when contextual information becomes readable from the residual stream. The raw embedding of "of" does not expose the geographic answer. Early and middle layers route and combine context, while the final layers reorganize probability mass around plausible city names. The raw logit lens shows this progression, but it does not localize the responsible computation.

Not all prompts produce clean, progressive convergence. Some produce what researchers call "sudden jumps": the prediction is flat and uninformative across many layers, then abruptly shifts to the correct answer in a single layer. This suggests that a specific computation at a specific layer is responsible for the answer, perhaps a particular attention head that routes a key fact from the position where it was mentioned to the prediction position.

Other prompts produce oscillating patterns where the top prediction changes back and forth between candidates across layers. This can indicate that the model is performing a computation over competing hypotheses and that multiple layers are contributing to the resolution.

Consider the prompt "The capital of France is". Here the answer "Paris" should come from world knowledge rather than surface statistics. In models with strong factual recall, the logit lens often shows an abrupt transition in a specific middle layer, consistent with the view that a dedicated circuit performs the factual lookup. In weaker models, the transition is messier and the answer may only emerge in the final few layers.

Let's implement the logit lens from scratch and apply it to GPT-2. We'll use HuggingFace Transformers to load the model and manually extract hidden states at each layer.

First, install the required packages and load GPT-2 with hooks to extract intermediate hidden states.

GPT-2 Small has 12 transformer layers, a hidden dimension of 768, and a vocabulary of 50,257 tokens. The unembedding matrix has shape $(V, d) = (50257, 768)$, mapping each 768-dimensional hidden state to a score for every token in the vocabulary. Applying this matrix to an intermediate hidden state is the heart of the logit lens.

To apply the logit lens, we need the hidden state after every layer. HuggingFace's GPT2Model returns all hidden states when output_hidden_states=True. This is a built-in capability that requires no hooks or model modifications.

The hidden_states tuple contains $L + 1$ tensors: the embedding layer output at index 0 (before any transformer layers) plus one entry per transformer layer at indices 1 through $L$. Each tensor has shape $(1, \text{seq\_len}, d)$, where the first dimension is the batch size (always 1 here).

The embedding layer output at index 0 is the sum of the token embedding and the positional encoding. It contains no contextual information from attention or feed-forward layers, making it a useful baseline for understanding how much the transformer layers contribute.

Now we implement the core logit lens computation. For each hidden state, we apply LayerNorm using the model's final layer norm weights, then multiply by the unembedding matrix.

Let's run the logit lens on our example prompt and examine the predictions at the last token position.

The output shows how GPT-2's prediction at the final position evolves through the layers. Early layers are often dominated by generic tokens. As layers increase, city names emerge, and by the final layer, "Paris" (tokenized as " Paris" with a leading-space marker in GPT-2's byte-pair encoding) ranks first without monopolizing the probability mass.

Notice that top-token probability is not a monotonic proxy for semantic progress. An early layer can be confidently generic, whereas a later layer can spread probability across several contextually appropriate city names. Reading rank, token identity, and distribution shape together is more informative than treating confidence alone as the signal.

Looking at one position at a time is useful for detailed analysis, but the logit lens becomes especially powerful when visualized across all token positions simultaneously. A heatmap showing the top-1 predicted token and its probability at every (layer, position) cell gives a complete picture of how the model processes the entire input.

The heatmap reveals a position-dependent pattern rather than a uniform rise in confidence. Some positions begin with a concentrated token-identity signal, while others change predictions repeatedly across layers. The rightmost column, where the model predicts the next word, is particularly informative for tracing how the factual completion emerges.

Notice also that positions in the middle of the input (like "Tower" or "located") show relatively high confidence early, because the model's "prediction" for those positions is largely determined by the token itself before much contextual processing has occurred. The prediction for the last position is harder and takes more layers to resolve, because it requires integrating information from the entire context.

For detailed analysis, it helps to track the evolution of probabilities for specific candidate tokens across all layers at a single position.

This line plot shows the probability trajectory for each candidate token across all layers. The key insight is that the model's confidence in the correct answer, "Paris", typically starts low and rises sharply in later layers as the model integrates contextual evidence from the entire prompt. Competing tokens like "London" or "Berlin" may register brief peaks in early or middle layers before being overtaken.

The relative ordering is more informative than any one probability in isolation. Paris ultimately ranks above the other tracked candidates, but the nonzero alternatives show that GPT-2 Small retains uncertainty among related locations. This fine-grained trajectory view is something you cannot get from just looking at the final argmax; the logit lens makes the competition visible.

The standard logit lens has a known limitation: it uses the final-layer unembedding matrix to decode intermediate representations that were never trained to be decoded that way. Intermediate hidden states are in the "processing domain" rather than the "output domain". They encode information in geometric arrangements optimized for subsequent transformer layers to process, not for direct vocabulary projection. Applying the output decoder to these states may therefore produce noisy or miscalibrated probability distributions, especially in early layers.

The tuned lens, introduced by Belrose et al. (2023), addresses this by training a small affine transformation per layer that bridges the gap between intermediate representations and the output domain. Conceptually, the tuned lens learns a "translation function" for each layer that converts the intermediate representation into a form the unembedding matrix was designed to decode.

For each layer $\ell$, the tuned lens learns a linear transformation parameterized by a matrix $A^{(\ell)}$ and a bias $\mathbf{b}^{(\ell)}$:

where:

• $A^{(\ell)} \in \mathbb{R}^{d \times d}$ is a learned linear transformation (the "translator") for layer $\ell$
• $\mathbf{b}^{(\ell)} \in \mathbb{R}^d$ is a learned bias for layer $\ell$
• $W_U$ is the fixed unembedding matrix, shared across all layers and not updated during tuned lens training

These per-layer affine transformations are trained to minimize the Kullback-Leibler divergence between the tuned lens prediction at layer $\ell$ and the final model prediction:

where:

• $p^{(L)}(x)$ is the final layer's probability distribution over the vocabulary given input $x$ (the model's actual output)
• $\hat{p}^{(\ell)}(x)$ is the probability distribution produced by the tuned lens at layer $\ell$
• $D_{\text{KL}}(p \| q) = \sum_{v=1}^{V} p_v \log(p_v / q_v)$ is the KL divergence from $q$ to $p$
• The expectation is taken over a corpus of training text $\mathcal{D}$

The training objective asks: given the intermediate hidden state at layer $\ell$, can we linearly transform it into a form from which the final model output is recoverable? If the tuned lens achieves low KL divergence at layer $\ell$, it means the information needed to make the final prediction was already present at layer $\ell$ in a linearly accessible form.

In practice, the $A^{(\ell)}$ matrices are initialized to the identity and the $\mathbf{b}^{(\ell)}$ biases are initialized to zero. This means the tuned lens starts from the logit lens (identity transformation) and learns corrections. The training uses the same text data used to pretrain the language model, running forward passes with frozen model weights and updating only the translator parameters.

An important constraint is that $A^{(\ell)}$ is often regularized to stay close to the identity matrix, preventing it from learning arbitrary transformations that would overfit to the training distribution. The tuned lens is intended to be a minimal correction that brings intermediate representations into the output domain, not a full re-learning of the unembedding.

The parameter overhead is modest. For a model with $L$ layers and hidden dimension $d$, the tuned lens adds $L \cdot (d^2 + d)$ parameters. For GPT-2 Small with $L = 12$ and $d = 768$, this is approximately $12 \times (768^2 + 768) \approx 7.1$ million parameters, a small fraction of the model's 117 million parameters.

The tuned lens tends to show smoother probability trajectories than the raw logit lens, with less noise in early layers. Critically, it separates two questions that the logit lens conflates:

1. When is information first present? At which layer does the model first have enough information to make the correct prediction?
2. When is that information used? At which layer does the model commit to using that information for the final output?

These can differ significantly. A fact might be recoverable from middle-layer representations (the information is "present") but only become the dominant prediction in later layers (the information is "used"). The tuned lens helps distinguish these cases by providing calibrated probabilities throughout the network, independent of the specific geometric orientation of intermediate representations.

The tuned lens has been applied to study factual recall, arithmetic reasoning, and code generation in large models. In each domain, it has revealed that the information required for correct outputs is often established in middle layers, with final layers primarily performing refinement and confidence calibration rather than introducing new information.

The logit lens reveals more than debugging information: it shows systematic patterns about how transformers process information. Researchers have identified several recurring patterns that appear across models and tasks.

The most common pattern is gradual refinement: early layers produce plausible but generic predictions, and later layers progressively sharpen toward the correct answer. This is consistent with the theoretical view of transformers as iterative refiners, where each layer builds on the previous one rather than performing independent computations.

The speed of convergence varies with difficulty. For "easy" predictions where the answer is clear from surface-level context (common collocations, syntactic completions like "I went to the" predicting "store"), the model converges quickly, often by layer 4 or 5 out of 12. For "hard" predictions requiring world knowledge or multi-step inference, convergence happens later and sometimes more abruptly.

The shape of the probability distribution can change dramatically from early to late layers, but raw logit-lens probabilities are not guaranteed to sharpen monotonically. In this GPT-2 Small example, an early layer is concentrated on a few generic continuations, while the final layer spreads probability across several plausible city names and ranks Paris first. That behavior illustrates both progressive semantic refinement and the raw lens's calibration limitations.

The early layer concentrates on generic continuations that fit the local phrase structure. By the final layer, the candidates become geographically specific and Paris ranks first, although the distribution remains broad. The contrast captures an important logit-lens lesson: semantic refinement does not necessarily mean a monotonic increase in confidence.

Across many prompts and models, the logit lens reveals that different types of information are committed at different layers. Research has consistently found that:

• Syntactic information (part of speech, subject-verb agreement, morphological form) tends to be incorporated early, in roughly the first third of layers (layers 1-4 in GPT-2 Small).
• Semantic information (word sense disambiguation, entity type, semantic role) is processed in middle layers (approximately layers 5-8).
• World knowledge (facts about named entities, geographic associations, historical events) tends to emerge in later layers (layers 8-12).

This layering of information types aligns with findings from probing classifier studies covered in the Probing Layers chapter. The logit lens provides a complementary, probe-free perspective: rather than asking what a trained probe can classify, we directly observe when the model's implicit prediction shifts to reflect each type of knowledge. The alignment is not perfect; individual attention heads and feed-forward layers can contribute knowledge at unexpected positions. But the broad pattern of syntactic-early, semantic-middle, factual-late is reliable enough to be practically useful.

Some prompts show a sudden jump pattern: the prediction is flat and uninformative across many layers, then abruptly changes to the correct answer in a single layer. This "phase transition" suggests that a specific computation at a specific layer is responsible for the answer, often associated with a particular attention head that routes a key fact from the position where it was mentioned to the prediction position.

These sudden jumps are of particular interest for mechanistic interpretability research, explored in the Mechanistic Interpretability chapter. When a single layer shows a large, discrete shift in the logit lens prediction, it is a strong signal that a circuit worth studying exists at that layer. Researchers have used logit lens "jump layers" as starting points for identifying induction heads and factual recall circuits.

The contrast between gradual refinement and sudden jumps also reveals differences in how the model represents different types of knowledge. Simple associations (collocations, syntactic patterns) tend to produce gradual refinement, while complex factual retrievals (capital cities, historical dates) often produce sudden jumps. This suggests the two types of processing are implemented differently in the network.

Transformers often appear to "know" the answer many layers before the final layer. The logit lens frequently shows the correct answer appearing as the top prediction by layer 8 or 9 out of 12 in GPT-2 Small, suggesting that the final three or four layers are doing refinement and confidence adjustment rather than introducing new information.

This raises an interesting question: what are the final few layers doing if the answer is already determined? Research suggests several functions. Final layers may be calibrating confidence by sharpening the probability distribution without changing which token is predicted. They may be resolving ambiguity between closely competing candidates. They may also be ensuring consistency with the broader discourse context in ways that are harder to measure than simple next-token prediction. The late layer effect is a reproducible observation, but its mechanistic explanation remains an active research question.

The heatmap view of the logit lens is particularly revealing because it shows how processing differs across token positions. Positions in the middle of a well-understood phrase converge quickly; positions at the boundary of ambiguous or knowledge-intensive structures converge slowly. A sentence like "Marie Curie was born in" will show rapid convergence for "Marie", "Curie", "was", and "born", because these are straightforward to predict in context, but slow convergence at "in", because the birthplace requires specific factual recall.

This position-dependent convergence pattern helps identify which parts of a text the model finds "easy" versus "hard." Positions where the logit lens converges slowly or shows oscillating predictions are exactly the positions where the model is doing its most complex processing, which makes them natural targets for deeper mechanistic analysis.

Beyond tracking specific tokens, we can analyze the entropy of the logit lens distribution at each layer. Entropy is a principled measure of uncertainty: high entropy means the probability mass is spread across many tokens, while low entropy means it is concentrated on a few.

For a probability distribution $p$ over a vocabulary of size $V$, the entropy in nats is:

where:

• $p_v$ is the probability of token $v$ under the logit lens distribution at a given layer
• The maximum possible entropy is $\log V$ (uniform distribution over all tokens)
• The minimum entropy is 0 (all probability on a single token)

For GPT-2's vocabulary of 50,257 tokens, the maximum entropy is $\log(50257) \approx 10.8$ nats.

The entropy curve is non-monotonic for this prompt. Drops identify layers that concentrate probability, while rises show layers that reopen competition among plausible continuations. The final distribution can therefore be semantically better even when its entropy is higher than an earlier, confidently generic prediction.

The entropy at the embedding layer (index 0) is also informative as a baseline. It tells you how much information the model can extract from just the token identity and position, before any contextual processing. For common words in predictable positions, this baseline entropy can be surprisingly low: the word "the" is almost always followed by a noun phrase, so even the embedding-level entropy for predictions after "the" is not at the theoretical maximum.

Entropy analysis across layers is useful for comparing how "hard" different prompts are for the model. A prompt where entropy drops sharply in layer 4 is "easy": the model commits early. A prompt where entropy stays high until layer 11 is "hard": the model keeps processing until the last moment. This comparison can reveal patterns in where the model's knowledge and uncertainty are concentrated, and can help identify which types of inputs are most taxing for the model's processing pipeline.

The logit lens implementation involves several important design choices:

• output_hidden_states=True: Must be set in the HuggingFace model call to extract intermediate representations.
• LayerNorm application: Using the final layer's LayerNorm parameters is standard for the logit lens. Some implementations experiment with using no normalization or the layer's own normalization; results vary.
• Weight tying: When the embedding and unembedding matrices are tied (as in GPT-2), the logit lens has a clean geometric interpretation as dot product with token embeddings. In models without weight tying, the unembedding matrix must be accessed separately.
• Top-k vs. argmax: Looking at top-5 predictions captures more of the probability mass in early layers, where the distribution is flat. Argmax alone can be misleading when the model is uncertain.

The logit lens is not only an academic interpretability tool. Practitioners use it to diagnose model behavior, debug training, and understand failure modes in deployed systems.

When a language model produces a wrong answer, the logit lens can pinpoint where the error originates. If the correct answer appears in the top-1 prediction at layers 5-11 but is displaced in the final layer, this suggests that late-layer processing is introducing an error (perhaps calibration or context integration). If the correct answer never appears in any layer, this suggests the model lacks the relevant knowledge entirely. These two failure modes require different interventions.

For factual recall specifically, the logit lens can distinguish between "the model knows the fact but ranks it wrong at the output" and "the model never learned the fact." This distinction is practically important for deciding whether fine-tuning on factual data is likely to help.

During language model pretraining, the layer at which the correct prediction first appears in the logit lens provides a view of how knowledge develops over training. Early in training, even the final layer may not produce correct predictions. As training proceeds, the correct answer appears first in the final layer, then in progressively earlier layers. This progression suggests that the model is first learning to make correct final predictions and then learning to compute them more efficiently, a form of emergent specialization that the logit lens makes visible.

The logit lens can reveal structural differences between model families. A model that converges to the correct answer by layer 6 of 24 may be processing information very differently from one that requires all 24 layers. Models trained with different objectives, different amounts of data, or different architectural choices can show dramatically different logit lens profiles for the same prompt.

This makes the logit lens a useful tool for model comparison beyond simple benchmark scores. Two models may produce the same final output for a given prompt but show very different intermediate trajectories, suggesting different internal representations and computation paths.

By running the logit lens on a diverse set of factual prompts and noting which layers show the most dramatic probability shifts, researchers have identified "knowledge-intensive" layers where factual information is preferentially encoded and retrieved. These layers are natural targets for knowledge editing techniques (like ROME and MEMIT, covered in the Knowledge Editing chapter), which aim to modify specific factual associations without affecting the rest of the model.

The logit lens provides a lightweight way to identify candidate layers for such interventions without needing to run the full suite of causal tracing experiments.

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.

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.

Ready to test your understanding? Take this quick quiz to reinforce what you've learned about the logit lens.