Module 21: Interpretability

Introduction

Every module so far taught you to build a language model. This one teaches you to look inside one while it runs. Interpretability is the study of what a trained network is actually computing — turning a pile of matrices into mechanisms you can name.

We start with the single most buildable interpretability tool: the logit lens (nostalgebraist, 2020). Recall that a decoder-only transformer keeps one vector per position — the residual stream — and refines it block by block, only reading it out to vocabulary logits at the very end:

logits = lm_head(ln_final(x))        # the model's own read-out head

The logit lens applies that same read-out head to the residual stream at every layer, not just the last. Each intermediate state becomes a distribution over the vocabulary — what would the model predict if it had to answer at this depth? Watching that guess sharpen from layer to layer is the closest thing to watching the model think.

Why it matters for LLMs:

  • The model already predicts, early. On many models the answer is legible in the residual stream several layers before the output — the later layers refine, not decide. The lens makes that visible.
  • It reuses machinery you already built. No new training, no model surgery: the m06 GPTModel already exposes its per-layer hidden states and ties its read-out to the token embedding. The lens is ten lines on top.
  • It is the honest start of a big field. The tuned lens, direct logit attribution, and circuit analysis all begin from “decode the residual stream.” We build the first two from scratch here — the logit lens, then direct logit attribution — and the rest have a foothold.

What You’ll Learn

After this module, you can:

  • Explain the residual stream as a running prediction the model refines and reads out once.
  • Apply the model’s own ln_final + lm_head to any layer to get a distribution — and see why the last layer’s lens equals the model’s output exactly.
  • Measure a prediction’s rank, entropy, and commit depth across layers, and track the residual-stream norm.
  • Train a tiny model to memorize a sentence and watch its guess crystallize with depth.
  • State honestly what the logit lens can and cannot tell you (the bias the tuned lens fixes).
  • Decompose a logit into an exact sum of per-component contributions (direct logit attribution) and see which attention head or FFN wrote the answer.
  • Use the logit difference to isolate what separates two candidate tokens, and know why frozen-LN attribution is exact bookkeeping rather than a causal claim.

Prerequisites

This module requires familiarity with:

Intuition: The Residual Stream as a Running Guess

Picture the model as an assembly line. A token enters as an embedding vector. Each transformer block adds something to that vector — attention pulls in context, the FFN transforms it — but crucially, every block writes back into the same running vector. That vector, carried from the embedding all the way to the output, is the residual stream. It is the model’s working memory for that position.

At the end of the line, one operation converts the vector into a prediction: normalize it (ln_final), then compare it against every token’s embedding (lm_head, whose weights are the token embeddings — that’s weight tying from m06). Tokens whose embedding points the same way as the residual stream get high logits.

Here is the move. That converter — normalize, then compare to token embeddings — doesn’t care which residual-stream vector you hand it. So hand it the half-finished vector after block 1, after block 2, and so on. Each gives a distribution: the model’s best guess so far. Early layers are vague; later layers sharpen. The logit lens is exactly this — reading out the residual stream before it’s done.

NoteKey Insight

The logit lens adds nothing to the model. It reuses the model’s own final normalization and unembedding, applied to states those layers never actually see. That reuse is the whole trick — and, as we’ll be careful to say later, also its main limitation.

The Math: Reading Out Any Layer

Let \(h_k\) be the residual stream after block \(k\) (with \(h_0\) the raw token+position embedding, and \(L\) blocks total, so \(h_L\) is the final state). The model’s output is

\[ \text{logits} = \text{lm\_head}\big(\text{ln\_final}(h_L)\big). \]

The logit lens at layer \(k\) applies the identical read-out to \(h_k\):

\[ \text{lens}_k = \text{lm\_head}\big(\text{ln\_final}(h_k)\big), \qquad p_k = \text{softmax}(\text{lens}_k). \]

