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

“How to train your ViT? Data, Augmentation, and Regularization in Vision Transformers” (AugReg) is a large-scale empirical study of ViT pre-training.

Paper: Andreas Steiner, Alexander Kolesnikov, Xiaohua Zhai, Ross Wightman, Jakob Uszkoreit and Lucas Beyer (Google Research, Brain Team), How to train your ViT? Data, Augmentation, and Regularization in Vision Transformers, TMLR 05/2022, arXiv:2106.10270.
Approach: no new architecture; a controlled grid search of more than 50,000 ViT training runs over pre-training dataset size, data augmentation, model regularization, model size and compute budget.
Question: for a practitioner with a fixed compute budget, is it better to spend it on more data or on more/better “AugReg” (the paper’s shorthand for the combination of data Augmentation and model Regularization)?
Headline finding: judicious AugReg substitutes for roughly an order of magnitude more pre-training data. ViT models pre-trained with AugReg on the public ImageNet-21k dataset match or beat the same architectures pre-trained on Google’s private 300M-image JFT-300M.
Why it matters: this closes a gap that had made the original ViT paper’s best results effectively unreproducible outside Google.
Used by: the authors release the resulting collection of pre-trained checkpoints, which became the de facto standard AugReg-initialized ViT backbones, reused by works such as SigLIP, SigLIP2, Image Captioners Are Scalable Vision Learners Too and PaliGemma.

Adding the right amount of AugReg has a similar effect on final accuracy as increasing pre-training dataset size by 10x
Fig. 1: ViT-B/32, ViT-B/16 and ViT-L/16 fine-tuned ImageNet top-1 accuracy as pre-training data grows from 1.28M (ImageNet-1k) through AugReg-augmented 1.28M and 13M (ImageNet-21k) to AugReg-augmented 13M and 300M (JFT-300M). Each “+AugReg” step recovers most of the gain of the next 10x data step.

Motivation

The original Vision Transformer paper showed ViT beats CNNs of similar cost only once pre-trained on very large corpora (ImageNet-21k, 13M images, or JFT-300M, 300M images); DeiT showed the opposite fix — heavy regularization and augmentation lets a ViT train competitively from scratch on the comparatively tiny ImageNet-1k (1.28M images). Neither paper studied the two knobs together, and neither varied training-data scale alongside AugReg strength in a controlled way. Practitioners were left without a systematic answer to a question that determines where to spend a fixed budget: given a model size and a fixed amount of pre-training compute, does adding data or adding AugReg buy more downstream accuracy, and how much compute does either cost to reach a given accuracy?

The paper explicitly separates three notions of “cost” that prior work conflated: the pre-training cost researchers pay to build the checkpoint, the practitioner’s cost of fine-tuning/adapting an already-pretrained checkpoint (which dominates when good checkpoints are freely downloadable), and the deployment cost of inference at scale. The study is designed around the middle case — practitioners who fine-tune an existing checkpoint rather than pre-train from scratch — while still reporting from-scratch results for comparison.

Experimental setup

All training and transfer runs use one JAX/Flax codebase on TPUs (inference-speed numbers alone come from PyTorch/timm on V100 GPUs), and all datasets are consumed through TensorFlow Datasets, specifically to eliminate the implementation inconsistencies (augmentation details, learning-rate schedules, dataset splits) that the authors argue make prior cross-paper comparisons noisy.

Pre-training datasets. ImageNet-1k (1.28M images, 1000 classes) and ImageNet-21k (~14M images, ~21k classes, de-duplicated against all downstream test sets); ImageNet-21k is roughly 10x larger than ImageNet-1k, which lets the study probe a full order of magnitude of data scale while other factors stay controlled.

Models (Table 1/2 of the paper).

Model	Layers	Width	MLP	Heads	Params
ViT-Ti	12	192	768	3	5.8M
ViT-S	12	384	1536	6	22.2M
ViT-B	12	768	3072	12	86M
ViT-L	24	1024	4096	16	307M

Patch size 16 is used for all configurations, plus patch size 32 for S and B. Hybrid ResNet+ViT models are also trained — a ResNet stem ( $7 \times 7$ conv + BN + ReLU + max-pool) followed by a variable number of bottleneck blocks feeds the spatial output into the ViT as patch embeddings, notated Rn+{Ti,S,L}/p where $n$ is the stem’s convolution count and $p$ the effective patch size in input pixels (e.g. R+Ti/16 halves resolution in the stem, then forms $8 \times 8$ patches, giving an effective patch size of 16). One deliberate deviation from the original ViT: the hidden layer in the classification head is dropped, since keeping it gives no accuracy benefit and can cause optimization instability when pre-training on ImageNet-1k.

