NCCDA: Neuron-wise Class-Conditional Distribution Alignment for Deep Neural Network Repair

Abstract

Neural network repair aims to correct prediction failures caused by multiple security threats—such as backdoor attacks, natural corruptions, and safety property violations—through limited adjustments to model parameters. However, most existing repair methods rely on single-sample, point-to-point correction strategies, overlooking the statistical regularities of the feature space. As a result, they are highly sensitive to the scale of faulty samples and struggle to simultaneously achieve repair generalization and original performance preservation under small-sample settings. To address these limitations, we propose a novel general neural network repair paradigm termed NCCDA (Neuron-wise Class-Conditional Distribution Alignment). The method is grounded in a key insight: prediction failures fundamentally arise from neuron-level internal representations deviating from the high-likelihood regions corresponding to their true classes. NCCDA constructs neuron-wise class-conditional distribution references and formulates the repair process as a joint optimization of distribution alignment and structure preservation. By guiding abnormal representations back to high-likelihood regions while anchoring the structure of normal samples, the method enables efficient and adaptive repair without explicit neuron localization. We theoretically prove a generalization error bound under small-sample settings based on Rademacher complexity, providing formal guarantees. Extensive experiments across 7 benchmark datasets and 38 models, covering three categories of repair tasks, demonstrate that NCCDA consistently outperforms existing methods in repair effectiveness, generalization repair capability (Gene), and original accuracy preservation.