Inspect Transformer Residual Streams with the Logit Lens

A language model returns token probabilities only after its final transformer block, but developers often need to inspect what happens before that point. A prompt may produce an unexpected completion, a fine-tuned checkpoint may behave differently from its base model, or a model modification may alter predictions several blocks before the final output.

The logit lens is a simple diagnostic for these cases. It takes an intermediate residual-stream vector, maps it through the model’s output machinery, and ranks vocabulary tokens as if that intermediate state were already ready for prediction.

That view can expose useful patterns across depth, but it is not a literal transcript of the model’s internal reasoning. The practical skill is knowing what the lens measures, what comparisons are meaningful, and where the method can mislead.

Start from the model’s normal output path

Consider an autoregressive transformer with vocabulary size (V) and hidden width (d). At a token position, the final transformer block produces a residual-stream vector (h_L \in \mathbb{R}^d), where (L) denotes the final block.

A simplified final output path is:

[ z = W_U N(h_L) + b_U ]

where:

• (N) is the model’s final normalization operation when one is present,
• (W_U) is the output projection, often called the unembedding matrix,
• (b_U) is an optional output bias,
• (z \in \mathbb{R}^V) contains one logit for each vocabulary token.

Applying softmax to (z) gives the next-token probability distribution.

The exact output path depends on the architecture. Some models tie the output projection to the token embedding matrix, some use a separate matrix, and bias handling varies. A diagnostic should use the actual model components rather than assume a particular arrangement.

The logit lens reuses this output path at an earlier block. If (h_8) is the residual stream after block 8, for example, the lens computes something like:

[ z_8 = W_U N(h_8) + b_U. ]

The resulting token ranking answers a narrow question: if this intermediate vector were interpreted through the model’s final output mapping, which vocabulary items would receive high scores?

Inspect one token position first

Suppose a model receives this prompt:

The capital of France is

Focus on the final prompt position, the position whose state will be used to predict the next token. A logit-lens inspection might produce a table such as this:

These token names are only a teaching example, not a claim about a specific model.

The useful observation is the progression. Earlier states can map to broad or unstable token rankings. Later states may place a task-relevant token near the top and keep it there. If the final answer is surprising, comparing this progression can help locate the depth range where the ranking changed.

Do not interpret a token appearing at block 10 as proof that the model had already made a fixed decision at block 10. Later attention and feed-forward blocks still transform the residual stream. The lens is a projection of the current state, not a guarantee about the eventual output.

Residual streams accumulate contributions across blocks

A transformer’s residual connections make the lens especially convenient. In simplified form, a block updates a residual state by adding new contributions:

[ h_{l+1} = h_l + a_l + m_l, ]

where (a_l) represents an attention contribution and (m_l) represents a feed-forward contribution. Real architectures differ in normalization placement and exact ordering, but the residual stream remains a shared representation carried through depth.

Because each block writes into that shared stream, applying the same output projection at several depths gives a common vocabulary-space view. You can compare token ranks from block to block without inventing a separate decoder for each state.

That convenience also creates the main limitation: intermediate states were not necessarily optimized to look like final states under the final output mapping. A strange early ranking can reflect a mismatch between the intermediate representation and the final decoder rather than a meaningful intermediate prediction.

Apply the model’s final normalization consistently

Normalization is an easy place to build a misleading implementation.

If the model normally applies a final LayerNorm or RMSNorm before the unembedding, a basic logit lens commonly applies that same final normalization to each inspected residual state. Omitting it changes the vector presented to the output projection and can change token rankings.

A compact implementation pattern is:

for block_index, residual in captured_residuals:
    normalized = final_norm(residual)
    logits = unembed(normalized)
    scores = logits[target_position]
    record_top_tokens(block_index, scores)

This pseudocode assumes the architecture has a final normalization stage. It also assumes captured_residuals correspond to compatible points in the block sequence. Mixing pre-attention states, post-attention states, and full block outputs without labeling them makes comparisons difficult to interpret.

For a model without that final normalization operation, do not add one solely to imitate another architecture. Match the model’s real output path.

Token rank is often more useful than raw probability

It is tempting to compare softmax probabilities across every block. That can be useful, but the numbers need care.

Softmax depends on the complete logit vector. If the scale of intermediate logits changes across depth, the distribution can become sharper or flatter even when the relative ordering of interesting tokens changes little. A probability of 0.4 at one block is therefore not automatically comparable to 0.4 at another block as a measure of internal certainty.

For initial diagnosis, inspect several signals together:

• rank of a target token,
• top-k token identities,
• target logit or logit margin against a relevant alternative,
• the same quantities across adjacent blocks.

Suppose Paris moves from rank 900 to rank 40, then rank 3, then rank 1. That trajectory can be more informative than a single intermediate softmax value. It shows that the output mapping becomes increasingly compatible with that token across depth.

