Rethinking robust generalization with energy minimization.

Abstract

Adversarial training has emerged as a leading strategy for enhancing the robustness of machine learning models against adversarial attacks. Its effectiveness often wanes when faced with unseen adversarial examples, resulting in suboptimal robust generalization. To address this issue, we introduce a novel energy-based optimization strategy to improve the robust generalization by incorporating the principles of energy-based models. Our framework models the energy of natural and adversarial examples, where natural samples are assigned to lower energy and adversarial samples to higher energy. During the inference phase, the influence of adversarial perturbation can be alleviated by energy minimization. Theoretically, we show that the proposed energy-based optimization strategy yields a tighter robust-generalization bound through an explicit energy-discrepancy term; this analysis provides an explanatory bound and should not be interpreted as certified robustness. Empirically, a series of evaluations provide evidence for the efficacy of the proposed methodology under the specified threat models and evaluation protocols, showing strong and competitive performance across three extensively utilized datasets. Specifically, EM-AT achieves 77.71% standard-AA robustness on CIFAR-10 and remains highly competitive under comparable lightweight settings. The source codes are available at https://github.com/LitterQ/EM-AT.