Incorruptible Neural Networks: Training Models that can Generalize to Large Internal Perturbations

Flat regions of the neural network loss landscape have long been hypothesized to correlate with better generalization properties. A closely related but distinct problem is training models that are robust to internal perturbations to their weights, which may be an important need for future low-power hardware platforms. In this paper, we explore the usage of two methods, sharpness-aware minimization (SAM) and random-weight perturbation (RWP), to find minima robust to a variety of random corruptions to weights. We consider the problem from two angles: generalization (how do we reduce the noise-robust generalization gap) and optimization (how do we maximize performance from optimizers when subject to strong perturbations). First, we establish, both theoretically and empirically, that an over-regularized RWP training objective is optimal for noise-robust generalization. For small-magnitude noise, we find that SAM’s adversarial objective further improves performance over any RWP configuration, but performs poorly for large-magnitude noise. We link the cause of this to a vanishing-gradient effect, caused by unevenness in the loss landscape, affecting both SAM and RWP. Lastly, we demonstrate that dynamically adjusting the perturbation strength to match the evolution of the loss landscape improves optimizing for these perturbed objectives.

💡 Research Summary

The paper investigates how to train neural networks that remain accurate when their weights are perturbed, a scenario increasingly relevant for low‑power analog hardware such as in‑memory computing accelerators. Two prominent techniques for finding flat minima are examined: Sharpness‑Aware Minimization (SAM), which performs an adversarial ascent step within an ℓ₂‑ball of radius ρ before descent, and Random‑Weight Perturbation (RWP), which adds isotropic Gaussian noise scaled by the maximum weight magnitude during training.

From a theoretical standpoint, the authors derive a PAC‑Bayes generalization bound (Theorem 4.1) that relates the expected test loss under test‑time noise σ_test to the expected training loss under training‑time noise σ_train and a complexity term proportional to ‖w‖²/σ_train². The bound shows an anti‑correlation: increasing σ_train raises the training loss term but reduces the complexity term, implying that “over‑regularizing” (σ_train > σ_test) can yield a tighter bound and therefore better generalization.

Empirically, the authors train ResNet‑18 on CIFAR‑100, Tiny‑ImageNet, and ImageNet‑100 with a range of σ_train values. Results confirm the theory: for each σ_test, the best test accuracy is achieved when σ_train exceeds σ_test, and the advantage grows with larger σ_test. Under‑regularization (σ_train < σ_test) never outperforms the matched case, while over‑regularization consistently improves robustness to weight noise.

When applying SAM, the picture changes. For small noise magnitudes (ρ ≤ 0.05) SAM outperforms any RWP configuration, delivering higher noisy‑test accuracy. However, as the perturbation radius grows, the loss surface becomes locally flat in the ascent direction, causing the gradient of the inner maximization to vanish. This “vanishing‑gradient effect” stalls SAM’s inner ascent step, leading to poor convergence and degraded robustness for large σ_test. RWP also suffers when σ_train is too large, but its random direction of perturbation mitigates the gradient collapse to a lesser extent.

To address these issues, the authors propose a dynamic perturbation schedule: start training with a small σ_train (or ρ) and gradually increase it as training progresses, aligning the perturbation strength with the evolving loss landscape. This schedule preserves meaningful gradients early on, allowing the optimizer to locate a flat region, and later introduces stronger noise to enforce robustness. Experiments show that both SAM and RWP benefit from the schedule, achieving 1.2–2.5 % higher noisy‑test accuracy compared to static settings. Moreover, simulations of analog in‑memory computing hardware demonstrate that dynamically‑trained models retain accuracy while being tolerant to realistic conductance variations, offering a practical path toward energy‑efficient, error‑robust deployment.

In summary, the paper makes four key contributions: (1) a novel PAC‑Bayes bound that predicts over‑regularized RWP yields tighter generalization guarantees; (2) empirical validation that over‑regularized RWP indeed finds more noise‑robust minima; (3) identification of a vanishing‑gradient limitation of SAM under large weight perturbations and demonstration that SAM still excels for small‑noise regimes; (4) a dynamic perturbation schedule that mitigates gradient collapse and improves convergence for both methods, with corroborating hardware‑level experiments. The findings suggest that matching training noise to test noise is not optimal; instead, deliberately training with stronger noise and adapting its magnitude over time yields models that are both flat and resilient, a valuable insight for future analog AI hardware design.