AugReg sweep (§3.3). Regularization: dropout on intermediate ViT activations, and stochastic depth with linearly increasing layer-drop probability. Augmentation: a combination of Mixup (parameterized by mixing strength $α$ , 0 = disabled) and RandAugment (parameterized by number of augmentation layers $l$ and magnitude $m$ ; the paper uses the original TensorFlow RandAugment implementation, which differs numerically from the timm reimplementation). The full sweep is a 28-way cross-product:

Regularization: off (no dropout, no stochastic depth) or on (dropout 0.1, max stochastic-depth drop probability 0.1) — 2 settings.
Data augmentation, 7 presets over $(l, m, α)$ : none $(0, 0, 0)$ , light1 $(2, 0, 0)$ , light2 $(2, 10, 0.2)$ , medium1 $(2, 15, 0.2)$ , medium2 $(2, 15, 0.5)$ , strong1 $(2, 20, 0.5)$ , strong2 $(2, 20, 0.8)$ .
Weight decay: $0.1$ or $0.03$ (decoupled à la AdamW, i.e. scaled by the peak learning rate of $0.001$ ) — 2 settings.

Pre-training (§3.4). Adam, $β_{1} = 0.9$ , $β_{2} = 0.999$ , batch size 4096, cosine LR schedule with 10k-step linear warmup, gradients clipped at global norm 1, Inception-style random-resized-crop + horizontal flip preprocessing. ImageNet-1k models train for 300 epochs; ImageNet-21k models train for both 30 and 300 epochs, so that (given IN-21k is ~10x larger) the 30-epoch IN-21k runs use roughly the same total compute as the 300-epoch IN-1k runs, isolating the effect of data scale at fixed compute.

Fine-tuning (§3.5). SGD with momentum 0.9 (optimizer state stored in bfloat16), 2–3 learning rates and 1–2 durations swept per dataset (Table 4, Appendix B), fixed batch size 512, gradient clipping at global norm 1, cosine decay with linear warmup, at both the original 224px resolution and a higher 384px resolution.

Key findings

1. AugReg trades off against roughly 10x more data

The central result (Fig. 1 above, §4.1): models trained with the right amount of AugReg on ImageNet-1k reach about the same fine-tuned accuracy as the same architecture trained without AugReg on the 10x-larger ImageNet-21k. Applying AugReg to ImageNet-21k similarly (given a correspondingly larger compute budget) lets models match or beat the same architectures pre-trained on the 25x-larger, non-public JFT-300M without AugReg — this is the result that lets the paper claim it reproduces private JFT-300M-level results using only the public ImageNet-21k. Regardless of whether the accuracy gain comes from more data or from better AugReg, the compute required to reach a given accuracy level is roughly the same either way — AugReg is not a free lunch, it substitutes compute-for-compute against data.

The effect has a floor: for genuinely small datasets, heavier augmentation on a subset does not recover the accuracy of training on the full dataset (citing a ResNet-50 result on 10%-of-ImageNet-1k), so AugReg extends but does not eliminate the value of scale.

2. Transfer beats from-scratch training for practically-sized datasets

For two realistic downstream datasets — Pets37 (~3,000 training images, similar modality to ImageNet-1k) and Resisc45 (~30,000 satellite images, a modality neither ImageNet-1k nor -21k covers well) — the authors ran a large from-scratch hyperparameter search (varying epochs, learning rate, weight decay, dropout, stochastic depth and AugReg strength) and compared it against fine-tuning a pre-trained checkpoint. On Pets37, no amount of from-scratch training time reaches the accuracy that a fine-tuned pre-trained model attains almost trivially; on the larger Resisc45, two orders of magnitude more from-scratch compute comes close but still does not match transfer. Since good checkpoints are freely downloadable, their pre-training cost is amortized to zero from a practitioner’s perspective — so transfer is simultaneously cheaper and better whenever a usable checkpoint exists, which the paper calls the standout practical recommendation of the whole study.

3. More upstream data yields more generic transfer, even when upstream accuracy alone does not show it