Two consequences fall straight out:

  1. The last layer is exact. By definition \(\text{lens}_L\) is the model’s own output — not an approximation, the same tensor. This is the anchor: any correct lens implementation reproduces model(tokens) at the final layer bit for bit.
  2. Weight tying makes the lens a similarity. Since lm_head.weight is the token-embedding matrix \(E\), \(\text{lens}_k = \text{ln\_final}(h_k)\,E^\top\) — the logit for token \(t\) is the (normalized) residual stream’s dot product with token \(t\)’s embedding. The lens asks: which token does this vector look like?

From the per-layer distribution \(p_k\) we read three numbers:

  • Rank of a target token \(t\): how many tokens outrank it, plus one. The shallowest layer where the rank hits 1 is the model’s commit depth for that answer.
  • Entropy \(H(p_k) = -\sum_t p_k(t)\log_2 p_k(t)\) (bits): high = unsure, and it typically falls with depth as the guess sharpens.
  • Residual-stream norm \(\lVert h_k\rVert\): it grows with depth, because every block adds to the stream.

Step by Step

Step through the read-out the lens performs at each layer:

Code: The Logit Lens from Scratch

The whole lens lives in interpretability.py. Its heart is one function — unembed — the model’s read-out head applied to any hidden state:

def unembed(model, hidden):
    return model.lm_head(model.ln_final(hidden))

Load a model and confirm the anchor: the lens at the final layer is the model’s output.

import importlib.util
import sys
from pathlib import Path

import torch

spec = importlib.util.spec_from_file_location(
    "interpretability", Path("interpretability.py").resolve()
)
I = importlib.util.module_from_spec(spec)
sys.modules["interpretability"] = I
spec.loader.exec_module(I)

torch.manual_seed(0)
model = I.GPTModel(vocab_size=48, embed_dim=64, num_heads=4, num_layers=4, max_seq_len=32)
tokens = torch.randint(0, 48, (1, 6))

lens = I.layer_logits(model, tokens)          # (num_layers+1, batch, seq, vocab)
model.eval()
with torch.no_grad():
    output = model(tokens)

print(f"lens tensor shape : {tuple(lens.shape)}")
print(f"final lens == model output : {torch.equal(lens[-1], output)}")
lens tensor shape : (5, 1, 6, 48)
final lens == model output : True

layer_logits runs one forward pass with return_hidden_states=True, then unembeds each of the num_layers + 1 residual-stream checkpoints (the embeddings plus one per block). The last slab equals model(tokens) exactly — that is the correctness anchor from the math above, made a torch.equal.

Now the reader-facing view: the top predictions at the last position, layer by layer. (This model is untrained, so the tokens are arbitrary ids — we’re checking the mechanism; the meaning comes next.)

rows = I.logit_lens(model, tokens, position=-1, top_k=3)
for layer, row in zip(["emb", "L1", "L2", "L3", "L4"], rows):
    shown = ", ".join(f"{tok}({p:.2f})" for tok, p in row)
    print(f"{layer:>4}: {shown}")
 emb: 38(0.04), 17(0.03), 18(0.03)
  L1: 32(0.03), 12(0.03), 15(0.03)
  L2: 32(0.03), 2(0.03), 39(0.03)
  L3: 32(0.03), 6(0.03), 4(0.03)
  L4: 32(0.03), 33(0.03), 39(0.02)

And the three diagnostics — rank of a chosen token, entropy per layer, and the residual-stream norm:

target = int(lens[-1, 0, -1].argmax())        # the model's final top token
print("rank of final top token :", I.prediction_rank(lens, target, position=-1))
print("entropy (bits) per layer :", [round(e, 2) for e in I.prediction_entropy(lens, position=-1)])

model.eval()
with torch.no_grad():
    _, hidden = model(tokens, return_hidden_states=True)
print("residual norm per layer  :", [round(n, 1) for n in I.residual_norms(hidden)])
rank of final top token : [7, 1, 1, 1, 1]
entropy (bits) per layer : [5.56, 5.57, 5.57, 5.57, 5.57]
residual norm per layer  : [0.2, 6.3, 8.6, 11.0, 14.2]

