Neural Network Repair With Shapley-Guided Search

Abstract

The deployment of deep neural networks (DNNs) in safety-critical domains is critically hampered by their vulnerability to defects, which can arise from malicious attacks or low-quality data. Therefore, precisely locating the network components responsible for these defects, and subsequently repairing them without compromising overall model performance, presents a significant challenge. To address this, this paper introduces NSRepair, a framework that combines interpretable fault localisation with multi-objective optimisation. Specifically, to accurately attribute blame for a defect, we employ Shapley values to quantify the contribution of each neuron. To systematically manage the trade-off between defect correction and performance preservation, we formulate the repair task as a multi-objective optimisation problem. We conducted extensive experiments across four distinct repair tasks, validating NSRepair on diverse model architectures against seven specialised state-of-the-art methods. The results demonstrate that our unified framework effectively repairs a wide range of defects, demonstrating its potential as a versatile and practical solution for improving DNN dependability. Our code is publicly available at https://doi.org/10.5281/zenodo.17494304.

Keywords

defect localization / multi-objective genetic algorithm / neural network repair / Shapley

Acknowledgements

This work was supported by the National Natural Science Foundation of China under Grant 62562009, in part by the Guangxi Natural Science Foundation of China under Grant 2026GXNSFAA00640946, and in part by Guangxi Key Lab of Multi-Source Information Mining and Security, and Guangxi Collaborative Innovation Center of Multi-source Information Integration and Intelligent Processing.

Conflicts of Interest

The authors declare no conflicts of interest.

Data Availability Statement

