A transformer produces its next-token distribution only after the final block, yet every block updates the residual state that eventually feeds that prediction. Logit lens examines those intermediate states by mapping them through the model’s output path into vocabulary logits. The result is a sequence of provisional token distributions across model depth.
The method is attractive because it reuses components already present in the model. Its output also needs careful interpretation. An intermediate residual state was not necessarily optimized to behave like a final residual state, so a readable token ranking is a diagnostic projection rather than a direct transcript of internal computation.
The projection reuses the model output space
For a decoder-only transformer, let h_l denote the residual state after block l. A simplified final output path can be written as:
logits = W_U * norm(h_L) + bHere, L is the final block, norm is the model’s final normalization operation, W_U is the vocabulary projection, and b is an optional output bias. Exact architecture details differ, including whether the output matrix is tied to the token embedding matrix.
Logit lens substitutes an earlier state for h_L:
lens_logits_l = W_U * norm(h_l) + bApplying softmax gives a vocabulary distribution for that depth. Comparing top tokens or selected token logits across blocks can reveal when a candidate continuation becomes prominent, disappears, or competes with another token.
The normalization step is not cosmetic. If the architecture applies a final normalization before unembedding, projecting an unnormalized intermediate state uses a different output path. A diagnostic implementation should make that choice explicit rather than treating all residual vectors as immediately comparable logits.
A token ranking is a projection, not a decoded thought
The vocabulary projection provides a convenient coordinate system because the final model output already uses it. That convenience does not establish that each intermediate block stores its state as an intended probability distribution over tokens.
An early block may encode features that later blocks transform before they become useful to the output projection. A low rank for the eventual token at an early depth does not imply that the model lacks relevant information there. The information may be present in a representation that the final unembedding matrix does not expose directly.
The reverse limitation also applies. A token can rank highly under the lens and then be suppressed by later computation. The projection shows compatibility between the current residual state and the output vocabulary basis. It does not guarantee that the token is committed as the final continuation.
This boundary keeps logit lens useful without assigning it more semantic authority than the operation supports.
Residual placement changes the observation point
A transformer block contains several updates rather than one indivisible operation. Attention and feed-forward components each contribute to the residual stream, with normalization placement depending on the architecture.
Capturing a state before attention, after attention, after the feed-forward update, or after the complete block answers different questions. Two tools can both report a value for block 12 and still be observing different tensors.
For comparisons across implementations, the capture point should therefore be part of the definition. Useful metadata includes whether the state is taken before or after the block, whether the final normalization is applied during projection, and which token position is inspected.
This is especially relevant for causal language models because each sequence position has its own residual state. A next-token analysis normally projects the state at the position whose output predicts the next token. Projecting every position can be informative, but it produces a different object from a single next-token trace.
Vocabulary probabilities can exaggerate small logit changes
Softmax converts logit differences into normalized probabilities. A modest shift in a few competing logits can produce a visually large probability change, especially when the distribution is concentrated.
For depth comparisons, raw logits or logit differences can therefore complement probabilities. Suppose tokens A and B have lens logits 8.0 and 7.8 at one block, then 8.1 and 7.2 at the next. The probability gap can expand substantially even though the absolute logit movement is limited.
Rank has a related limitation. Moving from rank 20 to rank 5 sounds large, but the underlying logits may be tightly clustered. Conversely, a token can keep the same rank while its margin over alternatives changes sharply.
A trace is more informative when it records the quantity that matches the analysis: rank for ordering, logit margin for pairwise competition, or probability when normalized mass itself matters.
Tokenization constrains what the lens can display
The output vocabulary contains tokens, not arbitrary words or concepts. A concept represented by several tokens may not appear as one readable entry in a top-token list. Leading spaces, punctuation, byte fragments, and subword pieces can also make intermediate rankings look less intuitive than plain text.
For a multi-token continuation, inspecting only the first token answers a narrow question: how the model scores that immediate next token across depth. It does not provide the probability of the full continuation, since later token probabilities depend on earlier generated tokens and their resulting states.
Comparisons should also use token IDs rather than rendered strings when exact identity matters. Distinct vocabulary entries can decode to text that appears similar after whitespace handling or display normalization.
Architecture details set the valid output path
A generic logit-lens implementation can silently become inaccurate when it assumes every transformer exposes the same final operations. Some models use RMS normalization, others use layer normalization. Output weights may be tied or untied. An output bias may exist or be absent. Model wrappers may expose hidden states before or after a final normalization.
The defensible implementation is architecture-aware: identify the tensor that normally enters the final output head, then reproduce that path for the selected intermediate state with only the depth substitution changed.
This also makes cross-model comparisons less direct. A rank trajectory from one architecture and a trajectory from another do not share a calibrated depth axis. Block 10 in a 12-block model and block 10 in a 40-block model occupy very different positions in their respective computations.
The strongest use is comparative diagnosis
Logit lens is most informative when a concrete comparison constrains interpretation. A developer can inspect the same prompt before and after a model change, compare two prompt variants, trace a target token against a competing token, or locate the depth range where a final preference first becomes visible under the output projection.
Those comparisons do not require treating the lens as a complete explanation of model behavior. They use it as an instrument: a consistent projection applied at multiple internal points.
When the projection itself appears to distort intermediate states, more specialized methods can fit mappings between intermediate representations and the output space. That changes the diagnostic question because the probe introduces additional fitted parameters. Plain logit lens remains distinctive precisely because its mapping comes from the model’s existing output path.
The practical limit is also its value: logit lens shows what intermediate residual states look like through the final vocabulary projection. Keeping that statement narrow makes the resulting traces easier to compare, reproduce, and interpret without mistaking a convenient view for the full internal computation.