Discrete Representations Strengthen Vision Transformer Robustness

Abstract

Vision Transformer (ViT) is emerging as the state-of-the-art architecture for image recognition. While recent studies suggest that ViTs are more robust than their convolutional counterparts, our experiments find that ViTs trained on ImageNet are overly reliant on local textures and fail to make adequate use of shape information. ViTs thus have difficulties generalizing to out-of-distribution, real-world data. To address this deficiency, we present a simple and effective architecture modification to ViT's input layer by adding discrete tokens produced by a vector-quantized encoder. Different from the standard continuous pixel tokens, discrete tokens are invariant under small perturbations and contain less information individually, which promote ViTs to learn global information that is invariant. Experimental results demonstrate that adding discrete representation on four architecture variants strengthens ViT robustness by up to 12% across seven ImageNet robustness benchmarks while maintaining the performance on ImageNet.

References in corpus (10)

• MLP-Mixer: An all-MLP Architecture for Vision
• Zero-Shot Text-to-Image Generation
• BEiT: BERT Pre-Training of Image Transformers
• Do ImageNet Classifiers Generalize to ImageNet?
• Intriguing Properties of Vision Transformers
• When Vision Transformers Outperform ResNets without Pre-training or Strong Data Augmentations
• Are we done with ImageNet?
• How to train your ViT? Data, Augmentation, and Regularization in Vision Transformers
• Self-Challenging Improves Cross-Domain Generalization
• Towards Robust Vision Transformer