Better plain ViT baselines for ImageNet-1k
Lucas Beyer, Xiaohua Zhai, Alexander Kolesnikov
Introduction
The ViT paper focused solely on the aspect of large-scale pre-training, where ViT models outshine well tuned ResNet (BiT ) models. The addition of results when pre-training only on ImageNet-1k was an afterthought, mostly to ablate the effect of data scale. Nevertheless, ImageNet-1k remains a key testbed in the computer vision research and it is highly beneficial to have as simple and effective a baseline as possible.
Thus, coupled with the release of the big vision codebase used to develop ViT , MLP-Mixer , ViT-G , LiT , and a variety of other research projects, we now provide a new baseline that stays true to the original ViT’s simplicity while reaching results competitive with similar approaches and concurrent , which also strives for simplification.
Experimental setup
We focus entirely on the ImageNet-1k dataset (ILSVRC-2012) for both (pre)training and evaluation. We stick to the original ViT model architecture due to its widespread acceptance , simplicity and scalability, and revisit only few very minor details, none of which are novel. We choose to focus on the smaller ViT-S/16 variant introduced by as we believe it provides a good tradeoff between iteration velocity with commonly available hardware and final accuracy. However, when more compute and data is available, we highly recommend iterating with ViT-B/32 or ViT-B/16 instead , and note that increasing patch-size is almost equivalent to reducing image resolution.
All experiments use “inception crop” at 224px² resolution, random horizontal flips, RandAugment , and Mixup augmentations. We train on the first 99% of the training data, and keep 1% for minival to encourage the community to stop selecting design choices on the validation (de-facto test) set. The full setup is shown in Appendix A.
Results
The results for our improved setup are shown in Figure 1, along with a few related important baselines. It is clear that a simple, standard ViT trained this way can match both the seminal ResNet50 at 90 epochs baseline, as well as more modern ResNet and ViT training setups. Furthermore, on a small TPUv3-8 node, the 90 epoch run takes only 6h30, and one can reach 80% accuracy in less than a day when training for 300 epochs.
The main differences from are a batch-size of 1024 instead of 4096, the use of global average-pooling (GAP) instead of a class token , fixed 2D sin-cos position embeddings , and the introduction of a small amount of RandAugment and Mixup (level 10 and probability 0.2 respectively, which is less than ). These small changes lead to significantly better performance than that originally reported in .
Notably absent from this baseline are further architectural changes, regularizers such as dropout or stochastic depth , advanced optimization schemes such as SAM , extra augmentations such as CutMix , repeated augmentations , or blurring, “tricks” such as high-resolution fine-tuning or checkpoint averaging, as well as supervision from a strong teacher via knowledge distillation.
Table 1 shows an ablation of the various minor changes we propose. It exemplifies how a collection of almost trivial changes can accumulate to an important overall improvement. The only change which makes no significant difference in classification accuracy is whether the classification head is a single linear layer, or an MLP with one hidden layer as in the original Transformer formulation.
Conclusion
It is always worth striving for simplicity.
Acknowledgements. We thank Daniel Suo and Naman Agarwal for nudging for 90 epochs and feedback on the report, as well as the Google Brain team for a supportive research environment.