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.

On this page 5
  1. Why it exists
  2. How it works
  3. Where you have already seen it
  4. Remember this
  5. What to learn next

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 confident

The 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

Developer — Code and libraries.

This uses plain gpt2, a small model that runs comfortably on CPU. Nothing beyond transformers and torch is needed.

Setup

bash
pip install transformers torch

The first run downloads gpt2, roughly 500 MB.

Reading out a guess at five different layers

logit_lens.py
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}")
Output
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

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:

text
logit_lens(l) = LN_f(h_l) @ W_U,   for l = 0, ..., L
  • h_l is the residual stream value after block l.
  • 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:

text
tuned_lens(l) = LN_f(A_l @ h_l + b_l) @ W_U

This 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

What to learn next

These follow on from what you just read.

  • Looking Inside a Trained Model

    Induction heads

    Induction heads let a model complete a repeated pattern it has only seen once before in the same conversation, the mechanism behind in-context learning.

  • Looking Inside a Trained Model

    Activation patching

    Activation patching swaps one internal number from a correct run into a broken run, to test which layer is actually causing a prediction rather than only correlated with it.

  • Looking Inside a Trained Model

    Probing hidden states

    A probe is a small classifier trained on a model's hidden states, testing whether specific information is present at a given layer, before the model even finishes reading.