The final rank is always 1 (that token is the argmax by construction), and the residual norm climbs with depth because each block adds to the stream. On an untrained model the earlier layers are noise — to see the lens do something meaningful, we need a model that has actually learned.

Watch a Prediction Crystallize

An untrained model’s intermediate lens is noise, so “watch the guess sharpen” would be a name-drop. Instead we make it real: train a tiny GPT for a couple of seconds to memorize one sentence, then lens its prediction for the missing final word. demonstrate_logit_lens does exactly this — build a tiny model, overfit it on

the cat sat on mat and the dog ran to the park

feed it every word but the last, and read out the final position, where the answer is park.

trace = I.demonstrate_logit_lens(seed=0)
print("prompt :", " ".join(trace["prompt"]))
print("answer :", trace["target"]["token"], f"(loss after training: {trace['final_loss']:.4f})")
print()
print(f"{'layer':>5} {'top token':>10} {'p(top)':>7} {'rank(park)':>11} {'H bits':>7}")
for lab, row, rank, ent, p in zip(
    trace["layers"], trace["top_k"], trace["target_rank"],
    trace["entropy_bits"], trace["target_prob"]
):
    print(f"{lab:>5} {row[0]['token']:>10} {row[0]['prob']:>7.2f} {rank:>11} {ent:>7.2f}")
prompt : <bos> the cat sat on mat and the dog ran to the
answer : park (loss after training: 0.0011)

layer  top token  p(top)  rank(park)  H bits
  emb        the    1.00           5    0.02
   L1       park    0.94           1    0.50
   L2       park    1.00           1    0.06
   L3       park    1.00           1    0.02
   L4       park    1.00           1    0.01

Read that table top to bottom. At the embeddings (emb), the lens just echoes the current input token, “the” — before any block has run, the residual stream is essentially the token embedding, and weight tying reads it back as itself, so “park” sits far down the ranking. Then a single block is enough: by L1 the model has already surfaced park as its top guess with probability past 0.9. The remaining blocks don’t change the answer — they just make it certain, driving the probability to essentially 1.0 and the entropy toward 0. The model decided early and committed late.

NoteKey Insight

The interesting computation happens in the first block or two; the rest is refinement. This “predict early, sharpen late” shape is exactly what the logit lens was invented to reveal — and why looking only at the output hides where the model actually did the work.

The heatmap makes the whole trajectory legible at once: each row is a layer, each column a token the model considered, and the cell brightness is the probability. The final answer lights up as a bright column that switches on after the first block; the rank and entropy lines confirm the “decide early, commit late” story.

Interactive Exploration

Drive the layer index yourself. Pick a depth and see that layer’s full top-k distribution as a bar chart — the answer token highlighted — alongside its rank, probability, entropy, and residual-stream norm. Step from the embeddings to the output and watch park climb from nowhere to certainty.

TipTry This
  1. Find the commit depth. Step until rank("park") first hits 1 — that’s the layer where the model committed. How many blocks did it take?
  2. Watch entropy collapse. Note the entropy at emb versus the output. Where does the biggest drop happen?
  3. Watch the norm grow. The residual-stream norm climbs every layer even after the answer is locked in — the later blocks keep writing, they just stop changing the ranking.

Direct Logit Attribution: Who Wrote the Answer?

The logit lens told us when the model’s guess for park appeared: it jumped to rank 1 after the first block. But it never told us which part of the model put it there. Was it an attention head pulling “park” in from earlier in the sentence? An FFN? The lens shows the running total; it cannot break that total into the individual writes that produced it.

Direct logit attribution (DLA) does exactly that. It splits a token’s logit into one number per component — the embeddings, and each block’s attention and FFN write — so you can point at the head or FFN that pushed the answer up. Where the lens is a cumulative read-out, DLA is the marginal contribution of each write, attributed to attention versus FFN. Two views of the same event: the lens showed the answer crystallize; DLA shows which write made it win.

