[Paper Review] A Light Recipe to Train Robust Vision Transformers

Edoardo Debenedetti, Vikash Sehwag+1 more|arXiv (Cornell University)|Sep 15, 2022

Adversarial Robustness in Machine Learning4 citations

TL;DR

This paper proposes a lightweight training recipe that significantly improves the adversarial robustness of Vision Transformers (ViTs) by replacing heavy data augmentation with stronger weight decay and epsilon-warmup, achieving state-of-the-art robust accuracy on ImageNet-1k and downstream tasks, outperforming even ResNet-50 under adversarial training.

ABSTRACT

In this paper, we ask whether Vision Transformers (ViTs) can serve as an underlying architecture for improving the adversarial robustness of machine learning models against evasion attacks. While earlier works have focused on improving Convolutional Neural Networks, we show that also ViTs are highly suitable for adversarial training to achieve competitive performance. We achieve this objective using a custom adversarial training recipe, discovered using rigorous ablation studies on a subset of the ImageNet dataset. The canonical training recipe for ViTs recommends strong data augmentation, in part to compensate for the lack of vision inductive bias of attention modules, when compared to convolutions. We show that this recipe achieves suboptimal performance when used for adversarial training. In contrast, we find that omitting all heavy data augmentation, and adding some additional bag-of-tricks ($\varepsilon$-warmup and larger weight decay), significantly boosts the performance of robust ViTs. We show that our recipe generalizes to different classes of ViT architectures and large-scale models on full ImageNet-1k. Additionally, investigating the reasons for the robustness of our models, we show that it is easier to generate strong attacks during training when using our recipe and that this leads to better robustness at test time. Finally, we further study one consequence of adversarial training by proposing a way to quantify the semantic nature of adversarial perturbations and highlight its correlation with the robustness of the model. Overall, we recommend that the community should avoid translating the canonical training recipes in ViTs to robust training and rethink common training choices in the context of adversarial training.

Motivation & Objective

• To investigate whether Vision Transformers (ViTs) can achieve superior adversarial robustness compared to CNNs when trained with adversarial training.
• To identify a training recipe tailored specifically for adversarial training of ViTs, distinct from canonical ViT training recipes.
• To understand why standard ViT training recipes—especially heavy data augmentation—harm robustness during adversarial training.
• To quantify the semantic nature of adversarial perturbations and link it to model robustness.
• To demonstrate that the proposed recipe generalizes across ViT architectures and enables efficient fine-tuning on smaller datasets.

Proposed method

• Conducting systematic ablation studies on a subset of ImageNet-1k to evaluate the impact of data augmentation, weight decay, and adversarial training hyperparameters on ViT robustness.
• Replacing strong data augmentation (e.g., mixup, CutMix) with minimal augmentation to improve robustness, contrary to canonical ViT training practices.
• Introducing epsilon-warmup during adversarial training to gradually increase the perturbation magnitude, stabilizing training.
• Applying larger weight decay values to improve generalization and robustness, especially in the absence of strong data augmentation.
• Using the Cross-Covariance Image Transformer (XCiT) as the primary architecture to validate the recipe, then extending to other ViT variants and ConvNeXt.
• Quantifying adversarial perturbations via gradient visualization and semantic alignment metrics to analyze their perceptual plausibility.

Experimental results

Research questions

• RQ1 Does a ViT trained with a custom adversarial training recipe outperform standard CNNs in adversarial robustness?
• RQ2 Why does heavy data augmentation, standard in ViT training, degrade performance in adversarial training?
• RQ3 Can a simpler data augmentation strategy combined with hyperparameter tuning (e.g., weight decay, epsilon-warmup) significantly improve ViT robustness?
• RQ4 How do the semantic properties of adversarial perturbations differ between robust ViTs and standard models like ResNet?
• RQ5 Does the proposed recipe generalize across different ViT architectures and downstream fine-tuning tasks?

Key findings

• The proposed recipe improves robust accuracy on ImageNet-1k by 13.1% compared to the canonical ViT training recipe.
• The recipe achieves state-of-the-art robust accuracy on ImageNet-1k and up to 5.5% higher robust accuracy on downstream fine-tuning tasks (e.g., CIFAR-100, Caltech-101).
• Omitting heavy data augmentation while increasing weight decay and using epsilon-warmup leads to better robustness without sacrificing clean accuracy.
• Models trained with the new recipe exhibit more effective 1-step attacks during training, indicating stronger adversarial signal learning.
• Adversarial perturbations for the robust XCiT model are more semantically aligned with human perception than those for ResNet, suggesting improved robustness through more meaningful perturbations.
• The recipe generalizes to other ViT architectures and even to ConvNeXt, enabling state-of-the-art performance with minimal changes.