PaperlayerView exact revision on arXiv

Source-linked AI summary

When Vision Transformers Outperform ResNets without Pre-training or Strong Data Augmentations

Xiangning Chen, Cho-Jui Hsieh, Boqing Gong

arXiv:2106.01548v3cs.CVcs.LG

TL;DR

Convolution-free ViTs and MLP-Mixers often require massive data or strong augmentations and converge to sharp minima. This paper applies SAM to smooth their loss landscapes, improving generalization across tasks and enabling ViTs to outperform comparable ResNets without pre-training or strong augmentations.

  • Problem

    Convolution-free ViTs and MLP-Mixers require massive pre-training or strong augmentations, while first-order optimizers overlook loss-geometry flatness linked to generalization.

  • Method

    The paper analyzes loss landscapes and trains ViTs and MLP-Mixers with sharpness-aware minimization (SAM) to explicitly smooth the loss geometry.

  • Results

    +5.3% and +11.0% top-1 ImageNet accuracy for ViT-B/16 and Mixer-B/16, respectively, while improving accuracy and robustness across supervised, adversarial, contrastive, and transfer tasks.

  • Takeaways & Limitations

    SAM enables ViTs to outperform ResNets of comparable size and throughput when trained from scratch without large-scale pre-training or strong augmentations.

  • Takeaways & Limitations

    SAM roughly doubles computational cost per update, and its effect diminishes as the training dataset becomes larger.

Abstract

from arXiv · show

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}.

1 INTRODUCTION

The paper argues that convolution-free ViTs and MLP-Mixers underperform partly because they converge to extremely sharp minima, and shows that sharpness-aware optimization improves their smoothness, accuracy, and robustness without large-scale pre-training or strong data augmentations.

  • Motivation: ViTs and MLPs offer general-purpose, end-to-end vision architectures but lack convolution-like inductive biases that support training.With conventional Inception-style preprocessing on ImageNet, Transformers achieve accuracies a few percentage points below comparable ResNets.
  • Contribution: ViTs can outperform larger ResNets in accuracy and robustness without large-scale pre-training or strong data augmentations, while MLP-Mixers become comparable to ResNets.The paper attributes these results to a principled optimizer.
  • Loss geometry: ViTs and MLP-Mixers converge at extremely sharp local minima whose largest principal curvatures are almost an order of magnitude larger than ResNets’.The sharpness accumulates as gradients backpropagate from the last layer to the first, with the initial embedding layer suffering the largest effect.
  • Sharpness-aware optimization: SAM smooths loss geometry by seeking parameter solutions whose entire neighborhood has low losses rather than optimizing only a singleton point.The resulting models have smoother loss landscapes and substantially improved generalization across supervised, adversarial, contrastive, and transfer learning tasks.
  • Mechanism: SAM increases active-neuron sparsity, especially in the first few layers, contributing to reduced Hessian eigenvalues.SAM also increases weight norms, suggesting weight decay alone may not be effective regularization; ViTs exhibit extremely sparse active neurons relative to ResNets and MLP-Mixers.

2 BACKGROUND AND RELATED WORK

This section introduces ViTs and MLP-Mixers as image architectures and frames the paper around their optimization and generalization behavior. It contrasts their sharp convergence with the smoothing effect of SAM and notes the pre-training context of prior ViT results.

  • Vision Transformers: ViT achieves state-of-the-art image-classification accuracy using a pure Transformer architecture pre-trained on large datasets.Its residual blocks contain multi-head self-attention, layer normalization, and an MLP layer.
  • Optimization behavior: ViTs and MLP-Mixers converge in sharp regions, while SAM smooths their loss landscapes and improves generalization.The reported comparison includes divergent NTK condition numbers and sharp convergence for ViT and MLP-Mixer, with SAM rescuing both behaviors.
  • MLP-Mixers: MLP-Mixer replaces self-attention with alternating token-mixing and channel-mixing MLPs while sharing ViT’s patch-based input layer.It partitions images into sequences of nonoverlapping patches or tokens.
  • Scope and notation: The paper focuses on ViTs and MLP-Mixers, using “S” and “B” for small and base model sizes and patch resolution to identify models.For example, ViT-B/16 is a base ViT receiving a sequence of 16×16 patches.

3 VITS AND MLP-MIXERS CONVERGE AT SHARP LOCAL MINIMA