The Math: A Logit is a Sum of Contributions

The whole method rests on one fact you built in m06: the residual stream is a sum. A pre-norm block adds to the stream twice — \(x \leftarrow x + \text{attn}(\text{ln}_1(x))\), then \(x \leftarrow x + \text{ffn}(\text{ln}_2(x))\) — so the final state unrolls into

\[ x_\text{final} = \underbrace{x_0}_{\text{embed}} + \sum_{k=1}^{L}\Big(\underbrace{a_k}_{\text{attn}_k} + \underbrace{f_k}_{\text{ffn}_k}\Big). \]

The model reads that out with \(\text{logits} = \text{lm\_head}(\text{ln\_final}(x_\text{final}))\). The read-out is almost linear — the only nonlinearity is ln_final, which divides by the per-position standard deviation \(\sigma\) of its input. So we freeze \(\sigma\) (and, for LayerNorm, the mean) at the value the true \(x_\text{final}\) produces. Frozen, ln_final is an affine map, and an affine map distributes over a sum. Writing \(W_U\) for the unembedding (\(=E^\top\), weight tying), \(\gamma,\beta\) for the LayerNorm gain and bias, and \(\bar\mu_c\) for a component’s own feature-mean:

\[ \text{logit}[t] \;=\; \sum_{c}\; \underbrace{W_U[t]\cdot\!\Big(\gamma \odot \tfrac{c-\bar\mu_c}{\sigma}\Big)}_{\text{contribution of component }c} \;+\; \underbrace{W_U[t]\cdot\beta}_{\text{bias}}. \]

Each bracketed term is the direct contribution of one write to token \(t\)’s logit, and — this is the anchor — they sum exactly back to the model’s real logit. It is the DLA analogue of the lens’s layer_logits[-1] == model(x): an identity a correct implementation reproduces to floating-point rounding.

NoteKey Insight

The residual stream being a sum, and the read-out being (once the norm is frozen) linear, is the entire reason attribution works. Every interpretability method that “reads a direction out of the residual stream” — DLA, the tuned lens, activation patching — leans on this same linearity.

Step by Step

Step through how a single logit is decomposed into component contributions:

Code: Attribution from Scratch

Everything lives in dla.py. The first function decomposes the residual stream into its additive writes; the second reads any one of them out with the norm frozen. Load the module and confirm the two structural facts the math promised — the writes sum to the final stream, and the per-component contributions sum to the model’s true logit.

import importlib.util
import sys
from pathlib import Path

import torch

spec = importlib.util.spec_from_file_location("dla", Path("dla.py").resolve())
DLA = importlib.util.module_from_spec(spec)
sys.modules["dla"] = DLA
spec.loader.exec_module(DLA)

torch.manual_seed(0)
model = DLA.GPTModel(vocab_size=48, embed_dim=64, num_heads=4, num_layers=3, max_seq_len=32)
tokens = torch.randint(0, 48, (1, 6))

labels, comps, x_final = DLA.residual_components(model, tokens)
print("components      :", labels)
print("writes sum to x :", torch.allclose(comps.sum(0), x_final, atol=1e-5))
components      : ['embed', 'L1.attn', 'L1.ffn', 'L2.attn', 'L2.ffn', 'L3.attn', 'L3.ffn']
writes sum to x : True

residual_components re-runs the model’s own sublayers, capturing each write — embed, then L1.attn, L1.ffn, L2.attn, … — and they sum back to the exact residual stream ln_final reads out. Now attribute a token’s logit:

out = DLA.direct_logit_attribution(model, tokens, position=-1)
print(f"target token id : {out['target']['id']}")
for lab, c in zip(out["labels"], out["contributions"]):
    print(f"  {lab:>8}: {c:+.3f}")
