When Vision Transformers Outperform ResNets without Pretraining or Strong Data Augmentations

By Xiangning Chen, Cho-Jui Hsieh, Boqing Gong. 2021

396 citations recorded by the source corpus.

Vision Transformers (ViTs) and MLPs signal further efforts on replacing hand-wired features or inductive biases with general-purpose neural architectures. Existing works empower the models by massive data, such as large-scale pre-training and/or repeated strong data augmentations, and still report optimization-related problems (e.g., sensitivity to initialization and learning rates). Hence, this paper investigates ViTs and MLP-Mixers from the lens of loss geometry, intending to improve the models' data efficiency at training and generalization at inference. Visualization and Hessian reveal extremely sharp local minima of converged models. By promoting smoothness with a recently proposed sharpness-aware optimizer, we substantially improve the accuracy and robustness of ViTs and MLP-Mixers on various tasks spanning supervised, adversarial, contrastive, and transfer learning (e.g., +5.3\% and +11.0\% top-1 accuracy on ImageNet for ViT-B/16 and Mixer-B/16, respectively, with the simple Inception-style preprocessing). We show that the improved smoothness attributes to sparser active neurons in the first few layers. The resultant ViTs outperform ResNets of similar size and throughput when trained from scratch on ImageNet without large-scale pre-training or strong data augmentations. Model checkpoints are available at \url{https://github.com/google-research/vision_transformer}.

Subfields: Artificial Intelligence, Computer Vision and Pattern Recognition

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

Citations in the corpus

Simpler is Better: off-the-shelf Continual Learning Through Pretrained Backbones (2022, background): “But, as can be seen by some recent works, Vision Transformers (ViT) are asserting themselves as a valuable alternative to CNNs for computer vision tasks, sometimes, achieving better performances with respect to CNNs [2].”
Transformers in Vision: A Survey (2021, neutral): “Another approach to data eicient training of ViTs is proposed in et al.[37].”
Questions for Flat-Minima Optimization of Modern Neural Networks (2022, methodology): “…adapt the hyper-parameter values from Hu et al. (2020) for GCN (Kipf & Welling, 2017), SAGE (Hamilton et al., 2017), and GIN (Xu et al., 2019), from Chen et al. (2021b) for (Lacroix et al., 2018) and ComplEx (Trouillon et al., 2016), and from Li et al. (2020) for DGCN. Due to high standard errors,…”
Attention mechanisms in computer vision: A survey (2021, background): “[184] significantly improved visual transformers by using a novel optimizer, the sharpness-aware minimizer (SAM) [185].”
Towards Stable Test-Time Adaptation in Dynamic Wild World (2023, background): “SAM and its variants (Kwon et al., 2021; Zheng et al., 2021; Du et al., 2022; Chen et al., 2022b) have shown outstanding performance on several deep learning benchmarks.”
A Simple Single-Scale Vision Transformer for Object Localization and Instance Segmentation (2021, background): “It is worth noting that, although recent works try to analyze vision transformer’s generalization [30], loss landscapes [9], and patterns of learned representations [31,38], they mostly focus on image classiﬁcation tasks.”
Artificial intelligence in histopathology: enhancing cancer research and clinical oncology (2022, background): “Historically, very large pretraining datasets were thought to be a prerequisite for using transformers; however, recently developed training strategies were shown to partially eliminate this dependence 122 .”
Revealing Occlusions with 4D Neural Fields (2022, background): “The attention mechanism introduced in [4, 58] has been applied with great success to computer vision [10,14,53,65].”
Stabilizing Transformer Training by Preventing Attention Entropy Collapse (2023, background): “From a loss geometry perspective, we hypothesize that these regions of instability are caused when the weights enter a region of high curvature, a hypothesis supported by (Chen et al., 2022), which”
Vision Transformers Need Registers (2023, methodology): “Other works have since reported interesting attention maps using various techniques: by modifying the optimisation procedure (Chen et al., 2022), by steering the attention scores towards useful image parts (Shi et al., 2023), by modifying the architecture of the transformer layers (Yu et al.,…”
How Well Do Sparse ImageNet Models Transfer? (2021, background): “Training restarts either from intialization [14], or by “rewinding” to an earlier point during training of the dense model [6, 7, 15].”
How Do Vision Transformers Work? (2022, background): “The flatter the loss landscape, the better the performance and generalization (Li et al., 2018; Keskar et al., 2017; Santurkar et al., 2018; Foret et al., 2021; Chen et al., 2022).”
Symbolic Discovery of Optimization Algorithms (2023, background): “Intuitively, the sign operation adds noise to the updates, which acts as a form of regularization and helps with generalization (Chen et al., 2022; Foret et al., 2021; Neelakantan et al., 2017).”
Adaptive Gradient Methods at the Edge of Stability (2022, methodology): “Meanwhile, [17] empirically observed that a VGG trained using an adaptive optimizer had 40x the sharpness than when trained using SGD, and [6] found that MLP-mixers and ViTs trained with Adam converged to far sharper solutions than a ResNet trained with SGD.”
PanGu-Coder: Program Synthesis with Function-Level Language Modeling (2022, background): “Vision [65, 18, 15, 13, 64, 25], Proteins [11], Time Series [71, 72, 55] and Code [42, 70, 26, 31, 54].”
Sharp-MAML: Sharpness-Aware Model-Agnostic Meta Learning (2022, methodology): “However, existing efforts have been focusing on improving generalization performance in single-level problems such as ERM (Bahri et al., 2021; Chen et al., 2021b).”
Sharpness-aware Quantization for Deep Neural Networks (2021, methodology): “To measure the loss curvature, we report the largest eigenvalue λmax of the Hessian of the converged quantized models following [8, 20].”