The Logit Lens — watching a transformer think, layer by layer

A transformer doesn't decide its next word all at once. It refines a guess across depth. The logit lens is an interpretability technique that decodes the residual stream at every layer — not just the last — so you can watch the prediction crystallize.

logit-lens(layer ℓ) = softmax( LayerNorm( h_ℓ ) · W_U )

Each token's representation flows up the network as a single vector — the residual stream. Every attention block and MLP block adds its output into that stream (the residual connection). Normally we only read the prediction after the final layer. The logit lens instead applies the model's own unembedding matrix W_U to the stream at each depth, asking: “if we stopped here, what would the model predict?” In this worked example you'll see the early layers simply echo the current token, a designed grammar block mid-stack flip the guess to a plausible next word, and the top layers commit. Everything below is computed live with genuine matrix arithmetic — no random placeholders.

1 Pick a prompt & a token position

Choose a sentence. The model predicts the next token after the position you select. The residual stream for that position is what the logit lens decodes.

Prompt

Read-out temperature 1.0

Position being decoded: — click any token above to move the read-out point.

2 The logit lens across depth

Each row decodes the residual stream after one block. Bar length = probability of that layer's top token. Watch the guess change as depth increases. Click a row to inspect it.

echoes current token forming confident & matches final

3 Distribution at the selected layer

Full softmax over the vocabulary, read from layer via the logit lens.

4 Confidence trajectory — how belief in the final answer grows

Probability the logit lens assigns to the model's final predicted token, plotted at every depth. The prediction stays pinned to the current token until a block injects the grammatical signal — then it jumps. Hover any point.

P(final token) — logit lens entropy of the layer's distribution (uncertainty)

5 Where did the update come from? Direct logit attribution

At each layer the stream h_ℓ = h_ℓ-1 + Attn_ℓ + MLP_ℓ. We project each added component through the final LayerNorm onto the predicted token's unembedding direction. Positive bars push the model toward its answer — and they sum exactly to the final logit.

attention block MLP block token + position embedding

How to read it — and why it's exact. The final logit is W_U[token] · LayerNorm(h). LayerNorm subtracts the mean and divides by the standard deviation — both linear-friendly operations — so each component c contributes exactly W_U[token] · (c − mean c) / std(h) logits. Sum the bars and you recover the final logit to machine precision. This is the same direct logit attribution used in mechanistic interpretability to pinpoint which blocks caused a prediction.

How the worked example is built

1. Embed Each token → a d=8 vector whose dimensions are interpretable concept axes (determiner, noun, verb, adjective…); a sinusoidal position vector is added.

2. 6 blocks Every block runs real single-head causal self-attention then a 2-layer ReLU MLP, each wrapped in pre-LayerNorm + a residual add.

3. Grammar circuit Block 4's MLP is hand-built: ReLU units detect “a determiner / verb is present” and write the matching next-word concept — a real, legible circuit.

4. Lens A fixed W_U maps any residual vector to vocabulary logits. Apply it after every block, not just the last — the logit lens.

Is the math real? Yes — every operation is genuine. Attention uses true scaled dot-product scores with a causal mask and softmax; the MLPs are real affine–ReLU–affine maps; LayerNorm truly normalizes to zero mean / unit variance. Most weights are random-but-fixed (deterministic seed); the embeddings and one grammar block are deliberately designed so the tiny model behaves sensibly — the same way mechanistic-interpretability researchers hand-construct toy models. The qualitative arc you see (echo → grammar flip → commitment) is exactly what the logit lens reveals on real GPT-style models.