Attribution-Guided Model Rectification of Unreliable Neural Network Behaviors

Peiyu Yang; Naveed Akhtar; Jiantong Jiang; Ajmal Mian

cs.LGMar 8, 2026

Attribution-Guided Model Rectification of Unreliable Neural Network Behaviors

Authors: Peiyu Yang, Naveed Akhtar, Jiantong Jiang, Ajmal Mian

Organizations: The University of Melbourne · The University of Western Australia

Abstract

The performance of neural network models deteriorates due to their unreliable behavior on non-robust features of corrupted samples. Owing to their opaque nature, rectifying models to address this problem often necessitates arduous data cleaning and model retraining, resulting in huge computational and manual overhead. In this work, we leverage rank-one model editing to establish an attribution-guided model rectification framework that effectively locates and corrects model unreliable behaviors. We first distinguish our rectification setting from existing model editing, yielding a formulation that corrects unreliable behavior while preserving model performance and reducing reliance on large budgets of cleansed samples. We further reveal a bottleneck of model rectifying arising from heterogeneous editability across layers. To target the primary source of misbehavior, we introduce an attribution-guided layer localization method that quantifies layer-wise editability and identifies the layer most responsible for unreliabilities. Extensive experiments demonstrate the effectiveness of our method in correcting unreliabilities observed for neural Trojans, spurious correlations and feature leakage. Our method shows remarkable performance by achieving its editing objective with as few as a single cleansed sample, which makes it appealing for practice.

Figures & tables

Figure 1: Given the original sample labeled as Agama , i.e., class y , the Trojaned model can correctly classify this sample. However, it misclassifies the poisoned sample containing a trigger as Tench , i.e., class $\tilde{\textbf{y}}$ . Attribution maps with Pearson Correlation Coefficients (PCCs) and predictive confidence for the vanilla model, fine-tuned model, and model rectified with our approach are provided. Our method restores the correct label by assigning appropriate attributions to the correct object.

Figure 2: False confidence reduction rank after individually rectifying different layers of ResNet-18. A lower value indicates better results.

Figure 3: Model rectifying workflow. Step 1: Given a pair of clean and corrupted samples, map their prediction attributions for all internal layers. Step 2: Transformed attributions to emphasize editable parameters and locate the suspect layer. Step 3: Apply Rank-one model editing to establish a new key-value association in the located layer for behavior correction.

Methods	CIFAR-10		ImageNet
Methods	OA $\uparrow$	ASR $\downarrow$	OA $\uparrow$	ASR $\downarrow$
Trojaned model	93.67	99.94	69.05	87.24
Fine-tuned model ( $n$ =1)	90.83	73.07	65.95	79.91
Fine-tuned model ( $n$ =10)	91.57	30.14	68.66	33.73
Fine-tuned model ( $n$ =20)	91.58	13.22	68.42	21.86
Patched model ( $n=20$ )	89.70	12.19	65.59	13.81

Table 1: Rectifying backdoor vulnerability. Overall accuracy (OA) and attack success rate (ASR) are reported for varying number (

n

) of samples. Ours are highlighted and the best metrics are in bold (with Trojaned model in gray for reference).

Panel of Figure 4 — Figure 4: Comparison of model performance between fine-tuned models (FT) and rectified models (Ours). (a) The mitigation of false confidence changes with the number of used samples. (b) The mitigation of false confidence changes with the overall accuracy degradation (%) during model rectifying and fine-tuning. Results are computed for ResNet-18 on CIFAR-10 dataset.

Methods	OA $\uparrow$	$\mathbf{\Gamma}_{\textbf{0.5}}$ $\downarrow$	$\Gamma_{0.3}$ $\downarrow$	$\Gamma_{0.7}$ $\downarrow$	$\Gamma_{1.0}$ $\downarrow$
Benign model	92.85	95.29	95.15	97.81	99.21
Patched model	89.61	26.86	30.84	32.42	37.19
Dyn. rectifying	91.21	5.17	6.84	7.65	7.91

Table 2: Generalization comparison for trigger in different visibility levels (0.3-1.0). A single corrupted sample is used for model patching and rectifying. ASR at visibility

\varphi

is denoted as

\Gamma_{\varphi}

Methods	OA $\uparrow$	$\mathbf{\Gamma}_{\text{{BR}}}$ $\downarrow$	$\Gamma_{\text{TL}}$ $\downarrow$	$\Gamma_{\text{C}}$ $\downarrow$	$\Gamma_{\text{BL}}$ $\downarrow$
Benign model	91.23	99.74	99.57	99.76	99.90
Patched model	89.22	29.31	34.42	34.58	34.88
Dyn. rectifying	90.85	6.36	9.24	9.47	8.95

Table 3: Generalization comparison for trigger located at top-left (TL), Center (C) and bottom-left (BL). One corrupted sample with trigger at bottom-right ( BR ) used for model patching and rectifying. ASR at location

\eta

is denoted as

\Gamma_{\eta}

