Activation patching / causal mediation

Mechanistic family · the core causal tool

Learn the method — a beginner's guide

New to explainable AI? This block explains the method as it is used on neural networks, then how we run the very same method on an Atari 2600, where — unlike for a network — we know the right answer. Every implementation step links to the code. See also the XAI primer for the shared vocabulary.

The idea in plain words

Gradient methods ask “what if I nudged this a tiny bit?”. Activation patching asks a bolder, causal question: what if this internal value had been the one from a different input? You run the network twice. The clean run uses the input you care about (“The Eiffel Tower is in” → “Paris”). The corrupted run uses a slightly different input that gives a different answer. Then you run a third time, and at one chosen place inside the network you overwrite (“patch”) the internal value with the one recorded in the other run. If the answer flips, that place carries the information that decides the answer.

An internal value of a network is called an activation: the output of one neuron, one layer or one attention head. Statisticians call a variable that sits on the causal path from cause to effect a mediator, so the method is also called causal mediation analysis. It is the workhorse of mechanistic interpretability, the attempt to reverse-engineer the algorithm a network runs. Unlike a correlation, a patch is an intervention: it sets a value and watches what happens. On the VCS we can do exactly the same thing to the machine's memory, and it turns out to be identical to our ground-truth oracle.

How it works on a neural network

Let \(x\) be the clean input and \(\tilde x\) the corrupted one. Let \(a_\ell(x)\) be the activation of component \(\ell\) (a neuron, an attention head, or a position in the residual stream: the running vector that every transformer layer reads from and adds to). Write \(y(x \mid a_\ell := v)\) for the output when we run on \(x\) but force \(a_\ell\) to the value \(v\). The patch effect of component \(\ell\) is

\[ \Delta_\ell \;=\; y\big(x \,\big|\, a_\ell := a_\ell(\tilde x)\big) \;-\; y(x). \]

This direction (corrupted value into the clean run) is called noising; the reverse (clean value into the corrupted run) is denoising. In ROME’s “causal tracing” the corruption is Gaussian noise on the subject tokens, and the patched-back clean state shows where a fact is recalled. Each \(\Delta_\ell\) costs one extra forward pass.

Toy example. Take a two-step network \(h = 2x_1\), \(g = x_2 + 1\), \(y = h + 3g\). Clean input \(x = (1, 0)\) gives \(h = 2\), \(g = 1\), \(y = 5\). Corrupted input \(\tilde x = (4, 0)\) gives \(h = 8\), \(g = 1\), \(y = 11\). Patch \(h\): run on \(x\) but set \(h := 8\). Then \(y = 8 + 3 = 11\), so \(\Delta_h = 6\): \(h\) mediates the whole difference. Patch \(g\): its corrupted value is also 1, so \(\Delta_g = 0\). The patch tells us which component carries the change between these two inputs, not which components matter in general. A different \(\tilde x\) (say, changing \(x_2\)) would light up \(g\) instead.

From a neural network to the Atari

The network and its forward pass	The ROM program running on the bit-exact emulated VCS (jutari). One “layer” is one video frame of 6507 CPU work.
The residual stream at layer \(\ell\)	The whole machine state at a frame boundary: the 128 bytes of RAM plus TIA video registers. Every frame reads it and writes it back.
A component (neuron / head)	One state site: a single RAM byte that the game uses as a variable (a candidate concept cell such as a ball coordinate), or one TIA register (the background colour COLUBK).
Clean run / corrupted run	Clean: the seeded random-action gameplay stream, checkpointed at frame \(f^*{=}90\). Donor: the same stream but with LEFT or RIGHT pressed at the last frame before \(f^*\).
The logit / metric	Outputs read 15 frames later: the values of the two most causally active candidate bytes, and the number of screen pixels that differ from the clean frame.
Ground truth	The intervention oracle, which pokes a value and re-runs the real ROM. On the VCS that is a patch, so activation patching and the truth coincide.

How we apply it to the VCS

The runner sits on the shared testbed that every method uses: a seeded stream of random joystick actions, a checkpoint at frame \(f^*\), and a horizon of 15 frames. Two fresh boots are first asserted to be byte-identical in RAM and screen, so any later difference is caused by our patch and nothing else.

