🤖 machine learning

When Attribution Patching Lies: Diagnosis and a Second-Order Correction

This paper identifies non-linearities in downstream networks as the primary source of error in attribution patching and introduces a computationally efficient Hessian-vector-product correction that significantly improves the accuracy and reliability of mechanistic interpretability circuits across various model scales.

Original authors: Luyang Zhang, Jialu Wang

Published 2026-06-10

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Luyang Zhang, Jialu Wang

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Big Picture: Trying to Find the "Brain Cells" of AI

Imagine you have a giant, complex machine (a Large Language Model) that can write stories, solve math problems, and chat with you. Scientists want to know exactly which specific parts of this machine are responsible for specific behaviors. For example, which "neuron" decides to use the word "apple" instead of "banana"?

To find these parts, researchers use a technique called Activation Patching. Think of this as a "surgery" on the machine:

You run the machine with a normal sentence (e.g., "The cat sat on the mat").
You run it again with a corrupted sentence (e.g., "The dog sat on the mat").
You swap the internal "brain activity" from the first run into the second run.
If the machine suddenly starts acting like it's talking about a cat again, you know that specific part is crucial.

The Problem: Doing this surgery on every single part of a modern AI is like trying to test every single brick in a skyscraper to see which one holds the roof up. It takes too much time and computer power.

The Shortcut: "Attribution Patching"

Because the full surgery is too slow, researchers use a shortcut called Attribution Patching. Instead of actually swapping the parts, they use a mathematical guess (a "first-order approximation") to predict which parts would matter.

Think of it like this:

Real Surgery (Activation Patching): You actually swap the engine of a car and see if it runs better. Accurate, but takes hours.
The Shortcut (Attribution Patching): You look at the engine, do a quick calculation, and guess, "Yeah, this engine is probably the problem." It's fast, but sometimes the guess is wrong.

The Discovery: Why the Shortcut Lies

The authors of this paper asked: "When does this shortcut fail, and why?"

They discovered that the shortcut fails not because of the part itself, but because of what happens after that part.

The Analogy of the Domino Chain:
Imagine a line of dominoes.

The Shortcut's Mistake: The shortcut looks at the first domino you push and assumes, "If I push this, the whole line falls." It assumes the push travels in a straight, predictable line.
The Reality: In a complex AI, the "push" travels through a winding path of other dominoes. Sometimes, the path curves, or the force gets amplified, or it gets canceled out by other forces. The shortcut doesn't see these curves; it only sees the immediate push.

The paper proves that the biggest errors happen because the shortcut ignores the curvature of the path the signal takes through the rest of the network. It's like trying to drive a car by only looking at the steering wheel, ignoring the fact that the road ahead is full of sharp turns.

The Solution: The "Screen-Flag-Fix" Pipeline

The authors propose a three-step workflow to fix this without slowing everything down. They call it Screen-Flag-Fix.

1. Screen (The Quick Guess)

First, run the fast, cheap shortcut (Attribution Patching) on every part of the AI. This gives you a rough list of who is important.

Analogy: You quickly scan a crowd of people to guess who is the VIP. You get a list of 100 suspects.

2. Flag (The Reliability Check)

Next, use a new "Reliability Score" to check your list. This score asks: "Is the road ahead likely to be curvy?"

If the score is low, the shortcut was probably right.
If the score is high, the shortcut is likely lying because the path is too complex.
Analogy: You look at your list of 100 suspects. You realize that for 90 of them, the path is straight, so your guess is fine. But for 10 of them, the path is full of sharp turns. You flag those 10 as "Unreliable."

3. Fix (The Targeted Surgery)

Finally, you only perform the expensive, accurate "surgery" (using a Hessian-Vector Product or HVP) on the flagged items.

Analogy: You don't check the whole crowd again. You only do a deep background check on the 10 suspicious people you flagged. This saves you 90% of the time but still catches the real VIPs.

Why This Matters

The paper shows that this method works incredibly well:

It's Accurate: It fixes the errors caused by the "curvy roads" in the AI's brain. In some tests, it reduced errors by over 80%.
It's Fast: Because it only fixes the parts that are actually broken, it's much cheaper than doing the full surgery on everything.
It Scales: Other methods that try to fix these errors (like "Integrated Gradients") become impossible to use on huge AI models (like 8 billion parameters). This new method works even on those massive models.

The Bottom Line

The paper teaches us that the common shortcut for understanding AI brains is often unreliable because it ignores the complex paths signals take later in the network. By using a smart "checklist" to find the unreliable guesses and only fixing those specific ones, we can get the accuracy of a full surgery with the speed of a quick guess. This helps scientists build a more accurate map of how AI models actually think.

Technical Summary: When Attribution Patching Lies

Problem Statement

Mechanistic interpretability relies on causal localization to identify which internal components (neurons, attention heads) drive specific model behaviors. The gold standard for this is activation patching, which measures the causal effect of replacing a component's activation under a corrupted input with its value from a clean input. However, activation patching is computationally prohibitive at scale because its cost scales linearly with the number of components.

Consequently, practitioners rely on attribution patching, a first-order, gradient-based approximation that linearizes the effect of an intervention around a corrupted activation. While efficient, the reliability of attribution patching is poorly understood. Systematic errors in these estimates can lead to the misidentification of underlying circuits. Prior work has identified specific failure modes (e.g., activation-region mismatch, softmax redistribution) and proposed partial remedies, but a fundamental question remains: When is attribution patching reliable, how large can its error be, and how should it be corrected?