print(f"  {'bias':>8}: {out['bias']:+.3f}")
print(f"\nΣ contributions + bias = {out['reconstructed']:.4f}")
print(f"model's true logit     = {out['model_logit']:.4f}")
print(f"exact reconstruction   : {abs(out['reconstructed'] - out['model_logit']) < 1e-3}")
target token id : 11
     embed: -0.004
   L1.attn: +0.072
    L1.ffn: +0.005
   L2.attn: +0.084
    L2.ffn: -0.001
   L3.attn: +0.240
    L3.ffn: +0.000
      bias: +0.000

Σ contributions + bias = 0.3957
model's true logit     = 0.3957
exact reconstruction   : True

That final line is the correctness anchor: the decomposition is lossless. Every logit the model produces is, exactly, the sum of what each component wrote plus the constant bias. (This model is untrained, so the contributions are arbitrary — the mechanism is what we’re checking. The meaning comes when we attribute a trained model’s real prediction below.)

The logit difference: what separates two candidates

A raw logit mixes two things: how much the model likes this token, and a shared “confidence” push that raises every token. To isolate the first, compare the answer against a distractor — attribute the logit difference \(\text{logit}[\text{target}] - \text{logit}[\text{distractor}]\). This projects each write onto the direction \(W_U[\text{target}] - W_U[\text{distractor}]\), so any part of a write that lifts both tokens equally drops out.

diff = DLA.logit_diff_attribution(model, tokens, target_id=3, distractor_id=7)
print("contribution to logit[3] - logit[7]:")
for lab, c in zip(diff["labels"], diff["contributions"]):
    print(f"  {lab:>8}: {c:+.3f}")
print(f"  {'bias':>8}: {diff['bias']:+.3f}")
print(f"reconstructed = {diff['reconstructed']:.4f}  vs true diff = {diff['model_logit_diff']:.4f}")
contribution to logit[3] - logit[7]:
     embed: +0.011
   L1.attn: -0.000
    L1.ffn: -0.009
   L2.attn: +0.024
    L2.ffn: -0.002
   L3.attn: +0.040
    L3.ffn: +0.002
      bias: +0.000
reconstructed = 0.0670  vs true diff = 0.0670

The bias term is the LayerNorm \(\beta\)’s per-token offset, \(W_U[\text{target}]\cdot\beta - W_U[\text{distractor}]\cdot\beta\). On this fresh model it reads 0.000 — because nn.LayerNorm initializes \(\beta\) to zero — which is a handy sanity check. But \(\beta\) is a learned parameter: once the model trains it drifts, and the term becomes genuinely nonzero (you’ll see it below). It is per-token, not a global constant, so it does not cancel in the difference. We report it so the reconstruction stays exact — honesty over a tidy-but-wrong “it cancels.”

Which Head Wrote It?

Attention is not one write — it is a sum over heads. A block’s attention output is out_proj(concat(head₀, …, head_{H-1})), and because concatenation lays the heads side by side and out_proj is linear, head \(h\)’s contribution is out_proj applied to only head \(h\)’s slice of the concatenation. So we can split a block’s attention contribution across its heads and ask the sharpest interpretability question there is: which head wrote the answer?

heads = DLA.attention_head_attribution(model, tokens, layer=0, target_id=int(out["target"]["id"]))
print(f"layer {heads['layer']} attention, per-head contribution to the target logit:")
for h, c in enumerate(heads["head_contributions"]):
    print(f"  head {h}: {c:+.3f}")
print(f"  bias   : {heads['bias']:+.3f}")
print(f"\nΣ heads + bias      = {heads['total']:.4f}")
print(f"whole L1.attn write = {heads['attn_contribution']:.4f}")
print(f"match               : {abs(heads['total'] - heads['attn_contribution']) < 1e-3}")
layer 0 attention, per-head contribution to the target logit:
  head 0: -0.039
  head 1: +0.042
  head 2: +0.021
  head 3: +0.048
  bias   : +0.000

Σ heads + bias      = 0.0719
whole L1.attn write = 0.0719
match               : True

The head contributions plus the shared out_proj bias sum to the block’s whole attention contribution — the same per-head decomposition Elhage et al. use to name the heads inside a circuit.

