Logit Lens › A Worked Example
Language AI HandbookAll chapters Logit Lens › A Worked ExampleOne section · Sections in this chapter4 / 12 ↓
Reading tools Definitions & inline links
Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.
Article links
Make inline references clickable
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.
When we inspect the logit lens, we typically look at the argmax token (the single highest-scoring token) or the top-k tokens at each layer. The argmax provides a clean narrative of how the prediction evolves, while the top-k view captures uncertainty in early layers where many tokens compete. A full probability distribution at each layer is most informative but harder to visualize compactly. Most logit lens visualizations show the argmax token annotated with its probability, which gives both the "answer" and the confidence at each step.
When the Pattern Differs
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.