Methodology and Core Insight

The authors analyze the structure of attribution patching error using a second-order Taylor expansion. They decompose the error $E$ between the true causal effect ( $\Delta$ ) and the first-order estimate ( $\hat{\Delta}$ ) into two components:
$E = \underbrace{\frac{1}{2} \delta^\top H \delta}_{\text{Dominant Quadratic Term}} + \underbrace{\Phi(\delta)}_{\text{Cubic Remainder}}$
where $\delta$ is the perturbation direction and $H$ is the Hessian of the scalar output metric with respect to the activation.

Key Insight: Downstream vs. Local Curvature

A critical finding is that the dominant error term stems from the network-level response to the intervention (how the perturbation propagates through the rest of the model), rather than the local curvature at the patched component itself.

Failure of Prior Fixes: Previous methods (e.g., AtP*, GIM) attempt to correct errors using local activation curvature (e.g., the second derivative of the activation function). The authors demonstrate that local curvature is a poor predictor of the true error, showing a near-zero correlation ( $r=0.04$ ) and underestimating the true error source by factors of 2–66 $\times$ .
Network-Level Correction: The true error depends on the Hessian-vector product (HVP) of the entire computational graph.

The Proposed Solution: Screen–Flag–Fix

Based on this decomposition, the authors propose a scalable workflow:

Screen: Compute standard attribution patching for all components (one backward pass).
Flag: Calculate a Reliability Score ( $\tilde{R}$ ) defined as the magnitude of the quadratic term relative to the first-order estimate:
$\tilde{R} = \frac{|\delta^\top H \delta|}{2 |\hat{\Delta}|}$
A high $\tilde{R}$ indicates that the omitted quadratic term is significant, flagging the estimate as unreliable.
Fix: Apply a HVP-based second-order correction only to the flagged components.
- Standard HVP: Adds the term $\frac{1}{2} \delta^\top H \delta$ to the estimate. This requires two backward passes (one to compute the gradient, one to compute the HVP).
- Multi-Step HVP (MS-HVP): For large perturbations where a single Taylor step might overshoot (violating the local smoothness assumption), the correction is split into $K$ sub-steps. This reduces the residual error from $O(\|\delta\|^2)$ to $O(\|\delta\|^3/K^2)$ .

Key Contributions

Reframing Attribution Reliability: The authors trace the dominant error source to the downstream network response rather than local nonlinearity, explaining why prior local fixes are structurally incomplete.
Scalable Second-Order Correction: They introduce HVP and MS-HVP as the first second-order corrections demonstrated to be feasible at large scales (up to 9B parameters). Unlike Integrated Gradients (IG), which becomes computationally infeasible at this scale (requiring $\sim$ 25 GPU-days per task), HVP remains tractable.
Design Space and Workflow: They organize prior refinement methods by their goals and propose a "Screen-Flag-Fix" pipeline that targets computational effort only toward unreliable components, balancing accuracy and efficiency.

Experimental Results

The authors evaluated their method across five model families (124M–9B parameters) and multiple tasks (IOI, Greater-Than, Factual Completion).

Error Reduction: Standard HVP correction reduces top-5 relative error by 72–90% on attention heads and 56–67% on pre-activation MLP neurons in low-to-moderate curvature regimes. In pathological high-curvature regimes (e.g., Pythia-410M IOI), single-step HVP can overshoot, but MS-HVP (K=5) recovers stability, reducing error by up to 82%.
Comparison to Baselines:
- vs. Integrated Gradients (IG): At matched computational costs, MS-HVP matches or exceeds the accuracy of IG. On GPT-2 IOI, MS-HVP (cost 10) significantly outperforms IG (cost 35) with a 3.5 $\times$ efficiency gain.
- vs. Integrated Hessians (IH): IH consistently underperforms MS-HVP because its weighting scheme (designed for feature interactions) systematically under-corrects the specific error term relevant to attribution patching.
- vs. GIM/AtP:* These methods provide modest gains or degrade performance, particularly in large models where higher-order network curvature dominates.
Circuit Recovery: The per-component accuracy gains translate to improved circuit recovery. On the GPT-2 Greater-Than benchmark, MS-HVP (K=5) improved top-5 head overlap with ground truth from 70.1% to 83.2%, outperforming both IG and GIM.
Scalability: At 8B+ parameters, IG is infeasible, while HVP remains practical. MS-HVP is the only second-order correction demonstrated at this scale.

Significance and Claims

The paper claims that attribution patching is not inherently unreliable but suffers from a specific, correctable source of error: the neglect of downstream network curvature. By shifting focus from local activation properties to network-level Hessian-vector products, the authors provide a general account of when attribution patching fails and how to correct it.

The significance lies in enabling high-fidelity circuit recovery at scales previously inaccessible to rigorous causal testing. The proposed Screen-Flag-Fix workflow allows practitioners to retain the speed of first-order approximations for reliable components while applying targeted, computationally efficient second-order corrections only where necessary. This approach supports the discovery of mechanistic circuits in modern Large Language Models without the prohibitive cost of full activation patching or path-integration methods like Integrated Gradients.