Watch the Components Vote

Now the payoff. demonstrate_dla trains the tiny GPT to memorize the same sentence the logit lens used, then attributes its prediction for the missing final word park — against the distractor sun. Where the lens showed park climb to rank 1 after block 1, DLA shows the individual writes that lifted it.

trace = DLA.demonstrate_dla(seed=0)
print("prompt :", " ".join(trace["prompt"]))
print("answer :", trace["target"]["token"],
      f"(loss after training: {trace['final_loss']:.4f})")
print()
print(f"{'component':>9} {'→ park':>9}")
for lab, c in zip(trace["labels"], trace["target_logit"]["contributions"]):
    print(f"{lab:>9} {c:>+9.3f}")
tl = trace["target_logit"]
print(f"{'bias':>9} {tl['bias']:>+9.3f}")
print(f"\nreconstructed = {tl['reconstructed']:.3f}  ==  model logit = {tl['model_logit']:.3f}")
prompt : <bos> the cat sat on mat and the dog ran to the
answer : park (loss after training: 0.0011)

component    → park
    embed    -0.015
  L1.attn    +1.688
   L1.ffn    +0.247
  L2.attn    +1.559
   L2.ffn    +0.991
  L3.attn    +1.315
   L3.ffn    +1.206
  L4.attn    +0.601
   L4.ffn    +1.687
     bias    -0.035

reconstructed = 9.243  ==  model logit = 9.243

Two things jump out of that column. First, the embeddings contribute almost nothing to “park”: the raw input at the final position is the word “the”, whose embedding points at “the”, not “park” — the model has to compute the answer. Second, every block writes a positive push, but the biggest ones come from attention in the early layers (which pull “park” in from context) and the FFN in the last layer (which sharpens it) — not from any single place. The bar chart below makes the votes visible; green pushes “park” up, red would push it down, and the biggest contributor is ringed.

And the per-head split of each layer’s attention — the finest grain DLA reaches here. Each row is a layer; each cell is one head’s push on “park”, so a bright cell is a head that did real work.

Interactive Exploration: Attribution

Switch between the two views. Target logit shows what pushed “park” up in absolute terms; logit difference shows what pushed “park” over the distractor “sun” — the same bars, but with the shared “raise everything” component removed. Watch which writes stay large (they genuinely separate the candidates) and which shrink (they were just raising confidence).

TipTry This
  1. Find the author. In “target logit” view, which single component has the tallest green bar? Is it attention or an FFN, and in which layer?
  2. Switch to the difference. Flip to “logit difference (park − sun)”. Which bars barely move (they separate the candidates) and which shrink (they were just raising every token)?
  3. Check the books. Read the header line: Σ contributions ± bias always equals the model’s true value. Attribution here is exact, not approximate.

Common Pitfalls

WarningThe logit lens is a biased probe, not ground truth

The lens applies the final ln_final + lm_head to intermediate states those layers were never trained to be read by. On GPT-2-style models it works well, but on others it can be misleading — an early layer’s “prediction” may be an artifact of forcing the final head onto an unfinished vector. Treat the lens as suggestive. The tuned lens (Belrose et al., 2023) fixes this by learning a small affine probe per layer; it is the rigorous successor.

WarningNorm growth is not meaning

The residual-stream norm rising with depth is mostly bookkeeping — blocks add to the stream, so it grows. A larger norm does not mean a layer is “more important.” Use rank and entropy for what the model predicts; use norm only as a sanity signal.

WarningSoftmax before you compare probabilities

Raw lens logits are not comparable across layers — their scale drifts. Always softmax first (as logit_lens and the diagnostics do) before reading a probability or an entropy.

WarningRead the right position

Next-token predictions live at the last position of the prompt. Lensing an interior position tells you what the model would predict there, which is a different question. position=-1 is the usual choice.

WarningFrozen-LN attribution is exact, but it is a convention, not a counterfactual