The section argues that ViTs and MLP-Mixers face optimization problems and converge to extremely sharp local minima, motivating loss-landscape analysis to reduce reliance on massive pre-training and strong augmentations. Perturbation-based flatness and Hessian curvature measurements show sharper landscapes than ResNet, while Mixer behavior suggests greater overfitting and poorer trainability.

  • Motivation: ViTs and MLP-Mixers are studied from an optimization perspective because their training recipes depend heavily on massive pre-training or strong augmentations and exhibit instability.Reported issues include gradient spikes, sudden accuracy dips, and sensitivity to initialization and hyperparameters.
  • Sharp local minima: ViTs and MLP-Mixers converge at extremely sharp local minima, whereas convergence to flatter regions is associated with better neural-network generalization.The paper visualizes their loss landscapes and evaluates curvature to characterize this behavior.
  • Sharp local minima: ViT-B/16 and Mixer-B/16 achieve lower training error than ResNet-152 but incur much higher loss after random weight perturbations.The perturbation-based training loss measures average flatness and indicates sharper solutions for the convolution-free models.
  • Sharp local minima: The dominant Hessian eigenvalues of ViTs and MLP-Mixers are orders of magnitude larger than ResNet’s, indicating substantially greater worst-case landscape curvature.The Hessian’s λmax is used as a mathematical evaluation of worst-case curvature.
  • Overfitting and trainability: Mixer-B/16 has fewer parameters than ViT-B/16 but lower training error and much worse test accuracy, suggesting cross-token MLPs are more prone to overfitting than softmax-restricted self-attention.The comparison is 59M versus 87M parameters.
  • Overfitting and trainability: ViTs and MLP-Mixers also suffer poor trainability, defined as how effectively gradient descent optimizes a network and characterized here using the neural tangent kernel’s condition number.The associated NTK is defined as Θ(x, x′) = J(x)J(x′)T, where J is the Jacobian matrix.

4 A PRINCIPLED OPTIMIZER FOR CONVOLUTION-FREE ARCHITECTURES

This section presents SAM as a sharpness-aware optimizer that seeks parameters whose neighborhoods have low training loss, using an efficient gradient-based approximation. SAM smooths ViT and MLP-Mixer loss landscapes, improving accuracy, robustness, data efficiency, and active-neuron sparsity, while enabling ViTs to outperform comparable ResNets when trained from scratch.

  • Sharpness-Aware Minimization: SAM minimizes worst-case training loss in a parameter neighborhood, approximating the inner maximization with a scaled gradient under the l2 norm.The neighborhood size is controlled by ρ, and updates use the gradient evaluated at the perturbed weights.
  • Improved Loss Geometry and Accuracy: SAM makes ViT and MLP-Mixer minima substantially smoother and improves ImageNet top-1 accuracy from 74.6% to 79.9% for ViT-B/16 and from 66.4% to 77.4% for Mixer-B/16.A similarly sized ResNet-152 improves by only 0.8%.
  • Robustness: 9.9% and 15.0% increases in ImageNet-C accuracy are reported for ViT-B/16 and Mixer-B/16, respectively, showing larger robustness gains under image corruptions.The reported ImageNet-C evaluation averages accuracy across 19 corruptions and five severity levels.
  • Comparison with ResNets: When trained from scratch on ImageNet, ViTs with SAM outperform similarly sized and larger ResNets in clean accuracy and robustness across ImageNet, ImageNet-ReaL, ImageNet V2, ImageNet-R, and ImageNet-C.ViT-B/16 reaches 79.9%, 26.4%, and 56.6% top-1 accuracy on ImageNet, ImageNet-R, and ImageNet-C, respectively.
  • Mechanisms of Smoothing: SAM reduces curvature by promoting sparser active neurons, especially in early layers, where the first linear embedding layer has the sharpest geometry.In ViTs, fewer than 10% of neurons are active in most layers for a given input image; MLP-Mixers also show decreased active GELU neurons.
  • Data Efficiency and Augmentation: With only 1/10 of ImageNet training images, ViT-B/16-SAM outperforms ViT-B/16-AUG by 7.6%, indicating stronger data-efficiency gains from SAM than strong augmentation.The section also reports that strong augmentations can enlarge λmax, whereas SAM reduces the largest curvature through a worst-case minimax objective.

5 ABLATION STUDIES

