Attribution / edge patching

Mechanistic family · patching at gradient cost

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

Activation patching is the gold standard of mechanistic interpretability, but it is slow. You need one extra forward pass for every component you test. A large language model has millions of neurons and thousands of attention heads at every token position, so exact patching of all of them is out of reach. Attribution patching is the shortcut. It predicts the effect of every patch at once from a single forward and backward pass.

The trick is the oldest one in calculus: near a point, a smooth function looks like a straight line. If you know the slope of the output with respect to an activation (the gradient), and you know how far the patch would move that activation, then slope × distance is a guess for how far the output moves. That guess is a first-order Taylor approximation. It is exact when the output is linear in the activation, and wrong when it is not. On the VCS we can measure the exact patch effect for every site, so we can see precisely how much accuracy the shortcut trades for its speed.

How it works on a neural network

Use the notation of activation patching: clean input \(x\), corrupted input \(\tilde x\), component \(\ell\) with clean activation \(a = a_\ell(x)\) and corrupted activation \(a' = a_\ell(\tilde x)\). The exact patch effect is \(\Delta_\ell = y(x \mid a_\ell := a') - y(x)\). Attribution patching replaces it with

\[ \widehat{\Delta}_\ell \;=\; \left.\frac{\partial y}{\partial a_\ell}\right|_{x} \cdot \,(a' - a). \]

The gradient for all components comes from one backward pass on the clean run; the \(a' - a\) terms come from one cached corrupted run. That is three passes in total, instead of one per component. Edge attribution patching applies the same estimate to a connection between two components (the sender’s change times the receiver’s gradient), so it can score every edge of the computational graph cheaply.

Toy example. Let \(y = a^2\) with clean \(a = 1\). The slope is \(2a = 2\).

• Small patch \(a' = 1.1\): estimate \(2 \times 0.1 = 0.2\); truth \(1.21 - 1 = 0.21\). Close.
• Large patch \(a' = 4\): estimate \(2 \times 3 = 6\); truth \(16 - 1 = 15\). Badly off in size, but the sign is right.
• Patch at a flat point, clean \(a = 0\), \(a' = 3\): slope \(0\), estimate \(0\); truth \(9\). The estimate misses the effect completely.

The first two cases are why attribution patching is good at which components matter but poor at how much. The third is the dangerous case: a real effect hidden behind a zero gradient, just like saturated attention or a dead ReLU.

From a neural network to the Atari

Activation \(a_\ell\)	One state site at frame \(f^*\): a candidate RAM byte or the COLUBK TIA register.
Corrupted activation \(a'\)	The value from a donor run (LEFT or RIGHT pressed before \(f^*\)), or the directed value base+17.
Gradient \(\partial y / \partial a\)	The local slope of the true 15-frame output, measured by poking the byte to \(a \pm 1\) and re-running the real ROM (a central finite difference).
Metric \(y\)	The two most causally active candidate bytes and the count of changed screen pixels, read 15 frames later.
Exact patching (the reference)	The full re-run with the patched value, which on the VCS equals the oracle.
An edge	A (patched site, output) pair: one cell of the effect matrix.

How we apply it to the VCS

The runner is built on top of the activation-patching harness. It uses the same shared gameplay checkpoint, the same outputs and the same patch list: donor patches from LEFT/RIGHT runs, directed base+17 patches on every candidate byte, and one background-colour patch.

1. The slope. For each site the runner writes \(a+1\) and \(a-1\), re-runs 15 frames each, and divides the difference by 2. At 0 or 255 it uses a one-sided difference. This is the program’s own local derivative, measured on the real machine. It is not a Zygote backward pass. For a single-frame content read the two coincide, but the finite difference also works across 15 frames of discrete code, where autodiff gives no useful answer.
2. The estimate. \(\widehat\Delta = \text{slope} \times (a' - a)\) for every (site, output) pair.
3. The truth. The exact patch: write \(a'\), re-run, read the outputs.
4. Scores. Pearson correlation of estimate vs truth over all cells, mean and relative error, edge precision/recall (does the estimate fire where the truth fires?), and the fraction of active edges where estimate and truth agree exactly (the “linear regime”).
5. Positive control. A byte read back at the patch frame itself is perfectly linear, \(y(a) = a\). For it the estimate must equal the truth for steps of 1 to 32. This proves the wiring is right, so any error elsewhere is the program’s nonlinearity.
6. Position gradient. Because this is a gradient method, the runner also reports the naive Zygote gradient of the screen-position output (identically zero) next to the bilinear-sampler gradient, which restores a non-zero value.

Note one honest consequence. On a network the gradient is nearly free. Here each slope costs two re-runs, so the VCS version is not cheaper than exact patching. It is run to measure the approximation error, not to save compute.

Read the code

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

• XAIattribution_patching.jl:136 — The slope \(\partial y/\partial a\): central difference of two real re-runs at \(a\pm1\).
• XAIattribution_patching.jl:351 — The attribution-patching estimate: slope times the patch distance.
• XAIattribution_patching.jl:382 — Edge precision/recall: the estimate’s firing pattern vs the exact data-flow.
• XAIattribution_patching.jl:481 — Positive control: a strictly linear site where the estimate must be exact.
• XAIattribution_patching.jl:427 — Naive (zero) vs bilinear-sampler Zygote gradient for the screen-position output.
• XAIactivation_patching.jl:282 — The exact patch both the slope probes and the reference are built from.
• substrateshared_testbed_impl.jl:357 — The bilinear sampler that gives the position output a usable gradient.
• substrateSoftStep.jl:106 — The differentiable one-hot RAM read the sampler gradient flows through.
• substratejutari_oracle.jl:176 — The byte write used for \(a\pm1\) probes and for the exact patch.
• substrateStellaEnvironment.jl:259 — One bit-exact frame of the real machine.

Reading the result

Expect the estimate to agree with the truth on which edges exist far better than on their size. Small donor patches and single-frame reads sit in the linear regime and match closely. Large directed patches, bytes that wrap at 255, and bytes the game recomputes each frame break linearity, and the scatter of estimate vs truth spreads out. The correlation is therefore good but clearly below the exact method, and the positive control shows the gap is the program’s fault, not the harness’s.

Pitfalls

• Large patches are not small. Taylor expansion holds near the clean point. A patch that moves an activation far (here, base+17) can land in a region the slope knows nothing about.
• Zero gradient, real effect. Saturated softmax, dead ReLUs or a threshold in the program give slope 0 and hide a large true effect. AtP* documents this for attention in LLMs.
• Wrap-around and thresholds. VCS bytes are 8-bit: 250 + 17 wraps to 11. Comparisons and branches make the output piecewise constant. Both break the linear model.
• Cancellation. Summing estimates over many components or positions can cancel positive and negative terms that are each large.
• Use it as a filter. The method is designed to shortlist candidates for exact patching, not to replace it.

Try it yourself

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

Writes attribution_patching_<game>.json/.npz and a core summary into tools/xai_study/phaseC_mechanistic/out/. Three re-runs per patch (two slope probes plus the exact patch); minutes per game on one CPU core.

Original papers & further reading

• Nanda (2023). Attribution Patching: Activation Patching At Industrial Scale. Blog post. — www.neelnanda.io/mechanistic-interpretability/attribution-patching
• Syed, Rager & Conmy (2023). Attribution Patching Outperforms Automated Circuit Discovery (edge attribution patching). — arxiv.org/abs/2310.10348
• Kramár, Lieberum, Shah & Nanda (2024). AtP*: An efficient and scalable method for localizing LLM behaviour to components. (Documents where the linear estimate fails, e.g. attention saturation, and proposes fixes.) — arxiv.org/abs/2403.00745

Read next: Activation patching / causal mediation · Vanilla gradient (saliency) · ACDC — automatic circuit discovery · Path patching / IOI circuit