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 headThe 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
GPTModelalready 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_headto 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).
- Build a tuned lens from scratch — an identity-initialized per-layer affine translator trained (model frozen) to distill the final distribution — and prove it starts equal to the logit lens and, by convexity, provably cannot do worse at any layer.
- 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.
- Build an induction head from scratch — a previous-token head plus a matching head that composes across two layers — and prove it continues
[A][B]…[A]→[B]. - Explain in-context learning as the induction bump, and why a fixed-period attention stripe can’t yet separate a real induction head from a positional shortcut.
- Run activation patching (causal tracing) from scratch: cache a clean run, corrupt the prompt, splice one activation back, and measure the recovery — then localize the head that causally carries the answer, and see it disagree with the correlational stripe score.
- Sharpen that to path patching: decompose the residual stream into its additive wires, restore one edge at a time, and separate a component’s direct effect on the logits from its total effect — revealing a head that node patching calls critical yet writes nothing straight to the output.
- Explain superposition — why a width-d stream holds far more than d features, and why that makes every neuron polysemantic.
- Build a sparse autoencoder from scratch (encoder/decoder, MSE + L1, unit-norm decoder) and prove it works by planting a dictionary and recovering it (MMCS → 1), then read a sparse feature code off the book’s own
GPTModel. - Derive the L1 shrinkage bias (f^\star = \max(0, a - \lambda/2)) and build the modern SAEs that cure it — TopK (L_0 = k by construction) and JumpReLU (z\cdot H(z-\theta), full magnitude above a learned threshold).
- Build the in-context-learning-as-gradient-descent construction from scratch — drop the softmax, and prove one linear-attention layer equals one GD step, that depth is the number of steps (both to machine precision), and that a trained blank layer rediscovers gradient descent on its own.
- Turn the logit lens into a decoder with DoLa — select a premature layer by Jensen–Shannon divergence, gate tokens with the adaptive plausibility constraint, and decode the log-ratio \log q_N - \log q_M; prove the exact margin-shift identity and watch the contrast demote a shallow frequency guess in favor of a late-earned fact.
Prerequisites
This module requires familiarity with:
- Module 5: Attention — query/key/value, the attention pattern, and heads, which the induction circuit composes.
- Module 6: Transformer — the residual stream,
ln_final,lm_head, and weight tying we read out. - Module 7: Training — we train a tiny model to memorize one sentence, and to grow an induction head, for the demonstrations.
- Module 8: Generation — logits → softmax → a next-token prediction, applied here layer by layer.
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:
- 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. - Weight tying makes the lens a similarity. Since
lm_head.weightis 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 : Truelayer_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.01Read 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
- 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? - Watch entropy collapse. Note the entropy at
embversus the output. Where does the biggest drop happen? - 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.
DoLa: Decoding by Contrasting Layers
Every tool so far has been diagnostic — the lens reads the stream, it never changes what the model says. This section turns the lens into a decoder. Every decoding strategy in Module 8 — greedy, top-p, typical — reshapes the model’s final distribution; DoLa is the one that reaches inside the stack and decodes the difference between two depths. Its premise is the exact phenomenon you just watched: the answer crystallizes with depth. DoLa (Chuang et al., ICLR 2024) asks which tokens the deeper layers made the model more confident about — and decodes those, because factual tokens tend to be earned late while fluent-but-wrong tokens are often confident early. Contrasting a mature layer against a premature one keeps what depth added, which on real LLaMA models raises TruthfulQA by 12–17 points with no retrieval and no fine-tuning.
Intuition: Decode the Layers’ Disagreement
Picture two ways the model can assign a token high probability. A token can be a shallow reflex — a frequent, fluent continuation an early layer already commits to before it has integrated the context. Or it can be a late retrieval — a specific fact the upper layers pull in, low early and high only at the end. The logit lens sees both as “high probability at the output.” DoLa separates them by looking at how the probability got there.
The move is a contrast. Take the final layer’s distribution q_N (the mature layer) and an earlier layer’s distribution q_M (a premature layer), and score each token by the log-ratio \log q_N(x) - \log q_M(x). A token the premature layer was already sure about (large q_M) is pushed down; a token the deep layers added (large q_N, small q_M) is pushed up. You decode the model’s late thoughts, net of its early reflexes.
Two questions remain, and DoLa answers each with one idea: which early layer do we contrast against (the one that disagrees most, by Jensen–Shannon divergence), and how do we stop the ratio from promoting rare garbage (a plausibility gate that keeps only tokens the mature layer already finds credible).
The Math: Pick a Layer, Then Contrast
1 · Select the premature layer. Run the logit lens at every candidate early layer j and pick the one whose distribution is most divergent from the final layer’s:
M \;=\; \arg\max_{j \in \mathcal{J}} \; \mathrm{JSD}\!\left(q_N \,\|\, q_j\right).
The Jensen–Shannon divergence \mathrm{JSD}(p\|q) = \tfrac12\mathrm{KL}(p\|m)
+ \tfrac12\mathrm{KL}(q\|m) with m=\tfrac12(p+q) is symmetric and — in bits — bounded in [0,1]. The most divergent layer is the one with the most to contrast against. (At scale the candidates are grouped into a few buckets and one bucket is fixed on a validation set; layer_buckets builds the partition.)
2 · Gate the tokens. The adaptive plausibility constraint (borrowed from Contrastive Decoding, Li et al., 2022) keeps only the mature layer’s credible tokens, \mathcal{V}_{\text{head}} \;=\; \bigl\{\, x : q_N(x) \ge \alpha \cdot \max_w q_N(w) \,\bigr\}, \qquad \alpha = 0.1, and scores everything else -\infty. Without it, a rare token with a tiny q_M would have a huge ratio and win — the gate makes the contrast safe. Note one fact you can read straight off the definition: the mature argmax always survives (\max \ge \alpha\cdot\max for any \alpha\le1), so DoLa can never be forced off a token the final layer is certain of.
3 · Contrast and decode. On the surviving set, score F(x) \;=\; \log q_N(x) - \log q_M(x), and take the argmax (greedy) or sample. A fixed premature layer is DoLa-static; selecting one per token is dynamic DoLa.
NoteKey Insight
DoLa doesn’t trust probability, it trusts the change in probability across depth. The contrast subtracts off whatever an early layer already believed, so what remains is the evidence the upper layers actually contributed — which is where factual recall lives.
Code: DoLa from Scratch
Everything is built on layer_logits from the logit-lens section — DoLa is that per-layer read-out plus arithmetic on the resulting distributions. The reference implementation is dola.py.
import importlib.util
import sys
from pathlib import Path
import torch
spec = importlib.util.spec_from_file_location("dola", Path("dola.py").resolve())
DL = importlib.util.module_from_spec(spec)
sys.modules["dola"] = DL
spec.loader.exec_module(DL)
# Jensen–Shannon divergence: symmetric and bounded in [0, 1] bits.
p = torch.tensor([1.0, 0.0]); q = torch.tensor([0.0, 1.0])
print(f"JSD(p, p) = {float(DL.js_divergence(p, p)):.3f} (identical → 0)")
print(f"JSD(p, q) = {float(DL.js_divergence(p, q)):.3f} (disjoint → 1 bit)")JSD(p, p) = 0.000 (identical → 0)
JSD(p, q) = 1.000 (disjoint → 1 bit)The contrast score, and its cleanest special case — contrast against a uniform premature layer (one that believes nothing) is just greedy decoding on the mature layer, since F(x) = \log q_N(x) - \log(1/V) = \log q_N(x) + \text{const}:
q_N = torch.tensor([0.45, 0.45, 0.10]) # mature: torn between tokens 0 and 1
q_M = torch.tensor([0.80, 0.10, 0.10]) # premature was already sure of token 0
scores = DL.dola_scores(q_N, q_M, alpha=0.1)
print("greedy would tie 0/1; DoLa picks token", int(scores.argmax()),
"(token 0 is demoted — the premature layer already had it)")
unif = torch.full((3,), 1 / 3)
print("uniform premature → argmax F =", int(DL.dola_scores(q_N, unif, alpha=0.0).argmax()),
"== argmax q_N =", int(q_N.argmax()))greedy would tie 0/1; DoLa picks token 1 (token 0 is demoted — the premature layer already had it)
uniform premature → argmax F = 0 == argmax q_N = 0The whole method collapses to one exact identity. The DoLa margin between a target t and a distractor d, minus the greedy margin, equals the premature layer’s own log-odds between them: \bigl[F(t)-F(d)\bigr]-\bigl[\log q_N(t)-\log q_N(d)\bigr] \;=\; \log q_M(d) - \log q_M(t). So if the premature layer prefers the distractor, DoLa shifts the decision toward the target by exactly that much — the method in one line:
shift = DL.margin_shift(q_N, q_M, target_id=1, distractor_id=0)
print(f"margin shift toward target = {shift:+.3f} (= log q_M(0) − log q_M(1))")margin shift toward target = +2.079 (= log q_M(0) − log q_M(1))Watch the Contrast Demote a Shallow Guess
A memorized single sentence is too easy — the model is certain everywhere, so there is nothing to contrast. DoLa’s real setting is a frequency trap. We train a tiny GPT on many copies of “hot the sky shows sun.” and one copy of “cold the sky shows moon.” Because hot→sun is far more frequent, a shallow layer that has not yet integrated the cue “cold” falls back on the frequency prior and predicts sun — a fluent hallucination. The deep layers integrate the cue and retrieve the rare-but-correct moon. demonstrate_dola trains this and decodes “cold the sky shows ___“:
trace = DL.demonstrate_dola(seed=0)
print("prompt :", " ".join(trace["prompt"]), "___")
print("answer :", trace["target"]["token"], f"(train loss {trace['final_loss']:.3f})")
print()
V = trace["vocab"]
print(f"{'layer':>5} {'p(sun)':>7} {'p(moon)':>7} {'JSD‖qN':>7}")
for lab, ps, pm, j in zip(
trace["layers"], trace["distractor_prob_by_layer"],
trace["target_prob_by_layer"], trace["jsd_bits"]
):
mark = " ← premature (max JSD)" if lab == trace["layers"][trace["premature_layer"]] else ""
print(f"{lab:>5} {ps:>7.3f} {pm:>7.3f} {j:>7.3f}{mark}")
print()
print("premature layer's top token :", trace["distractor"]["token"], "(the shallow guess)")
print("mature layer's top token :", trace["greedy_token"]["token"])
print("DoLa decodes :", trace["dola_token"]["token"])
print(f"margin shift toward 'moon' : {trace['margin_shift']:+.3f} nats "
f"(= how hard DoLa demotes 'sun')")prompt : <bos> cold the sky shows ___
answer : moon (train loss 0.059)
layer p(sun) p(moon) JSD‖qN
emb 0.000 0.000 0.999
L1 0.909 0.070 0.809 ← premature (max JSD)
L2 0.116 0.876 0.062
L3 0.003 0.995 0.001
L4 0.000 0.999 0.000
premature layer's top token : sun (the shallow guess)
mature layer's top token : moon
DoLa decodes : moon
margin shift toward 'moon' : +2.566 nats (= how hard DoLa demotes 'sun')Read the table top to bottom: sun is the confident guess at the shallow premature layer (L1), then collapses as depth integrates “cold”, while moon climbs from nothing to near-certainty. The premature layer is the one that disagrees most with the final answer — exactly where JSD points. DoLa decodes the difference, so it inherits the deep layers’ correctness and actively suppresses the shallow layer’s frequency reflex.
The shaded column is the premature layer DoLa picks — the JSD peak, where the shallow “sun” reflex is loudest. Now compare what three different read-outs would emit at the decision step: the shallow premature layer, the mature layer, and the DoLa contrast.
TipTry This
- Find the crossover. In the divergence chart, note the layer where
p("moon")overtakesp("sun"). Everything before it is the shallow reflex; everything after is the retrieved fact. - Read the margin shift. The printed shift equals \log q_M(\text{sun}) -
\log q_M(\text{moon}) at the premature layer — the shallow layer’s own log-odds. That is exactly how many nats DoLa moves the decision toward the fact. Confirm it against
p(sun)/p(moon)in theL1row.
Interactive Exploration: The α and Layer Dials
Drive the two knobs yourself. α sets the plausibility floor — how much of the mature distribution survives the gate — and the premature layer chooses what to contrast against. Watch the head set grow as α shrinks, and the decoded token change as you contrast against different depths.
TipTry This
- Open the gate. Drag α to 0 and watch the head set grow to every token. On this over-confident toy the contrast still decodes “moon” — but in a less certain model, α=0 lets the ratio promote a rare token with a tiny premature probability. That failure mode is why the gate exists.
- Contrast against the wrong depth. Set the premature layer to
L3. It has already agreed with the answer, so the JSD is near zero and the contrast does almost nothing — DoLa deliberately picks the most divergent layer, not the deepest one.
WarningDoLa contrasts distributions, not logits
The score is a difference of log-probabilities (\log q_N - \log q_M), where each q is a softmax of a lens read-out — not a difference of raw logits. Skip the softmax and the constant per-layer shift no longer cancels, and the “ratio” stops meaning “how much more probable.” Contrast normalized distributions.
WarningThe premature layer is chosen per token, and it moves
DoLa is not “layer N minus layer k” for a fixed k. The JSD-selected layer changes from token to token — a factual token may disagree most with an early layer while a syntactic token disagrees most with a mid layer. Fixing one layer (DoLa-static) is a weaker baseline the paper reports precisely to show the dynamic choice matters.
The Tuned Lens: A Probe That Learns
The logit lens has a crack we glossed over. The read-out head — ln_final + lm_head — was trained to decode exactly one thing: the final residual. We then aim it at every intermediate layer and trust what comes back. On GPT-2-style models that mostly works, but it is a biased probe: an early state lives in a slightly different basis than the final head expects, so the lens can under-read a prediction the model has, in fact, already made — exactly what we saw at the embeddings, where “park” sat far down the ranking even though a single block would surface it.
The tuned lens (Belrose et al., 2023) removes the bias by learning the missing piece: a small per-layer affine translator that maps an intermediate state into the basis the final head expects, before the same frozen read-out. It is the rigorous successor the earlier caveats kept promising — and, satisfyingly, it contains the logit lens as its own starting point.
Intuition: Correct the Basis, Then Read
Picture the residual stream as a sentence being progressively rewritten toward the “language” of the output head. The logit lens reads a half-rewritten sentence with the output head’s dictionary and hopes for the best. The tuned lens first applies a learned translation — one matrix and one bias per layer — that finishes the rewrite into the output basis, then reads with the very same frozen dictionary. It adds nothing the model didn’t already compute (the read-out head is untouched); it only corrects where the intermediate vector sits.
Two facts make this honest rather than magic:
- The translator is initialized to the identity (
A = I,b = 0), so before any training the tuned lens is the logit lens, exactly. It can only improve. - We never train a translator for the last layer — that state already is the final residual, so there is nothing to translate. Both lenses agree with the model there, bit-for-bit.
NoteKey Insight
The tuned lens keeps the model’s read-out head frozen and inserts a learned affine “translator” in front of it. Because the translator starts as the identity, the tuned lens is a strict refinement of the logit lens — same tool, one learned correction bolted on — not a different, incomparable probe.
The Math: An Affine Translator
Write the model’s frozen read-out as the logit lens itself:
\text{LogitLens}(h) = \text{lm\_head}(\text{ln\_final}(h))
The tuned lens inserts a learned affine map — the translator (A_\ell, b_\ell) — in front of it, one per layer \ell (Belrose et al., Eq. 8):
\text{TunedLens}_\ell(h_\ell) = \text{LogitLens}(A_\ell\, h_\ell + b_\ell)
The translators are trained with the transformer frozen, by distillation onto the model’s own final prediction: make each layer’s read-out match the final distribution p_\text{final} by minimizing the KL divergence to it (Eq. 9), averaged over a corpus of positions x:
\min_{A_\ell,\, b_\ell}\ \ \mathbb{E}_x\!\left[\, D_{\mathrm{KL}}\!\big(p_\text{final}(x)\ \big\|\ \text{TunedLens}_\ell(h_\ell)\big) \,\right]
Nothing about the model changes — only A_\ell and b_\ell are learned. Two properties fall straight out of this parameterization:
- Identity initialization. Set A_\ell = I, b_\ell = 0 and \text{TunedLens}_\ell \equiv \text{LogitLens} exactly. Training starts at the logit lens.
- It cannot do worse. For a fixed layer, (A,b)\mapsto logits is linear and the objective is a softmax cross-entropy (KL to a fixed target) on top of it — convex in (A,b). The identity is a feasible point, so the optimum’s KL is \le the logit lens’s KL. Gradient descent can only match or beat it.
That second point is the anchor we get to prove below: on our toy model the tuned lens’s KL-to-final drops at or below the logit lens’s at every layer.
Code: A Tuned Lens from Scratch
The whole thing lives in tuned_lens.py. The translator is one linear layer, parameterized as identity-plus-correction so it starts as the identity:
class Translator(nn.Module):
def __init__(self, embed_dim):
super().__init__()
self.delta = nn.Linear(embed_dim, embed_dim)
nn.init.zeros_(self.delta.weight) # Δ ≡ 0 ⇒ A = I + Δ = I
nn.init.zeros_(self.delta.bias) # b = 0
def forward(self, h):
return h + self.delta(h) # (I + Δ) h + b_Δ — an affine mapTunedLens holds one translator per intermediate checkpoint; tuned_lens_logits applies each translator, then the frozen unembed, exactly like layer_logits but with the learned correction in the middle. Load it and confirm the two exact anchors.
tl_spec = importlib.util.spec_from_file_location("tuned_lens", Path("tuned_lens.py").resolve())
TL = importlib.util.module_from_spec(tl_spec)
sys.modules["tuned_lens"] = TL
tl_spec.loader.exec_module(TL)
torch.manual_seed(0)
tl_model = TL.GPTModel(vocab_size=48, embed_dim=64, num_heads=4, num_layers=4, max_seq_len=32)
tl_ids = torch.randint(0, 48, (1, 6))
tuned = TL.TunedLens(tl_model) # identity-initialized translators
# Anchor 1: at init, the tuned lens IS the logit lens, everywhere.
tuned_logits = TL.tuned_lens_logits(tl_model, tuned, tl_ids)
logit_logits = TL.layer_logits(tl_model, tl_ids)
print("max |tuned − logit| at init :", float((tuned_logits - logit_logits).abs().max()))
# Anchor 2: the last checkpoint has no translator → tuned lens == model output.
tl_model.eval()
with torch.no_grad():
print("final tuned == model output :", torch.equal(tuned_logits[-1], tl_model(tl_ids)))max |tuned − logit| at init : 0.0
final tuned == model output : TrueBoth anchors hold: the difference at initialization is 0.0, and the final layer reproduces the model exactly. The tuned lens genuinely starts as the logit lens — training moves it from there.
Watch the Lens Learn
Now make the win real. We reuse the memorized-sentence setup from the logit-lens demo — overfit a tiny GPT on the cat sat on mat and the dog ran to the park — then freeze it and fit the per-layer translators by distillation onto the model’s own final prediction. demonstrate_tuned_lens does exactly this and returns both lenses’ per-layer diagnostics for the final position (the answer is park).
tl_trace = TL.demonstrate_tuned_lens(seed=0)
print("prompt :", " ".join(tl_trace["prompt"]))
print("answer :", tl_trace["target"]["token"])
print()
print(f"{'layer':>5} {'logit KL':>9} {'tuned KL':>9} {'logit rank':>10} {'tuned rank':>10}")
for lab, lkl, tkl, lr, tr in zip(
tl_trace["layers"], tl_trace["logit"]["kl_bits"], tl_trace["tuned"]["kl_bits"],
tl_trace["logit"]["rank"], tl_trace["tuned"]["rank"],
):
print(f"{lab:>5} {lkl:>9.3f} {tkl:>9.3f} {lr:>10} {tr:>10}")prompt : <bos> the cat sat on mat and the dog ran to the
answer : park
layer logit KL tuned KL logit rank tuned rank
emb 12.644 0.001 5 1
L1 0.079 0.000 1 1
L2 0.004 0.000 1 1
L3 0.000 0.000 1 1
L4 0.000 0.000 1 1Read the KL columns — the distance (in bits) from each layer’s read-out to the model’s final answer, the tuned lens’s own training target. At the embeddings the logit lens is effectively blind (KL of several bits, “park” far down the ranking), because the raw token embedding lives nowhere near the output basis. The tuned lens’s translator carries that same vector into the basis and the KL collapses toward zero — it recovers the answer at a depth where the logit lens sees nothing. And by the convexity argument, the tuned KL sits at or below the logit KL on every row, never above.
WarningWe fit and read on the same toy sentence
As with the logit-lens demo, this tiny experiment trains and evaluates the translators on the same memorized sentence — it shows the mechanism (a learned basis correction lowers the KL), not held-out generalization. Belrose et al. fit translators on a large corpus and evaluate on held-out text, where the tuned lens’s lower bias and better calibration are the headline. Here the point is narrower and exact: identity-init + a convex objective ⇒ the tuned lens is never worse.
The curve below plots both lenses layer by layer. Switch the metric between KL-to-final and cross-entropy to “park” — the tuned line (solid) rides at or under the logit line (dashed) throughout.
Interactive Exploration: Logit vs Tuned
Finally, drive the layer index and compare the two lenses head to head at that depth. The left column is the logit lens’s top predictions; the right column is the tuned lens’s, with the answer park highlighted and each lens’s rank and KL-to-final shown above. Step from the embeddings to the output and watch the tuned lens lock onto park layers earlier than the logit lens.
TipTry This
- Race the commit depth. Step until each lens first ranks park at 1. The tuned lens gets there at least as early — often at the embeddings, where the logit lens is still guessing.
- Flip the metric on the curve above between KL and cross-entropy. The tuned (solid) line never crosses above the logit (dashed) line — the convexity guarantee, seen.
- Read the final layer. Both lenses’ last column is identical and matches the model — there is no translator there, and nothing left to fix.
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 : Trueresidual_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 : TrueThat 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.0670The 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 : TrueThe 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.243Two 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
- 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?
- 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)?
- Check the books. Read the header line: Σ contributions ± bias always equals the model’s true value. Attribution here is exact, not approximate.
Intuition: Learning Inside a Single Prompt
DLA named which write pushed a logit up. But it stopped at a ranking — it never said what algorithm those writes implement. This section builds the first named algorithm, and the reason transformers can learn from their own context.
Show a model a brand-new pattern inside the prompt and it continues it:
… Dr. Lidenbrock … Dr.→Lidenbrock
Nobody trained the model on “Lidenbrock”; it appeared moments ago in the same prompt. Continuing a pattern seen earlier in the same sequence — with no weight update — is in-context learning (ICL), and its simplest form is literal: find where this token appeared before, and predict whatever came next. That is an induction head (Elhage et al., 2021; Olsson et al., 2022): a two-step lookup spread across two attention layers.
NoteKey Insight
An induction head runs the rule [A][B] … [A] → [B]: “I am at A; last time I saw A, the next token was B; so predict B.” It needs two things working together — a way to know what token preceded each position, and a way to match the current token against those predecessors and copy the token that followed. Those are two heads in two layers, composing.
The Math: Two Heads That Compose
The circuit is two attention heads reading one residual stream:
Previous-token head (layer 0). A purely positional head: position i attends to position i-1 and copies that token into a dedicated “prev-token” slot. After it runs, position j carries an annotation of \text{token}[j-1].
Induction head (layer 1). Its query is the current token \text{token}[i]; its key at position j is that prev-token annotation \text{token}[j-1]. The attention score fires exactly when
\text{token}[j-1] = \text{token}[i],
i.e. “j is one step after an earlier occurrence of my current token.” Its OV copies \text{token}[j] — the token that followed mine last time. That copied token is the prediction.
The second head can only work because it reads the first head’s write — its keys are the previous-token annotations. This reading-of-an-earlier-head is K-composition, and it is why induction needs (at least) two layers: one to label each position with its predecessor, one to match and copy.
Step by Step
Walk the lookup on the prompt A B C A, predicting what follows the second A. Step through the two hops: read the current token, use the previous-token annotations to find where it occurred before, then copy the token after it.
NoteKey Insight
Neither head alone is enough. The previous-token head knows what preceded each position but never looks at the current token; the induction head knows the current token but would have nothing to match against without the first head’s labels. Induction is the composition — the canonical two-layer circuit.
Code: An Induction Circuit from Scratch
Everything lives in induction.py. Rather than hope the mechanism emerges from training, we build it by hand — an explicit previous-token head and induction head over one-hot tokens — so the algorithm is transparent and provably correct.
To make it exact, we probe with a distinct-token repeated sequence: a permutation of distinct tokens, repeated once ([perm; perm]). Because each token value occurs exactly twice, every second-half position has a unique earlier match — no ties, no averaging.
import importlib.util
import sys
from pathlib import Path
import torch
spec = importlib.util.spec_from_file_location("induction", Path("induction.py").resolve())
IND = importlib.util.module_from_spec(spec)
sys.modules["induction"] = IND
spec.loader.exec_module(IND)
period = 6
seq = IND.distinct_repeated_sequence(period, seed=0)
print("sequence:", seq[0].tolist(), " (a distinct block, repeated once)")
# Layer 0: annotate each position with its predecessor's token.
feats = IND.prev_token_features(seq, vocab_size=period)
print("pos 7 predecessor token:", int(feats[7].argmax()), "== token at pos 6:", seq[0, 6].item())sequence: [2, 5, 3, 0, 1, 4, 2, 5, 3, 0, 1, 4] (a distinct block, repeated once)
pos 7 predecessor token: 2 == token at pos 6: 2The previous-token head is just previous_token_pattern @ one_hot(tokens) — attend one step back, copy. Now the induction head matches the current token against those predecessor labels and copies what followed:
attn = IND.induction_pattern(seq, temperature=0.02) # (2P, 2P) attention
preds = IND.induction_predictions(seq) # argmax next-token per position
# On the second copy, the induction head predicts the true next token EXACTLY.
for i in range(period, 2 * period - 1):
j = int(attn[i].argmax())
print(f"pos {i} (token {seq[0,i].item()}) → attends pos {j} → predicts {preds[i].item()}"
f" (true next: {seq[0,i+1].item()})")pos 6 (token 2) → attends pos 1 → predicts 5 (true next: 5)
pos 7 (token 5) → attends pos 2 → predicts 3 (true next: 3)
pos 8 (token 3) → attends pos 3 → predicts 0 (true next: 0)
pos 9 (token 0) → attends pos 4 → predicts 1 (true next: 1)
pos 10 (token 1) → attends pos 5 → predicts 4 (true next: 4)Every second-half prediction matches the true next token — the circuit copies token[j] from the unique position j = i - period + 1, one step after the earlier occurrence. That exactness is the correctness anchor; the tests in tests/test_induction.py assert it across seeds.
The Induction Stripe
An induction head has a signature you can see. Plot its attention as a matrix (row = query position, column = the position it reads) and the second copy lights up a diagonal stripe offset by the period: every position attends to the one just after its earlier twin. induction_stripe_score measures exactly this — the average attention paid along that offset (≈1 for a sharp induction head, ≈0 for a previous-token head, which sits on the first off-diagonal).
The bright stripe is the induction head at work: no attention in the first copy (nothing has been seen twice yet), then a clean off-period diagonal in the second.
In-Context Learning Is the Induction Bump
The circuit’s payoff is a measurable jump in prediction quality within one sequence. On the distinct repeat, the first copy is unguessable (the tokens are random), so accuracy hovers at chance; the moment the second copy begins, the induction head kicks in and accuracy snaps to 1.0. That step up — better predictions later in the same prompt, with no weight change — is in-context learning, and in_context_learning_score quantifies it.
Interactive Exploration: Do Induction Heads Emerge?
We built the circuit — but Olsson et al.’s result is that gradient descent grows one on its own. Train the book’s own m06 GPTModel (two layers) on repeated random sequences — the second copy is only predictable by copying — and a head develops the induction stripe while in-context loss collapses. demonstrate_induction runs this end to end; here are the heads it found, scored by how much each attends along the induction offset:
demo = _trace["emergent"] # _trace = IND.induction_trace(...), computed just above
print(f"in-context learning score (loss drop on the repeat): {demo['icl_loss_score']:.2f}")
best = demo["best_head"]
print(f"top induction head: layer {best['layer']}, head {best['head']} "
f"(stripe score {best['score']:.2f})")in-context learning score (loss drop on the repeat): 2.76
top induction head: layer 0, head 0 (stripe score 0.81)One head stands out with a strong stripe — the emergent induction head — and its attention pattern is the same diagonal we constructed by hand:
TipTry This
- Read the bars. Which (layer, head) owns the strongest stripe? On this fixed-period probe it often sits in layer 0 — read the pitfall below on why that is not yet proof of a content-based induction head.
- Compare the two heatmaps. The trained stripe and the hand-built stripe trace the same diagonal. The mechanism you constructed is the one training discovered.
- Break it. In
induction.py, raise the induction head’stemperaturetoward 1.0 and re-plotinduction_pattern— the sharp stripe blurs into a soft band, and the second-copy accuracy falls below 1.0. A crisp induction head is a low-entropy head.
Intuition: Break It, Then Fix One Wire
Everything so far reads activations the model already produced — the lens, DLA, the stripe score. All three are correlational: they attribute after the fact. We even had to admit it twice — DLA’s push is “given the normalization the whole model produced,” and a fixed-period stripe can’t separate a content head from a positional shortcut. To make a causal claim you have to intervene.
Activation patching (Meng et al.’s causal tracing, 2022) does exactly that. Run the model twice:
- a clean run on a prompt where it does the task (predicts the right token), and
- a corrupted run where one detail is changed so it fails.
Then run the corrupted prompt again, but at one internal site splice the clean activation back in and let the forward continue. If restoring that one wire brings the answer back, that site causally carries the information. It is the difference between “this component’s write correlated with the answer” and “put this component’s clean value back and the answer returns.”
The whole method rests on two exact primitives, both built here from scratch on the book’s own GPTModel: a faithful cache (run_with_cache) and a re-runnable, patchable forward (patched_forward). “Faithful” and “exact” are not adjectives — they are tested identities, below.
The Math: The Recovery Metric
Patching needs a number for “did the answer come back.” We read the logit difference at the answer position — the same quantity DLA preferred, the target token minus a rival:
\Delta \;=\; \text{logit}[\text{correct}] \;-\; \text{logit}[\text{wrong}].
Call \Delta_{\text{clean}} its value on the clean run (large and positive) and \Delta_{\text{corrupt}} its value on the broken run (small or negative). For any patched run we report a normalized recovery:
\text{recovery} \;=\; \frac{\Delta_{\text{patched}} - \Delta_{\text{corrupt}}} {\Delta_{\text{clean}} - \Delta_{\text{corrupt}}}.
It reads off a clean scale: 0 means the patch changed nothing (still broken), 1 means it fully restored the clean behavior. A site whose patch scores near 1 is causally responsible for the answer. (This is the denoising direction — patch clean into corrupt; the mirror-image noising direction patches corrupt into clean and asks what breaks it.)
Code: Cache and Patch from Scratch
The cache re-runs the pre-norm forward reusing the model’s own sublayers, stashing the residual stream after each block and — the piece head-patching needs — every attention head’s write into the stream (the same additive split DLA used, now kept per position). It is faithful by construction, and we check it:
import importlib.util
import sys
from pathlib import Path
import torch
spec = importlib.util.spec_from_file_location("activation_patching", Path("activation_patching.py").resolve())
AP = importlib.util.module_from_spec(spec)
sys.modules["activation_patching"] = AP
spec.loader.exec_module(AP)
torch.manual_seed(0)
model = AP.GPTModel(vocab_size=20, embed_dim=32, num_heads=4, num_layers=2, max_seq_len=16).eval()
ids = torch.randint(0, 20, (1, 7))
logits, cache = AP.run_with_cache(model, ids)
# Anchor 1: the cache is a faithful trace of the real forward.
print("cache logits == model(ids):", bool(torch.allclose(logits, model(ids), atol=1e-5)))
readout = model.lm_head(model.ln_final(cache["resid_post"][-1]))
print("final residual reads out to logits:", bool(torch.allclose(readout, logits, atol=1e-5)))
print("head_out shape (layers, batch, heads, seq, dim):", tuple(cache["head_out"].shape))cache logits == model(ids): True
final residual reads out to logits: True
head_out shape (layers, batch, heads, seq, dim): (2, 1, 4, 7, 32)patched_forward runs that same forward but overwrites chosen activations before continuing, so a change propagates downstream — the essence of a causal intervention. Two identities pin it down: patching nothing is the model, and patching a run’s own cache back into itself is a no-op.
# Anchor 2: empty patch == the model; self-patch == no-op.
print("empty patch == model:",
bool(torch.allclose(AP.patched_forward(model, ids), model(ids), atol=1e-5)))
S = ids.shape[1]
self_patch = {1: {p: cache["resid_post"][1, 0, p] for p in range(S)}}
print("self-patch resid is a no-op:",
bool(torch.allclose(AP.patched_forward(model, ids, resid_patches=self_patch), logits, atol=1e-5)))empty patch == model: True
self-patch resid is a no-op: TrueAnd the identity that bounds the whole metric: splice the entire clean final residual into a different, corrupted run and you have run clean — full patch ⇒ recovery 1.
torch.manual_seed(1)
clean_ids = torch.randint(0, 20, (1, 8))
corrupt_ids = torch.randint(0, 20, (1, 8))
logits_clean, cache_clean = AP.run_with_cache(model, clean_ids)
last = len(model.blocks) - 1
full_patch = {last: {p: cache_clean["resid_post"][last, 0, p] for p in range(8)}}
patched = AP.patched_forward(model, corrupt_ids, resid_patches=full_patch)
print("patch the whole final residual → clean logits:",
bool(torch.allclose(patched, logits_clean, atol=1e-5)))patch the whole final residual → clean logits: TrueThis mirrors the exactness anchors elsewhere in the module — the lens’s layer_logits[-1] == model(x) and DLA’s Σ contributions + bias == logit. The production code lives in activation_patching.py.
Tracing the Flow: Residual Patching
Now the payoff. We reuse the induction setup: train_induction_model grew a copy head on [r; r] repeats. Build a clean repeated prompt (the correct next token is the copied one) and a corrupt twin that changes one source token’s identity — positions untouched, only what sits there. Then denoise every (layer, position) in turn: patch the clean residual at that site into the corrupt run and record the recovery. demonstrate_patching runs the whole thing:
pt = AP.demonstrate_patching(period=8, vocab_size=12, steps=400)
ojs_define(patchTrace = pt)print(f"clean Δ = {pt['clean_diff']:.2f} corrupt Δ = {pt['corrupt_diff']:.2f}")
print(f"corrupted the source token at position {pt['source']} (content only)")
print(f"causal top head: L{pt['causal_best']['layer']+1}·H{pt['causal_best']['head']}"
f" (recovers {pt['causal_best']['recovery']:.2f} of the clean Δ alone)")
print(f"stripe top head: L{pt['stripe_best']['layer']+1}·H{pt['stripe_best']['head']}"
f" (score {pt['stripe_best']['score']:.2f})")
print(f"do the two metrics agree on the head? {pt['metrics_agree']}")clean Δ = 8.81 corrupt Δ = -9.44
corrupted the source token at position 7 (content only)
causal top head: L1·H1 (recovers 0.81 of the clean Δ alone)
stripe top head: L1·H2 (score 0.78)
do the two metrics agree on the head? FalseThe grid below is the causal trace. Step through the depth: at the input (embed) the only bright cell is the corrupted source token — that is where the answer’s information enters. By the copy layer the bright cell has moved to the query position — the model has carried the answer forward to where it reads out. That jump is the copy, localized in space and depth.
NoteKey Insight
Denoising traces where the answer lives at each depth. Here it enters at the corrupted source token (bright at embed) and is readable at the query position by the copy layer — a two-cell picture of the copy, drawn by intervention, not inspection.
Which Head Carries It: Head Patching
The residual grid localizes depth and position; to name the head we patch one head’s clean write at a time and score the recovery. One head recovers most of the clean logit difference on its own — a causal fingerprint you cannot read off an attention pattern. Flip the metric below between causal recovery (patch and measure) and the induction-stripe score (eyeball the attention) and watch the winner change:
The two metrics need not agree — and here they don’t. The stripe score (how induction-shaped a head’s attention looks) crowns one head; causal patching (does restoring it actually recover the answer) crowns a different one. When a correlational signal and an intervention disagree, the intervention wins: the stripe is a hint, the patch is evidence. And notice where the causal head sits — often layer 0, the cheap positional copy the fixed-period pitfall warned about. Patching doesn’t just find the mechanism; it confirms which mechanism the model actually used.
TipTry This
- Flip the metric above and note the two winners. Which head is induction-shaped, and which one causally moves the answer?
- Follow the flow. In the residual grid, step
embed → L1 → L2. At which depth does the bright cell jump from the source column to the query column? That is the layer that does the copy. - Raise
stepsindemonstrate_patching(e.g. 200 → 600). Does the causal head sharpen (recover a larger fraction), and do the two metrics ever agree?
Intuition: One Wire, Not the Whole Node
Head patching told you which head carries the answer. But it lumps together everything that head causes: what it writes straight to the output, and what it sets in motion in every later head and MLP. The recovery is the head’s total effect. It cannot tell you how the head matters — whether it speaks to the logits directly, or works entirely by feeding a component downstream.
Path patching is the finer instrument (Wang et al., the IOI circuit, 2022; formalized by Goldowsky-Dill et al., 2023). Picture the model as a wiring diagram: the embedding, every head, and every MLP are boxes, and each writes into the residual stream, which fans out to everything after it. Activation patching flips a whole box to clean and lets the change flood all its wires. Path patching flips just one wire — one box’s contribution as it reaches one receiver — and holds every other wire at its corrupt value.
Point that wire at the logits and you read a component’s direct effect: what it writes straight into the final read-out, excluding whatever it triggers in later layers. Below, restore one wire at a time and watch how much of the answer each one recovers on its own:
Most wires recover almost nothing on their own — the answer is not written by any single direct edge. That is the door path patching opens: a component can be critical and yet write nothing straight to the logits.
The Math: The Residual Stream Is a Sum of Wires
Path patching works because the residual stream is additive. Every sublayer adds its output back through a residual connection, so the vector that reaches the read-out is a plain sum — the embedding, plus, for each block, every head’s write, the shared attention bias, and the MLP write:
x_{\text{final}} \;=\; x_{\text{embed}} \;+\; \sum_{\ell}\Big( \sum_{h} \text{head}_{\ell,h} \;+\; b^{O}_{\ell} \;+\; \text{mlp}_{\ell} \Big)
This is the same decomposition direct logit attribution used to split a logit into per-component votes (Elhage et al., 2021). A component’s direct path to the logits is exactly its own term in that sum. So path-patching one wire is one subtraction: take the corrupt final residual, swap the sender’s term for its clean value, and read out through the model’s own \text{ln\_final} + \text{lm\_head}.
Two identities make the tool exact and tie it back to activation patching:
- Restore no wires → the residual is unchanged → the corrupt logits.
- Restore every wire → each term becomes clean → the clean residual, bit-for-bit → the clean logits. Patching all of a node’s paths is just ordinary node (activation) patching. Path patching is the strict generalization.
NoteKey Insight
Activation patching asks “does this component matter?” Path patching asks “which of its wires matters?” The residual stream being a literal sum is what lets you answer the second question by swapping a single additive term — no re-run, no approximation.
Code: Path Patching from Scratch
The one new primitive derives the additive decomposition from the faithful cache we already built — the MLP write is recovered as resid_post − resid_pre − attention, so nothing is re-run. Its sum equals the final residual exactly:
import importlib.util
import sys
from pathlib import Path
import torch
spec = importlib.util.spec_from_file_location("path_patching", Path("path_patching.py").resolve())
PP = importlib.util.module_from_spec(spec)
sys.modules["path_patching"] = PP
spec.loader.exec_module(PP)
torch.manual_seed(0)
model = PP.GPTModel(vocab_size=20, embed_dim=32, num_heads=4, num_layers=2, max_seq_len=16).eval()
ids = torch.randint(0, 20, (1, 7))
c = PP.node_contributions(model, ids)
total = (c["embed"] + c["head_out"].sum(dim=(0, 2)) + c["mlp"].sum(dim=0) + c["attn_bias"].sum(dim=0))
print("Σ node contributions == final residual:", bool(torch.allclose(total, c["final"], atol=1e-5)))
print("final reads out to the model's logits:", bool(torch.allclose(PP.readout(model, c["final"]), c["logits"], atol=1e-5)))Σ node contributions == final residual: True
final reads out to the model's logits: TrueDirect-path patching is then a single subtraction, and the two identities hold exactly — restore nothing and you have the corrupt run; restore every wire and you have the clean run (path patching all paths = node patching):
clean, corrupt, meta = PP.make_induction_pair(period=6, vocab_size=20, seed=7)
cc = PP.node_contributions(model, clean)
xc = PP.node_contributions(model, corrupt)
empty = PP.direct_path_patch(model, cc, xc, []) # restore no wires
allw = PP.direct_path_patch(model, cc, xc, PP.all_nodes(cc)) # restore every wire
print("restore nothing → corrupt logits:", bool(torch.allclose(empty, xc["logits"], atol=1e-5)))
print("restore all wires → clean logits: ", bool(torch.allclose(allw, cc["logits"], atol=1e-5)))restore nothing → corrupt logits: True
restore all wires → clean logits: TrueThe production code lives in path_patching.py. It reuses run_with_cache, head_patching, and the induction pair from activation_patching.py by reference — path patching is a lens over the same machinery, not a second copy of it.
pt = PP.demonstrate_path_patching(period=8, vocab_size=12, steps=400)
ojs_define(pathTrace = pt)Direct vs. Total: A Critical Head That Writes Nothing
Now put the two grids side by side. For every head, total recovery is node patching (its whole effect) and direct recovery is path patching to the logits (only what it writes straight out). Where the bars agree, a head speaks to the output directly. Where they diverge, a head does its work downstream:
The tallest total bar is the layer-0 previous-token head — node patching crowns it the most important head in the circuit. Yet its direct bar is almost flat: it writes essentially nothing to the logits. Its entire causal role is upstream — it stamps each position with “the token before me,” building exactly the residual the layer-1 induction head reads to copy the answer. Activation patching sees a critical head; path patching sees a critical head that never touches the output. Both are true, and only the finer tool tells them apart — the same upstream-vs-output split that lets the IOI circuit separate its “S-inhibition” heads from its “name-mover” heads.
NoteKey Insight
A large total effect with a near-zero direct effect is the signature of an upstream component: it matters only because of what reads it later. This is the one thing head patching structurally cannot show you — and the reason circuit work reaches for path patching.
Cutting the Stream Higher: Path Patching Through a Layer
The logits are not the only receiver. Aim the wire at the residual entering block L instead, and only the part of a sender’s signal that travels through blocks L{\ldots}\text{end} is let clean. The same two identities hold at the boundary — restore nothing and you keep the corrupt run; restore every upstream wire and the whole entering residual is clean, so the rest of the network reproduces the clean run:
for L in (0, 1, 2):
none_through = PP.path_patch_through(model, cc, xc, [], receiver_layer=L)
all_through = PP.path_patch_through(model, cc, xc, PP.upstream_nodes(cc, L), receiver_layer=L)
print(f"block {L}: restore none → corrupt:", bool(torch.allclose(none_through, xc["logits"], atol=1e-5)),
"| restore all upstream → clean:", bool(torch.allclose(all_through, cc["logits"], atol=1e-5)))block 0: restore none → corrupt: True | restore all upstream → clean: True
block 1: restore none → corrupt: True | restore all upstream → clean: True
block 2: restore none → corrupt: True | restore all upstream → clean: TrueChoosing the receiver is choosing where in the network you let a component’s signal flow. The logits give you its direct effect; a later block gives you its effect routed through that block. Sweep the receiver and you trace a signal’s path through the model — which is how a full circuit gets drawn.
TipTry This
- Scrub the wires. Step the wire selector through every sender. Is there any single direct wire that recovers most of the answer? (There isn’t — that is the point.)
- Read the gap. In the direct-vs-total chart, find the head with the biggest gap between its bars. That head is doing its job through other components, not at the output.
- Move the receiver. In
path_patch_through, restore only the previous-token head into block 1’s input vs. into the logits. Where does it recover more — confirming its signal reaches the output through the later block, not directly?
Intuition: Below the Component — Features in Superposition
Every tool so far stopped at a component: a layer (the lens), a write (DLA), a head (the induction circuit and head patching). But look closely at any single component and it is polysemantic — one neuron, one residual direction, fires for a jumble of unrelated things: DNA sequences and HTTP headers and a name in Korean. The grain we have been reading is not the grain the model thinks in.
The reason is superposition (Elhage et al., Toy Models of Superposition, 2022). A residual stream of width d does not store d features in d neurons. It stores many more than d features by packing each one along its own nearly orthogonal direction and relying on sparsity — only a few fire at once — so the interference between them stays small. A d-dimensional space has room for only d orthogonal axes, but exponentially many almost-orthogonal ones. The model exploits every last one. That is why a neuron is polysemantic: it is one axis, and many feature directions have a component along it.
So the honest unit of meaning is a feature: a direction in activation space, not a neuron. To read the model’s own vocabulary we need to decompose each activation into a sparse sum of feature directions — recover the dictionary the model packed. That is exactly what a sparse autoencoder learns.
NoteKey Insight
A neuron is an axis; a feature is a direction. Because the model stores more features than it has neurons, the two never line up — every neuron is a blend. A sparse autoencoder gives up on neurons and learns the directions directly.
The Math: A Sparse, Overcomplete Dictionary
An SAE is the simplest autoencoder that could untangle superposition: one hidden layer, wider than the input, with a sparsity penalty. For an activation \mathbf{x} \in \mathbb{R}^{d} it computes a non-negative code \mathbf{f} \in \mathbb{R}^{m} with m \gg d (an overcomplete dictionary), then reconstructs:
\bar{\mathbf{x}} = \mathbf{x} - \mathbf{b}_{\text{dec}}, \qquad \mathbf{f} = \mathrm{ReLU}\!\left(\mathbf{W}_{\text{enc}}\,\bar{\mathbf{x}} + \mathbf{b}_{\text{enc}}\right), \qquad \hat{\mathbf{x}} = \mathbf{W}_{\text{dec}}\,\mathbf{f} + \mathbf{b}_{\text{dec}}.
Each row of \mathbf{W}_{\text{dec}} is one atom — a feature direction in activation space — and \hat{\mathbf{x}} is the sum of the active atoms, each weighted by its code entry. The \mathrm{ReLU} makes the code non-negative (a feature is present or absent, never “anti-present”), and the pre-encoder bias \mathbf{b}_{\text{dec}} is subtracted before encoding and added back after, so the empty code reconstructs to \mathbf{b}_{\text{dec}} — the data mean, the part no feature needs to explain.
Training minimizes reconstruction error plus an L1 penalty on the code:
\mathcal{L} \;=\; \mathbb{E}_{\mathbf{x}}\!\left[\; \underbrace{\lVert \mathbf{x} - \hat{\mathbf{x}} \rVert_2^2}_{\text{reconstruction}} \;+\; \lambda \underbrace{\textstyle\sum_i f_i(\mathbf{x})}_{\text{sparsity}} \;\right].
The L1 term is what forces monosemanticity: to pay the least sparsity cost, the encoder learns to fire one atom per underlying feature instead of smearing it across many. There is one catch — the network could cheat by shrinking \mathbf{f} (cheap L1) and scaling \mathbf{W}_{\text{dec}} up to keep the reconstruction. So after every step we renormalize each decoder row to unit norm, pinning the atoms’ scale and leaving \lambda in real control of sparsity. The single knob \lambda trades reconstruction against sparsity — the frontier we chart below.
Code: A Sparse Autoencoder from Scratch
sae.py builds exactly the three equations above. Load it and watch the two identities that pin the math down — the loss really is recon + λ·sparsity, and the empty code really reconstructs to b_dec:
import importlib.util
import sys
from pathlib import Path
import torch
spec = importlib.util.spec_from_file_location("sae", Path("sae.py").resolve())
SAE = importlib.util.module_from_spec(spec)
sys.modules["sae"] = SAE
spec.loader.exec_module(SAE)
torch.manual_seed(0)
sae = SAE.SparseAutoencoder(input_dim=8, dict_size=32, seed=1) # 4× overcomplete
x = torch.randn(6, 8)
x_hat, f = sae(x)
print(f"activation x: {tuple(x.shape)} feature code f: {tuple(f.shape)} (dict is 4× wider)")
print(f"code is non-negative (ReLU): {bool((f >= 0).all())}")
parts = SAE.sae_loss(x, x_hat, f, l1_coeff=0.2)
print(f"loss = recon + λ·sparsity exactly: "
f"{bool(torch.allclose(parts['total'], parts['reconstruction'] + 0.2 * parts['sparsity']))}")
# Anchor: with every atom off, the reconstruction is the decoder bias — nothing else.
empty = sae.decode(torch.zeros(1, 32))
print(f"empty code → b_dec exactly: {bool(torch.allclose(empty[0], sae.b_dec))}")activation x: (6, 8) feature code f: (6, 32) (dict is 4× wider)
code is non-negative (ReLU): True
loss = recon + λ·sparsity exactly: True
empty code → b_dec exactly: TrueThe production model lives in sae.py (SparseAutoencoder, sae_loss, l0_norm, dead_features, …). Next we make it earn its atoms.
Recover Planted Features
On a real model “the atoms look interpretable” is a judgment call. So we test the SAE where we can check it: plant a known dictionary of feature directions, generate data from it in superposition, and see whether the SAE rediscovers the planted atoms. The recovery metric is mean max cosine similarity (MMCS) — for each true atom, the best cosine match among the learned atoms, averaged. MMCS = 1 means every planted feature was found, up to permutation and scale.
summary = SAE.demonstrate_superposition_recovery(
num_features=64, dim=32, active=2, l1_coeff=0.2, steps=500, seed=0
)
hist = summary["history"]
ojs_define(
saeStepsX = hist["steps"],
saeMMCS = hist["mmcs"],
saeL0 = hist["l0"],
saeRecon = hist["reconstruction"],
saeTrueActive = summary["true_active"],
saeInitMMCS = summary["initial_mmcs"],
saeFinalMMCS = summary["final_mmcs"],
)print("64 features planted in a 32-dim space (2× overcomplete); 2 active per input")
print(f"MMCS recovery: {summary['initial_mmcs']:.3f} → {summary['final_mmcs']:.3f}")
print(f"code sparsity L0: {summary['final_l0']:.2f} (true active per input = {summary['true_active']})")
print(f"dead atoms: {summary['final_dead']} explained variance: {summary['explained_variance']:.3f}")64 features planted in a 32-dim space (2× overcomplete); 2 active per input
MMCS recovery: 0.393 → 0.978
code sparsity L0: 4.40 (true active per input = 2)
dead atoms: 0 explained variance: 0.945At initialization MMCS sits near the chance level for random directions in 32 dimensions (~0.4); as training proceeds it climbs toward 1.0 — the SAE has recovered the dictionary the data was built from. That is the honest version of “sparse autoencoders find features”: here we planted them, so we can measure that they came back.
The Sparsity–Reconstruction Tradeoff
The whole art of training an SAE is the single knob \lambda. Turn it up and the code gets sparser (fewer atoms per input) but the reconstruction gets worse; turn it down and you reconstruct almost perfectly with a dense, uninterpretable code. Sweeping \lambda traces a Pareto frontier — each point an SAE, plotting how many atoms it uses against how much variance it leaves unexplained. Somewhere on that curve is the sweet spot where the learned atoms line up with real features (highest MMCS).
lambdas = [0.05, 0.1, 0.2, 0.35, 0.5]
true_dict = SAE.make_feature_dictionary(64, 32, seed=0)
X_pareto, _ = SAE.sample_sparse_activations(true_dict, n=3072, active=2, seed=0)
pareto = SAE.sae_pareto(X_pareto, lambdas, dict_size=64, steps=300, seed=7, true_dictionary=true_dict)
ojs_define(
paretoLambda = pareto["l1_coeffs"],
paretoL0 = pareto["l0"],
paretoRecon = pareto["reconstruction"],
paretoMMCS = pareto["mmcs"],
paretoDead = pareto["dead"],
)
TipTry This
- Sweep λ with the slider. Watch L0 and the reconstruction error trade off, and note where MMCS peaks — the best-recovering SAE is neither the densest nor the sparsest.
- Push λ high (right slider positions). Dead atoms start to appear: the L1 pressure kills features that never fire. That is the failure mode dead-feature resampling (
resample_dead_features) exists to fix.
Interactive Exploration: Features on a Real Model
Now the same SAE on the book’s own GPTModel. We gather the residual stream at the last block over a handful of prompts (reusing m21’s faithful run_with_cache — no new tracing), train an overcomplete SAE on those vectors, and look at the learned code. The heatmap below is atoms × token positions: each column is one token’s feature code. A well-trained SAE fires few atoms per token — the columns are nearly one-hot, the visual signature of a sparse, monosemantic decomposition.
# Load the sibling induction module for its GPTModel + a quick trained model.
ap_spec = importlib.util.spec_from_file_location("activation_patching", Path("activation_patching.py").resolve())
AP = importlib.util.module_from_spec(ap_spec)
sys.modules["activation_patching"] = AP
ap_spec.loader.exec_module(AP)
torch.manual_seed(0)
gpt = AP.GPTModel(vocab_size=24, embed_dim=32, num_heads=4, num_layers=2, max_seq_len=16).eval()
prompts = torch.randint(0, 24, (8, 12))
acts = SAE.collect_residual_activations(gpt, prompts, layer=-1) # (8*12, 32)
feat_sae = SAE.SparseAutoencoder(input_dim=32, dict_size=128, seed=1)
SAE.train_sae(feat_sae, acts, l1_coeff=2.0, steps=300, lr=1e-2, seed=1)
# One prompt's code, atoms × positions; keep the atoms that ever fire for a compact heatmap.
one = SAE.collect_residual_activations(gpt, prompts[:1], layer=-1) # (12, 32)
with torch.no_grad():
code = feat_sae.encode(one) # (12, 128)
live = (code > 1e-6).any(dim=0)
code_live = code[:, live].t() # (n_live_atoms, 12)
ojs_define(
saeHeat = code_live.detach().numpy().tolist(),
saeHeatL0 = SAE.l0_norm(code),
)Only a few atoms light per column — most of the 128-atom dictionary stays dark for any given token. That sparsity is the point: the SAE has re-expressed a dense, polysemantic 32-dim residual as a short list of features, the grain the earlier tools could not reach.
Modern SAEs: Fixing L1’s Shrinkage
The SAE you just built is the vanilla one — a ReLU encoder trained against an L1 penalty. It works, but it pays for its sparsity with a bias the whole field has since moved away from. Two symptoms you already met — the shrunk magnitudes the unit-norm decoder had to guard against, and the dead atoms that appeared as \lambda climbed — are the same disease: the L1 term punishes the size of a feature, not just its presence. The 2024 generation of SAEs cures it by making sparsity a property of the activation instead of a penalty. We’ll build the two that matter — TopK and JumpReLU — from scratch.
Intuition: L1 Buys Sparsity With a Magnitude Tax
Picture one feature that genuinely fires with strength a. The reconstruction wants the code f to equal a. But the loss also carries \lambda f — a toll that grows with f. So the optimizer settles for something smaller than a: it trades a little reconstruction error for a lower toll. Every active feature comes out shrunk, and features whose honest strength is below the toll are pushed all the way to zero — a dead atom. The sparsity is real, but so is the distortion: the numbers an L1 SAE reports for “how strongly did this feature fire?” are systematically too small.
NoteKey Insight
An L1 penalty conflates two questions — which features are active and how strongly. It answers the first by damaging the second. TopK and JumpReLU separate them: decide activity by a hard rule, then let the survivors report their full magnitude.
The Math: Shrinkage Is Soft-Thresholding
We can make the bias exact. Take a single atom whose decoder direction d is a unit vector, reconstructing an input x = a\,d. Because \lVert d\rVert = 1, the reconstruction term collapses to (a - f)^2, and the code that minimizes the SAE objective for this feature is
f^\star \;=\; \arg\min_{f \ge 0}\;\big[(a - f)^2 + \lambda f\big] \;=\; \max\!\Big(0,\; a - \tfrac{\lambda}{2}\Big).
That is the soft-threshold: the recovered magnitude is the truth minus a constant \lambda/2, no matter how large a is. Below \lambda/2 the feature dies; above it, it survives but reads low by exactly \lambda/2. This isn’t a tuning artifact — it’s the closed-form optimum. We don’t have to take it on faith; shrinkage_demo minimizes the real objective numerically and checks it against the formula:
shr = SAE.shrinkage_demo(l1_coeff=0.4, verify=True) # numerically minimize, then compare
ojs_define(
shrinkOffset = shr["offset"],
shrinkError = shr["max_abs_error"],
)
print(f"offset λ/2 = {shr['offset']:.3f}")
print(f"max |numeric − analytic| across magnitudes = {shr['max_abs_error']:.2e}")The numeric minimum lands on \max(0, a - \lambda/2) to within roundoff — the shrinkage is provable. Drive \lambda below and watch the L1 recovery line slide a constant \lambda/2 under the diagonal, while TopK and JumpReLU stay on it.
Code: The TopK SAE From Scratch
The simplest cure is the bluntest: after the encoder’s ReLU, keep the k largest features and zero the rest. There is no penalty at all — sparsity is enforced by the activation — so nothing shrinks the survivors, and the code’s L_0 is k by construction (Gao et al., 2024; the idea is Makhzani & Frey’s 2013 k-sparse autoencoder). You stop tuning toward a sparsity and simply set it.
z = torch.tensor([[0.1, 0.9, 0.4, 0.7, 0.2]])
kept = SAE.topk_activation(z, k=2) # keep the 2 largest, zero the rest
print("pre-activations :", z.tolist()[0])
print("after TopK(k=2) :", kept.tolist()[0]) # 0.9 and 0.7 survive, at full valueTopKSAE is the vanilla autoencoder with exactly this one line swapped into its encoder — everything else (the unit-norm decoder, the transpose init, the training loop) is inherited unchanged. Because sparsity is structural, we train it with l1_coeff=0: pure reconstruction, nothing to tune.
tk = SAE.TopKSAE(input_dim=32, dict_size=128, k=4, seed=0)
_, f_tk = tk(torch.randn(6, 32))
print("L0 per input:", SAE.l0_norm(f_tk), " (exactly k = 4, every input)")Watch TopK Pin the Sparsity
Train an L1 SAE and a TopK SAE on the same planted data. The L1 SAE’s L_0 drifts as the optimizer negotiates the penalty; the TopK SAE’s L_0 is a flat line nailed to k from step 0 — and it reaches the same recovery (MMCS) without a \lambda to sweep.
dict64 = SAE.make_feature_dictionary(64, 32, seed=0)
Xtk, _ = SAE.sample_sparse_activations(dict64, n=3072, active=3, seed=0)
l1_sae = SAE.SparseAutoencoder(32, 64, seed=7)
h_l1 = SAE.train_sae(l1_sae, Xtk, l1_coeff=0.3, steps=400, lr=1e-2, true_dictionary=dict64, record_every=20)
tk_sae = SAE.TopKSAE(32, 64, k=3, seed=7)
h_tk = SAE.train_sae(tk_sae, Xtk, l1_coeff=0.0, steps=400, lr=1e-2, true_dictionary=dict64, record_every=20)
ojs_define(
tkSteps = h_tk["steps"],
tkL0L1 = h_l1["l0"], tkL0Tk = h_tk["l0"],
tkMMCSL1 = h_l1["mmcs"], tkMMCSTk = h_tk["mmcs"],
tkK = 3,
)
TipTry This
- Read the two lines. The TopK curve is dead flat at k; the L1 curve wanders because L_0 is a side effect of \lambda, not a setting. Same recovery, one fewer hyperparameter.
- Change
kin the code cell (say to 2 or 6) and re-run. The flat line moves to the new k exactly — you dial sparsity directly, in the units you care about.
Code: JumpReLU — Full Magnitude Above a Threshold
TopK fixes the count. JumpReLU (Rajamanoharan et al., 2024, the SAE behind DeepMind’s Gemma Scope) fixes the magnitude with a continuous switch: keep a learned per-feature threshold \theta, and let a feature pass its full pre-activation the instant it clears \theta —
\text{JumpReLU}_\theta(z) \;=\; z \cdot H(z - \theta),
where H is the Heaviside step. The discontinuity at \theta is the “jump”: unlike a soft-threshold, which would subtract \theta and shrink the survivor, JumpReLU passes z untouched. Below \theta it is exactly zero.
z = torch.tensor([[0.2, 0.9, 0.5, 0.35]])
gated = SAE.jumprelu_activation(z, theta=0.5) # z · H(z − 0.5)
print("pre-activations :", z.tolist()[0])
print("after JumpReLU :", gated.tolist()[0]) # 0.9 and 0.5 survive at full value; 0.2, 0.35 → 0JumpReLUSAE stores the threshold as \log\theta so \theta > 0 always. Training \theta faithfully is the one subtlety: the step function has no gradient, so the paper pushes one through with straight-through estimators and penalizes L_0 (the count) directly instead of an L1 proxy — that’s what lets it dodge shrinkage entirely. We keep that honest here: rather than fake the STE, we drive \theta by hand and watch the gate, using TopK as our trained example.
TipTry This
Raise \theta and watch features drop out — but notice the survivors never shrink: each bar keeps its full height right up until it crosses \theta and vanishes. That faithful magnitude, plus a sparsity you control, is why JumpReLU and TopK have replaced the L1 SAE in most current interpretability work.
Gated SAEs: The Idea That Led Here
JumpReLU didn’t arrive alone. Its direct ancestor, the Gated SAE (Rajamanoharan et al., 2024), was the first to name the fix out loud: split the encoder into a gate that decides which features fire and a magnitude path that decides how strongly, and apply the L1 penalty only to the gate. The magnitude path never sees the toll, so it never shrinks — the gate reported half as many firing features for the same reconstruction fidelity as a vanilla SAE. The JumpReLU you just built is the same insight compressed into a single thresholded activation, which is why it superseded the two-tower Gated design.
NoteKey Insight
Every fix here is the same move: decouple activity from magnitude. L1 couples them and pays with shrinkage; Gated SAEs separate them with two towers; JumpReLU separates them with a threshold; TopK separates them with a hard count. The trend line of SAE research is exactly this decoupling, made ever simpler.
Feature Steering: Turn a Direction Into a Dial
Every tool in this module so far has read the residual stream. The logit lens read out a layer; DLA split a logit into per-component pushes; activation patching localized a computation causally; the SAE named the directions the stream packs in superposition. Reading is where interpretability earns its trust — but it is only half the loop. The other half is the one everyone remembers: once you hold a direction, you can write it back in and causally change what the model does.
That is steering, and it is one line:
h \;\leftarrow\; h + \alpha \, d
Add a direction d to the residual stream at some layer, at strength \alpha, and let it propagate. Turn \alpha up and the model leans harder into whatever d means. This single move is behind ActAdd (Turner et al., 2023 — steering with no fine-tuning, just a vector) and the feature clamping of Scaling Monosemanticity (Templeton et al., 2024) that produced “Golden Gate Claude” by pinning one SAE feature on. We will build it from scratch on the book’s own GPTModel, prove it is an exact perturbation of the real forward, and then turn the dial and watch a confident prediction flip.
Intuition: Two Ways to Get a Steering Direction
The mechanism is direction-agnostic — h \leftarrow h + \alpha d works for any d. What makes steering useful is having a d that means something. Two constructions dominate, and this module has already built the ingredients for both:
- A contrast of two contexts (ActAdd / CAA). Run the model on a “toward” prompt and an “away-from” prompt, read the residual stream of each at some layer, and subtract: d = h_{\text{pos}} - h_{\text{neg}}. This points from what the model computes on one context toward the other — “Love” minus “Hate” (Turner et al., 2023), or a mean over many contrastive pairs (Rimsky et al., 2023). It needs no SAE and is the strongest, most reliable lever.
- A single SAE feature (Golden Gate). Take one unit decoder atom d_i = W_{\text{dec},i} from the SAE you just trained. Now \alpha is a dial on that named feature. On a frontier model, clamping one such feature is enough to make it talk about the Golden Gate Bridge in every reply (Templeton et al., 2024). It is less about raw strength than interpretability: you know exactly which concept you turned up.
We build the contrast direction as the headline (it flips a decision cleanly) and then use the SAE to name which feature points the same way.
The Math: Add a Direction, Read the Effect
Steering is a perturbation of the model’s own forward — nothing is re-derived. Recall the pre-norm block from m06: the residual after block \ell is a running sum, and the model reads the final state out with \text{lm\_head}(\text{ln\_final} (x)). Steering inserts one extra summand after block \ell:
x^{(\ell)} \;\leftarrow\; x^{(\ell)} + \alpha\, d \qquad\text{then continue blocks } \ell{+}1, \ell{+}2, \dots
Two facts pin it down and make it trustworthy, not just plausible:
- \alpha = 0 is the model, exactly. Adding a zero field changes nothing — the steered logits equal \text{model}(x) bit-for-bit. Steering is a controlled nudge to the real computation.
- The stream really is x^{(\ell)} + \alpha d. Read the residual back after the steered block: at the steered positions it is the un-steered residual plus exactly \alpha d. Everything downstream is the model faithfully propagating that one write — which is what makes the effect causal.
Why the direct logit push is linear in \alpha. If we freeze the final LayerNorm’s per-position scale (the DLA convention from earlier in this module), the read-out is affine, so the direction’s first-order contribution to any token’s logit is \alpha \cdot W_U\!\left(\gamma \odot d / \sigma\right) — a straight line in \alpha. The full effect bends once \alpha is large enough to change the LayerNorm statistics and the attention pattern downstream, but near the base point the dial is linear, which is exactly why a small \alpha nudges and a larger one commits.
Clamping. Instead of a fixed \alpha, set a feature’s activation to a target value v. Along the unit atom d_i, the feature’s current contribution to the reconstruction is f_i(x)\,d_i; to move it to v\,d_i you add
\big(v - f_i(x)\big)\, d_i .
Ask for the value the feature already has and the field is zero — clamping to the current value is a no-op. Clamp high and you have the Golden-Gate move.
Code: Steering From Scratch
steering.py is a faithful forward — it replays the model’s own sub-layers (ln1/attention/ln2/ffn/ln_final/lm_head) exactly as run_with_cache does — and only adds a field to the residual stream after a chosen block. Load it and watch the three anchors hold live:
import importlib.util
from pathlib import Path
import torch
spec = importlib.util.spec_from_file_location("steering", Path("steering.py").resolve())
STEER = importlib.util.module_from_spec(spec)
spec.loader.exec_module(STEER)
torch.manual_seed(0)
model = STEER.GPTModel(vocab_size=20, embed_dim=32, num_heads=4, num_layers=3, max_seq_len=16, dropout=0.0).eval()
ids = torch.randint(0, 20, (1, 7))
# Anchor 1: α = 0 (a zero field) reproduces the model bit-for-bit.
zero_steer = STEER.steered_forward(model, ids, layer=1, vector=torch.zeros(32))
print(f"α=0 steer == model: {bool(torch.allclose(zero_steer, model(ids), atol=1e-5))}")
# Anchor 2: the steered residual after the layer is base + α·d, propagated.
ap = STEER._load_sibling("activation_patching")
d = torch.randn(32); d = d / d.norm(); alpha = 3.0
_, cache = ap.run_with_cache(model, ids)
base_resid = cache["resid_post"][1][0] # (seq, dim) after block 1
patch = {1: {p: base_resid[p] + alpha * d for p in range(ids.shape[1])}}
expected = ap.patched_forward(model, ids, resid_patches=patch) # base+α·d, propagated
got = STEER.steered_forward(model, ids, layer=1, vector=alpha * d)
print(f"steered resid == base + α·d: {bool(torch.allclose(got, expected, atol=1e-5))}")
# Anchor 3: clamping a feature to its own current value changes nothing.
sae = STEER._load_sibling("sae").SparseAutoencoder(32, 64, seed=3)
STEER._load_sibling("sae").train_sae(
sae, STEER._load_sibling("sae").collect_residual_activations(model, ids, layer=0),
l1_coeff=0.05, steps=60, lr=1e-2, batch_size=7, seed=0,
)
cur = float(STEER.feature_activation(sae, model, ids, layer=0)[0, -1, 7])
noop = STEER.clamp_feature(model, ids, layer=0, sae=sae, feature=7, value=cur, positions=[-1])
print(f"clamp to current value == no-op: {bool(torch.allclose(noop[0, -1], model(ids)[0, -1], atol=1e-4))}")α=0 steer == model: True
steered resid == base + α·d: True
clamp to current value == no-op: Truesteered_forward is the primitive; steer_with_feature(model, ids, layer, sae, feature, alpha) is just it with vector = α·d_feature, and clamp_feature computes the (v - f_i)\,d_i field for you. Next we point a meaningful direction at a real decision.
Turn the Dial: Flip a Prediction
Here is the whole loop, from scratch and deterministic. We train a tiny GPTModel on a two-branch corpus where a single marker token decides the ending:
<bos> dog . the→ park<bos> cat . the→ sun
The model must route the marker (“dog”/“cat”) to the last position, so the decision-position residual genuinely encodes which branch it is on. We build the contrastive steering vector — the cat-context residual minus the dog-context residual at layer 1 — then feed the dog prompt (the model is 99.9% sure the next word is “park”) and steer along that direction. Sweep \alpha:
demo = STEER.demonstrate_steering()
ojs_define(
stAlphas = demo["alphas"],
stRise = demo["rise_prob"], # p(sun) — steered toward
stFall = demo["fall_prob"], # p(park) — the original prediction
stRank = demo["rise_rank"], # rank of "sun"
stFlip = demo["flip_alpha"],
stPosTarget = demo["pos_target"], # "park"
stNegTarget = demo["neg_target"], # "sun"
stContrastNorm = demo["contrast_norm"],
stReadout = demo["readout_top"],
stFeature = demo["sae_feature"],
stCosine = demo["sae_cosine"],
)
print(f'dog-prompt base: p("{demo["pos_target"]}") = {demo["base_prob_pos"]:.3f}')
print(f'flip at α ≈ {demo["flip_alpha"]:.2f} → p("{demo["neg_target"]}") = {demo["final_prob_neg"]:.3f} at α_max')
print(f'||contrast vector|| = {demo["contrast_norm"]:.1f} (steer strength is measured in residual-norm units)')Step the dial below. At \alpha = 0 the model is committed to park; as you turn it up, park collapses and sun rises, crossing over near \alpha \approx the flip strength — the model changes its mind, with no weights touched. Steering strength is measured in residual-norm units; the contrast vector has norm \approx 21.6, so the crossover sits at roughly half the distance between the two branches.
TipTry This
Step from \alpha = 0 upward. Nothing much happens at first — then, right around the flip line, park falls off a cliff and sun shoots up: the model commits to the other branch. Notice you never edited a weight. You reached into the residual stream mid-forward, added one vector, and the whole downstream computation re-routed. That is the entire idea behind test-time control of a model’s behavior.
What the Feature Promotes
The contrast vector is strong but anonymous. The SAE gives it a name: which of its learned atoms points the same way? We take the atom with the highest cosine to the steering direction and read out what it promotes with the frozen-LayerNorm lens from the DLA section — its first-order push on each vocabulary token. If the SAE found a real “sun-branch” feature, that atom should push “sun” up:
The top-pushed token is exactly the branch we steered toward — the SAE’s atom is the interpretable stand-in for the anonymous contrast vector. (In a tiny toy model a single atom is a weak lever, which is why we steered along the full contrast direction; on a frontier model one such atom is strong enough to be Golden Gate Claude.)
NoteKey Insight
Steering is the write half of interpretability: reading tools find directions, and the same additive structure that let DLA decompose the stream lets us inject into it — h \leftarrow h + \alpha d. The exact anchors (\alpha = 0 is the model; the stream really is h + \alpha d) are what make an edit a piece of evidence rather than a party trick. And the two ways to get d — a contrast of contexts, or one named SAE feature — are ActAdd and Golden Gate Claude, built from the same one line.
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.
WarningA fixed-period stripe can’t tell positional from content
On a probe with a constant period, “attend period-1 back” (a purely positional head) and “attend after my current token’s last occurrence” (a content-based induction head) draw the same diagonal — so a high stripe score alone does not prove induction. That is exactly why the emergent top head often lands in layer 0: gradient descent found the cheaper positional shortcut. To distinguish them you need variable offsets (the earlier occurrence at an unpredictable distance), where only content matching generalizes. The hand-constructed circuit is content-based by construction — that is why it is the lesson’s proof, and the trained model its (softer) illustration.
WarningInduction needs two layers — and the right composition
An induction head is useless without the previous-token head feeding its keys (K-composition). A one-layer model cannot do induction: there is no earlier head to write the “what preceded me” labels the induction query matches against. If you probe a 1-layer model for an induction stripe, you will not find one.
WarningPatching depends on your corruption and your metric
A patch’s recovery is only as meaningful as the clean/corrupt pair it denoises. Corrupt too much and every site restores something; corrupt too little and nothing breaks to recover. Denoising (patch clean → corrupt) and noising (patch corrupt → clean) can even disagree, because models keep backup components that step in only when a primary one is ablated — so a site can look unimportant under denoising yet matter under noising. Report the corruption and metric you used, and prefer the logit difference (Zhang & Nanda, 2023). Our result is robust because the corruption is a single content-only token swap and the metric is the normalized logit-difference recovery.
Warning“Direct” means straight to the logits — a head still feeds its own MLP
Direct-path patching restores a component’s own additive term and reads it out — so “direct” bypasses every MLP and every later block, including the head’s own block’s MLP (which reads the post-attention residual). A head can therefore have a large total effect and a small direct effect even in a single layer, not because it acts through a later head but because its signal is reshaped by the MLP beside it. Don’t read a low direct recovery as “this head is unimportant” — read it as “this head’s effect is mediated.” And because the read-out LayerNorm is nonlinear, direct and indirect effects are additive in the residual, not in the scalar recovery metric — so don’t expect direct + indirect to sum to the total number.
WarningWithout the unit-norm decoder, L1 is free
An SAE can drive the sparsity term to zero without becoming sparse: shrink every f_i toward 0 (cheap L1) and scale the matching decoder atom up to keep the reconstruction. The code looks sparse by L1 but every atom still fires. The fix is the constraint the lesson applies every step — renormalize each decoder row to unit norm — so \lambda, not a scale trick, controls sparsity. Drop normalize_decoder() from the loop and watch L0 stay high while the L1 loss collapses.
WarningDead features and the sparsity–reconstruction trade
Push \lambda too high and atoms die — they never fire, wasting dictionary capacity, and reconstruction suffers. That is not a bug in the SAE; it is the frontier. Revive dead atoms with resampling (resample_dead_features, pointing them at poorly-reconstructed inputs) and pick \lambda on the Pareto curve where MMCS peaks — neither the densest nor the sparsest SAE recovers features best. And remember an SAE is lossy: it explains most, not all, of the variance, and the atoms it finds are a hypothesis about the model’s features, not ground truth (which is why we validated it against a planted dictionary first).
WarningAn L1 SAE’s feature magnitudes are shrunk — don’t read them as strengths
The soft-threshold f^\star = \max(0, a - \lambda/2) means every active feature an L1 SAE reports is low by a constant \lambda/2. So “feature 17 fired at 0.8” is not 0.8 — it’s 0.8 plus the toll. If your analysis depends on absolute activation magnitudes (steering vectors, activation scaling, thresholds on strength), use a TopK or JumpReLU SAE, whose survivors keep their true magnitude; or at least compare features only relatively, where the constant offset cancels.
WarningSteering is causal, but it is not surgical
Adding \alpha d does change behavior — but a direction is rarely a single clean concept. Steering a “sun” feature can also drag along whatever else correlates with it in the training distribution, and pushing \alpha far enough to force one outcome degrades fluency everywhere else (watch the readout’s other tokens move too). Two honest habits: report the whole effect, not just the token you hoped to move; and prefer the smallest \alpha that achieves the effect over the largest one that guarantees it. Steering demonstrates a direction matters; it does not prove the direction is only that concept.
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 = 0Exercise 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.04Exercise 6: The induction bump
Run the hand-built induction circuit on a distinct repeat and confirm the in-context-learning jump: near-chance accuracy on the first copy, exact on the second.
seq = IND.distinct_repeated_sequence(10, seed=4)
preds = IND.induction_predictions(seq)
acc = IND.per_position_accuracy(preds[:-1], seq[0, 1:])
# Your turn: split `acc` at the period and print the mean of each half.
first_half = acc[:10].mean().item()
second_half = acc[10:].mean().item()
print(f"accuracy — first copy: {first_half:.2f} second copy: {second_half:.2f}")
print(f"ICL score: {IND.in_context_learning_score(acc, 10):+.2f}")accuracy — first copy: 0.10 second copy: 1.00
ICL score: +0.90Exercise 7: Patch a single head and measure recovery
head_patching sweeps every head; here, do one by hand. Cache a clean run, take the clean/corrupt induction pair, and splice one head’s clean write into the corrupt forward — then confirm the head found by the full sweep recovers far more of the clean logit difference than an arbitrary other head.
model_ap = AP.train_induction_model(period=8, vocab_size=12, embed_dim=48,
num_heads=4, num_layers=2, steps=400, seed=0)
clean_p, corrupt_p, meta = AP.make_induction_pair(period=8, vocab_size=12, seed=7)
_, cache_c = AP.run_with_cache(model_ap, clean_p)
grid = AP.head_patching(model_ap, clean_p, corrupt_p,
correct_id=int(AP.run_with_cache(model_ap, clean_p)[0][0, meta["read"]].argmax()),
wrong_id=int(AP.run_with_cache(model_ap, corrupt_p)[0][0, meta["read"]].argmax()))
best = grid["best"]
# Your turn: read off the best head's recovery vs. the mean over all heads.
flat = [v for row in grid["grid"] for v in row]
print(f"top causal head L{best['layer']+1}·H{best['head']}: recovery {best['recovery']:.2f}")
print(f"mean over all heads: {sum(flat)/len(flat):.2f} (the copy head stands out)")top causal head L1·H1: recovery 0.81
mean over all heads: 0.14 (the copy head stands out)Exercise 8: Direct vs. total effect of a head
Take the head the sweep just crowned and split its effect. Its total recovery is the node-patch number above; its direct recovery is a path patch to the logits. Confirm the pattern the section found — a head can be critical by total effect yet write almost nothing directly (it acts upstream).
correct_id = int(PP.node_contributions(model_ap, clean_p)["logits"][0, meta["read"]].argmax())
wrong_id = int(PP.node_contributions(model_ap, corrupt_p)["logits"][0, meta["read"]].argmax())
total = best["recovery"]
direct = PP.direct_path_effect(model_ap, clean_p, corrupt_p,
[("head", best["layer"], best["head"])],
correct_id, wrong_id, meta["read"])
print(f"top head total (node patch): {total:.2f}")
print(f" direct (path patch): {direct:.2f}")
# Your turn: which components carry the missing effect? Try restoring the head
# together with its own block's MLP: PP.direct_path_effect(..., [("head", L, H), ("mlp", L)], ...)top head total (node patch): 0.81
direct (path patch): 0.04Exercise 9: The sparsity knob
Train the SAE at two very different \lambda values and watch the code sparsity (L0) and the reconstruction trade off. Which \lambda leaves more atoms active, and which reconstructs the planted data better?
import importlib.util, sys
from pathlib import Path
spec = importlib.util.spec_from_file_location("sae", Path("sae.py").resolve())
SAE = importlib.util.module_from_spec(spec); sys.modules["sae"] = SAE; spec.loader.exec_module(SAE)
D = SAE.make_feature_dictionary(64, 32, seed=0)
X, _ = SAE.sample_sparse_activations(D, n=2048, active=2, seed=0)
for lam in (0.05, 0.5):
s = SAE.SparseAutoencoder(32, 64, seed=7)
h = SAE.train_sae(s, X, l1_coeff=lam, steps=300, lr=1e-2, true_dictionary=D)
print(f"λ={lam}: L0={h['l0'][-1]:.2f} recon={h['reconstruction'][-1]:.3f} MMCS={h['mmcs'][-1]:.3f}")
# Your turn: which λ recovers the planted features best (highest MMCS)?λ=0.05: L0=23.49 recon=0.056 MMCS=0.691
λ=0.5: L0=3.05 recon=0.278 MMCS=0.962Exercise 10: Prove recovery against ground truth
MMCS only means something because we planted the dictionary. Confirm the SAE’s atoms line up with the true atoms up to permutation: for a few true features, print the best-matching learned atom and its cosine.
D = SAE.make_feature_dictionary(64, 32, seed=0)
X, _ = SAE.sample_sparse_activations(D, n=4096, active=2, seed=0)
sae = SAE.SparseAutoencoder(32, 64, seed=1234)
SAE.train_sae(sae, X, l1_coeff=0.2, steps=500, lr=1e-2, batch_size=256, true_dictionary=D)
rec = SAE.feature_recovery(sae.dictionary(), D)
for t in range(4):
print(f"true feature {t} → learned atom {rec['match_index'][t].item():2d} cos={rec['cosine'][t]:.3f}")
print(f"MMCS over all 64 features: {rec['mmcs']:.3f}")
# Your turn: raise `active` to 4 in the sampler — does recovery get harder?true feature 0 → learned atom 57 cos=0.994
true feature 1 → learned atom 40 cos=0.994
true feature 2 → learned atom 47 cos=0.996
true feature 3 → learned atom 0 cos=0.995
MMCS over all 64 features: 0.978Exercise 11: TopK recovers magnitudes the L1 SAE shrinks
Train an L1 SAE and a TopK SAE on the same planted data, then, for a feature that fires, compare the recovered code value against the true amplitude. The L1 SAE reads low by about \lambda/2; the TopK SAE reads the true value. This is the shrinkage bias, measured on a real model rather than derived.
D = SAE.make_feature_dictionary(64, 32, seed=0)
X, S = SAE.sample_sparse_activations(D, n=4096, active=3, seed=0)
l1 = SAE.SparseAutoencoder(32, 64, seed=7)
SAE.train_sae(l1, X, l1_coeff=0.2, steps=400, lr=1e-2, batch_size=256, true_dictionary=D)
tk = SAE.TopKSAE(32, 64, k=3, seed=7)
SAE.train_sae(tk, X, l1_coeff=0.0, steps=400, lr=1e-2, batch_size=256, true_dictionary=D)
# For each SAE, average the active code value over all firing entries.
for name, sae in (("L1", l1), ("TopK", tk)):
with __import__("torch").no_grad():
f = sae.encode(X)
active = f > 1e-6
print(f"{name:4s}: mean active code = {f[active].mean():.3f} L0 = {SAE.l0_norm(f):.2f}")
print(f"true mean amplitude ≈ {S[S > 0].mean():.3f} (λ/2 = 0.10 is the L1 shortfall)")
# Your turn: does the L1 gap ≈ λ/2 grow when you raise l1_coeff to 0.4?L1 : mean active code = 0.291 L0 = 9.42
TopK: mean active code = 0.903 L0 = 3.00
true mean amplitude ≈ 1.002 (λ/2 = 0.10 is the L1 shortfall)Exercise 12: Steer the other way — suppress a token
Steering with a negative \alpha (or clamping a feature to a negative value) is just as valid: it points the write the other way. Take the demo’s steering direction (cat − dog) and sweep \alpha from negative to positive, tracking “park”, the branch the model started on. Confirm the dial runs both ways: negative \alpha (steering away from the cat branch) leaves the model firmly on “park”, while positive \alpha walks it off the cliff and flips to “sun” — the sign of \alpha chooses the direction of the nudge.
import importlib.util as _ilu
from pathlib import Path as _Path
import torch as _torch
_spec = _ilu.spec_from_file_location("steering", _Path("steering.py").resolve())
_ST = _ilu.module_from_spec(_spec); _spec.loader.exec_module(_ST)
# Rebuild the demo's model + contrast direction (deterministic).
_torch.manual_seed(0)
_V = len(_ST._load_sibling("interpretability").TOY_VOCAB)
_model = _ST.GPTModel(vocab_size=_V, embed_dim=64, num_heads=4, num_layers=4, max_seq_len=32, dropout=0.0)
_corpus = _torch.tensor([_ST._POS_SENTENCE, _ST._NEG_SENTENCE])
_opt = _torch.optim.Adam(_model.parameters(), lr=3e-3); _lf = _torch.nn.CrossEntropyLoss()
_model.train()
for _ in range(400):
_opt.zero_grad()
_lo = _model(_corpus[:, :-1])
_lf(_lo.reshape(-1, _V), _corpus[:, 1:].reshape(-1)).backward(); _opt.step()
_model.eval()
_pos = _corpus[0:1, :-1]; _neg = _corpus[1:2, :-1]
_dir = _ST.contrast_direction(_model, _neg, _pos, layer=1) # cat − dog
for _a in (-8.0, -4.0, 0.0, 6.0, 12.0):
_r = _ST.steering_response(_model, _pos, 1, _dir, [_a], target_token=9) # p("park")
print(f"α = {_a:6.1f} p(park) = {_r['prob'][0]:.3f}")
# Your turn: at which α does clamp_feature (on the matched SAE atom) match this curve?α = -8.0 p(park) = 0.999
α = -4.0 p(park) = 0.999
α = 0.0 p(park) = 0.999
α = 6.0 p(park) = 0.995
α = 12.0 p(park) = 0.009Exercise 13: DoLa-static vs. dynamic
Dynamic DoLa selects the premature layer per token by JSD; DoLa-static fixes one. Decode the frequency-trap prompt with dola_decode twice — once with the full candidate bucket (candidate_layers=None, dynamic) and once pinned to a single deep layer that already agrees with the answer (e.g. candidate_layers=[3]). Contrasting against a layer that has nothing to disagree about should barely change the output, while the dynamic choice targets the layer where the shallow guess is loudest. Confirm the dynamic decoder still recovers “moon”.
import importlib.util as _ilu
from pathlib import Path as _Path
import torch as _torch
_spec = _ilu.spec_from_file_location("dola", _Path("dola.py").resolve())
_D = _ilu.module_from_spec(_spec); _spec.loader.exec_module(_D)
_torch.manual_seed(0)
_Vt = _D.TRAP_VOCAB
_batch = _torch.tensor([_D.TRAP_COMMON] * 8 + [_D.TRAP_RARE])
_m = _D.GPTModel(vocab_size=len(_Vt), embed_dim=64, num_heads=4, num_layers=4, max_seq_len=16, dropout=0.0)
_opt = _torch.optim.Adam(_m.parameters(), lr=3e-3); _lf = _torch.nn.CrossEntropyLoss(); _m.train()
for _ in range(600):
_opt.zero_grad()
_lf(_m(_batch[:, :-1]).reshape(-1, len(_Vt)), _batch[:, 1:].reshape(-1)).backward(); _opt.step()
_m.eval()
_prompt = _torch.tensor([_D.TRAP_PROMPT])
_dyn = _D.dola_decode(_m, _prompt, max_new_tokens=1) # dynamic (all early layers)
_stat = _D.dola_decode(_m, _prompt, max_new_tokens=1, candidate_layers=[3]) # static, deep layer
print("dynamic DoLa →", _Vt[_dyn[0]])
print("static (L3) →", _Vt[_stat[0]])
# Your turn: pin candidate_layers to the shallow [1] and compare.dynamic DoLa → moon
static (L3) → moonIn-Context Learning as Gradient Descent
Back at the induction head, we watched in-context learning do one thing: copy. On the cat sat … the cat the model finished the repeat because a two-head circuit matched the current token to an earlier occurrence and echoed what came next. That is ICL as pattern-matching — real, mechanical, and only half the story.
Here is the other half. Give a model a handful of (x, y) pairs drawn from a function it has never seen — say a noisy line — then a fresh x, and it predicts the right y. Nothing was copied; there is no earlier occurrence of this x. The model fit a regressor from the prompt alone, with every weight frozen. That looks less like matching and more like learning — so which learning algorithm is running, and where?
The answer is one of the most beautiful results in interpretability (von Oswald et al., 2023; Akyürek et al., 2023): a single linear-attention layer performs one step of gradient descent on an in-context least-squares problem. A forward pass through a deep linear transformer is secretly running an optimizer — its layers are training iterations, its depth is the number of steps. We will build the exact construction from scratch, prove one layer equals one GD step to machine precision, stack it so depth becomes training, watch the fixed point land on ordinary least squares, and finally train a blank attention layer and catch it rediscovering gradient descent on its own.
Intuition: The Prompt Is a Training Set
Lay the prompt out as a dataset. Each in-context example becomes one token by stacking its input and its label:
e_j = \begin{pmatrix} x_j \\ y_j \end{pmatrix}, \qquad j = 1,\dots,N, \qquad e_{\text{query}} = \begin{pmatrix} x_\text{test} \\ 0 \end{pmatrix}.
Now recall what attention does: every token builds a weighted sum over all the others. The query token — the one holding x_\text{test} with an empty answer slot — gets to look at every labelled pair at once. What if that “looking” were arranged so the number it writes into its own empty slot is exactly the correction a gradient step would make to a prediction? Then one attention layer would be one training step, and the empty slot would fill in with the model’s fitted guess for y_\text{test} — no weights touched, all of it inside the forward pass.
That is precisely what the construction below does. The only ingredient we must give up is the softmax.
NoteKey Insight
Drop the softmax from attention and its scores become raw inner products x_i\cdot x_j — a linear map of the data. With the right (fixed) projection matrices, that linear map computes the gradient of an in-context least-squares loss and writes one descent step into the query’s answer slot. In-context learning, for linear regression, is implicit gradient descent — and the depth of the network is the number of steps.
The Math: One Linear-Attention Layer = One GD Step
Take the least-squares loss a linear model W pays on the in-context data (init W_0, and we take W_0=0 — “start from a blank model”):
L(W) = \frac{1}{2N}\sum_{i=1}^{N}\big\|W x_i - y_i\big\|^2, \qquad \Delta W = -\eta\,\nabla_W L(W_0) = -\frac{\eta}{N}\sum_{i=1}^N (W_0 x_i - y_i)\,x_i^\top.
Linear self-attention is ordinary attention with the softmax removed. With tokens as the columns of E and four weight matrices (W_Q, W_K, W_V, P),
\text{LSA}(E) = E + P\,(W_V E)(W_K E)^\top (W_Q E).
The middle factor (W_K E)^\top (W_Q E) is the score matrix — here the unnormalized inner products, no softmax, no 1/\sqrt{d_k}. Now choose the blocks so the layer reads x, forms residuals, and writes a scaled sum back into the y-slot (von Oswald et al., Prop. 1):
W_K = W_Q = \begin{pmatrix} I_x & 0 \\ 0 & 0 \end{pmatrix},\quad W_V = \begin{pmatrix} 0 & 0 \\ W_0 & -I_y \end{pmatrix},\quad P = \frac{\eta}{N}\,I.
Read it block by block: W_Q,W_K zero the y-rows, so the score of example i for token j is exactly x_i\cdot x_j. W_V maps e_i=(x_i,y_i) to (0,\;W_0 x_i - y_i) — the residual sits in the y-slot. P=\tfrac{\eta}{N}I supplies the learning rate and the 1/N average. Summing over i writes
\Delta y_j = \frac{\eta}{N}\sum_i (W_0 x_i - y_i)(x_i\cdot x_j) = -\,\Delta W\,x_j
into every token’s y-slot — the exact change one GD step makes to the prediction W x_j (the sign is the read-out convention; we read the query’s prediction as the negation of its slot). The query’s empty slot fills with W_1 x_\text{test}: a fitted prediction, computed by attention.
Code: The Construction from Scratch
The whole thing lives in icl_gradient_descent.py, and none of it needs a GPTModel — the point is a bare linear-attention layer. Its heart is the softmax-free update:
def linear_self_attention(E, W_Q, W_K, W_V, P):
Q, K, V = W_Q @ E, W_K @ E, W_V @ E
scores = K.transpose(-1, -2) @ Q # raw inner products — no softmax
return E + P @ (V @ scores)Load the module, build the Prop. 1 weights, run one layer, and read the query’s prediction. The anchor: it equals one explicit gradient step, to machine precision.
import importlib.util
import sys
from pathlib import Path
import torch
icl_spec = importlib.util.spec_from_file_location(
"icl_gradient_descent", Path("icl_gradient_descent.py").resolve()
)
ICL = importlib.util.module_from_spec(icl_spec)
sys.modules["icl_gradient_descent"] = ICL
icl_spec.loader.exec_module(ICL)
torch.manual_seed(0)
X = torch.randn(8, 3) # 8 in-context inputs (n_x = 3)
Y = torch.randn(8, 1) # their labels
x_test = torch.randn(3) # the query
eta = 0.2
# One linear-attention layer, using the hand-built construction.
W_Q, W_K, W_V, P = ICL.gd_construction_weights(X, Y, eta)
E = ICL.pair_tokens(X, Y, x_test)
E1 = ICL.linear_self_attention(E, W_Q, W_K, W_V, P)
lsa_pred = ICL.read_prediction(E1, n_x=3, n_y=1)
# One explicit GD step from a blank model, read out on x_test.
W1 = ICL.gd_step(torch.zeros(1, 3), X, Y, eta)
gd_pred = (W1 @ x_test).reshape(-1)
print("linear-attention prediction :", lsa_pred.tolist())
print("one-GD-step prediction :", gd_pred.tolist())
print("max |difference| :", float((lsa_pred - gd_pred).abs().max()))linear-attention prediction : [0.02213941141963005]
one-GD-step prediction : [0.02213941514492035]
max |difference| : 3.725290298461914e-09The difference is zero to floating-point precision (~10^{-9}): the attention layer did not approximate a gradient step, it is one. Every number the layer produced came from inner products of the data and a fixed set of weights — the softmax was the only thing standing between attention and an optimizer.
Depth Is Training Steps
One step is a taste; the payoff is stacking. If one layer is one GD step, then a model of L identical layers runs L steps of gradient descent — depth is the training loop. The stackable form keeps a running prediction in each token’s y-slot (initialized to -y_i so the slot holds the negative residual) and lets the query attend only to the context tokens, so their growing residuals never leak into each other. Applying lsa_gd_layer L times reproduces L explicit GD steps exactly:
steps = 15
stacked = ICL.lsa_gd_predict(X, Y, x_test, eta, steps)["predictions"]
gd = ICL.explicit_gd(X, Y, eta, steps)["weights"]
gap = max(float((s - (W @ x_test).reshape(-1)).abs().max())
for s, W in zip(stacked, gd))
print(f"max |stacked-LSA − explicit-GD| over all {steps} depths : {gap:.2e}")max |stacked-LSA − explicit-GD| over all 15 depths : 2.98e-08The two agree to floating-point noise at every depth. So watching the fitted line settle onto the data, layer by layer, is quite literally watching gradient descent run inside a forward pass. Drive the depth below: at depth 0 the model predicts a flat zero line; each layer is one GD step, rotating and shifting the line onto the context points while the in-context loss (right) falls toward the least-squares floor (dashed).
TipTry This
- Step the depth from 0 to the end. The flat zero line pivots onto the points and the loss marker walks down to the least-squares floor — each click is one gradient step, executed by stacked attention.
- Watch the last few steps barely move. Gradient descent’s diminishing returns are visible: the big rotations happen early, then the line only settles. That is why a handful of layers already fits well.
- Notice the floor. The loss flattens at the dashed line, not at zero — with noisy labels the best any linear fit can do is the least-squares residual, the subject of the next part.
The Fixed Point Is Least Squares
Run the descent long enough and it stops moving — it has reached the minimizer of L(W), which for a linear model is the ordinary-least-squares solution in closed form, W^\star = (X^\top X)^{-1}X^\top Y. So a deep-enough linear transformer doesn’t just take a step toward a good regressor; in the limit it computes the exact least-squares fit the data admits.
gd_far = ICL.explicit_gd(X, Y, eta, 5000)["weights"][-1] # many GD steps
W_ols = ICL.ols_solution(X, Y) # closed form
print("max |GD(5000) − OLS| :", float((gd_far - W_ols).abs().max()))max |GD(5000) − OLS| : 2.384185791015625e-07They coincide. This is the bridge to Akyürek et al. (2023), who came at the same question from the other side: they proved a transformer can implement not just one GD step but the full closed-form ridge/OLS solver, and found that trained in-context learners behave like few-step GD early and the exact minimizer as depth grows — exactly the trajectory the widget above traces.
NoteKey Insight
Two papers, two directions, one picture. von Oswald et al. pin the tight identity — one linear-attention layer = one GD step, with an explicit weight construction. Akyürek et al. show the broader menu — a transformer can also realize closed-form ridge regression (OLS as the no-penalty limit), and depth is what carries a trained model from a few GD steps to the exact minimizer.
Trained Transformers Discover GD
We built the construction. The stunning empirical result is that you don’t have to: train a blank linear-attention layer to solve in-context regression, tell it nothing about gradients, and it rediscovers the construction on its own. We can watch that happen from scratch. train_lsa_layer fits a single layer’s weights (Adam, a few hundred steps) across many random regression tasks; learned_vs_gd then compares its predictions to a single GD step on held-out tasks it never trained on.
trained = ICL.train_lsa_layer(seed=0) # ~a few seconds, deterministic
cmp = ICL.learned_vs_gd(trained["model"], eta=0.5)
print(f"training loss : {trained['loss_curve'][0]:.2f} → {trained['loss_curve'][-1]:.2f}")
print(f"cosine(learned predictions, one-GD-step predictions) : {cmp['cosine']:.3f}")training loss : 6.12 → 1.32
cosine(learned predictions, one-GD-step predictions) : 0.984The learned layer’s held-out predictions line up with gradient descent at cosine \approx 0.98 — it converged to the same computation we constructed by hand. The scatter below plots each held-out task’s GD prediction (x-axis) against the trained layer’s prediction (y-axis); the points hug the diagonal. Toggle to the training curve to see the loss it descended to get there.
WarningThe weights never changed
It is worth saying plainly what did and did not happen. In the trained demo we ran a real optimizer to shape one attention layer — that is ordinary training. But when that finished layer (or the hand-built one) then solves a fresh regression task in a single forward pass, no weight moves at all. The “learning” is entirely in the activations: the prompt is the training set, the layers are the steps, and the fitted model exists only as numbers flowing through the residual stream. That is what makes in-context learning feel like magic — and why “it’s just gradient descent in the forward pass” is such a satisfying answer.
In-Context Learning as Preconditioned Descent: GD++
The construction above is exact, but it leaves a loose end. When you actually train a transformer on these in-context regression tasks, it does not settle for plain gradient descent — at the same depth it fits better than K GD steps ever could. von Oswald et al. found what the trained layer is really doing and named it GD++: before each gradient step the layer transforms its own inputs, and that transform is a curvature correction — a preconditioner discovered in the forward pass. Same descent, better-shaped landscape, far faster convergence.
Intuition: Straightening the Ravine
Least-squares gradient descent is slow for one reason: the loss surface is a ravine. When the input features have very different scales, the bowl L(W)=\tfrac{1}{2N}\sum\lVert Wx_i-y_i\rVert^2 is stretched — steep across the ravine, almost flat along it. Gradient descent can only step as far as the steepest direction tolerates, so it zig-zags across the ravine while barely crawling toward the minimum along the floor. The stretch is measured by the condition number \kappa of the input Gram X^\top X (largest eigenvalue over smallest); the bigger \kappa, the slower GD.
GD++ tilts the floor. Its input transform x\leftarrow(I-\gamma X^\top X)\,x pulls in the loud, steep directions harder than the quiet, flat ones — rounding out the ravine so a single step makes progress everywhere at once. Crucially it is an invertible change of coordinates: it moves the path descent takes, not the destination it reaches. The fit at the bottom is the same least-squares fit; you just get there in a fraction of the steps.
NoteKey Insight
A trained transformer is not limited to running an optimizer — it can improve one. GD++ is a learned preconditioner: the attention layer reshapes the problem’s curvature in the forward pass, so its later “GD-step” layers count for far more. Depth is still training steps; GD++ just makes each step worth more.
The Math: Reshaping the Hessian’s Spectrum
The full GD++ layer updates every token e_j=(x_j;y_j) — context and query — at once (von Oswald et al., Eq. 93):
x_j \;\leftarrow\; x_j - \gamma\,X^\top X\,x_j \;=\; (I-\gamma X^\top X)\,x_j, \qquad y_j \;\leftarrow\; y_j - \Delta W\,x_j ,
where the y-slot update \Delta W=-\tfrac{\eta}{N}\sum_i (Wx_i-y_i)x_i^\top is the same gradient step as before. The only new piece is the x-slot map. Setting \gamma=0 leaves the inputs untouched and recovers plain GD exactly — GD++ contains the previous section.
Why does the map accelerate descent? Because it acts directly on the loss Hessian. The curvature of L is set by the spectrum of X^\top X; after the transform, each eigenvalue \lambda becomes, exactly,
\lambda \;\longmapsto\; \lambda\,(1-\gamma\lambda)^2 .
This is a polynomial reshaping: for a well-chosen \gamma it squeezes the large eigenvalues toward the small ones, cutting the condition number. The paper notes the transform “closely resembles a heavily truncated Neumann series approximation of the inverse X^\top X” — i.e. an approximate preconditioner (X^\top X)^{-1}, the matrix that would make the bowl perfectly round in one shot. And since GD’s per-step contraction is governed by \tfrac{\kappa-1}{\kappa+1}, halving \kappa pushes that factor away from 1 and multiplies the progress per layer.
Code: GD++ from Scratch
The new code lives in gd_plus_plus.py and reuses the previous section’s layer by reference — GD++ is a superset, not a rewrite. Load it and check the three facts the math promised, each to machine precision.
gdpp_spec = importlib.util.spec_from_file_location(
"gd_plus_plus", Path("gd_plus_plus.py").resolve()
)
GDPP = importlib.util.module_from_spec(gdpp_spec)
sys.modules["gd_plus_plus"] = GDPP
gdpp_spec.loader.exec_module(GDPP)
task = GDPP.ill_conditioned_task(seed=0) # deliberately stretched features
Xi, Yi = task["X"], task["Y"]
n_x, N = Xi.shape[1], Xi.shape[0]
# Anchor 1 — γ = 0 makes the GD++ layer the plain-GD layer, bit-for-bit.
E = ICL.build_tokens(Xi, Yi, torch.ones(n_x))
same = torch.equal(GDPP.gd_plus_plus_layer(E, n_x, N, 0.1, 0.0),
ICL.lsa_gd_layer(E, n_x, N, 0.1))
print("γ=0 GD++ layer == plain-GD layer :", same)
# Anchor 2 — the transform maps the Hessian spectrum by exactly λ(1−γλ)².
eig = torch.linalg.eigvalsh(GDPP.feature_gram(Xi))
gamma = 0.5 / float(eig[-1])
Xt, _ = GDPP.curvature_transform(Xi, gamma)
direct = torch.sort(torch.linalg.eigvalsh(GDPP.feature_gram(Xt))).values
formula = torch.sort(GDPP.spectrum_map(eig, gamma)).values
print("spectrum map matches eigendecomposition, max err :",
float((direct - formula).abs().max()))γ=0 GD++ layer == plain-GD layer : True
spectrum map matches eigendecomposition, max err : 6.103515625e-05The \gamma=0 reduction is exact (True), and the eigenvalue map agrees with a direct eigendecomposition of the transformed Gram to float precision. Now find the \gamma that best rounds the ravine, and confirm the block-matrix attention weights (von Oswald et al., Eq. 78) really do write (I-\gamma X^\top X)x into the x-slot:
opt = GDPP.optimal_gamma(Xi)
print(f"condition number κ : {opt['kappa_before']:.1f} → {opt['kappa_after']:.1f}"
f" ({opt['reduction']:.1f}× tighter at γ* = {opt['gamma']:.5f})")
# Anchor 3 — the Eq. 78 block matrices realize the curvature transform.
W_Q, W_K, W_V, P = GDPP.gd_plus_plus_weights(n_x, 1, 0.1, gamma, N)
out = ICL.linear_self_attention(E, W_Q, W_K, W_V, P)
x_all = E[:n_x, :]
expected = x_all - gamma * ((x_all @ x_all.T) @ x_all)
print("block-matrix x-slot == (I − γ XXᵀ) x, max err :",
float((out[:n_x, :] - expected).abs().max()))condition number κ : 51.8 → 8.6 (6.0× tighter at γ* = 0.00090)
block-matrix x-slot == (I − γ XXᵀ) x, max err : 7.152557373046875e-07On this stretched task the single best curvature step cuts \kappa from ~52 to ~9 — about a 6× rounder bowl — and the attention weights produce the transform exactly. Drive \gamma yourself below and watch the spectrum compress.
TipTry This
- Slide γ up from 0. Every bar starts at its ghost height (plain GD) and the tall ones shrink fastest — the ravine rounding out. The \kappa(\gamma) curve on the right dives to a minimum at \gamma^\*.
- Push past γ*. Keep going and \kappa climbs again: overshoot and the map \lambda(1-\gamma\lambda)^2 starts sending the biggest eigenvalue below the smallest, re-stretching the bowl the other way. There is a sweet spot, not a “more is better.”
- Find where a bar hits zero. At \gamma=1/\lambda a direction is annihilated (1-\gamma\lambda=0) — the transform is no longer invertible there, which is why the search stays below 1/\lambda_{\max}.
The Payoff: Same Fit, Far Fewer Steps
Rounder curvature is only worth something if it actually speeds descent up without changing the answer. It does both. Preconditioning the data once (transform, then run ordinary GD) reaches the same least-squares floor as plain GD — but because the well-conditioned problem admits a larger stable learning rate, it gets there in a small fraction of the steps.
race = GDPP.descent_race(Xi, Yi, opt["gamma"], steps=60)
print(f"stable learning rate η : plain {race['plain']['eta']:.4f}"
f" → preconditioned {race['precond']['eta']:.4f}")
print(f"same least-squares floor : plain {race['plain']['floor']:.5f}"
f" preconditioned {race['precond']['floor']:.5f}")
print(f"loss after 30 steps : plain {race['plain']['losses'][30]:.4f}"
f" preconditioned {race['precond']['losses'][30]:.5f}")
print(f"steps to reach the floor : plain {race['plain_steps']}"
f" preconditioned {race['precond_steps']}")stable learning rate η : plain 0.0471 → preconditioned 0.2927
same least-squares floor : plain 0.00445 preconditioned 0.00445
loss after 30 steps : plain 0.1384 preconditioned 0.00450
steps to reach the floor : plain None preconditioned 20The two runs share an identical floor to five decimals — the transform is a change of coordinates, not a change of problem — yet preconditioned GD converges in about 20 steps while plain GD is still ~30× higher at step 30 and never reaches the floor inside the budget. That gap is the trained transformer’s edge over a fixed-step optimizer, and the widget below shows both descents at once.
NoteKey Insight
The transform changes the path, not the destination. Both descents minimize the same loss and land on the same least-squares fit; GD++ simply travels a rounder valley, so a bigger step is safe and the tail no longer crawls. That is the whole trick behind “trained transformers beat gradient descent” — they are not doing something other than GD, they are doing GD on a better-conditioned problem they conditioned themselves.
Common Pitfalls
WarningThe softmax is not optional in real models
The clean “one layer = one GD step” identity needs linear attention — the softmax dropped. Real transformers keep the softmax, so this is a mechanistic model of ICL, not a claim that GPT literally runs this exact update. Its value is the existence proof: attention can implement gradient-based learning, and trained linear versions provably do. Softmax attention is understood to run a related, kernelized descent (von Oswald et al. and follow-ups), but the bare linear case is the one you can build and verify to machine precision.
WarningMind where the 1/N and the sign live
Two bookkeeping choices trip up a from-scratch build. The 1/N averaging in the loss must live in exactly one place — here, inside P=\tfrac{\eta}{N}I — or your “learning rate” is off by a factor of N. And the construction writes -\Delta W\,x into the answer slot, so the prediction is read as its negation. Put the 1/N in twice, or forget the sign, and the layer will look broken while the algebra is fine. Our tests pin both by comparing to an explicit gd_step.
WarningGD++ needs the sweet spot, and its transform is a reparameterization
Two ways to misread GD++. First, \gamma is not “bigger is better”: past \gamma^\* the map \lambda(1-\gamma\lambda)^2 pushes the largest eigenvalue below the smallest and re-stretches the ravine, and at \gamma=1/\lambda it annihilates a direction (the transform stops being invertible). Second, the curvature transform changes the coordinates, not the fit — it converges to the same least-squares solution, so don’t read the accelerated loss as “GD++ finds a better model.” It finds the same model faster. The tests pin the invertibility (det(H) ≠ 0) and the shared OLS floor.
Exercises
- Change the learning rate. In the descent widget’s task, re-run
explicit_gdwith a largereta. Find the value where the loss stops decreasing monotonically — the same divergence GD shows on any quadratic when the step exceeds 2/\lambda_{\max}. - Ridge from the prompt. Compare
ols_solution(X, Y, ridge=0.0)withridge=5.0and re-fit. Which generalizes better when you add labelnoiseto the task? This is the Akyürek “in-context ridge” knob. - Break the sign. Flip the sign in
read_predictionand watch the anchor test fail — then explain, from the block matrices, exactly why the slot holds -\Delta W\,x. - Two heads, two learning rates. The construction uses one head. Sketch how two linear-attention heads with different \eta could implement a momentum step, and check your idea against
explicit_gd. - Sweep γ past the sweet spot. Call
GDPP.spectrum_map(eig, γ)over a grid and plot the condition number. Confirm it bottoms out atGDPP.optimal_gamma(X)["gamma"]and climbs again — then find the \gamma where the smallest mapped eigenvalue first hits zero, and relate it to 1/\lambda_{\max}. - Precondition a well-conditioned task. Run
descent_raceon the well-conditionedICL.regression_taskinstead ofill_conditioned_task. With \kappa already near 1, how much does GD++ buy? Explain why curvature correction helps most exactly when plain GD hurts most.
Summary
Key takeaways:
- 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.
- The logit lens reads it out early. Apply the model’s own
ln_final+lm_headto any layer’s hidden state to get that layer’s guess — no new training, no model surgery. - The last layer’s lens is the model’s output, exactly. That identity (
layer_logits(model, x)[-1] == model(x)) is the correctness anchor. - 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?
- 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.
- 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).
- The lens is biased — and the tuned lens fixes it, provably. The logit lens forces the final head onto states it never saw. The tuned lens inserts a learned per-layer affine translator
A_ℓ h_ℓ + b_ℓbefore that frozen head, fit (model frozen) to distill the final distribution. It is identity-initialized, so it starts as the logit lens, and its per-layer objective is convex, so it provably cannot do worse — on our demo it recovers “park” at the embeddings, where the logit lens is still blind. - 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. - 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”.
- Use the logit difference to isolate a decision. Attributing
logit[target] − logit[distractor]projects ontoW_U[target] − W_U[distractor], removing the shared “raise everything” push and keeping only what separated the candidates. - The induction head is your first named circuit. Two heads compose across two layers — a previous-token head labels each position with its predecessor, then an induction head matches the current token to those labels and copies what followed:
[A][B] … [A] → [B]. Built by hand, it predicts the repeat exactly. - In-context learning is the induction bump. On a repeated sequence, prediction accuracy jumps from chance to exact the moment the second copy starts — learning within one prompt, no weights changed. It is what training grows a head to do, though a fixed-period stripe alone can’t separate a content-based induction head from a positional shortcut.
- Activation patching makes it causal. Cache a clean run, break the prompt, then splice one clean activation back into the corrupt run and measure the recovery of the logit difference (0 = still broken, 1 = fully restored). Where the lens, DLA, and the stripe read activations, patching intervenes — the difference between correlation and cause. Built from scratch on the m06 model with exact anchors: the cache reproduces the logits, and patching the whole clean final residual reproduces the clean output.
- The intervention outranks the pattern. Head patching localizes the copy to a single head that recovers most of the answer alone — and it need not be the head with the top induction-stripe score. When a correlational signal and a causal one disagree, trust the patch; on the fixed-period probe it lands on the layer-0 positional shortcut, confirming the pitfall by intervention.
- Path patching isolates one wire. The residual stream is a literal sum, so a component’s direct effect on the logits is its own contribution term. Restore one wire at a time (∅ ⇒ corrupt, all ⇒ clean — path patching all paths is node patching) and you split a head’s total effect from its direct one. The finding: the layer-0 previous-token head is the most important head by node patching, yet writes almost nothing to the logits directly — it acts entirely upstream, a split only path patching can see.
- Features live in superposition, below the neuron. A width-d residual stream stores far more than d features by packing each along its own nearly-orthogonal direction and relying on sparsity — so every neuron is polysemantic, and the honest unit of meaning is a direction, not an axis.
- A sparse autoencoder recovers the dictionary. An overcomplete encoder/decoder with a ReLU code, an L1 penalty, and unit-norm atoms decomposes each activation into a short sum of feature directions. We proved it works where we could check: plant a dictionary, superpose data from it, and the SAE’s MMCS recovery climbs from chance (~0.4) toward 1.0 — the atoms come back up to permutation — before we read a sparse feature code off the book’s own
GPTModel. - In-context learning is implicit gradient descent. ICL has a second face beyond the induction head’s copying: shown fresh
(x, y)pairs, a model fits a regressor from the prompt alone. Drop the softmax and one linear-attention layer, with a fixed weight construction, computes the least-squares gradient and writes one GD step into the query’s answer slot — proven equal to an explicitgd_stepto machine precision. - Depth is the training loop. Stacking that layer L times runs L GD steps (bit-for-bit versus
explicit_gdat every depth), and in the limit the fitted line lands on the exact ordinary-least-squares solution — the Akyürek bridge from few-step GD to the closed-form minimizer. - Trained transformers rediscover GD. Fit a blank linear-attention layer on random regression tasks, tell it nothing about gradients, and its held-out predictions align with a single GD step at cosine ≈ 0.98. Crucially, when it solves a new task the weights never move — the learning lives entirely in the activations of one forward pass.
- GD++ improves on GD by correcting curvature. A trained transformer beats fixed-step GD because each layer also transforms its inputs by (I-\gamma X^\top X) — a preconditioner that remaps the Hessian spectrum by exactly \lambda\mapsto\lambda(1-\gamma\lambda)^2, cutting the condition number (~52 → ~9 on our stress task). \gamma=0 recovers plain GD exactly; the block matrices realize the transform to machine precision.
- The correction is a reparameterization, not a better fit. Preconditioning changes the path, not the destination: it reaches the same least-squares floor, but the rounder valley admits a larger stable step, so it converges in ~20 steps where plain GD is still ~30× higher and never lands inside the same budget.
- The logit lens can decode, not just diagnose — that is DoLa. Contrast the mature layer against the most-divergent (max-JSD) premature layer, gate to the plausible tokens (q_N(x)\ge\alpha\max), and score \log q_N(x)-\log q_M(x). It keeps what depth earned and demotes what a shallow layer already believed — on our frequency trap the shallow layer hallucinates the frequent “sun” while the contrast recovers the rare-but-correct “moon”, by exactly the premature layer’s own log-odds. Because factual recall lives in the upper layers, this lifts real LLaMA models 12–17 points on TruthfulQA with no retrieval and no fine-tuning.
What’s Next
You have gone from what the residual stream predicts (the logit lens) to which component wrote it (direct logit attribution) to naming your first circuit — the induction head — to causally localizing that circuit with activation patching, and finally below the component to the features in superposition a sparse autoencoder recovers. Every step decoded the same running vector; the frontier from here sharpens the feature dictionary and reads structure from the weights:
- Better SAEs — built above. The Modern SAEs section constructs TopK (Gao et al., 2024), which sets the active count directly and sidesteps the L1-shrinkage bias, and JumpReLU / Gated (Gemma Scope, 2024), which learns a per-atom threshold. From here, the frontier scales them to production dictionaries (millions of atoms) and studies the features they surface.
- Feature steering & circuits-with-features — once activations are features, you can set one and watch generation change (Scaling Monosemanticity, 2024), and wire DLA and patching through the features to draw circuits in the model’s own vocabulary rather than in neurons.
- The tuned lens (Belrose et al., 2023) — a learned per-layer probe that fixes the logit lens’s bias.
- Full circuit recovery (Wang et al., 2022) — this section built path patching to the logits and through a layer; the frontier aims each edge at a specific head’s query, key, or value input, then chains the edges into a complete circuit graph (the IOI circuit’s name-mover, S-inhibition, and induction heads), validated end-to-end with causal scrubbing (Chan et al., 2022).
- Noising, and separating content from position — patch with a corruption that breaks the fixed period (a variable-offset repeat) to force the causal head to be content-based, closing the gap our fixed-period probe left open.
- QK/OV circuits (Elhage et al., 2021) — analyze a head’s behavior from its weights (
W_Q Wᵀ_K,W_O W_V) rather than one run’s activations, the next rung above the induction stripe. - Preconditioned in-context GD (GD++) — von Oswald et al. show trained multi-layer transformers beat plain GD by learning a curvature correction H=(I-\gamma XX^\top) and descending on the whitened data; the natural sequel to the vanilla step this section builds, and the bridge to how softmax attention runs a kernelized descent.
Each begins exactly where this module ends — reading structure out of the residual stream, or — as in the last section — recognizing that the forward pass was an optimizer all along.
Going Deeper
Core Ideas:
- interpreting GPT: the logit lens — nostalgebraist (2020), the original logit lens.
- A Mathematical Framework for Transformer Circuits — Elhage et al. (2021), the residual-stream view, the unembedding as a read-out, and direct logit attribution.
- Eliciting Latent Predictions from Transformers with the Tuned Lens — Belrose et al. (2023), the learned per-layer probe that fixes the logit lens’s bias.
- In-context Learning and Induction Heads — Olsson et al. (2022), the mechanism behind in-context learning.
- Transformers Learn In-Context by Gradient Descent — von Oswald et al. (ICML 2023), the exact construction: one linear-attention layer = one GD step, and depth = training steps. Appendix A.10 defines GD++ — the input transform (I-\gamma X^\top X), the eigenvalue map \lambda(1-\gamma\lambda)^2, and why the trained model beats plain GD.
- What Learning Algorithm Is In-Context Learning? Investigations with Linear Models — Akyürek et al. (ICLR 2023), transformers implement GD and closed-form ridge/OLS in-context; depth selects the algorithm.
- DoLa: Decoding by Contrasting Layers Improves Factuality in Large Language Models — Chuang et al. (ICLR 2024), the logit lens as a decoder: contrast the mature layer against the max-JSD premature layer, gated by the adaptive plausibility constraint, for +12–17 points on TruthfulQA with no retrieval.
- Contrastive Decoding: Open-ended Text Generation as Optimization — Li et al. (ACL 2023), the expert-minus-amateur log-ratio and the adaptive plausibility constraint DoLa borrows to gate the contrast.
- What Can Transformers Learn In-Context? A Case Study of Simple Function Classes — Garg et al. (NeurIPS 2022), the empirical setup: transformers learn linear (and richer) function classes purely in-context.
- Why Can GPT Learn In-Context? — Dai et al. (ACL 2023 Findings), the “attention as meta-optimizer” dual view.
- Locating and Editing Factual Associations in GPT — Meng et al. (2022), causal tracing: corrupt the input, restore hidden states, and localize where a fact is stored.
- Interpretability in the Wild: the IOI circuit — Wang et al. (2022), where path patching was introduced: freeze every component to its clean value except one sender, let the counterfactual reach a chosen receiver, and recover a full circuit in GPT-2.
- Localizing Model Behavior with Path Patching — Goldowsky-Dill, MacLeod, Sato & Arora (2023), the formalization this section builds: patch along a chosen set of paths, with the identities that patching no paths is the clean run and patching all of a node’s paths is ordinary node patching.
- Causal Scrubbing — Chan et al. (2022), a systematic way to test a whole circuit hypothesis by resampling everything not on the claimed paths.
- Toy Models of Superposition — Elhage et al. (2022), why a network stores more features than dimensions, and the sparsity that makes it work.
- Sparse Autoencoders Find Highly Interpretable Features in Language Models — Cunningham et al. (2023), the SAE recipe and the MMCS recovery evaluation.
- Towards Monosemanticity: Decomposing Language Models With Dictionary Learning — Bricken et al. (2023), the MSE + L1, unit-norm-decoder SAE and dead-feature resampling this section builds.
- Scaling Monosemanticity — Templeton et al. (2024), SAE features on a production model, and the feature clamping (“Golden Gate Claude”) this section’s steering builds toward.
- Activation Addition: Steering Language Models Without Optimization — Turner et al. (2023), ActAdd: steer with a single contrast vector
h ← h + α·vat inference, no fine-tuning — the mechanism this section builds. - Steering Llama 2 via Contrastive Activation Addition — Rimsky et al. (2023), CAA: build the steering vector from the mean difference over many contrastive pairs, the robust version of the two-context contrast used here.
- Scaling and evaluating sparse autoencoders — Gao et al. (2024), the TopK SAE built in this section — set the active count directly, no L1.
- Jumping Ahead: Improving Reconstruction Fidelity with JumpReLU Sparse Autoencoders — Rajamanoharan et al. (2024), the JumpReLU activation
z·H(z−θ)and the straight-through training that penalizes L0 directly. - Improving Dictionary Learning with Gated Sparse Autoencoders — Rajamanoharan et al. (2024), the Gated SAE that first split gate from magnitude to solve shrinkage.
- Gemma Scope: Open Sparse Autoencoders Everywhere All At Once — Lieberum et al. (2024), JumpReLU SAEs trained across every layer of Gemma 2.
- k-Sparse Autoencoders — Makhzani & Frey (2013), the original keep-the-top-k idea TopK revives.
Practical Resources:
- A Comprehensive Mechanistic Interpretability Explainer & Glossary — Neel Nanda, on the logit-difference view and direct logit attribution in practice.
- Towards Best Practices of Activation Patching — Zhang & Nanda (2023), how corruption and metric choices change the answer (and why the logit difference is the safe default).