Adversarial Feature Alignment: Balancing Robustness and Accuracy in Deep Learning via Adversarial Training

Deep learning models are continually improving in accuracy, but they remain vulnerable to adversarial attacks, often resulting in the misclassification of adversarial examples. Adversarial training can mitigate this problem by enhancing the model's robust accuracy on adversarial examples. However, such training typically compromises the model's standard accuracy on clean samples. The necessity for deep learning models to balance both robustness and accuracy for security is evident, but achieving this balance remains challenging, and the underlying reasons are yet to be clarified. This paper proposes an innovative pre-training method called Adversarial Feature Alignment (AFA) to address these problems. Our approach involves identifying the trade-off in the model's feature space and fine-tuning it to achieve accuracy on both standard and adversarial examples concurrently. Our research unveils an intriguing insight: misalignment within the feature space often leads to misclassification, regardless of whether the samples are benign or adversarial. AFA mitigates this risk by employing a novel optimization algorithm based on contrastive learning to alleviate potential feature misalignment. Through our evaluations, we demonstrate the superior performance of AFA. Our method delivers state-of-the-art robust accuracy while minimizing the drop in clean accuracy to 1.86% and 8.91% on CIFAR10 and CIFAR100, respectively, compared to cross-entropy. We also show that joint optimization of AFA and TRADES, accompanied by data augmentation using a recent diffusion model, achieves state-of-the-art accuracy and robustness. Through AFA, we expect to enhance security while preserving accuracy in deep learning models through adversarial training.