Pyramid Adversarial Training Improves ViT Performance
This work addresses a key limitation in adversarial training for ViTs, offering a method that enhances both accuracy and robustness without extra data, which is significant for improving reliability in vision tasks.
The paper tackles the problem of adversarial training often reducing clean accuracy in Vision Transformers (ViT), and presents Pyramid Adversarial Training (PyramidAT) which improves both in-distribution accuracy and out-of-distribution robustness, achieving a 1.82% absolute improvement on ImageNet clean accuracy and setting new state-of-the-art results on robustness benchmarks like ImageNet-C (41.42 mCE).
Aggressive data augmentation is a key component of the strong generalization capabilities of Vision Transformer (ViT). One such data augmentation technique is adversarial training (AT); however, many prior works have shown that this often results in poor clean accuracy. In this work, we present pyramid adversarial training (PyramidAT), a simple and effective technique to improve ViT's overall performance. We pair it with a "matched" Dropout and stochastic depth regularization, which adopts the same Dropout and stochastic depth configuration for the clean and adversarial samples. Similar to the improvements on CNNs by AdvProp (not directly applicable to ViT), our pyramid adversarial training breaks the trade-off between in-distribution accuracy and out-of-distribution robustness for ViT and related architectures. It leads to 1.82% absolute improvement on ImageNet clean accuracy for the ViT-B model when trained only on ImageNet-1K data, while simultaneously boosting performance on 7 ImageNet robustness metrics, by absolute numbers ranging from 1.76% to 15.68%. We set a new state-of-the-art for ImageNet-C (41.42 mCE), ImageNet-R (53.92%), and ImageNet-Sketch (41.04%) without extra data, using only the ViT-B/16 backbone and our pyramid adversarial training. Our code is publicly available at pyramidat.github.io.