The ablation studies examine SAM across vision models and training setups, including reduced-data ImageNet training. SAM substantially mitigates the sharper accuracy decline of ViTs and MLP-Mixers on smaller training sets relative to their lack of convolutional inductive bias.

  • Ablation scope: The section evaluates SAM’s effects across vision models and training setups, with adversarial, contrastive, and transfer-learning results deferred to Appendices B–D.The ablation scope explicitly includes multiple training paradigms.
  • Training-data scale: The reduced-data experiments construct ImageNet subsets using one-quarter and one-half of each class, containing 320,291 and 640,583 images, respectively.The subsets are named i1k (1/4) and i1k (1/2).
  • Training-data scale: With fewer training examples, ViT and MLP-Mixer accuracies decline faster than ResNet accuracies, reflecting their lack of convolutional inductive bias.The comparison is reported for the smaller ImageNet training sets.
  • Training-data scale: SAM drastically rescues ViT and MLP-Mixer performance on smaller training sets, and its improvement over vanilla SGD is proportional to the number of training images.The relationship is shown in Figure 4 (right).

6 CONCLUSIONS AND LIMITATIONS

The paper analyzes convolution-free ViTs and MLP-Mixers through loss-landscape geometry, using SAM to reduce reliance on massive pre-training and strong data augmentations. SAM improves loss smoothness and generalization but roughly doubles update cost, with diminishing effects as datasets grow.

  • The study examines ViTs and MLP-Mixers through loss-landscape geometry to reduce dependence on massive pre-training and strong data augmentations.
  • SAM regularizes loss geometry, producing flatter loss landscapes and improved generalization.
  • Around 2x computational cost per update arises because SAM requires an additional forward and backward propagation to update ϵ.
  • SAM’s effect diminishes as the training dataset becomes larger, motivating algorithms that improve or accelerate large-scale pre-training.

ETHICS STATEMENT

The authors identify no immediate ethical issues and expect the work to inform efficient convolution-free architectures. They note that potential negative societal impacts depend mainly on how such architectures are applied.

  • ETHICS STATEMENT: The authors report no immediate ethical issues and hope the work benefits future data- and computation-efficient neural architectures.Potential negative societal impacts are described as mainly hinging on applications of convolution-free architectures.

APPENDICES · A ARCHITECTURES

The appendix defines the ViT and MLP-Mixer architectures used in the paper, including their scale and patch-size notation and fixed 224 × 224 input resolution. It also reports that ViTs trained with SAM outperform similarly sized and similarly fast ResNets under adversarial training in clean accuracy and robustness.

  • A ARCHITECTURES: Table 8 specifies the ViT and MLP-Mixer architectures used throughout the paper.These architectures are attributed to prior ViT and MLP-Mixer work.
  • A ARCHITECTURES: “S” and “B” denote small and base model scales, respectively.The notation follows conventions from prior ViT and MLP-Mixer studies.
  • A ARCHITECTURES: “B/16” denotes a base-scale model using non-overlapping 16 × 16 image patches.The notation combines model scale with patch resolution.
  • A ARCHITECTURES: ViTs-SAM outperform similarly sized and similarly fast ResNets-SAM in clean accuracy and adversarial robustness.This comparison is made under the ImageNet adversarial training framework.
  • A ARCHITECTURES: The adversarial-training comparison reports improvements over standard adversarial training without SAM in parentheses.The table caption defines the parenthetical values as those improvements.
  • A ARCHITECTURES: 224 × 224 is the input image resolution used throughout the paper.Table 8 likewise states that all architectures are trained at this resolution.

B WHEN SAM MEETS ADVERSARIAL TRAINING … G NTK CONDITION NUMBER

The paper evaluates SAM across adversarial, contrastive, and transfer learning, while documenting its training settings and analyzing loss geometry, attention maps, Hessian eigenvalues, and NTK condition numbers. These studies connect smoother geometry with improved robustness, accuracy, transferability, and architectural properties.

  • B WHEN SAM MEETS ADVERSARIAL TRAINING: SAM and adversarial training are both minimax formulations, differing in whether the inner maximization perturbs network weights or input images.The paper studies ViTs and MLP-Mixers with fast adversarial training using FGSM with random start, the l∞ norm, and a maximum per-pixel change of 2/255.
  • B WHEN SAM MEETS ADVERSARIAL TRAINING: Landscape smoothing improves convolution-free architectures’ clean and adversarial accuracy, while ResNets trade slightly lower clean accuracy for improved robustness.The evaluation uses clean accuracy, real-world robustness, and adversarial robustness under a 10-step PGD attack; ViTs surpass similar-size ResNets under adversarial training.
  • C WHEN SAM MEETS CONTRASTIVE LEARNING: SAM is coupled with supervised contrastive learning for 350 epochs, followed by 90 epochs of classification-head fine-tuning for ViT-S/16 and ViT-B/16.The method uses a small SAM perturbation strength, ρ = 0.02, during contrastive training.
  • D WHEN SAM MEETS TRANSFER LEARNING: SAM improves ViTs and MLP-Mixers’ transferability, and ViTs transfer better than similarly sized ResNets.Transfer experiments cover CIFAR-10/100, Oxford-IIIT Pets, and Oxford Flowers-102, using 224 × 224 resolution without SAM during fine-tuning.
  • E.1 LOSS LANDSCAPE: Loss-landscape visualizations use filter normalization, cross-entropy loss, 2,500 sampled points, and 10% of ImageNet training images.Cross-entropy is used for fair architectural comparison, although ViTs and MLP-Mixers were originally trained with sigmoid loss, whose geometry is sharper.
  • E.2 ATTENTION MAP: Attention visualization averages classification-token self-attention from the last MSA layer, reshapes it over image patches, and upsamples it to image resolution.The resulting attention matrix has dimensions H/P × W/P before upsampling to H × W.
  • F HESSIAN EIGENVALUE: Hessian λmax is computed on 10% of ImageNet training images with 100 power-iteration steps, while NTK condition numbers are estimated from initialization using 48 × 48 block averaging.For NTK comparison, GELU is replaced by ReLU because smoother activations produce smaller condition numbers.