Pretraining on more data yields more transferable models, evaluated on the 19-task VTAB suite across three task categories
Fig. 3: VTAB natural / specialized / structured mean accuracy vs. inference speed, comparing IN-1k (300 epochs), IN-21k (30 epochs, same compute as IN-1k/300ep) and IN-21k (300 epochs). IN-21k pre-training wins on all three VTAB categories even at matched compute.

Comparing models pre-trained for equal compute (IN-1k/300ep vs. IN-21k/30ep), the IN-21k-pretrained model transfers noticeably better across all 19 VTAB tasks and all three of its task categories (natural, specialized, structured) — even though this compute-matched advantage is much less visible when looking only at ImageNet-1k’s own validation accuracy. This motivates the paper’s recommendation to select checkpoints by pre-training data scale (and upstream validation score) rather than by ImageNet-1k accuracy alone, since IN-1k accuracy is an inflated, narrow proxy for general-purpose representation quality.

4. AugReg helps on medium data, hurts on large data at fixed compute — unless compute also grows

Validation accuracy heatmaps across the 7x2 AugReg grid, for IN-1k/300ep, IN-21k/30ep and IN-21k/300ep
Fig. 4: rows are models (increasing capacity top-to-bottom), columns are augmentation strength (increasing left-to-right) split into no-regularization / regularization-0.1 halves; the rightmost thin column is the best-regularized-minus-best-unregularized delta. On IN-1k (left panel) almost every AugReg setting helps. On IN-21k at fixed 30-epoch budget (middle panel) almost every AugReg setting hurts. Only once the IN-21k budget grows to 300 epochs (right panel) does AugReg start helping again, and even then it still hurts the smallest models.

This is the paper’s most nuanced finding (§4.4): the intuition “more AugReg is always safer” is dataset- and compute-budget-dependent, not a universal rule. On the 10x larger ImageNet-21k at a fixed (i.e. proportionally reduced-epoch) compute budget, adding AugReg of almost any strength actively hurts almost every model — the regularization/augmentation eats into a compute budget that is already thin relative to the dataset size. It is only once total pre-training compute is also increased (300 vs. 30 epochs on IN-21k) that AugReg starts paying off again, and even then, mainly for the larger models — model regularization specifically (dropout + stochastic depth, isolated in Appendix D / Fig. 7) helps primarily the largest models trained the longest, and actively hurts smaller/shorter runs across the board. Augmentation is beneficial in more of the grid than regularization is; the paper’s blanket rule of thumb is therefore “prefer augmentation to regularization” (§4.4 title) when a choice must be made without further tuning.

5. Selecting checkpoints by upstream validation score is a good, cheap heuristic

Rather than fine-tuning every pre-trained checkpoint on every downstream task and picking the best result per task (expensive), the paper shows that picking a single checkpoint by its own upstream (pre-training) validation accuracy and using that one checkpoint everywhere works nearly as well across a large collection of pre-trained models and 5 downstream datasets — with the exception of a few outlier cases where fine-tuning several checkpoints and picking the best helps materially. One subtlety: for models pre-trained on ImageNet-21k and evaluated on the IN-21k minival split before transfer to ImageNet-1k, large models on long schedules partially memorize the ImageNet-1k images embedded inside ImageNet-21k, which biases minival accuracy as a model-selection signal. The fix the paper adopts and recommends is to use ImageNetV2 (an independently collected, non-overlapping test set) as the validation split for this specific model-selection step; other datasets did not show the issue.

6. Prefer a larger patch size over a thinner (“Tiny”) model at matched inference speed

ImageNet transfer results: validation accuracy vs. throughput for the S/B/L and Tiny/hybrid model families, and best model per pretraining dataset
Fig. 6: left — accuracy vs. images/sec/core for the best checkpoint per pretraining dataset (i1k, IN-21k@30ep, IN-21k@300ep). Right — at matched throughput, “/32” patch-size variants (B/32, S/32) clearly beat “Ti” (thin, small-width) variants on both IN-21k and IN-1k transfer.

At a fixed inference throughput, models using the “Ti” (Tiny, reduced-width) configuration are consistently and clearly worse than models of a similar speed obtained instead by using a larger patch size (e.g. B/32, S/32). Patch size controls the number of self-attention tokens for a given input resolution and therefore contributes to effective model capacity in a way that raw parameter count does not capture — so, per the paper, parameter count alone is not a reliable proxy for either speed or quality, and shrinking a model by widening its patch size is a better lever than shrinking its width.

Hyperparameters

Pre-training (all models).


