Logit lens: the model's running guess at every layer

Note

Generated from notebooks/applications/logit_lens.ipynb, which you can run yourself. The outputs below are the ones it produced.

A decoder-only transformer builds its prediction incrementally: each layer adds to the residual stream, and only the last one is read out. The logit lens cheats by taking an intermediate layer’s residual stream, pushing it through the final norm and unembedding, and asking “if we stopped here, what would you predict?”

Reading those predictions layer by layer shows when the answer appears.

Key questions

• At which layer does the model commit to the right next token?
• How confident is it before then?
• Does the answer arrive gradually, or snap into place?

Structure

1. Set up
2. Run the lens
3. Plot the layer-by-layer prediction
4. Read the prediction off each layer
5. Save and reload

Model and data. GPT-2 small in float32. It is small enough that bfloat16 rounding visibly degrades its predictions, so the notebooks load it explicitly in float32. One factual prompt. The step is architecture-agnostic, so swapping model_id for any causal LM nnterp can standardize (Llama, Mistral, Qwen, OPT) works unchanged.

Requirements. pip install murano-interp[plot]

1. Set up

import torch

from murano import MuranoModel, Pipeline, keys
from murano.steps import LoadPrompts, LogitLens, Plot, Save

OUTPUT_DIR = "murano_outputs/logit_lens"

# GPT-2 is small enough that bfloat16 rounding degrades its predictions.
model = MuranoModel("gpt2", dtype=torch.float32)
print(f"{model.model_id}: {model.n_layers} layers, {model.n_heads} heads")

PROMPT = "The Eiffel Tower is in the city of"

gpt2: 12 layers, 12 heads

2. Run the lens

LogitLens projects every layer’s residual stream through the unembedding and records, for each layer and each token position, the predicted token and its probability.

results = Pipeline([LoadPrompts([PROMPT]), LogitLens(model)]).run()

lens = results[keys.LOGIT_LENS]
print("layers x batch x positions:", tuple(lens.max_probs.shape))
print("input tokens:", lens.input_words[0])

layers x batch x positions: (12, 1, 10)
input tokens: ['The', ' E', 'iff', 'el', ' Tower', ' is', ' in', ' the', ' city', ' of']

3. Plot the layer-by-layer prediction

Rows are layers, columns are token positions. Each cell shows the token that layer would predict at that position, shaded by its probability. Read the rightmost column bottom-to-top to watch the final answer form.

Plot(output_dir=OUTPUT_DIR)(results);

4. Read the prediction off each layer

The heatmap is the overview. For the actual sequence of guesses at the final position, index the result directly.

final_position = -1
print(f"prompt: {PROMPT!r}\n")
print(f"{'layer':>16}  {'prediction':<16} probability")
for layer_index, address in enumerate(lens.addresses):
    word = lens.predicted_words[layer_index][0][final_position]
    prob = float(lens.max_probs[layer_index, 0, final_position])
    bar = "#" * int(prob * 30)
    print(f"{str(address):>16}  {word!r:<16} {prob:5.3f}  {bar}")

prompt: 'The Eiffel Tower is in the city of'

           layer  prediction       probability
   L0.resid_post  ' the'           0.690  ####################
   L1.resid_post  ' the'           0.813  ########################
   L2.resid_post  ' the'           0.720  #####################
   L3.resid_post  ' the'           0.688  ####################
   L4.resid_post  ' the'           0.414  ############
   L5.resid_post  ' the'           0.058  #
   L6.resid_post  ' East'          0.057  #
   L7.resid_post  ' Ing'           0.136  ####
   L8.resid_post  ' Rome'          0.192  #####
   L9.resid_post  ' London'        0.285  ########
  L10.resid_post  ' Paris'         0.183  #####
  L11.resid_post  ' Paris'         0.070  ##

5. Save and reload

Every artifact round-trips through disk, so an analysis can be split across sessions.

from murano import load_logit_lens

run_dir = Save(
    output_dir="murano_outputs", run_name="logit_lens", model_id=model.model_id
)(results)[keys.OUTPUT_DIR]

reloaded = load_logit_lens(f"{run_dir}/logit_lens/logit_lens.pt")
print("saved to:", run_dir)
print("reloaded shape:", tuple(reloaded.max_probs.shape))
print("matches:", torch.allclose(reloaded.max_probs, lens.max_probs))

saved to: murano_outputs/logit_lens
reloaded shape: (12, 1, 10)
matches: True

What next

• logit_attribution.ipynb decomposes the final logit into per-component contributions, answering which components put the answer there.