Logit lens¶
Where in depth a prediction emerges, on any family.
model.project_on_vocab(layer.layer_output) reads the residual stream leaving a block through the model's own final norm, its unembedding and whatever the model does after the head, so one loop over model.layers is the logit lens on every family. The method is the one in nnsight's logit lens tutorial. What nnterp changes is that the norm, the head and any softcap are no longer yours to find per architecture.
import torch
from nnterp import StandardizedTransformer
model = StandardizedTransformer("openai-community/gpt2", device="cpu")
prompt = "The Eiffel Tower is in the city of"
paris = model.tokenizer(" Paris", add_special_tokens=False).input_ids[0]
lens = [] # bound outside the trace: a list made inside it is gone afterwards
with model.trace(prompt):
for layer in model.layers:
lens.append(model.project_on_vocab(layer.layer_output)[0, -1].save()) # [vocab_size]
logits = model.logits[0, -1].save()
for i, row in enumerate(lens):
print(f"block {i:2d} top {model.tokenizer.decode(row.argmax())!r:>10} P(Paris) {row.softmax(-1)[paris]:.3f}")
print(torch.equal(lens[-1], logits))
# block 0 top ' the' P(Paris) 0.000
# block 1 top ' the' P(Paris) 0.000
# block 2 top ' the' P(Paris) 0.000
# block 3 top ' the' P(Paris) 0.000
# block 4 top ' the' P(Paris) 0.000
# block 5 top ' the' P(Paris) 0.000
# block 6 top ' East' P(Paris) 0.001
# block 7 top ' Ing' P(Paris) 0.003
# block 8 top ' Rome' P(Paris) 0.025
# block 9 top ' London' P(Paris) 0.248
# block 10 top ' Paris' P(Paris) 0.183
# block 11 top ' Paris' P(Paris) 0.070
# True
GPT-2 decodes the for the first six blocks, a city from block 8 (Rome, then London) and Paris from block 10. The probability says more than the argmax: P(Paris) peaks at block 9, where London is still the top token, and falls through the last two blocks to the 0.070 the model outputs. The loop reads the blocks in forward order and logits after them, which is the order one trace requires.
On many newer checkpoints the early blocks decode to unrelated tokens and the answer appears only in the last few blocks; Qwen/Qwen2.5-0.5B puts Paris on top from block 21 of 24 and nothing recognizable before it. That is not a wiring fault, and the wiring check is what decides.
The final True is the wiring check. At the last block the lens is the model's own computation, so it must equal model.logits exactly. Run it first on a new checkpoint. A lens that differs there is reading through the wrong norm or head, and every earlier block is then on the wrong scale. Decoding model.lm_head(hidden) without the norm is not a cheaper lens. A norm is not affine, so the curve comes out confident and wrong rather than as an error.
The step after the head¶
model.logits is what the model returns and model.lm_head.output the raw projection. On Gemma-2 the two differ, because the model applies cap * tanh(logits / cap) after the head (final_logit_softcapping, 30.0 on google/gemma-2-2b, where the raw projection and the logits differ by up to 7.8 at the last position of this prompt). project_on_vocab reads that key from the text config and applies the same cap, so the wiring check passes there too and every block is read on the model's scale rather than an uncapped one. Cohere's logit_scale and Granite's logits_scaling are the same kind of step, and those families define their own project_on_vocab with it.
On another family¶
The recipe runs unchanged. layer_output is a tensor on every family, including the families with a tuple-returning block (BLOOM, MPT, Falcon), and on DeepSeek-V4, where it carries an extra stream axis, the family's project_on_vocab collapses the streams the way the model does. The lens reads module boundaries, so it does not need attn_implementation="eager". The per-family work is all in project_on_vocab: the norm, the head and the step after it.
Contribution decomposition reads each sublayer's addition through the same head, in a form that sums.