Methods	CIFAR-10			ImageNet
Methods	Overall $\uparrow$	Clean $\uparrow$	Spurious	Overall $\uparrow$	Clean $\uparrow$	Spurious
Benign model	94.00	94.42	100.00 ${}_{\textbf{{\color[rgb]{0.8633,0,0}\,+\,5.58}}}$	69.04	81.25	91.66 ${}_{\textbf{{\color[rgb]{0.8633,0,0}\,+\,10.41}}}$
FT model (n=10)	93.32	88.22	99.66 ${}_{\textbf{{\color[rgb]{0.8633,0,0}\,+\,11.40}}}$	68.01	64.58	74.99 ${}_{\textbf{{\color[rgb]{0.8633,0,0}\,+\,10.41}}}$
FT model (n=20)	93.47	88.97	99.62 ${}_{\textbf{{\color[rgb]{0.8633,0,0}\,+\,10.65}}}$	68.18	64.58	74.99 ${}_{\textbf{{\color[rgb]{0.8633,0,0}\,+\,10.41}}}$
P-ClArC (n=20)	88.29	16.89	17.12 ${}_{\textbf{{\color[rgb]{0.8633,0,0}\,+\,0.23}}}$	66.84	8.32	10.91 ${}_{\textbf{{\color[rgb]{0.8633,0,0}\,+\,2.59}}}$
A-ClArC (n=20)	92.41	76.77	79.34 ${}_{\textbf{{\color[rgb]{0.8633,0,0}\,+\,2.57}}}$	67.01	75.66	82.25 ${}_{\textbf{{\color[rgb]{0.8633,0,0}\,+\,6.59}}}$

Table 4: Performance comparison for mitigating spurious correlation. Accuracy is reported for the overall testing, clean, and spurious sets. The erroneously increased accuracy on the spurious set, compared to the accuracy on the clean set, is indicated in red .

Figure 5: BlockMNIST and feature leakage. (a) Null block randomly appended at the top or bottom of MNIST samples. (b&c) Integrated gradients estimated on benign and our rectified models.

Methods	Accuracy $\uparrow$	Feature Leakage $\downarrow$
Benign model	99.17	3.597
Benign model + IG Reg.	94.14	3.417
FT model (n=20)	98.67	2.929
Stat. rectifying (n=1)	98.97	2.655
Dyn. rectifying (n=1)	99.03	2.417

Table 5: Accuracy (%) and leakage comparison on BlockMNIST.

Figure 6: ISIC samples and the inherent spurious patches. Samples containing malignant and benign lesions are presented, where benign samples are partly polluted with spurious colored patches.

Methods	Overall $\uparrow$	Clean $\uparrow$	Spurious
Benign model	79.00	61.50	87.50 ${}_{\textbf{{\color[rgb]{0.8633,0,0}\,+\,26.00}}}$
FT model (n=10)	79.50	62.00	83.00 ${}_{\textbf{{\color[rgb]{0.8633,0,0}\,+\,21.00}}}$
FT model (n=20)	80.50	53.00	64.50 ${}_{\textbf{{\color[rgb]{0.8633,0,0}\,+\,11.50}}}$
A-ClArC (n=20)	79.50	54.50	59.50 ${}_{\textbf{{\color[rgb]{0.8633,0,0}\,+\,5.00}}}$
Stat. rectifying	79.50	60.00	64.50 ${}_{\textbf{{\color[rgb]{0.8633,0,0}\,+\,4.50}}}$
Dyn. rectifying	80.00	61.00	62.50 ${}_{\textbf{{\color[rgb]{0.8633,0,0}\,+\,1.50}}}$

Table 6: Performance comparison for mitigating spurious correlation on ISIC. For our edited models, we use n=10. The erroneously increased accuracy on the spurious set is indicated in red .

Figure 7: Illustration of samples utilized for neural Trojans and spurious correlations. Two patterns serve as backdoor triggers and spurious features. Top row : For neural Trojans, original samples $x$ with label $y\neq\tilde{y}$ are attached with a trigger and changed its label to the target label $\tilde{y}$ . Bottom row : To induce spurious correlations, samples $x$ of a class $y$ are polluted with spurious features.

Figure 8: Illustration of cleansed samples on ISIC. For benign samples polluted with colored patches, we manually clean them by covering the patches with skin tissue from unaffected regions.

Panel of Figure 9 — Figure 9: Performance in reducing false confidence after individually editing different layers of ResNet-18. A lower value indicates better suppression of the model’s false confidence. Red arrows indicate the layer yielding the best results for a given dataset after model editing.

Method	CIFAR-10		ImageNet
Method	Overall Accu. $\uparrow$	ASR $\downarrow$	Overall Accu. $\uparrow$	ASR $\downarrow$
Trojaned model	94.01	99.79	68.95	78.24
Fine-tuned model (n=1)	91.59	69.07	65.45	77.45
Fine-tuned model (n=20)	92.85	9.70	68.63	20.23
Static edited model (n=1)	93.32	4.49	66.06	15.24
Dynamic edited model (n=1)	93.37	0.65	66.74	6.15

Table 7: Performance comparison of defending against the backdoor attack on Trojaned models trained with the Phoenix logo on CIFAR-10 and ImageNet. Overall accuracy (%) and attack success rate (ASR) are compared between fine-tuned models and models edited by our methods. Ours are highlighted and the best metrics are in bold (with Trojaned model in gray for reference).