H TRAINING DETAILS … I LONGER SCHEDULE OF VANILLA SGD

The paper details scratch-training protocols, SAM perturbation tuning, ImageNet-subset adjustments, strong-augmentation settings, and longer vanilla-SGD schedules. Across these settings, training generally follows established protocols, while perturbation strength and schedule are tuned empirically.

  • H TRAINING DETAILS: ImageNet scratch-training hyperparameters use 224 × 224 resolution with basic Inception-style preprocessing.These settings are summarized in Table 12 for ViT-B/16 and Mixer-B/16 training from scratch.
  • H.1 PERTURBATION STRENGTH IN SAM: ImageNet top-1 accuracy is evaluated for ViT-B/16 and Mixer-B/16 across different SAM perturbation strengths ρ.The corresponding results are reported in Table 13.
  • H TRAINING DETAILS: Models are generally trained from scratch on ImageNet with basic Inception-style preprocessing: random crops and 50% horizontal flips.This excludes strong-augmentation and contrastive-learning experiments and follows the original ResNet and ViT settings; Table 12 provides detailed hyperparameters.
  • H.1 PERTURBATION STRENGTH IN SAM: Different architecture species require different SAM perturbation strengths ρ, selected through grid search and reported as the best results.ResNets perform best with relatively small ρ ∈[0.02, 0.05], whereas ViTs and MLP-Mixers benefit from larger values.
  • H.2 TRAINING ON IMAGENET SUBSETS: For ImageNet subsets, only batch size changes to preserve similar total iterations: 2048 for i1k (1/2), 1024 for i1k (1/4), and 512 for i1k (1/10).All other settings remain unchanged, and learning-rate scaling is avoided because it harms performance.
  • H.3 TRAINING WITH STRONG AUGMENTATIONS: With strong augmentations, ViT uses a 1e-3 peak learning rate, 0.1 weight decay, 0.1 Dropout, and 0.1 stochastic depth.The augmentations are mixup with probability 0.5 and RandAugment with two layers and magnitude 15.
  • H.3 TRAINING WITH STRONG AUGMENTATIONS: MLP-Mixer uses the original strong-augmentation hyperparameters: 1e-3 peak learning rate, 0.1 weight decay, 0.0 Dropout, and 0.1 stochastic depth.These settings follow Tolstikhin et al. (2021).
  • I LONGER SCHEDULE OF VANILLA SGD: SAM incurs approximately 2× the training overhead of vanilla training, while extending vanilla training to 600 epochs does not improve clean accuracy or robustness.The results indicate that 300 epochs suffice for ViTs and MLP-Mixers to converge.

J VARYING WEIGHT DECAY STRANGTH

Varying weight decay improves ImageNet accuracy without SAM but drives models toward sharper regions. SAM instead improves accuracy and curvature smoothness across all tested decay strengths, while consistently increasing the weight norm.

  • Effects of weight decay: Weight decay improves ImageNet accuracy without SAM, while increasing its strength decreases the weight norm.The decrease in weight norm follows the expected effect of stronger decay.
  • Effects of weight decay: Stronger weight decay aggravates convergence to sharper regions, measured by both LN train and λmax.
  • Effects of SAM: SAM consistently increases ∥w∥2 across every weight decay strength while improving ImageNet accuracy and producing smoother landscape curvature.
Loading 2106.01548v3…