[Paper Review] Self-Healing Robust Neural Networks via Closed-Loop Control

Zhuotong Chen, Qianxiao Li+1 more|arXiv (Cornell University)|Jun 26, 2022

Adversarial Robustness in Machine Learning4 citations

TL;DR

This paper proposes a post-training, closed-loop control framework for self-healing neural networks that automatically detects and corrects errors caused by unforeseen perturbations or attacks without prior knowledge of the attack type. By leveraging Pontryagin’s Maximum Principle and margin-based analysis, the method enhances robustness across diverse perturbation types (ℓ₁, ℓ₂, ℓ∞) with consistent accuracy improvements of over 60% against AutoAttack on CIFAR-10.

ABSTRACT

Despite the wide applications of neural networks, there have been increasing concerns about their vulnerability issue. While numerous attack and defense techniques have been developed, this work investigates the robustness issue from a new angle: can we design a self-healing neural network that can automatically detect and fix the vulnerability issue by itself? A typical self-healing mechanism is the immune system of a human body. This biology-inspired idea has been used in many engineering designs but is rarely investigated in deep learning. This paper considers the post-training self-healing of a neural network, and proposes a closed-loop control formulation to automatically detect and fix the errors caused by various attacks or perturbations. We provide a margin-based analysis to explain how this formulation can improve the robustness of a classifier. To speed up the inference of the proposed self-healing network, we solve the control problem via improving the Pontryagin Maximum Principle-based solver. Lastly, we present an error estimation of the proposed framework for neural networks with nonlinear activation functions. We validate the performance on several network architectures against various perturbations. Since the self-healing method does not need a-priori information about data perturbations/attacks, it can handle a broad class of unforeseen perturbations.

Motivation & Objective

• Address the critical challenge of neural network vulnerability to unforeseen data perturbations and hardware faults in real-world deployment.
• Overcome limitations of existing defense methods that rely on prior assumptions about attack types or data distributions.
• Enable post-training robustness enhancement through autonomous error detection and correction using control-theoretic principles.
• Design a self-healing mechanism inspired by biological immune systems, applicable to both standard and robustly trained models.
• Ensure broad applicability across diverse, unknown perturbations by avoiding reliance on specific attack models or data characteristics.

Proposed method

• Formulate a closed-loop control system where embedding functions at input and hidden layers detect potential errors in real time.
• Implement a control process that adjusts neuron activations via a control loss function derived from Pontryagin’s Maximum Principle to correct detected errors.
• Use margin-based analysis to theoretically justify that the control mechanism increases the decision boundary margin, enhancing robustness.
• Speed up inference by improving the numerical solver for the optimal control problem using the method of successive approximations.
• Design embedding functions as denoising autoencoders trained to project inputs onto robust manifolds that preserve classification under perturbations.
• Integrate the control mechanism into pre-trained models without retraining, enabling post-hoc robustness improvement.

Experimental results

Research questions

• RQ1 Can a closed-loop control framework enable post-training self-healing in neural networks to correct errors caused by unknown or unforeseen perturbations?
• RQ2 How does the proposed control mechanism improve robustness without requiring prior knowledge of attack types or data distributions?
• RQ3 To what extent can the margin-based control formulation enhance the generalization and robustness of a classifier across different perturbation norms?
• RQ4 How effective is the method in a white-box setting where the attacker has full knowledge of the control mechanism?
• RQ5 Can the integration of optimal control via Pontryagin’s Maximum Principle be efficiently solved in inference time to maintain practical deployment feasibility?

Key findings

• The proposed self-healing framework improves robustness against AutoAttack with ℓ∞, ℓ₂, and ℓ₁ perturbations, achieving over 60% accuracy on CIFAR-10 for standard models and over 50% for robustly trained models.
• On ResNet-18, the method achieves 87.04% accuracy under ℓ∞-PGD attack (ε=8/255) and 73.04% under ℓ₁-attack (ε=12), significantly outperforming the baseline of 0% accuracy with no control.
• In the white-box setting, where the attacker knows the control mechanism, the method still improves average accuracy by 0.5%–1% over adversarially trained baselines under AutoAttack.
• The use of 2-layer denoising autoencoders as embedding functions enables effective manifold learning, allowing the control system to detect and correct perturbations in hidden representations.
• The method maintains high performance across multiple architectures (ResNet-18, -34, -50, WRN-28-8, WRN-34-8), demonstrating broad applicability and generalization.
• Theoretical analysis confirms that the control loss increases the margin of the classifier, providing a principled explanation for improved robustness.