[Paper Review] How to train your ViT? Data, Augmentation, and Regularization in Vision Transformers

Andreas Steiner, А. И. Колесников+4 more|arXiv (Cornell University)|Jun 18, 2021

Advanced Neural Network ApplicationsComputer Science228 citations

TL;DR

The paper systematically studies how data size, augmentation, regularization, model size, and compute budget interact for Vision Transformers, showing that AugReg can substitute for large data and that transfer learning from larger upstream datasets is often more cost-effective.

ABSTRACT

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.

Motivation & Objective

• Understand how training data size, augmentation, and regularization interact for ViTs across model sizes.
• Quantify the compute and data trade-offs for pre-training ViTs and transferring to downstream tasks.
• Evaluate transfer learning versus training from scratch with AugReg across diverse datasets.
• Identify practical recommendations for selecting pre-trained checkpoints and training strategies.

Proposed method

• Pre-train ViT models of multiple sizes (Ti, S, B, L) and ResNet+ViT hybrids on ImageNet-1k and ImageNet-21k with a unified setup on TPU/JAX/Flax.
• Apply regularization (Dropout, stochastic depth) and augmentation (Mixup, RandAugment) with a 28-configuration hyperparameter sweep.
• Fine-tune on downstream datasets with SGD and cosine schedules, comparing from-scratch versus transfer from upstream models.
• Systematically vary data scale, compute budget, and augmentation strength to map performance landscapes.

Experimental results

Research questions

• RQ1 How do data size, augmentation, and regularization interact to affect ViT performance across model sizes?
• RQ2 Is transfer learning from larger upstream datasets more cost-effective than training from scratch with AugReg for typical downstream tasks?
• RQ3 How does increasing pre-training data vs. augmentations affect transfer performance across VTAB tasks?
• RQ4 What practical guidance can be given for selecting pre-trained checkpoints and training strategies under compute constraints?

Key findings

• AugReg can match the gains of 10x increases in training data for ViTs on comparable tasks.
• Transfer learning from pre-trained models is generally more cost-efficient and yields better results than training from scratch for practical datasets.
• More pre-training data yields more generic models that transfer better across diverse VTAB tasks.
• Augmentation helps more on medium-sized datasets, but on very large upstream data (e.g., ImageNet-21k with fixed compute) AugReg often harms unless compute is increased.
• Choosing upstream models by upstream validation accuracy is usually effective for transfer; using ImageNet-21k checkpoints is recommended.
• Increasing patch size to 32 can outperform thinner models of similar speed, indicating patch size is a key capacity factor beyond parameter count.