Looking Inside a Trained Model
The logit lens
The logit lens reads out a model's current best guess at every intermediate layer, showing a prediction sharpen gradually rather than appear all at once.
- 9 min read
- 3 reading levels
- Published
Read these first
On this page 5
One lesson, three depths. Pick the one that fits you today — you can switch any time.
Beginner — No maths. Plain English.
The logit lens checks what a model would guess at every layer, not only at the end.
Picture a portrait artist at work. The first strokes are a rough oval. A few strokes later, you can guess it is a face. Only near the end can you tell whose face it is.
A language model builds its answer the same way, layer by layer. The logit lens peeks at the canvas after every stroke. It borrows the model's own final judgement to read out its guess so far.
Why it exists
A trained model looks like a black box from outside. You give it text, it gives you an answer, and the middle is hidden.
Researchers wanted to know: does the model decide its answer instantly, or build it up gradually? Without a way to look inside, that question had no answer.
The logit lens gives a cheap, honest peek. It reuses a piece the model already has: the part that turns internal numbers into a word. That same piece gets applied early, at every layer in turn.
How it works
Prompt: "The Eiffel Tower is located in the city of"
Peek after each layer, using the model's own
final step to read out a guess:
after layer 0: "destro" <- nonsense, too early
after layer 3: "the" <- still generic
after layer 6: "England" <- wrong country, but a country
after layer 9: "Rome" <- a European capital, closer
after layer 11: "Paris" <- correct, and now confidentThe guess does not appear out of nowhere at the last layer. It sharpens gradually, moving from nonsense, to a rough category, to the right kind of answer, to the correct one.
Where you have already seen it
- "Why did the model say that?" investigations. Researchers debugging a wrong or biased answer use this to find which layer first went wrong.
- AI safety research. Checking at which layer a model "decides" to refuse or comply with a request.
- Model distillation research. Comparing how quickly a small model's guess sharpens compared to a large one, layer for layer.
Remember this
- The logit lens reads out a guess at every layer, not only the final one.
- A model's prediction sharpens gradually across layers, rather than appearing instantly.
- It reuses the model's own final step early, rather than needing any extra training.
What to learn next
- Induction heads — a specific, well-studied mechanism that helps build these guesses.
- Residual stream — the running total each layer's guess is read from.
- Attention — how earlier tokens inform the guess at each layer.
Developer — Code and libraries.
This uses plain gpt2, a small model that runs comfortably on CPU. Nothing beyond transformers and torch is needed.
Setup
pip install transformers torchThe first run downloads gpt2, roughly 500 MB.
Reading out a guess at five different layers
import torch
from transformers import GPT2LMHeadModel, GPT2Tokenizer
tok = GPT2Tokenizer.from_pretrained("gpt2")
model = GPT2LMHeadModel.from_pretrained("gpt2")
model.eval()
text = "The Eiffel Tower is located in the city of"
inputs = tok(text, return_tensors="pt")
with torch.no_grad():
out = model(**inputs, output_hidden_states=True)
# hidden_states[i] is the residual stream after block i.
# hidden_states[12] (the last one) has ALREADY passed the model's own
# final layer norm; earlier ones have not, so we apply it ourselves.
hidden_states = out.hidden_states
ln_f = model.transformer.ln_f
unembed = model.lm_head
last_pos = inputs["input_ids"].shape[1] - 1
for layer_idx in [0, 3, 6, 9, 11, 12]:
h = hidden_states[layer_idx][0, last_pos]
logits = unembed(h) if layer_idx == 12 else unembed(ln_f(h))
top3 = torch.topk(logits, 3).indices
guesses = ", ".join(repr(tok.decode([t])) for t in top3)
print(f"after block {layer_idx:2}: top 3 guesses = {guesses}")after block 0: top 3 guesses = ' destro', ' mathemat', ' challeng' after block 3: top 3 guesses = ' the', ' course', ' a' after block 6: top 3 guesses = ' England', ' the', ' La' after block 9: top 3 guesses = ' Rome', ' San', ' La' after block 11: top 3 guesses = ' Paris', ' London', ' Amsterdam' after block 12: top 3 guesses = ' Paris', ' London', ' Amsterdam'
Block 12 (the real output) and block 11 already agree. The interesting part is everything before that: the guess visibly narrows from nonsense, to European capitals, to the correct one.
Line by line
ln_f is the model's final layer normalisation. GPT-2 applies this once, right before turning numbers into word predictions. Applying it at an earlier layer too is the entire trick: it reuses a piece the model already learned, instead of training a new one.
Layer 12 skips the extra ln_f call. output_hidden_states=True already returns the last hidden state post-normalisation. Applying ln_f a second time there would double-normalise it and silently break the result, a mistake worth watching for.
unembed, also called lm_head, is the matrix that turns a 768-number vector into a score for every word in the vocabulary. It exists once in the model and gets reused for every layer here.
Common mistakes
Applying the final layer norm twice, at the last layer. As noted above, this quietly produces wrong guesses at exactly the layer that should be most correct. Always check whether your hidden_states output is pre- or post-norm.
Reading too much into layer 0. The embedding layer alone has no attention, no context mixing between tokens. Guesses here are close to meaningless, and are not a fair test of the model's reasoning.
Assuming every fact converges this cleanly. Some facts sharpen gradually like this one. Others jump late, or waver between two answers until the very last layer. Do not generalise from a single example.
Try it yourself
Change the prompt to one about a fact the model likely knows less confidently, for example "The tallest mountain in the world is". Rerun the script and watch when, or if, the guess stabilises.
A less certain fact often keeps changing its mind for more layers, or settles later, than a well-known one like the Eiffel Tower's location.
What to learn next
- Activation patching — testing which layer's information is actually responsible for a prediction, not only displaying it.
- Residual stream — the object every layer reads from and writes back to.
- Where a model stores a fact — going one level deeper, into which sublayer adds the key information.
Researcher — Mathematics and papers.
Formal definition
For a decoder-only transformer with L blocks, residual stream states h_0, h_1, ..., h_L (post-embedding through post-final-block), final layer norm LN_f, and unembedding matrix W_U:
logit_lens(l) = LN_f(h_l) @ W_U, for l = 0, ..., Lh_lis the residual stream value after blockl.logit_lens(L)is, by construction, identical to the model's real output logits, since that is exactly how the model computes them.- For
l < L, this applies a transformation the model never explicitly trained for that intermediate point, an important caveat covered below.
nostalgebraist (2020) introduced this technique informally, observing that intermediate-layer predictions in GPT-2 often resemble plausible next-token distributions well before the final layer, particularly in the second half of the network.
Why it works at all
The residual stream is additive: each block adds its output to a running total, rather than replacing it, h_l = h_{l-1} + block_l(h_{l-1}). So W_U implicitly gets applied to a sum that already includes early contributions. The network has some incentive, during training, to keep that stream in a basis W_U can partially decode throughout, not only at the final step.
This is an empirical regularity, not a designed property, and it does not hold uniformly across all models, layers or token positions.
Known limitations
Belrose et al. (2023), Eliciting Latent Predictions from Transformers with the Tuned Lens, identify a systematic bias. Applying the model's final layer norm and unembedding to an intermediate layer is a mismatch, since that layer was never optimised against this specific readout.
Their fix, the tuned lens, learns a small per-layer affine probe (A_l, b_l) that translates intermediate hidden states into the final layer's basis before decoding:
tuned_lens(l) = LN_f(A_l @ h_l + b_l) @ W_UThis consistently produces more calibrated, less noisy intermediate predictions than the raw logit lens, at the cost of requiring a lightweight training step per model.
Complexity
Computing the logit lens at every layer costs L additional matrix multiplications of size d_model x |vocab|, applied to already-computed hidden states. This is cheap relative to the forward pass itself, since no new forward computation through the model is needed, only reuse of the existing unembedding matrix at each checkpoint.
Key references
- nostalgebraist (2020). interpreting GPT: the logit lens. LessWrong.
- Belrose, N. et al. (2023). Eliciting Latent Predictions from Transformers with the Tuned Lens. arXiv:2303.08112
- Geva, M. et al. (2022). Transformer Feed-Forward Layers Build Predictions by Promoting Concepts in the Vocabulary Space. arXiv:2203.14680 — related work reading MLP updates directly in vocabulary space.
Current state and open problems
The logit lens remains a first-pass diagnostic: cheap, requires no extra training, and often informative. The tuned lens is preferred when the intermediate predictions themselves are the object of study, rather than a rough sanity check.
Neither tells you why a layer produced a given intermediate guess, only what that guess was. That causal question is what activation patching, covered next, is built to answer. The two are typically used together: the logit lens to find an interesting layer, activation patching to confirm it is actually doing the work.
What to learn next
- Activation patching — moving from observing a layer's guess to testing its causal role.
- Where a model stores a fact — a specific, well-studied case of localising a computation to one sublayer.
- Residual stream — the additive structure that makes this technique possible at all.