Optimizer	Adam, $β_{1} = 0.9$ , $β_{2} = 0.999$
Batch size	4,096
LR schedule	cosine decay, 10,000-step linear warmup
Gradient clipping	global norm 1
Preprocessing	Inception-style random-resized crop, random horizontal flip
Epochs	300 (IN-1k); 30 and 300 (IN-21k, to match/exceed IN-1k compute)
Weight decay (AugReg sweep)	0.1 or 0.03 (decoupled, scaled by peak LR 0.001)
Dropout / stochastic depth (AugReg sweep)	off, or dropout 0.1 + max stochastic-depth drop-prob 0.1
RandAugment $(l, m)$ presets	none $(0, 0)$ , light1 $(2, 0)$ , light2 $(2, 10)$ , medium1/2 $(2, 15)$ , strong1/2 $(2, 20)$
Mixup $α$ presets	0, 0, 0.2, 0.2 / 0.5, 0.5 / 0.8 (paired with the RandAugment presets above)

Fine-tuning (recommended-checkpoint sweep, Table 4).

Dataset	Learning rates swept	(Total steps, warmup steps)
ImageNet-1k	0.01, 0.03	(20k, 500)
Pets37	1e-3, 3e-3, 0.01	(500, 100) or (2.5k, 200)
Kitti-distance	1e-3, 3e-3, 0.01	(500, 100) or (2.5k, 200)
CIFAR-100	1e-3, 3e-3, 0.01	(2.5k, 200) or (10k, 500)
Resisc45	1e-3, 3e-3, 0.01	(2.5k, 200) or (10k, 500)

Common to all fine-tuning runs: SGD, momentum 0.9, bfloat16 optimizer state, batch size 512, gradient clipping at global norm 1, cosine decay with linear warmup, evaluated at both 224px and 384px resolution. Full VTAB fine-tuning (Appendix C) additionally sweeps 4 learning rates ${0.001, 0.003, 0.01, 0.03}$ and 2 schedule lengths ${500, 2500}$ steps, selected on the VTAB validation split.

Results

On the recommended-checkpoint ImageNet-1k transfer table (Table 3), at 384px resolution the best AugReg-trained ImageNet-21k checkpoints match or exceed the paper’s own reproduction of the corresponding JFT-300M-pretrained models from the original ViT paper for several configurations — e.g. ViT-L/16 reaches 87.08% (i21k, 300 epochs, selected by ImageNetV2) against 87.12% for the JFT-300M-pretrained ViT-L/16, and ViT-B/32 reaches 83.59% against 80.73% for JFT-300M, i.e. exceeding the private-data result. Across the VTAB 19-task suite (Table 5), the best AugReg IN-21k/300-epoch checkpoints reach mean scores in the 73–81% range across natural/specialized/structured task groups, consistently ahead of the same architectures pre-trained without AugReg or on the smaller dataset. An additional experiment (Appendix E) contrastively fine-tunes a locked ImageNet-21k-pretrained image tower against a BERT-Base text tower on CC12M image–text pairs: the “recommended” (best upstream-validation) B/32 checkpoint yields consistently better zero-shot ImageNet and Coco/Flickr image–text retrieval accuracy than checkpoints using deliberately too-little or too-much AugReg, confirming that upstream validation accuracy remains predictive even for a downstream task (retrieval) far removed from classification.

Limitations

The study is restricted to the canonical ViT architecture (plus ResNet-ViT hybrids); it does not cover ResNets in the same sweep (already well studied elsewhere) or later ViT variants (hierarchical/windowed transformers, convolutional hybrids beyond the simple stem), though the authors expect the qualitative findings to transfer.
The compute cost of the study itself is very large (50,000+ runs); the authors frame the release of ready-made checkpoints as specifically intended to remove the need for others to repeat searches of this scale, rather than as encouragement to do so.
AugReg’s benefit is compute-budget-dependent and non-monotonic (finding 4 above) — there is no single AugReg setting that is safe to apply blindly; the “prefer augmentation to regularization” rule of thumb is a coarse simplification of a genuinely two-dimensional, model-size- and budget-dependent effect (Fig. 4/7).
AugReg substitutes for more data only down to a point — it does not rescue arbitrarily small datasets.
Self-supervised pre-training is explicitly out of scope; the paper only studies supervised pre-training and leaves systematic comparison to self-supervised objectives to future work.
All headline transfer results are classification-shaped (plus one small-scale retrieval experiment); no detection or dense-prediction transfer is evaluated.

Wiki

Explorer

How to train your ViT