A margin can also clarify close competition. For a target token (t) and alternative (a), define:

[ \Delta_l = z_{l,t} - z_{l,a}. ]

A positive (\Delta_l) means the lens ranks the target above that alternative at block (l). Tracking the margin across blocks can reveal a reversal that a top-1-only display would hide.

Compare models at matching semantic points

The logit lens is useful for checkpoint comparisons, but only when the states being compared mean roughly the same thing.

Imagine comparing a base model with a fine-tuned version. For the same prompt and token position, you can inspect the target-token rank after each corresponding block. If both models share the same architecture, the comparison may show that their output-space trajectories separate sharply after a particular depth range.

That result is a clue, not a causal explanation. It tells you where a visible difference emerges under this projection. It does not prove that a single block caused the final behavioral change, because later states depend on all preceding updates and residual interactions.

Comparisons become less direct when models have different numbers of blocks, hidden widths, tokenizers, output projections, or normalization schemes. Mapping “block 12” in one model to “block 12” in another can be arbitrary when their depths differ substantially.

Common mistakes make the lens look more certain than it is

The most serious mistake is treating intermediate token rankings as a faithful sequence of internal decisions. The model is not required to encode every intermediate computation in a form directly readable by its final unembedding.

Another mistake is inspecting the wrong token position. In a causal language model, next-token logits at each position are derived from that position’s hidden state. If the question concerns the token generated after the full prompt, inspect the final prompt position rather than averaging states across all positions.

Tokenization can also confuse interpretation. A human-readable word may correspond to several tokens, and a vocabulary item may include leading whitespace or other tokenizer-specific markers. Display decoded token strings carefully and keep token IDs available during debugging.

Top-k output alone can hide useful movement. A target can improve from rank 5,000 to rank 20 without appearing in a top-10 table. If you already have a token or small candidate set of interest, track its rank and logit directly.

Finally, avoid using one prompt as evidence for a general model property. A trajectory that looks clean for a factual completion may look very different for syntax, arithmetic, multilingual text, or a long-context task. Use a prompt set that represents the behavior under investigation.

The basic lens has a representation-mismatch limitation

The final output projection is trained to consume the representation produced at the end of the network. Earlier residual states can follow different distributions and can contain information in forms that later blocks still need to transform.

The basic logit lens ignores that mismatch. Its strength is simplicity: the same final decoder is reused everywhere. Its weakness comes from the same choice.

More elaborate interpretability methods can introduce mappings that better align intermediate states with the final output space. Those methods answer a somewhat different question and add parameters or assumptions. For many debugging tasks, starting with the plain lens is useful because its mechanics are transparent. If a conclusion depends heavily on subtle early-block rankings, a more careful method or a separate intervention-based test is appropriate.

Use the lens as a diagnostic, not a causal test

A logit-lens plot is observational. It shows how an intermediate residual state maps into vocabulary logits under a chosen decoder. It does not establish that a component is necessary for a prediction.

If you need a causal claim, pair inspection with an intervention. Depending on the research question, that might mean replacing, removing, patching, or otherwise altering a component and measuring the resulting output. Such experiments require their own controls because interventions can push activations away from normal model behavior.

For routine engineering diagnosis, the observational view is often enough to narrow the search. If a target token’s rank changes abruptly between adjacent blocks, those blocks become reasonable places for deeper inspection. If two checkpoints diverge from the earliest inspected states, focusing only on the final block is unlikely to explain the full difference.

A practical inspection workflow

Start with a small set of prompts that reliably exhibit the behavior you care about. Record the exact tokenization and identify the position used for next-token prediction. Capture residual states at consistent points after each transformer block.

Pass each state through the architecture’s real final normalization and output projection. Record top tokens plus explicit ranks or margins for any target tokens. Plotting rank or margin against block index often makes abrupt changes easier to spot than reading dozens of token tables.

Then check whether the pattern repeats across related prompts. A repeated transition is more useful for diagnosis than an isolated one. If the result will support a strong mechanistic claim, move beyond the lens and test the suspected component with a controlled intervention.

This workflow keeps the tool in its proper role: a cheap vocabulary-space probe that helps decide where to inspect next.

Treat intermediate logits as evidence with limits

The logit lens gives developers a direct way to project transformer residual states into a familiar space: vocabulary logits. It can reveal when target tokens rise or fall in rank, where checkpoints begin to diverge, and which depth ranges deserve closer inspection.

Its output remains a projection through the final decoder, not a privileged view of the model’s internal computation. Match the architecture’s actual output path, inspect the correct token position, compare ranks and margins carefully, and avoid causal claims from observation alone.

Used with those constraints, the logit lens is a compact diagnostic that turns an opaque stack of residual vectors into evidence you can inspect block by block.