Method	CIFAR-10			ImageNet
Method	Overall $\uparrow$	Clean $\uparrow$	Spurious	Overall $\uparrow$	Clean $\uparrow$	Spurious
Benign model	94.14	94.67	97.15 ${}_{\textbf{{\color[rgb]{0.8633,0,0}\,+\,2.48}}}$	69.14	77.08	95.83 ${}_{\textbf{{\color[rgb]{0.8633,0,0}\,+\,18.75}}}$
Fine-tuned model (n=10)	93.67	86.80	93.93 ${}_{\textbf{{\color[rgb]{0.8633,0,0}\,+\,7.13}}}$	67.41	65.99	89.24 ${}_{\textbf{{\color[rgb]{0.8633,0,0}\,+\,23.25}}}$
Fine-tuned model (n=20)	94.07	86.67	93.28 ${}_{\textbf{{\color[rgb]{0.8633,0,0}\,+\,6.61}}}$	67.83	68.32	85.72 ${}_{\textbf{{\color[rgb]{0.8633,0,0}\,+\,17.40}}}$
Dyn. edited model (n=1)	94.03	93.28	94.78 ${}_{\textbf{{\color[rgb]{0.8633,0,0}\,+\,1.50}}}$	66.19	93.35	86.42 ${}_{\textbf{{\color[rgb]{0.8633,0,0}\,+\,6.93}}}$
Dyn. edited model (n=20)	94.04	97.15	97.89 ${}_{\textbf{{\color[rgb]{0.8633,0,0}\,+\,0.74}}}$	67.60	81.25	84.08 ${}_{\textbf{{\color[rgb]{0.8633,0,0}\,+\,2.83}}}$

Table 8: Performance comparison of mitigating spurious correlation on susceptible models trained with the Phoenix logo on CIFAR-10 and ImageNet. Accuracy (%) is reported for the overall testing set, clean set and spurious set. To facilitate comparison, we present the increased accuracy on the spurious set relative to the accuracy on the clean set in red . Our results are highlighted .

Method	Worst-Group Accuracy	Overall Accuracy
Benign model	62.90	87.70
Group DRO	63.60	87.60
Fine-tuned model (n=10)	63.12	86.50
Static edited model (n=10)	66.84	87.64
Dyn. Edited model (n=10)	69.18	87.68

Table 9: Performance comparison for mitigating spurious correlation on Waterbirds dataset. The accuracy values (%) for both the worst group and the entire dataset are reported. Ours are highlighted and the best metrics are in bold (with benign model in gray for reference).

Method	Top-1 Recall $\uparrow$	Top-3 Recall $\uparrow$	Top-5 Recall $\uparrow$
ResNet-18	80%	100%	100%
ResNet-34	80%	80%	100%
EfficientNet-B4	50%	100%	100%

Table 10: Results of recall rate (%) in using the proposed susceptible layer localization technique on ResNet-18, ResNet-34 and EfficientNet-B4 models.

Figure 10: Attribution map comparisons on ImageNet among the vanilla model, fine-tuned model and dynamic edited model (Ours). When the model misclassifies poisoned samples containing triggers, our method effectively corrects this unreliable behavior, aligning the attribution maps with those derived from the original samples.

Figure 11: Comparisons of attribution maps on ImageNet among the vanilla model, fine-tuned model and dynamic edited model (Ours). Our method effectively mitigates the model’s reliance on spurious correlated features, aligning the attribution maps with those derived from the original samples.

Figure 12: Comparisons of attribution maps on ISIC dataset among the vanilla model, fine-tuned model and dynamic edited model (Ours). When the model relies on the spurious feature to make predictions, our method effectively corrects this unreliable behavior, aligning the attribution maps with those derived from the original samples.

Model Editing

Model editing changes targeted facts or behaviors in trained models while aiming to preserve their other capabilities.

See all 64 papers

Correcting Attribution Patching with a Few Exact Measurements
Oct 19, 2026Luyang Zhang, Minxing Zheng, Jialu WangLLM Interpretability Activation Patching
Carnegie Mellon University · University of California, Santa Cruz
Reinforcement Learning for Neural Model Editing
Jun 11, 2026Shaivi MalikReinforcement Learning Model Editing
Shortcut Mitigation via Spurious-Positive Samples
May 13, 2026Phuong Quynh Le, Jörg Schlötterer, Sari Sadiya +2Neural Network Robustness Spurious Correlation Robustness
University of Marburg · Goethe University Frankfurt

Oct 19, 2026J/K move · Enter open · S save

Correcting Attribution Patching with a Few Exact Measurements

Luyang Zhang, Minxing Zheng, Jialu Wang

Carnegie Mellon University · University of California, Santa Cruz

LLM Interpretability Activation Patching Gradient-Based Attribution

Open paper AskPDF arXiv

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.

Abstract

Model Editing

Explore similar work

Correcting Attribution Patching with a Few Exact Measurements

Reinforcement Learning for Neural Model Editing

Shortcut Mitigation via Spurious-Positive Samples