DLA is exactly additive only because we freeze ln_final’s scale at the true x_final. A component’s attributed push is therefore “its effect on the logit given the normalization the whole model produced.” It is not what you’d get by deleting that component and re-running — the norm would then rescale everything. Treat DLA as an honest bookkeeping of the final logit, and reach for activation patching when you need a true causal (delete-and-rerun) answer.

WarningAttribution is not causation

A component with a big positive contribution wrote toward the answer, but it may have done so because an earlier component set up the residual stream that way. DLA tells you the final tally, not the chain of cause. Naming a circuit means tracing those dependencies (which head reads which earlier write), not just ranking bars.

WarningPrefer the logit difference for a decision

A raw-logit attribution mixes “likes this token” with a shared “raise every token” push. When you care why the model picked A over B, attribute the logit difference — it removes the common component and keeps only what separated them.

Exercises

Exercise 1: Commit depth

The commit depth is the shallowest layer where a target token reaches rank 1. Compute it from a trace’s target_rank list.

trace = I.demonstrate_logit_lens(seed=0)

def commit_depth(ranks):
    # Your implementation here: return the first index where rank == 1 (else None).
    for i, r in enumerate(ranks):
        if r == 1:
            return i
    return None

depth = commit_depth(trace["target_rank"])
print("target ranks per layer:", trace["target_rank"])
print(f"'{trace['target']['token']}' commits at layer index {depth} "
      f"({trace['layers'][depth]})")
target ranks per layer: [5, 1, 1, 1, 1]
'park' commits at layer index 1 (L1)

Exercise 2: Prediction sharpening

Compute the entropy drop between consecutive layers and find the block that sharpens the prediction most. Does it line up with the commit depth?

ent = trace["entropy_bits"]
drops = [round(ent[i] - ent[i + 1], 3) for i in range(len(ent) - 1)]
print("entropy drop after each block:", drops)
# Your turn: which block index has the largest drop? Compare it to commit_depth above.
entropy drop after each block: [-0.477, 0.435, 0.039, 0.007]

Exercise 3: Reconstruct the logit from its parts

The DLA anchor is that the contributions plus the bias equal the model’s true logit. Verify it yourself, and find the single component that pushed the answer hardest.

trace = DLA.demonstrate_dla(seed=0)
contribs = trace["target_logit"]["contributions"]
bias = trace["target_logit"]["bias"]

# Your implementation: sum the contributions and the bias, and compare to the model.
total = sum(contribs) + bias
print(f"reconstructed = {total:.3f}  vs model = {trace['target_logit']['model_logit']:.3f}")

top = trace["labels"][contribs.index(max(contribs))]
print(f"biggest contributor to '{trace['target']['token']}': {top} (+{max(contribs):.2f})")
reconstructed = 9.243  vs model = 9.243
biggest contributor to 'park': L1.attn (+1.69)

Exercise 4: The most important head

Attribute the answer to a chosen layer’s attention heads and find the head that wrote the most. Does the same head dominate in every layer?

model = DLA._train_toy_model(seed=0, train_steps=400, lr=3e-3,
                             embed_dim=64, num_layers=4, num_heads=4)
prompt = torch.tensor([DLA.TOY_SENTENCE[:-1]])
answer = DLA.TOY_SENTENCE[-1]

for layer in range(4):
    h = DLA.attention_head_attribution(model, prompt, layer=layer, target_id=answer)
    hc = h["head_contributions"]
    print(f"L{layer+1}: heads = {[round(c, 2) for c in hc]}  →  top head = {hc.index(max(hc))}")
# Your turn: is it always the same head index, or does the "author" move by layer?
L1: heads = [0.29, 0.59, 0.59, 0.21]  →  top head = 2
L2: heads = [1.04, 0.24, 0.18, 0.1]  →  top head = 0
L3: heads = [0.24, 0.32, 0.68, 0.07]  →  top head = 2
L4: heads = [0.3, 0.18, -0.02, 0.15]  →  top head = 0

