Organizations: Key Laboratory of Ministry of Education for Brain Inspired Intelligent Perception and Cognition, University of Science and Technology of China, Hefei, Anhui 230026, China
Mechanistic interpretability seeks to localize model behavior to the internal components that causally realize it. Prior work has advanced activation-space localization and causal tracing, but modules that appear important in activation space may merely aggregate or amplify upstream signals rather than encode the target capability in their own parameters. To address this gap, we propose Weight Patching, a parameter-space intervention method for source-oriented analysis in paired same-architecture models that differ in how strongly they express a target capability under the inputs of interest. Given a base model and a behavior-specialized counterpart, Weight Patching replaces selected module weights from the specialized model into the base model under a fixed input. We instantiate the method on instruction following and introduce a framework centered on a vector-anchor behavioral interface that provides a shared internal criterion for whether a task-relevant control state has been formed or recovered in open-ended generation. Under this framework, the analysis reveals a hierarchy from shallow candidate source-side carriers to aggregation and routing modules, and further to downstream execution circuits. The recovered component scores can also guide mechanism-aware model merging, improving selective fusion across the evaluated expert combinations and providing additional external validation.
Figures & tables
Figure 1: Weight Patching versus standard Activation Patching. Standard activation patching intervenes on activations under clean/corrupted inputs, whereas Weight Patching intervenes on parameters under a fixed input. In our instruction-following analyses, the activation-side companion intervention uses a fixed-input cross-model form with the same restoration logic.
Figure 2: Mechanistic summary of instruction following under the proposed framework. The recovered hierarchy suggests that shallow components act as candidate source-side carriers of instruction-conditioned control features, mid-layer attention heads aggregate and route them into an instruction-following direction, and downstream instruction-specific modules help unfold the recovered control state into response generation.
Figure 3: Overview of the proposed framework. We first extract an instruction-direction anchor as a shared behavioral interface, then apply Weight Patching for source-level localization, and finally combine parameter-space and activation-space evidence to identify source-side carriers, aggregation/routing modules, and downstream execution circuits.
Figure 4: Reusing WP-recovered component scores for mechanism-aware model merging.
Figure 5: Layer-wise recovery under task-vector injection. Injecting the extracted task vector into different layers of Llama-3.2-3B-Base yields clear recovery peaks across six IFEval tasks, indicating that the anchor-level control representation emerges primarily in middle-to-late layers.
Figure 6: Component-importance heatmaps under Activation Patching and Weight Patching on the English Capital task.
Figure 7: Layer distributions and overlaps of top-ranked heads and neurons identified by Activation Patching and Weight Patching on the English Capital task.
Task
Max
Top-5 Avg
Top-20 Avg
Llama-3.2-3B
English Capital
0.775
0.734
0.507
Multiple Sections
0.405
0.371
0.264
Number Highlighted Sections
0.485
0.442
0.305
Title
0.435
0.422
0.321
No Comma
0.445
0.363
0.246
Table I: Overlap between neurons localized directly by Weight Patching and neurons recovered by tracing upstream from the Top-20 AP-ranked heads. We report the maximum overlap (Max), the mean overlap over the Top-5 heads (Top-5 Avg), and the mean overlap over all 20 heads (Top-20 Avg).
Figure 8: Causal validation of upstream (top row) and downstream (bottom row) modules on the English Capital task. The results support a source–convergence separation in upstream modules and a neuron-dominant pattern in downstream execution.
Figure 9: Fine-grained comparison of upstream and downstream neurons on the English Capital task. Left panels show ablation behavior, and right panels show output-space projections. Upstream neurons exhibit threshold-like collective effects and more abstract projections, whereas downstream neurons behave more additively and show more directly interpretable output-space features.
Figure 10: MLP submodule ablation for neuron localization. (a) Overlap between neurons localized from gate, up, down, or gate+up and the full-unit neuron ranking under WP. (b) Restoration completeness after globally restoring neurons selected under each setting, averaged over six IFEval tasks.
Figure 11: Selective attention of critical heads to instruction tokens. Critical heads in the instruction-tuned model allocate substantially more attention to instruction regions than matched heads in the base model or random heads, supporting their routing role.
Figure 12: Cross-task ablation and restoration across six IFEval tasks. The plots compare head-level and neuron-level knockout/restoration under Activation Patching, Weight Patching, and random selection.
Figure 13: Overlap of critical modules across tasks. Intersection counts are shown for top-ranked heads and neurons across six IFEval tasks.
Figure 14: Cross-task specificity of restored upstream neurons. Rows denote the restored task and columns the evaluation task; values are normalized restoration completeness. Left: AP-traced upstream neurons. Right: WP-localized upstream neurons.
Gradient
Intersection
Total Time
Peak VRAM
Model
Attribution
(%)
(Second)
(MiB)
Head-level
LLama-3.2-3B
✘
66.00
384.85
13150
✔
4.03
15760
LLama-3.1-8B
✘
67.50
1099.68
34458
✔
4.58
37747
Table II: Comparison between exact Weight Patching and its gradient-based approximation, averaged over six tasks, including the overlap of Top-20 heads and Top-1000 neurons, total analysis time, and peak VRAM on three RTX 4090 GPUs.
Figure 15: Layer-wise parameter changes induced by instruction tuning. Top: relative changes in attention submodules (q, k, v, o, and overall). Bottom: relative changes in MLP submodules (gate, up, down, and overall).
Figure 16: Instruction-following recovery via injection of localized critical components. Injecting the identified head and neuron subsets into base models substantially restores performance across model sizes, supporting their reuse as transferable capability carriers.
Model(s)
Method
HumanEval
MBPP
MMLU
MATH
GSM8K
IFEval
Avg.
Base Models
Base
-
18.90
28.20
52.34
0.80
3.00
23.59
21.14
Code
-
19.51
29.20
52.96
5.90
24.60
23.55
25.95
Instruction Tuned
-
30.49
32.20
52.23
7.30
42.10
35.51
33.31
Math
-
14.02
23.40
50.00
12.90
57.80
21.80
29.99
Merged Models
Table III: Performance comparison of the average baseline, prior merging methods, and precise fusion (ours) across four expert combinations.
Appendix figures & tables27 assets
Supplementary material from the paper’s appendix.
Appendix
Model
Layers
Hidden Dim
Attention
KV Heads
MLP Width
Vocab Size
( L )
( dmodel )
Heads
Llama-2-13B
40
5120
40
40
13824
32000
Llama-3.1-8B
32
4096
32
8
14336
128256
Llama-3.2-3B
28
3072
24
8
8192
128256
Appendix
Table IV: Architectural statistics of the model families used in our experiments.
Figure 17: Layer-wise steering correction rates for Llama-3.1-8B and Llama-2-13B. This figure complements Fig. 5 in the main text by showing the steering effects of task vectors across all layers for the 8B and 13B models on six representative IFEval tasks. Similar to the 3B model, optimal instruction-following recovery consistently peaks in the mid-to-late layers across both larger models. This demonstrates the cross-model stability of the vector-anchor interface and confirms that instruction-conditioned control representations are reliably formed after early token-level processing.
Figure 18: Heatmaps of component importance under activation patching and Weight Patching on the English Capital task on Llama-3.1-8B .
Figure 19: Layer distributions and overlaps of top-ranked components identified by activation patching and weight patching on Llama-3.1-8B .
Figure 20: Heatmaps of component importance under activation patching and Weight Patching on the English Capital task on Llama-2-13B .
Figure 21: Layer distributions and overlaps of top-ranked components identified by activation patching and weight patching on Llama-2-13B .
Task
Layer
Projection (Top-10 vocabulary)
No Comma
3
¡ wear ¿, ¡ adam ¿, ¡ Dipl ¿, ¡ cca ¿, ¡ McCorm ¿, ¡ ¡U+FFFD¿ ¿, ¡ Synd ¿, ¡ least ¿, ¡ dex ¿, ¡ ook ¿
¡ egy ¿, ¡ Gram ¿, ¡ firm ¿, ¡ clam ¿, ¡ haft ¿, ¡ rails ¿, ¡ Gal ¿, ¡ WF ¿, ¡ Wars ¿, ¡ Sequ ¿
15
¡ comb ¿, ¡ zero ¿, ¡ emb ¿, ¡ connected ¿, ¡ cord ¿, ¡ succession ¿, ¡ port ¿, ¡ God ¿, ¡ pair ¿, ¡ esch ¿
20
¡ zero ¿, ¡ neg ¿, ¡ without ¿, ¡ Jordan ¿, ¡ no ¿, ¡ ash ¿, ¡ sans ¿, ¡ Without ¿, ¡ single ¿
25
¡ single ¿, ¡ without ¿, ¡ un ¿, ¡ neg ¿, ¡ : ¿, ¡ , ¿, ¡ negative ¿, ¡ Single ¿
30
¡ none ¿, ¡ without ¿, ¡ single ¿, ¡ , ¿, ¡ neg ¿, ¡ None ¿, ¡ acc ¿, ¡ ats ¿, ¡ tall ¿, ¡ no ¿
Appendix
Table VII: Vocabulary projection results of task vectors from shallow to deep layers across six tasks in the Llama-2-13B model.
Figure 24: Causal validation of Llama-2-13B upstream and downstream modules on the Title task.
Figure 25: Causal validation of Gemma2-2B global modules on the English Capital task.
Figure 26: Causal validation of Mistral-7B-v0.3 global modules on the English Capital task.
Figure 27: Causal validation of Qwen2.5-3B global modules on the English Capital task.
Figure 28: Vocabulary projection results of downstream neurons in the Llama-3.2-3B model, identified via weight patching across six tasks.
Figure 29: Selective attention of critical heads to instruction tokens.
Figure 30: Overlap of critical modules across tasks on Llama-3.1-8B .
Figure 31: Overlap of critical modules across tasks on Llama-2-13B .
Figure 32: Layer-wise parameter changes induced by instruction tuning on Llama-3.1-8B .
Figure 33: Layer-wise parameter changes induced by instruction tuning on Llama-2-13B .
Figure 34: Absolute performance comparison (IFEval strict accuracy) across three model scales: Llama-3.2-3B, Llama-3.1-8B, and Llama-2-13B. The radar charts illustrate the performance of the Vanilla Base Model (brown), the Vanilla Instruct Model (blue), and the Base Model with optimal layer-wise activation steering (green). By injecting the extracted task vector, the base model successfully recovers, and in some cases slightly exceeds, the instruction-following capabilities of the fully fine-tuned instruct model across all six representative tasks.
Task
Max
Top-5 Avg
Top-20 Avg
Llama-3.2-3B
English Capital
0.761
0.703
0.480
Multiple Sections
0.403
0.380
0.288
Number Highlighted Sections
0.409
0.388
0.286
Title
0.495
0.446
0.337
No Comma
0.415
0.354
0.247
Appendix
Table VIII: The intersection rate of neurons found by directly applying WP versus tracing WP from the top-20 heads identified by AP (Top-5000 Neurons). Metrics include the maximum overlap ( Max ), the average of the top-5 overlaps ( Top-5 Avg ), and the average overlap across all 20 heads ( Top-20 Avg ).
Figure 35: Case study of Llama-3.2-3B on the Englist Capital task.
Figure 36: Case study of Llama-3.2-3B on the Title task.
Figure 37: Heatmaps comparing component localization across distinct metrics (anchor similarity vs. KL divergence) on Llama-3.1-8B .
Figure 38: Layer distributions and overlaps of top-ranked components identified across distinct metrics (anchor similarity vs. KL divergence) on Llama-3.1-8B .
Mechanistic interpretability asks which language model components matter for a behavior, and activation patching answers by changing one component's activation and measuring its effect, at a cost growing with components and prompts. Attribution patching estimates these effects from one gradient pass, but its estimates can be far off, and verifying each one restores the cost. Can the structure of these errors reduce this cost? We propose Guided-AP, which uses a few exact measurements to reveal attribution patching's errors and correct unmeasured estimates. On a single prompt, many components' errors share one sign, so a small exact sample yields one correction for the rest. Over many prompts of one task, components are misestimated similarly, so a few measured prompts show what to verify and how to correct each component's average effect. Our analysis shows that choosing what to verify needs error size, while correction also needs sign. On a single prompt, the correction lowers error by up to 26% across three behaviors at the compute of verifying the largest estimates and also on larger models. Over many prompts, learned corrections cut the error of average effects on factual completion by 75% to 79%, repaying their cost in long analyses.
Luyang Zhang, Minxing Zheng, Jialu Wang
Carnegie Mellon University · University of California, Santa Cruz
A central goal of mechanistic interpretability is to identify which internal components causally drive a language model's behavior. Because these importance estimates serve as the evidence for identifying circuits, systematic errors can lead to the misidentification of the underlying mechanisms. While activation patching provides a gold-standard causal metric, its computational cost is prohibitive at scale. Practitioners instead rely on attribution patching, a gradient-based, first-order approximation whose reliability remains poorly understood. In this work, we characterize the source of this unreliability, demonstrating that the dominant error stems from the non-linearities in the downstream network rather than local curvature at the patched component. This insight yields three practical tools: (i) a reliability score to detect untrustworthy estimates, (ii) error bounds quantifying potential attribution mis-specifications, and (iii) a Hessian-vector-product (HVP) correction that eliminates the leading-order error with only one additional backward pass. In evaluations across five model families (124M-9B parameters) and both random-token and naturalistic (name-swap) perturbations, HVP is the only second-order correction feasible at larger scale, where standard baselines like Integrated Gradients become computationally prohibitive. In comparative experiments, a multi-step HVP variant matches or exceeds the accuracy of Integrated Gradients at significantly lower compute, outperforming prior second-order baselines. These improvements lead to higher-fidelity circuit recovery on standard benchmarks and support a Screen-Flag-Fix workflow that targets computational effort only toward the components flagged as unreliable.
Luyang Zhang, Jialu Wang
1Carnegie Mellon University · University of California, Santa Cruz
Attributing language model outputs to their internal computations is an open problem in interpretability. Existing methods, which use causal interventions, gradients, or learnable masks, either are infeasibly expensive or struggle to identify actual causally-important internal computations. We propose framing attribution as the problem of identifying nested subsets of internal components which minimise a downstream loss. To learn this task, we introduce Matryoshka Attribution (MAttr), a mask learning method that parametrises the mask with a simple differentiable sigmoid top-k operator. We supervise training over all sparsities simultaneously by randomising k over training, resulting in a learned ordering of components by attribution score. MAttr achieves number 1 on the official leaderboard of the Mechanistic Interpretability Benchmark (Mueller et al., 2025); our method identifies sparse and task-transferrable circuits across varying circuit bases. As a practical application, we show that MAttr can be trained with reinforcement learning to identify weight changes responsible for downstream behaviours in LLM finetuning. We train MAttr on refusal judge scores and find that restoring 1% of Llama 3.1 8B Instruct's weights to their base model state is sufficient to remove refusals while maintaining capabilities. We view MAttr as a successful formulation of interpretability into a learnable objective that we can tackle with gradient descent, and encourage future work along these lines.