How to train your ViT? Data, Augmentation, and Regularization in Vision Transformers

By Andreas Steiner, Alexander Kolesnikov, Xiaohua Zhai, Ross Wightman, Jakob Uszkoreit, Lucas Beyer. 2021

854 citations recorded by the source corpus.

Vision Transformers (ViT) have been shown to attain highly competitive performance for a wide range of vision applications, such as image classification, object detection and semantic image segmentation. In comparison to convolutional neural networks, the Vision Transformer's weaker inductive bias is generally found to cause an increased reliance on model regularization or data augmentation ("AugReg"for short) when training on smaller training datasets. We conduct a systematic empirical study in order to better understand the interplay between the amount of training data, AugReg, model size and compute budget. As one result of this study we find that the combination of increased compute and AugReg can yield models with the same performance as models trained on an order of magnitude more training data: we train ViT models of various sizes on the public ImageNet-21k dataset which either match or outperform their counterparts trained on the larger, but not publicly available JFT-300M dataset.

Subfields: Artificial Intelligence, Computer Vision and Pattern Recognition, Radiology, Nuclear Medicine and Imaging

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Artificial Intelligence

In Computer Vision and Pattern Recognition

In Radiology, Nuclear Medicine and Imaging

Citations in the corpus

Scalable Diffusion Models with Transformers (2022, neutral): “Unlike much prior work with ViTs [57, 61], we did not ﬁnd learning rate warmup nor regularization necessary to train DiTs to high performance.”
LiT: Zero-Shot Transfer with Locked-image text Tuning (2021, methodology): “Specifically, we tune an ImageNet-21k pre-trained ViTL/16 model [54] on the union of the YFCC100m-CLIP and CC12M datasets.”
The Efficiency Misnomer (2021, background): “Values for accuracies and cost indicators are obtained from (Steiner et al., 2021; Liu et al., 2021).”
A Simple Baseline that Questions the Use of Pretrained-Models in Continual Learning (2022, methodology): “The use of large-scale datasets, such as ImageNet-1k[23] and ImageNet-21k [22] and pretraining [24] has led to significant advances.”
Segmenter: Transformer for Semantic Segmentation (2021, methodology): “The original ViT models [19] have been trained with random cropping only, whereas the training procedure proposed by [47] uses a combination of dropout [46] and stochastic depth [30] as regularization and Mixup [62] and RandAugment [15] as data augmentations.”
A ConvNet for the 2020s (2022, methodology): “(T)ViT results with 90-epoch AugReg [67] training, provided through personal communication with the authors.”
Simple Open-Vocabulary Object Detection with Vision Transformers (2022, background): “It is known that large Transformers trained on datasets of this size (such as ImageNet-1k) require carefully-tuned regularization and data augmentation to perform well [39,38,3].”
OpenEarthMap: A Benchmark Dataset for Global High-Resolution Land Cover Mapping (2022, background): “Moreover, data augmentation, regularization, and tuning of hyper-parameters still need to be explored when training on limited training data [40].”
Run, Don't Walk: Chasing Higher FLOPS for Faster Neural Networks (2023, background): “Many follow-up works have attempted to improve ViT in terms of training setting [51, 58, 59] and model design [12,34,35,65,77].”
Dual PatchNorm (2023, methodology): “The AugReg recipe provides the optimal mixup regularization (Zhang et al., 2017) and RandAugment (Cubuk et al., 2020) for each ViT backbone.”
Visformer: The Vision-friendly Transformer (2021, positive): “[37] further study data, augmentation and regularization in vision Transformer and finds that augmentation can yield the same performance as that trained on an order of magnitude more data.”
Token Merging: Your ViT But Faster (2022, methodology): “Both AugReg (Steiner et al., 2022) and SWAG (Singh et al., 2022) are ViT models pretrained on a large supervised (or weakly supervised) pretraining dataset and fine-tuned on ImageNet-1k.”
TiC-CLIP: Continual Training of CLIP Models (2023, positive): “When training with large batches, linear learning rate warm-up is typically employed to stabilize the start of the training when beginning from a random initialization (Goyal et al., 2017; Steiner et al., 2021).”
DINOv2: Learning Robust Visual Features without Supervision (2023, methodology): “For reference, using the Mask2Former pipeline (Steiner et al., 2021) with a ViT-Adapter (Chen et al., 2023b) on top of our frozen ViT-g/14 backbone gives 60.2 mIoU on ADE-20k.”
Vision Transformer Adapter for Dense Predictions (2022, methodology): “Using the ImageNet-22K pre-trained weights [65], our ViT-Adapter-B achieves 51.”
Meta Knowledge Distillation (2022, methodology): “Training large ViT models is nontrivial (Steiner et al., 2021), and previous works proposed different techniques (Touvron et al., 2021a; He et al., 2021) to achieved good results.”
muNet: Evolving Pretrained Deep Neural Networks into Scalable Auto-tuning Multitask Systems (2022, background): “This constrains hyperparameters to be changed incrementally and biases the search toward the initial values of the root model, that in our application are the result of an extensive study aimed to identify a fine-tuning configuration for ViT models that is as generic as possible [Steiner et al., 202”