The data that support the findings of this study are available in the following public repositories: MNIST [59] (http://yann.lecun.com/exdb/mnist), Fashion-MNIST [63] (https://github.com/zalandoresearch/fashion-mnist), LFW [66] (http://vis-www.cs.umass.edu/lfw), GTSRB [64] (https://benchmark.ini.rub.de/gtsrb_dataset.html), and CIFAR-10 [61] (https://www.cs.toronto.edu/~kriz/cifar.html). These datasets are derived from public-domain resources and are widely used in the machine learning community.

References

[1] L. Meng, H. Li, B.-C. Chen, et al., AdaViT: Adaptive Vision Transformers for Efficient Image Recognition,” in Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (IEEE, 2022), 12309-12318.

[2] Z. Tu, H. Talebi, H. Zhang, et al., MAXIM: Multi-Axis MLP for Image Processing,” in Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (IEEE, 2022), 5769-5780.

[3] Y. Almalioglu, M. Turan, N. Trigoni, and A. Markham, “Deep Learning-Based Robust Positioning for All-Weather Autonomous Driving,” Nature Machine Intelligence 4, no. 9 (2022): 749-760, https://doi.org/10.1038/s42256-022-00520-5.

[4] D. Wu, W. Han, Y. Liu, et al., Language Prompt for Autonomous Driving,” in Proceedings of the AAAI Conference on Artificial Intelligence, vol. 39, no. 8 (2025), 8359-8367, https://doi.org/10.1609/aaai.v39i8.32902.

[5] A. W. K. Yeung, A. Torkamani, A. J. Butte, et al., The Promise of Digital Healthcare Technologies,” Frontiers in Public Health 11 (2023): 1196596, https://doi.org/10.3389/fpubh.2023.1196596.

[6] B. Gomes and E. A. Ashley, “Artificial Intelligence in Molecular Medicine,” New England Journal of Medicine 388, no. 26 (2023): 2456-2465, https://doi.org/10.1056/nejmra2204787.

[7] J. Chen, H. Chen, K. Chen, Y. Zhang, Z. Zou, and Z. Shi, “Diffusion Models for Imperceptible and Transferable Adversarial Attack,” IEEE Transactions on Pattern Analysis and Machine Intelligence (2024), https://doi.org/10.1109/TPAMI.2024.3480519.

[8] H. Mohammadian, A. A. Ghorbani, and A. H. Lashkari , “A Gradient-Based Approach for Adversarial Attack on Deep Learning-based Network Intrusion Detection Systems,” Applied Soft Computing 137 (2023): 110173, https://doi.org/10.1016/j.asoc.2023.110173.

[9] Z. Yang, B. Xu, J. M. Zhang, et al., Stealthy Backdoor Attack for Code Models,” IEEE Transactions on Software Engineering 50, no. 4 (2024): 721-741, https://doi.org/10.1109/tse.2024.3361661.

[10] Y. Sun, X. Huang, D. Kroening, J. Sharp, M. Hill, and R. Ashmore, “Structural Test Coverage Criteria for Deep Neural Networks,” supplement, ACM Transactions on Embedded Computing Systems 18, no. S5 (2019): 1-23, https://doi.org/10.1145/3358233.

[11] Y. Yuan, Q. Pang, and S. Wang, “Provably Valid and Diverse Mutations of Real-World Media Data for DNN Testing,” IEEE Transactions on Software Engineering 50, no. 5 (2024): 1040-1064, https://doi.org/10.1109/tse.2024.3370807.

[12] H. Duong, D. Xu, T. Nguyen, and M. B. Dwyer, “Harnessing Neuron Stability to Improve DNN Verification,” in Proceedings of the ACM on Software Engineering, Proceedings of the ACM on Software Engineering, vol. 1 (FSE, 2024), 859-881, https://doi.org/10.1145/3643765.

[13] J. Liu, Y. Xing, X. Shi, F. Song, Z. Xu, and Z. Ming, “Abstraction and Refinement: Towards Scalable and Exact Verification of Neural Networks,” ACM Transactions on Software Engineering and Methodology 33, no. 5 (2024): 1-35, https://doi.org/10.1145/3644387.

[14] X. Yang, T. Yamaguchi, H.-D. Tran, B. Hoxha, T. T. Johnson, and D. Prokhorov, “Neural Network Repair With Reachability Analysis,” in International Conference on Formal Modeling and Analysis of Timed Systems (Springer, 2022), 221-236.

[15] F. Bauer-Marquart, D. Boetius, S. Leue, and C. Schilling, “SpecRepair: Counter-Example Guided Safety Repair of Deep Neural Networks,” in International Symposium on Model Checking Software (Springer, 2022), 79-96.

[16] A. Sinitsin, V. Plokhotnyuk, D. Pyrkin, S. Popov, and A. Babenko, “Editable Neural Networks,” in International Conference on Learning Representations (OpenReview.net, 2020), 1-12.

[17] M. Usman, D. Gopinath, Y. Sun, Y. Noller, and C. S. Păsăreanu, “NN_REPAIR: Constraint-Based Repair of Neural Network Classifiers ,” in International Conference on Computer Aided Verification (Springer, 2021), 3-25.

[18] P. Henriksen, F. Leofante, and A. Lomuscio, “Repairing Misclassifications in Neural Networks Using Limited Data,” in Proceedings of the 37th ACM/SIGAPP Symposium on Applied Computing (ACM, 2022), 1031-1038.

[19] G. Dong, J. Sun, X. Wang, X. Wang, and T. Dai, “Towards Repairing Neural Networks Correctly,” in International Conference on Software Quality, Reliability and Security (IEEE, 2021), 714-725.

[20] J. Sohn, S. Kang, and S. Yoo, “Arachne: Search-Based Repair of Deep Neural Networks,” ACM Transactions on Software Engineering and Methodology 32, no. 4 (2023): 1-26, https://doi.org/10.1145/3563210.

[21] B. Sun, J. Sun, L. H. Pham, and J. Shi, “Causality-Based Neural Network Repair,” in Proceedings of the 44th International Conference on Software Engineering (ACM/IEEE, 2022), 338-349.

[22] F. Liu, X. Du, H. Ding, and J. Qian, “Towards Robust Neural Networks: Exploring Counterfactual Causality-Based Repair,” Expert Systems with Applications 257 (2024): 125082, https://doi.org/10.1016/j.eswa.2024.125082.

[23] A. Ghorbani and J. Y. Zou, “Neuron Shapley: Discovering the Responsible Neurons,” in Advances in Neural Information Processing Systems, Vol. 33 (Curran Associates, 2020), 5922-5932.

[24] J. Chen, L. Song, M. J. Wainwright, and M. I. Jordan, “L-Shapley and C-Shapley: Efficient Model Interpretation for Structured Data,” in International Conference on Learning Representations (OpenReview.net, 2019), 1-33.

[25] S. M. Lundberg and S.-I. Lee, “A Unified Approach to Interpreting Model Predictions,” Advances in Neural Information Processing Systems 30 (2017): 1-10, https://dl.acm.org/doi/10.5555/3295222.3295230.

[26] B. Schölkopf, F. Locatello, S. Bauer, et al., Toward Causal Representation Learning,” Proceedings of the IEEE 109, no. 5 (2021): 612-634, https://doi.org/10.1109/jproc.2021.3058954.

[27] M. Sundararajan, A. Taly, and Q. Yan, “Axiomatic Attribution for Deep Networks,” in International Conference on Machine Learning (PMLR, 2017), 3319-3328.

[28] R. R. Selvaraju, M. Cogswell, A. Das, R. Vedantam, D. Parikh, and D. Batra, “Grad-CAM: Visual Explanations From Deep Networks via Gradient-Based Localization,” in Proceedings of the IEEE International Conference on Computer Vision (IEEE, 2017), 618-626.

[29] A. Chattopadhay, A. Sarkar, P. Howlader, and V. N. Balasubramanian, “Grad-CAM++: Generalized Gradient-Based Visual Explanations for Deep Convolutional Networks,” in 2018 IEEE Winter Conference on Applications of Computer Vision (WACV) (IEEE, 2018), 839-847.

[30] S. Bach, A. Binder, G. Montavon, F. Klauschen, K.-R. Müller, and W. Samek, “On Pixel-Wise Explanations for Non-linear Classifier Decisions by Layer-Wise Relevance Propagation,” PLoS One 10, no. 7 (2015): e0130140, https://doi.org/10.1371/journal.pone.0130140.

[31] A. Shrikumar, P. Greenside, and A. Kundaje, “Learning Important Features Through Propagating Activation Differences,” in International Conference on Machine Learning (PMlR, 2017), 3145-3153.

[32] K. Deb and H. Jain, “An Evolutionary Many-Objective Optimization Algorithm Using Reference-Point-Based Nondominated Sorting Approach,” IEEE Transactions on Evolutionary Computation 18, no. 4 (2014): 577-602, https://doi.org/10.1109/TEVC.2013.2281535.

[33] Bilal, M. Pant, H. Zaheer, L. Garcia-Hernandez, and A. Abraham, “ Differential Evolution: A Review of More Than Two Decades of Research,” Engineering Applications of Artificial Intelligence 90 (2020): 103479, https://doi.org/10.1016/j.engappai.2020.103479.

[34] X. Dong, M. Tan, A. W. Yu, D. Peng, B. Gabrys, and Q. V. Le, “AutoHAS: Differentiable Hyper-Parameter and Architecture Search,” arXiv preprint arXiv:2006.03656 4, no. 5 (2020), https://doi.org/10.48550/arXiv.2006.03656.

[35] K.-D. Lu, J.-C. Huang, G.-Q. Zeng, M.-R. Chen, G.-G. Geng, and J. Weng, “Multi-Objective Discrete Extremal Optimization of Variable-Length Blocks-Based CNN by Joint Nas and Hpo for Intrusion Detection in IIoT,” IEEE Transactions on Dependable and Secure Computing 22, no. 4 (2025): 4266-4283, https://doi.org/10.1109/tdsc.2025.3545363.

[36] G.-Q. Zeng, Y.-W. Yang, K.-D. Lu, G.-G. Geng, and J. Weng, “Evolutionary Adversarial Autoencoder for Unsupervised Anomaly Detection of Industrial Internet of Things,” IEEE Transactions on Reliability 74, no. 3 (2025): 3454-3468, https://doi.org/10.1109/tr.2025.3528256.

[37] X. Huang, M. Kwiatkowska, S. Wang, and M. Wu, “Safety Verification of Deep Neural Networks,” in International Conference on Computer Aided Verification (Springer, 2017), 3-29.

[38] I. J. Goodfellow, J. Shlens, and C. Szegedy , Explaining and Harnessing Adversarial Examples (ICLR, 2015), https://doi.org/10.48550/arXiv.1412.6572.

[39] X. Gao, R. K. Saha, M. R. Prasad, and A. Roychoudhury, “Fuzz Testing Based Data Augmentation to Improve Robustness of Deep Neural Networks,” in Proceedings of the ACM/IEEE 42nd International Conference on Software Engineering (ACM/IEEE, 2020), 1147-1158.

[40] E. Wong, L. Rice, and J. Z. Kolter, “Fast Is Better Than Free: Revisiting Adversarial Training,” in International Conference on Learning Representations (OpenReview.net, 2020), 1-17.

[41] C. Xie, J. Wang, Z. Zhang, Z. Ren, and A. Yuille, “Mitigating Adversarial Effects Through Randomization,” in International Conference on Learning Representations (OpenReview.net, 2018), 1-16.

[42] B. Sun, J. Sun, T. Dai, and L. Zhang, “Probabilistic Verification of Neural Networks Against Group Fairness,” in International Symposium on Formal Methods (Springer, 2021), 83-102.

[43] V. A. Dasu, A. Kumar, S. Tizpaz-Niari, and G. Tan, “NeuFair: Neural Network Fairness Repair With Dropout,” in Proceedings of the 33rd ACM SIGSOFT International Symposium on Software Testing and Analysis (ACM, 2024), 1541-1553.

[44] N. Srivastava, G. Hinton, A. Krizhevsky, I. Sutskever, and R. Salakhutdinov, “Dropout: A Simple Way to Prevent Neural Networks From Overfitting,” Journal of Machine Learning Research 15, no. 1 (2014): 1929-1958, https://dl.acm.org/doi/abs/10.5555/2627435.2670313.

[45] T. Gu, K. Liu, B. Dolan-Gavitt, and S. Garg, “BadNets: Evaluating Backdooring Attacks on Deep Neural Networks,” IEEE Access 7 (2019): 47230-47244, https://doi.org/10.1109/access.2019.2909068.

[46] D. Wu and Y. Wang, “Adversarial Neuron Pruning Purifies Backdoored Deep Models,” Advances in Neural Information Processing Systems 34 (2021): 16913-16925, https://doi.org/10.48550/arXiv.2110.14430.

[47] S. Chai and J. Chen, “One-Shot Neural Backdoor Erasing via Adversarial Weight Masking,” in Advances in Neural Information Processing Systems, ed. A. H. Oh, A. Agarwal, D. Belgrave, and K. Cho (Curran Associates, 2022), 1-15.

[48] Y. Zeng, S. Chen, W. Park, Z. Mao, M. Jin, and R. Jia, “Adversarial Unlearning of Backdoors via Implicit Hypergradient,” in International Conference on Learning Representations (OpenReview.net, 2022).

[49] R. Zhu, D. Tang, S. Tang, X. Wang, and H. Tang, “Selective Amnesia: On Efficient, High-Fidelity and Blind Suppression of Backdoor Effects in Trojaned Machine Learning Models,” in 2023 IEEE Symposium on Security and Privacy (SP) (IEEE, 2023), 1-19.

[50] L. H. Pham, J. Li, and J. Sun, “Socrates: Towards a Unified Platform for Neural Network Verification,” arXiv preprint arXiv:2007.11206 (2020), https://doi.org/10.48550/arXiv.2007.11206.

[51] L. S. Shapley, “A Value for n-Person Games,” Contributions to the Theory of Games 2 (1953): 307-317, https://doi.org/10.1017/CBO9780511528446.003.

[52] H. Chen, I. C. Covert, S. M. Lundberg, and S.-I. Lee, “Algorithms to Estimate Shapley Value Feature Attributions,” Nature Machine Intelligence 5, no. 6 (2023): 590-601, https://doi.org/10.1038/s42256-023-00657-x.

[53] S. M. Lundberg and S.-I. Lee, “A Unified Approach to Interpreting Model Predictions,” in Advances in Neural Information Processing Systems (Curran Associates, 2017), 4765-4774.

[54] S. Gupte and J. Paparrizos, “Understanding the Black Box: A Deep Empirical Dive Into Shapley Value Approximations for Tabular Data,” Proceedings of the ACM on Management of Data 3, no. 3 (2025): 1-31, https://doi.org/10.1145/3725420.

[55] S. M. Lundberg, G. Erion, H. Chen, et al., From Local Explanations to Global Understanding With Explainable AI for Trees,” Nature Machine Intelligence 2, no. 1 (2020): 56-67, https://doi.org/10.1038/s42256-019-0138-9.

[56] R. Jia, D. Dao, B. Wang, et al., Towards Efficient Data Valuation Based on the Shapley Value,” in The 22nd International Conference on Artificial Intelligence and Statistics (PMLR, 2019), 1167-1176.

[57] A. Ghorbani and J. Zou, “Data Shapley: Equitable Valuation of Data for Machine Learning,” in International Conference on Machine Learning (PMLR, 2019), 2242-2251.

[58] Y. Tian, X. Xiang, X. Zhang, R. Cheng, and Y. Jin, “Sampling Reference Points on the Pareto Fronts of Benchmark Multi-Objective Optimization Problems,” in 2018 IEEE Congress on Evolutionary Computation (CEC) (IEEE, 2018), 1-6, https://doi.org/10.1109/CEC.2018.8477730.

[59] Y. LeCun, C. Cortes, and C. Burges, “MNIST Handwritten Digit Database,” AT Labs [Online] (2010), http://yann.lecun.com/exdb/mnist/.

[60] Y. LeCun, L. Bottou, Y. Bengio, and P. Haffner, “Gradient-Based Learning Applied to Document Recognition,” in Proceedings of the IEEE, vol. 86, no. 11 (1998), 2278-2324, https://doi.org/10.1109/5.726791.

[61] A. Krizhevsky , Learning Multiple Layers of Features From Tiny Images (Technical Report, University of Toronto, 2009).

[62] K. He, X. Zhang, S. Ren, and J. Sun, “Deep Residual Learning for Image Recognition,” in Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition (IEEE, 2016), 770-778.

[63] H. Xiao, K. Rasul, and R. Vollgraf, “Fashion-MNIST: A Novel Image Dataset for Benchmarking Machine Learning Algorithms,” arXiv preprint arXiv:1708.07747 [Online] (2017), https://github.com/zalandoresearch/fashion-mnist.

[64] J. Stallkamp, M. Schlipsing, J. Salmen, and C. Igel, “The German Traffic Sign Recognition Benchmark: A Multi-Class Classification Competition,” in The 2011 International Joint Conference on Neural Networks [Online] (2011), 1453-1460, https://benchmark.ini.rub.de/gtsrb_dataset.html.

[65] K. D. Julian, M. J. Kochenderfer, and M. P. Owen, “Deep Neural Network Compression for Aircraft Collision Avoidance Systems,” Journal of Guidance, Control, and Dynamics 42, no. 3 (2019): 598-608, https://doi.org/10.2514/1.g003724.

[66] G. B. Huang, M. Mattar, T. Berg, and E. Learned-Miller, “Labeled Faces in the Wild: A Database for Studying Face Recognition in Unconstrained Environments,” in Workshop on Faces in ‘Real-Life’ Images: Detection, Alignment, and Recognition [Online] (2008), http://vis-www.cs.umass.edu/lfw.

[67] J. Blank and K. Deb, “Pymoo: Multi-Objective Optimization in Python,” IEEE Access 8 (2020): 89497-89509, https://doi.org/10.1109/access.2020.2990567.

[68] K. Deb, A. Pratap, S. Agarwal, and T. Meyarivan , “A Fast and Elitist Multiobjective Genetic Algorithm: NSGA-II,”IEEE Transactions on Evolutionary Computation 6, no. 2 (2002): 182-197, https://doi.org/10.1109/4235.996017.

[69] Q. Zhang and H. Li, “MOEA/D: A Multiobjective Evolutionary Algorithm Based on Decomposition,” IEEE Transactions on Evolutionary Computation 11, no. 6 (2007): 712-731, https://doi.org/10.1109/tevc.2007.892759.

[70] O. Sener and V. Koltun, “Multi-Task Learning as Multi-Objective Optimization,” Advances in Neural Information Processing Systems 31 (2018): 1-12, https://doi.org/10.48550/arXiv.1810.04650.

[71] X. Lin, H.-L. Zhen, Z. Li, Q.-F. Zhang, and S. Kwong, “Pareto Multi-Task Learning,” Advances in Neural Information Processing Systems 32 (2019): 1-11, https://dl.acm.org/doi/10.5555/3454287.3455367.