1. Outputs. The runner picks the two candidate RAM bytes whose own value moves most under the oracle’s interventions, and adds a whole-screen output: the count of pixels that differ from the clean final frame.
2. Donor patches. Two donor runs share the gameplay prefix and press LEFT or RIGHT just before \(f^*\). For every candidate byte the donor’s value is written into a copy of the clean checkpoint. The record notes whether the donor actually differed there; equal values are honest no-op patches.
3. Directed patches. Because a one-frame donor leaves many bytes unchanged, every candidate byte is also patched to base+17 (the same value the oracle uses), plus one TIA patch that sets the background colour to white.
4. Recovered vs exact. Each patch is re-run on the clean checkpoint (the method). The same write is then replayed from an independent fresh boot (the oracle). The two effect tables are compared cell by cell. They are equal, because both are the same single-site write followed by the same deterministic program.
5. Edge precision/recall. A (patch, output) pair “fires” if its effect is non-zero. The recovered firing pattern is scored against the exact one. An honest extra audit lists transient sites: bytes the program overwrites on the next frame, so their patch never reaches the outputs.

The difference from LLM practice is worth stating plainly. In a transformer, patching and “the truth” are different things, and we argue about which metric to trust. On the VCS, patching a byte and re-running is the oracle’s own definition of a cause. So this method is the reference the other Phase C methods approximate.

Read the code

XAI the interpretability method   substrate the bit-exact emulator it runs on (jutari, Julia).

• XAIactivation_patching.jl:282 — The patch itself: copy the checkpoint, write one RAM byte or TIA register, step the tail, snapshot.
• XAIactivation_patching.jl:417 — The donor (“corrupted”) runs: same gameplay prefix, LEFT or RIGHT at the last step.
• XAIactivation_patching.jl:460 — The recovered effect: patch into the clean checkpoint and re-run.
• XAIactivation_patching.jl:466 — The exact effect via an independent fresh boot, so equality is not a shared-code artefact.
• XAIactivation_patching.jl:487 — Site precision/recall of the recovered firing pattern against the exact data-flow.
• substratejutari_oracle.jl:176 — The write: one byte of RIOT RAM in the live emulator.
• substratejutari_oracle.jl:187 — The same write for a TIA video register (the background-colour patch).
• substrateshared_testbed_impl.jl:115 — The shared gameplay state: seeded actions, checkpoint at \(f^*\), causes, cause-density gate.
• substrateStellaEnvironment.jl:259 — One bit-exact video frame of the real machine; every re-run is a loop over this.
• substrateoracle_intervene.jl:242 — The oracle: apply a cause to a copy of the checkpoint and run the horizon. Same operation as a patch.

Reading the result

The headline number is the largest gap between the recovered and the exact effect table. Expect it to be zero on every game: this is a positive result by construction, not a discovery. The more informative parts are the firing pattern and the transient list. Many bytes show no effect because the game recomputes them every frame. Others move a score byte but not the screen, or the reverse. Read the page as a map of which variables mediate which outputs over 15 frames, not as a test the method could fail.

Pitfalls

• The answer depends on the corruption. A patch measures what carries the difference between these two runs. A donor that does not change a byte can never reveal that byte, which is why we add directed patches.
• Off-distribution values. A directed value like base+17, or Gaussian noise in ROME, may put the system in a state it never visits naturally. The effect is real, but it may not describe normal behaviour.
• Effect is not meaning. A perfect effect table says which byte matters, not what it represents or how bytes combine. Hase et al. also found that where patching localises a fact is a poor guide to where editing it works.
• Metric choice changes the picture. Zhang & Nanda show that logit difference, probability and noising vs denoising can rank components differently in LLMs.
• Cost. One full re-run per site and value. Fine for 128 bytes, expensive for billions of activations; this is what attribution patching tries to fix.

Try it yourself

julia --project=jutari tools/xai_study/phaseC_mechanistic/activation_patching.jl --games core

Writes activation_patching_<game>.json/.npz and a core summary into tools/xai_study/phaseC_mechanistic/out/. Each patch is replayed twice (clean checkpoint and fresh boot), so expect minutes per game on one CPU core.

Original papers & further reading

• Vig et al. (2020). Causal Mediation Analysis for Interpreting Neural NLP: The Case of Gender Bias. NeurIPS. — arxiv.org/abs/2004.12265
• Meng, Bau, Andonian & Belinkov (2022). Locating and Editing Factual Associations in GPT (ROME, causal tracing). NeurIPS. — arxiv.org/abs/2202.05262
• Zhang & Nanda (2024). Towards Best Practices of Activation Patching in Language Models: Metrics and Methods. ICLR. — arxiv.org/abs/2309.16042
• Hase, Bansal, Kim & Ghandeharioun (2023). Does Localization Inform Editing? NeurIPS. (A critique: where patching localises a fact is not where editing works best.) — arxiv.org/abs/2301.04213

Read next: Attribution / edge patching · Path patching / IOI circuit · Interchange interventions / DAS · Causal scrubbing