Exercise 5: Target logit vs logit difference

Compare a component’s raw contribution to its contribution after removing a distractor. Which components shrink most — the ones that were only raising overall confidence?

raw = DLA.direct_logit_attribution(model, prompt, target_id=answer)
diff = DLA.logit_diff_attribution(model, prompt, target_id=answer,
                                  distractor_id=DLA.TOY_VOCAB.index("sun"))
print(f"{'component':>9} {'raw':>8} {'diff':>8} {'shrunk by':>10}")
for lab, r, d in zip(raw["labels"], raw["contributions"], diff["contributions"]):
    print(f"{lab:>9} {r:>+8.2f} {d:>+8.2f} {r - d:>+10.2f}")
# Your turn: a large "shrunk by" means that write mostly raised confidence, not the choice.
component      raw     diff  shrunk by
    embed    -0.01    +0.04      -0.06
  L1.attn    +1.69    +1.60      +0.08
   L1.ffn    +0.25    +0.28      -0.03
  L2.attn    +1.56    +1.52      +0.04
   L2.ffn    +0.99    +1.05      -0.06
  L3.attn    +1.31    +1.25      +0.06
   L3.ffn    +1.21    +1.31      -0.10
  L4.attn    +0.60    +0.75      -0.15
   L4.ffn    +1.69    +1.73      -0.04

Summary

Key takeaways:

  1. The residual stream is a running prediction. A transformer refines one vector per position through its blocks and reads it out to logits only at the end.
  2. The logit lens reads it out early. Apply the model’s own ln_final + lm_head to any layer’s hidden state to get that layer’s guess — no new training, no model surgery.
  3. The last layer’s lens is the model’s output, exactly. That identity (layer_logits(model, x)[-1] == model(x)) is the correctness anchor.
  4. Weight tying makes the lens a similarity. A token’s lens logit is the normalized residual stream’s dot product with that token’s embedding — which token does this vector look like?
  5. Models decide early and commit late. On our memorizing model the answer reaches rank 1 after a single block, then the remaining blocks only sharpen probability toward 1 and drive entropy toward 0.
  6. Rank, entropy, and norm are the read-outs. Rank gives commit depth, entropy gives confidence, and the residual norm grows with depth (bookkeeping, not meaning).
  7. The lens is biased, and honest about it. It forces the final head onto states it never saw; the tuned lens learns a per-layer probe to do it right.
  8. A logit is a sum of component contributions. Because the residual stream is a sum and the read-out is affine (once ln_final’s scale is frozen), a token’s logit decomposes exactly into one contribution per write — embeddings, and each block’s attention and FFN. Σ contributions + bias == model logit.
  9. Direct logit attribution says who wrote the answer. DLA turns that decomposition into a signed bar per component, and — since attention is a sum over heads — down to which head did the work. On our memorizing model the embeddings contribute ~0, and early-layer attention plus the last FFN write “park”.
  10. Use the logit difference to isolate a decision. Attributing logit[target] − logit[distractor] projects onto W_U[target] − W_U[distractor], removing the shared “raise everything” push and keeping only what separated the candidates.

What’s Next

You have gone from what the residual stream predicts (the logit lens) to which component wrote it (direct logit attribution) — both by decoding the same running vector. The frontier from here trades exact bookkeeping for causal claims:

  • The tuned lens (Belrose et al., 2023) — a learned per-layer probe that fixes the logit lens’s bias.
  • Activation patching — replace one component’s activation with another run’s and measure the effect, the causal counterpart to DLA’s bookkeeping.
  • Induction heads and circuits (Olsson et al., 2022; Elhage et al., 2021) — trace which head reads which earlier write, naming the mechanism DLA only ranks.
  • Sparse autoencoders / features — decompose a single write into interpretable directions, going below the per-component grain this module reaches.

Each begins exactly where this module ends — reading structure out of the residual stream.

Going Deeper

Core Ideas:

Practical Resources: