Source-linked AI summary
How to train your ViT? Data, Augmentation, and Regularization in Vision Transformers
Andreas Steiner, Alexander Kolesnikov, Xiaohua Zhai, Ross Wightman, Jakob Uszkoreit, Lucas Beyer
TL;DR
ViT training lacked comprehensive evidence on the trade-offs among data size, AugReg, model size, and compute. The paper conducts a controlled empirical study across these factors and transfer settings, finding that selected AugReg can roughly substitute for a tenfold increase in training data, while comparable performance still requires similar compute.
Problem
The study addresses the missing comprehensive evidence on trade-offs among ViT regularization, augmentation, training-data size, model size, and compute.
Method
The authors systematically pre-train ViTs and hybrids across datasets, model sizes, AugReg settings, and compute budgets, then evaluate the resulting models through transfer learning.
Results
10x training-data equivalence: selected regularization and augmentation roughly match the accuracy of a tenfold larger training dataset, while similar performance requires roughly the same compute.
Takeaways & Limitations
Among similarly performing pre-trained models, the study suggests preferring more training data over more augmentation for transfer learning.
Takeaways & Limitations
The study restricts itself to the default ViT architecture and excludes ResNets and newer ViT variants.
Abstract
from arXiv · showhide
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.
1 Introduction
The study addresses the missing systematic evidence on how data size, AugReg, model size, and compute interact in ViT training. It uses controlled experiments to derive practitioner-oriented insights, including a roughly 10x data-equivalence effect from selected AugReg.
- ViTs often need large datasets or strong AugReg because they lack CNNs’ translational equivariance.
- The paper pre-trains diverse ViTs across dataset sizes, model configurations, and regularization and augmentation settings, then evaluates transfer learning.
- Consistent training and evaluation reduce noise from uncontrolled augmentation, optimization, preprocessing, and dataset choices.
- 10x training-data equivalence: carefully selected regularization and augmentations roughly match the accuracy gains of increasing training data tenfold.
- Similar performance from more data or better AugReg requires roughly the same compute expenditure.
- The study also compares transfer learning, training from scratch, and compute-versus-performance trade-offs for practitioners with limited budgets.
2 Scope of the study
The study distinguishes several ways to define training efficiency and emphasizes practitioner’s cost when pre-training is available or amortized. It also recognizes deployment cost as a separate efficiency perspective.
- Overall training cost can include both pre-training and subsequent fine-tuning, with pre-training often dominating by orders of magnitude.
- Practitioners commonly use publicly available parameters, making fine-tuning, adaptation, or task-specific training the more relevant cost.
- Training costs may be negligible when amortized over massive deployment, shifting attention toward eventual inference cost.
- The paper mainly studies practitioner’s cost: time and compute for finding an adaptation strategy or tuning a from-scratch setup when pre-training is free or amortized.
- The study touches all three cost viewpoints while concentrating on the practitioner-oriented setting.
3 Experimental setup
The experiments use a unified, reproducible setup spanning ImageNet pre-training, multiple ViT and hybrid architectures, AugReg configurations, and transfer-learning evaluations. Training schedules also enable comparison of dataset size under roughly constant pre-training compute.
- Unified setup: A single JAX/Flax codebase, TPUs, TensorFlow Datasets, and a separate V100-based inference measurement setup provide a unified experimental framework.
- Datasets and metrics: Pre-training uses ImageNet-1k with about 1.3M images and ImageNet-21k with approximately 14M images and about 21,000 categories.
- Datasets and metrics: Transfer evaluation covers natural, specialized, and geometric vision tasks through VTAB datasets including CIFAR-100, Pets37, Resisc45, and Kitti-distance.
- Datasets and metrics: Top-1 classification accuracy is the main metric, with validation splits used for hyper-parameter selection and test splits for final reporting.
- Models: The model suite spans ViT-Ti, ViT-S, ViT-B, ViT-L, and hybrids that feed ResNet spatial features into ViT patch embeddings.
- Regularization and augmentation: AugReg sweeps dropout, stochastic depth, Mixup, RandAugment, and two weight-decay values across 28 configurations.
- Pre-training: ImageNet-1k uses 300 epochs, while ImageNet-21k uses 30 or 300 epochs, allowing increased dataset size to be examined at roughly constant pre-training compute.
4 Findings
The study finds that data scale, AugReg, compute, and model choices jointly shape ViT performance and transferability. More data or carefully selected AugReg can improve results, but benefits depend on compute, model capacity, dataset size, and validation strategy.
- 4.1 Scaling datasets with AugReg and compute: AugReg can match the accuracy of training on roughly 10x more data, while achieving comparable performance through either route requires roughly similar compute.AugReg ImageNet-1k models perform about equally to plain ImageNet-21k models; AugReg ImageNet-21k models with increased compute match or outperform reported JFT-300M results.
- 4.2 Transfer is the better option: Pre-trained models are generally more cost-efficient and accurate than training ViTs from scratch on practical-sized downstream datasets.On tiny Pet37, scratch training remains far below transferred models regardless of training time; on Resisc45, even 100x more compute and extensive search may not reach them.
- 4.3 More data yields more generic models: More upstream data produces more generic models: ImageNet-21k pretraining significantly outperforms ImageNet-1k across natural, specialized, and structured VTAB tasks under the same compute budget.Longer ImageNet-21k schedules further improve most tasks, though gains are small on nearly solved tasks.
- 4.4 Prefer augmentation to regularization: AugReg helps on ImageNet-1k, but on ImageNet-21k it generally hurts at fixed compute and continues hurting smaller models even when compute increases.The study reports more cases where augmentation helps than where regularization helps, with regularization almost always harmful on ImageNet-21k.
- 4.5 Choosing which pre-trained model to transfer: Selecting one pre-trained model by upstream validation is usually as effective as adapting all models, while adapting additional checkpoints can help in some cases.The cheaper strategy is generally equally effective, but extra compute may improve adaptation performance for notable outliers.
- 4.5 Choosing which pre-trained model to transfer: ImageNet-21k-to-ImageNet-1k transfer requires ImageNetV2 validation because minival scores can be biased by training-data overlap and memorization.The issue is especially severe for large models trained with long schedules; the authors did not observe similar problems on other datasets.
- 4.6 Prefer increasing patch-size to shrinking model-size: For similarly fast small models, larger /32 patch sizes significantly outperform thinner Tiny variants because patch size contributes to capacity beyond parameter count.The paper notes that parameter count reflects neither speed nor capacity reliably.
5 Related work
Prior ViT work established strong performance with large datasets and addressed overfitting with augmentation and regularization, but did not comprehensively study their trade-offs. Related research also explores architectural inductive biases and self-supervised objectives, which this paper does not systematically compare.
- ViT data and AugReg studies: ViTs achieved competitive performance with comparatively large datasets, while separate work used strong regularization and augmentation to address overfitting on ImageNet-1k.Neither line of work analyzed stronger augmentation or regularization together with larger datasets.
- ViT data and AugReg studies: Earlier studies used ImageNet-21k for selected experiments, but did not thoroughly investigate its combined use with AugReg.This paper positions its systematic comparison as addressing that unresolved combination.
- Architectural inductive biases: Other ViT variants introduce convolutional backbones, hierarchical structures, or convolution-like initialization to add architectural inductive biases.These approaches are distinct from the paper’s focus on pretraining data, AugReg, and compute.
- Self-supervised learning: Self-supervised ViT methods use objectives such as reconstructing perturbed patches or related representation-learning procedures, but systematic comparison with supervised pretraining is left for future work.The paper treats self-supervised learning as an orthogonal approach to overfitting and transfer performance.
6 Discussion
The study acknowledges substantial compute use and limits its scope to the default ViT architecture. It aims to reduce future search costs through starting points and off-the-shelf checkpoints.
- The study uses substantial compute, but its stated aim is to reduce the need for repeated extensive searches through reusable checkpoints.The authors frame the checkpoints as practical starting points rather than encouraging every practitioner to reproduce the full study.
- The experiments restrict analysis to the default ViT architecture and exclude both ResNets and newer ViT variants.The authors anticipate that many findings may extend to other ViT-based architectures.
7 Summary of recommendations
The recommendations favor varied upstream data, carefully tuned AugReg, and upstream validation for selecting transfer models. They caution that AugReg choices depend on model capacity and training schedule.
- Prefer checkpoints pre-trained on more upstream data because varied data yields more widely applicable models.ImageNet-1k validation accuracy can be inflated when pre-training uses ImageNet-1k itself.
- Judicious AugReg improves models at a fixed dataset size, but its best settings vary with model capacity and training schedule.Applying AugReg to models that are too small or training runs that are too short can deteriorate model quality.
- For transfer, select the model with the best upstream validation performance when evaluating all checkpoints is impractical.Fine-tuning every checkpoint may yield slightly better results in some scenarios, but model selection requires care around ImageNet-1k and ImageNet-21k overlap.
8 Conclusion
The paper systematically studies how regularization, augmentation, model size, data size, and compute interact in ViT pre-training and transfer learning. It finds complex, setting-dependent AugReg effects, while transfer remains effective across diverse datasets.
- The study systematically evaluates regularization, augmentation, model size, training data size, compute requirements, and transfer learning for ViTs.Its experiments characterize a complex landscape of pre-training settings and identify when augmentation and regularization help or hurt.
- Augmentation and regularization produce setting-dependent benefits, with their usefulness varying across model sizes and training conditions.The paper emphasizes that its experiments reveal both beneficial and unfavorable situations rather than a single universal recipe.
- Transfer learning remains the best available option across a wide range of datasets, including datasets weakly related to pre-training data.Among similarly performing pre-trained models, the authors suggest preferring one trained with more data over one with more augmentation.
- The study is intended to guide practitioners optimizing final model performance under a given computational budget.The authors present the work as a source of effective training settings for future research and practice.
A From-scratch training details
The from-scratch training appendix documents parameter sweeps for models and fine-tuning, alongside detailed result tables. It covers epochs, learning rates, weight decay, AugReg, and transfer settings.
- A From-scratch training details: From-scratch experiments train B/32 and B/16 models on Resisc45 and Pets37 using grid searches over training hyperparameters.The listed sweeps include epochs, learning rates, and weight decays.
- A From-scratch training details: All from-scratch runs sweep dropout, stochastic depth, and data augmentation configurations.The tested dropout and stochastic-depth rates are paired across three settings, while augmentation parameters vary across eight listed combinations.
- A From-scratch training details: Fine-tuning uses the same hyperparameter sweep for all pre-trained models in the paper.The appendix identifies Table 4 as the source of the fine-tuning sweep details.
- A From-scratch training details: The appendix explains that its weight-decay values are decoupled and become coupled values after multiplication by the base learning rate.This convention follows the cited optimization approach and differs from the values specified earlier in the paper.
- A From-scratch training details: Table 5 reports detailed VTAB results grouped into natural, specialized, and structured dataset categories.The table includes the mean accuracy used in Figure 3.
C VTAB results
VTAB evaluation reports results across benchmark settings, while regularization effects vary by model size and training duration. Regularization mainly benefits larger models trained longer, but hurts all except the largest ImageNet-21k models.
- C VTAB results: VTAB results are reported as percentages for all models, using recommended checkpoints selected by upstream validation accuracy.The full VTAB score is reported only for the best pre-trained models, and hyper-parameters are selected on the VTAB validation split.
- C VTAB results: Regularization gains or losses are measured by the accuracy change from adding dropout and stochastic depth.Positive values indicate improved accuracy, while negative values indicate deterioration; absolute accuracies are shown elsewhere.
- C VTAB results: Regularization mainly helps larger models when they are trained for longer.The comparison combines dropout and stochastic depth, with peak drop probability 0.1 identified as the best setting in earlier experiments.
- C VTAB results: For ImageNet-21k pre-training, regularization hurts all but the largest models across the board.
E Using recommended checkpoints for other computer vision tasks
The study mainly evaluates classification, but it also tests a recommended ImageNet-21k B/32 checkpoint in image-text retrieval. Its upstream validation accuracy predicts zero-shot classification, and its representations perform better for retrieval than those from models with non-ideal AugReg.
- E Using recommended checkpoints for other computer vision tasks: The study’s main limitation is its focus on classification, with only a limited extension to multimodal image-text retrieval.Broader evaluation across detection, segmentation, tracking, and other tasks is described as prohibitive.
- E Using recommended checkpoints for other computer vision tasks: The recommended ImageNet-21k B/32 checkpoint is evaluated in a contrastive image-text training setup with a locked image tower.The text tower is initialized from BERT-Base and trained for 20 epochs on CC12M.
- E Using recommended checkpoints for other computer vision tasks: Upstream validation accuracy is a good predictor of zero-shot classification in the image-text setup.
- E Using recommended checkpoints for other computer vision tasks: Representations from the recommended checkpoint yield better image-text retrieval results than models without the ideal amount of AugReg.The comparison is reported across different retrieval datasets after contrastive training.