Mechanistic interpretability

Open the model up, read what it's thinking layer by layer, and run controlled experiments to find which parts actually cause an answer.

Source: Suraj Sharma, X post, stage 9

Going deeper: everything below is the full version, with some maths. On a first read, skim the headings and the "Why it matters for robotics" box, and come back for the rest later.

1. What this stage covers

• Learn: activation patching, the logit lens, circuit discovery, and the tooling (TransformerLens, Baukit) that makes hooking into a model easy.
• Practice: localise a factual recall such as the Eiffel Tower → Paris, and patch activations to change the output. On this page you do it in GPT-2 small on a laptop, with a patching heatmap over layers and token positions.
• Why (the post's claim): you can't fix hallucinations or bias without knowing which circuit produces them. That's partly true and partly overstated. See the reality check.

2. Concepts, explained

The residual stream and the logit lens

In a decoder transformer, each layer reads the residual stream and adds its output back to it. The final layer norm and unembedding turn the last residual vector into logits. The logit lens (nostalgebraist, 2020, LessWrong post) applies that same final step to the residual stream after every layer, to see what the model "would say" if it stopped there:

h_0 = embed(x) + pos
h_l = h_{l-1} + attn_l(h_{l-1}) + mlp_l(…)      # each block adds to the stream
lens_l = W_U · LN_final(h_l)                     # logits "as of" layer l

It's cheap and often revealing in GPT-2: the right answer climbs the ranking over the later layers. It's also biased, because middle layers weren't trained to be read by the final unembedding. The tuned lens (Belrose et al., 2023) learns a small affine map per layer to fix this. Treat the logit lens as a quick look, not a proof. It's correlational: it shows what's present, not what's used.

Activation patching (causal tracing)

Patching is an intervention. You need two inputs that differ in one controlled way, a metric, and a list of components to test:

clean:     "The Eiffel Tower is located in the city of"  → Paris
corrupted: same template, different subject              → Rome (say)
metric:    LD = logit(Paris) − logit(Rome) at the last position

for each component c (layer × position, or head):
    run corrupted, but replace c's activation with its clean value
    recovered(c) = (LD_patched − LD_corr) / (LD_clean − LD_corr)

This direction (clean into corrupted) asks "is this component sufficient to restore the behaviour?". The reverse (corrupted into clean, sometimes called noising) asks "is it necessary?". The two can disagree when information is stored redundantly. ROME's causal tracing (Meng, Bau, Andonian & Belinkov, 2022) used a variant: it corrupts the subject's token embeddings with Gaussian noise rather than swapping subjects, then restores one hidden state at a time. Practical details that change results: the choice of metric (logit difference is more stable than probability), whether the two prompts tokenise to the same length, and whether you patch a single position or all of them.

What ROME found, and the correction to the post

The post frames the project as finding "the exact attention head" for "The Eiffel Tower is in…". The paper that studied exactly this kind of fact found something different. Meng et al. traced factual predictions in GPT-2 XL and found a distinct set of steps in middle-layer feed-forward (MLP) modules while the model processes the last token of the subject ("Tower"). They then edited one MLP weight matrix with a rank-one update (ROME) to rewrite a fact, and showed the edit generalised to paraphrases without spilling onto unrelated subjects.

Follow-up work filled in the picture. Geva, Bastings, Filippova & Globerson (2023) found a three-step flow: early MLPs enrich the last-subject-token representation with many attributes, the relation ("is located in") is moved to the final position, and then attention heads at the final position extract the attribute from the enriched subject. So heads do matter, but as readers and movers spread over several layers, not as a single head that "contains" the fact. A good project result is a map with an MLP hot spot at the subject and an attention hot spot at the end, not one glowing head.

One more caution. Hase, Bansal, Kim & Ghandeharioun (2023) found that where causal tracing localises a fact doesn't predict which layer is best to edit. Localisation and control are separate questions.

Circuits and circuit discovery

A circuit is a subgraph of the model (heads, MLPs and the edges between them) that implements one behaviour. The framing comes from Anthropic's "A Mathematical Framework for Transformer Circuits" (Elhage et al., 2021), which treats each attention head as reading from and writing to the residual stream independently.

The landmark manual study is the IOI circuit (Wang, Variengien, Conmy, Shlegeris & Steinhardt, 2022). In "When Mary and John went to the store, John gave a drink to", GPT-2 small predicts "Mary". They explained it with 26 attention heads in 7 classes, found with causal interventions, and checked the explanation for faithfulness, completeness and minimality. Even then they reported gaps.

ACDC (Conmy, Mavor-Parker, Lynch, Heimersheim & Garriga-Alonso, 2023) automates the search: it prunes edges of the computational graph whose removal barely changes the metric. On GPT-2 small's greater-than task it kept 68 of 32,000 edges and recovered all 5 component types that manual work had found. Path patching and attribution patching (a gradient-based linear approximation of patching) are cheaper ways to rank components before confirming with real patches.

circuit ⊂ graph(model)
faithful: running only the circuit (rest ablated) ≈ full model on the task
minimal:  removing any node hurts the metric

Tools: TransformerLens and Baukit

TransformerLens (Neel Nanda and contributors) reimplements GPT-style models with a named hook at every activation: blocks.5.hook_mlp_out, blocks.9.attn.hook_z and so on. run_with_cache stores them all, and run_with_hooks lets you edit any of them during a forward pass. That makes patching a few lines. Baukit (David Bau's lab, the ROME authors) wraps an ordinary PyTorch model and lets you trace or edit any named module's output. It's lighter and works on models TransformerLens doesn't support. Under the hood both use the forward hooks you met in stage 01. For larger models, nnsight follows the same idea. The newer direction of the field, sparse autoencoders that split activations into interpretable features, is covered in the Transformer Circuits Thread linked below.

3. The practice project

(a) Laptop version: patch a fact in GPT-2 small

1. pip install transformer_lens. Load gpt2 (124M parameters, runs on CPU).
2. Check the model knows the fact first. Print the top 5 next tokens for your prompt. If " Paris" isn't top-1, try a rephrase, a different fact, or gpt2-medium/gpt2-xl on Colab. Patching a fact the model doesn't know measures noise.
3. Pick a corrupted subject with the same number of tokens (check with m.to_str_tokens) and a single-token answer.
4. Logit lens: for each layer, project the last-position residual through ln_final and the unembedding. Record the rank of " Paris".
5. Patching heatmaps: for resid_pre, mlp_out and attn_out, patch each (layer, position) from clean into corrupted and record the recovered fraction. Then patch individual heads (hook_z) at the final position.
6. Repeat on 20–50 facts (city, language, company), averaged. One prompt is an anecdote.

from transformer_lens import HookedTransformer, utils
import torch

m = HookedTransformer.from_pretrained("gpt2")
clean = m.to_tokens("The Eiffel Tower is located in the city of")
corr = m.to_tokens("The Colosseum is located in the city of")  # check lengths match
assert clean.shape == corr.shape, m.to_str_tokens(corr)
P, R = m.to_single_token(" Paris"), m.to_single_token(" Rome")
ld = lambda lg: (lg[0, -1, P] - lg[0, -1, R]).item()

lg_c, cache = m.run_with_cache(clean)
base_c, base_x = ld(lg_c), ld(m(corr))

def patch(act, hook, pos):
    act[:, pos] = cache[hook.name][:, pos]
    return act

L, T = m.cfg.n_layers, clean.shape[1]
heat = torch.zeros(L, T)
for l in range(L):
    for p in range(T):
        h = (utils.get_act_name("mlp_out", l), lambda a, hook, p=p: patch(a, hook, p))
        heat[l, p] = (ld(m.run_with_hooks(corr, fwd_hooks=[h])) - base_x) / (base_c - base_x)

# logit lens at the last position
for l in range(L):
    r = cache["resid_post", l][:, -1:]
    print(l, m.to_string(m.unembed(m.ln_final(r))[0, -1].argmax()))

What to measure: the recovered-fraction heatmap for MLP, attention and residual; the layer where " Paris" first reaches top-1 in the lens; the top 5 heads at the final position. What to look for, based on ROME and Geva et al. rather than a run I can vouch for in GPT-2 small: a band of MLP effect around the subject's last token in early-to-middle layers, and attention effect at the final position in later layers. If your map differs, report it. That's a finding.

"Change the behaviour": patching from a corrupted run already flips the answer in that run. To make a persistent change, add a steering vector (for example, mean activation difference between Paris and Rome prompts at the hot layer) or try a ROME edit with the authors' code. Then test on paraphrases and unrelated facts so you can see side effects.

(b) Full version from the post

The same experiment on a 1–8B model (Pythia, Gemma 2 2B or Llama-class), over a few hundred facts, with head-level patching, path patching between the key heads, and an edit or steering intervention evaluated for generalisation and specificity. A 2B model fits a free Colab T4 in half precision. An 8B model with caching of all activations needs a 24–48 GB GPU, about $0.50–2 per hour rented; head-level sweeps over hundreds of prompts take a few GPU-hours. Use attribution patching to shortlist components so you don't run every patch.

(c) Pitfalls

• Token misalignment. Different-length subjects shift every position. Match lengths or use ROME-style noise on the subject embeddings.
• Bad metric. Probabilities saturate. Use logit difference and normalise to the clean/corrupted range.
• One prompt. Average over many, and look at the variance.
• Sufficiency isn't necessity. A head that restores the answer when patched in may not be needed when knocked out, because of backup heads (the IOI paper found such backup name movers).
• Logit lens over-reading. Early-layer lens output is unreliable. Use the tuned lens if you need to compare layers.
• Calling it "the" circuit. Report what fraction of the effect your components explain.

4. Reality check

5. In this course

• 30u30 Day 13 · The Transformer and Day 14 · Annotated Transformer: the attention heads, MLPs and residual stream you'll be hooking into.
• 30u30 Day 01 · RNN Effectiveness: the early "interpretable neuron" result (cells tracking quotes and line position), the ancestor of this field.
• Stage 01: PyTorch hooks and autograd, the mechanics behind TransformerLens and Baukit.
• Module 02 · TensorTonic: lists an Interpretability project track for practice.
• Module 07 · Preference alignment: the behavioural route to fixing model behaviour, for contrast.

6. Why it matters for robotics

7. Self-check

8. Courses & content

Start with ARENA's transformer interpretability chapter: it walks you through the logit lens, patching and the IOI circuit in TransformerLens, which is exactly this stage's project.

Full course map for all 12 stages →

9. Papers & docs

• Meng, Bau, Andonian & Belinkov (2022). Locating and Editing Factual Associations in GPT (ROME). Project page: rome.baulab.info
• Geva, Bastings, Filippova & Globerson (2023). Dissecting Recall of Factual Associations in Auto-Regressive Language Models
• Hase, Bansal, Kim & Ghandeharioun (2023). Does Localization Inform Editing?
• Wang, Variengien, Conmy, Shlegeris & Steinhardt (2022). Interpretability in the Wild: a Circuit for Indirect Object Identification in GPT-2 small
• Conmy, Mavor-Parker, Lynch, Heimersheim & Garriga-Alonso (2023). Towards Automated Circuit Discovery for Mechanistic Interpretability (ACDC)
• Belrose et al. (2023). Eliciting Latent Predictions from Transformers with the Tuned Lens
• nostalgebraist (2020). interpreting GPT: the logit lens
• Elhage et al. (2021). A Mathematical Framework for Transformer Circuits
• Docs: TransformerLens · Baukit