PaperlayerView exact revision on arXiv

Source-linked AI summary

Early Convolutions Help Transformers See Better

Tete Xiao, Mannat Singh, Eric Mintun, Trevor Darrell, Piotr Dollár, Ross Girshick

arXiv:2106.14881v3cs.CV

TL;DR

ViTs are unusually sensitive to optimizer, hyperparameter, and training-schedule choices, motivating investigation of their atypical patchify stem. The paper replaces that stem with a lightweight convolutional stem while matching computational budgets, finding more stable optimization and approximately 1-2% ImageNet top-1 improvement across model and dataset scales.

  • Problem

    ViT models have substandard optimizability, including sensitivity to optimizer choice, hyperparameters, and training schedule length.

  • Method

    The paper replaces ViT’s stride-p p×p patchify stem with a lightweight convolutional stem and removes one transformer block to preserve computational parity.

  • Results

    ∼1-2% ImageNet top-1 improvement accompanies faster convergence, greater optimization stability, and consistent gains from 1G to 36G flops and ImageNet-1k to ImageNet-21k.

  • Takeaways & Limitations

    A standard, lightweight convolutional stem is recommended as a more robust and higher-performing ViT architecture within the analyzed model-complexity and dataset-scale regime.

Abstract

from arXiv · show

Vision transformer (ViT) models exhibit substandard optimizability. In particular, they are sensitive to the choice of optimizer (AdamW vs. SGD), optimizer hyperparameters, and training schedule length. In comparison, modern convolutional neural networks are easier to optimize. Why is this the case? In this work, we conjecture that the issue lies with the patchify stem of ViT models, which is implemented by a stride-p p*p convolution (p=16 by default) applied to the input image. This large-kernel plus large-stride convolution runs counter to typical design choices of convolutional layers in neural networks. To test whether this atypical design choice causes an issue, we analyze the optimization behavior of ViT models with their original patchify stem versus a simple counterpart where we replace the ViT stem by a small number of stacked stride-two 3*3 convolutions. While the vast majority of computation in the two ViT designs is identical, we find that this small change in early visual processing results in markedly different training behavior in terms of the sensitivity to optimization settings as well as the final model accuracy. Using a convolutional stem in ViT dramatically increases optimization stability and also improves peak performance (by ~1-2% top-1 accuracy on ImageNet-1k), while maintaining flops and runtime. The improvement can be observed across the wide spectrum of model complexities (from 1G to 36G flops) and dataset scales (from ImageNet-1k to ImageNet-21k). These findings lead us to recommend using a standard, lightweight convolutional stem for ViT models in this regime as a more robust architectural choice compared to the original ViT model design.

1 Introduction

The paper attributes ViT’s optimization difficulties primarily to its atypical patchify stem and tests a lightweight convolutional replacement. With matched computational budgets, the convolutional-stem variant improves optimization behavior and ImageNet performance across model and dataset scales.

  • ViTs are more sensitive than CNNs to optimizer choice, hyperparameters, and training schedule length.
  • The original ViT stem uses a non-overlapping stride-p p×p convolution, with p = 16 by default, unlike standard CNN stems.CNN practice favors a small stack of stride-two 3×3 convolutions.
  • The proposed convolutional stem replaces patchification with approximately five convolutions while removing one transformer block to maintain flops, parameters, and runtime parity.
  • The convolutional-stem variant converges faster, supports AdamW or SGD without a significant accuracy drop, and is more stable to learning-rate and weight-decay choices.
  • ∼1-2 percentage points of ImageNet top-1 error improvement is obtained with the convolutional stem across 1G to 36G flops and ImageNet-1k to ImageNet-21k.
  • The authors recommend a standard, lightweight convolutional stem as a more robust and higher-performing ViT design within the analyzed regime.

2 Related Work

Related work contrasts CNNs’ established optimization robustness with ViTs’ optimization and data requirements. The paper focuses on whether a minimal convolutional stem improves ViT optimizability without compromising its transformer-based design.

  • CNNs are widely considered easy to optimize and are typically trained with stochastic gradient descent.
  • Vision transformers use global multi-headed self-attention instead of convolutional local-processing inductive bias.
  • ViTs initially showed lackluster ImageNet-1k training results but strong transfer learning with large-scale pretraining data.
  • Prior ViT improvements include stronger regularization, data augmentation, multi-scale networks, increased depth, and locality priors.
  • This paper studies optimizability under minimal ViT modifications and reports that approximately five added stem convolutions enable both AdamW and SGD, reduce hyperparameter sensitivity, accelerate convergence, and improve ImageNet accuracy across scales.

3 Vision Transformer Architectures

The paper defines ViT’s patchify stem and constructs matched convolutional-stem variants for controlled comparisons. The replacement uses overlapping convolutions while preserving comparable computational complexity.

  • ViT partitions an image into non-overlapping p×p patches, projects them to d-dimensional vectors, adds positional and classification embeddings, and processes them with a transformer encoder.
  • Corresponding ViTP and ViTC models match closely on flops, parameters, activations, and ImageNet-1k epoch training time.
  • The patchify stem is a stride-p, p×p convolution, whereas a convolutional stem uses multiple overlapping convolutions with stride smaller than kernel size.
  • ViTP models use the original patchify stem and are configured near 1GF, 4GF, 18GF, and 36GF for comparison with CNNs.
  • ViTC replaces the patchify stem with stacked 3×3 convolutions followed by a 1×1 projection, then removes one transformer block to compensate for added stem flops.
  • A prior hybrid ViT used a deep partial BiT-ResNet-50 stem with 40 convolutional layers.

4 Measuring Optimizability

The paper operationalizes optimizability through training-length, optimizer, hyperparameter, and peak-performance comparisons. These measures quantify convergence speed, sensitivity to optimizer choice, sensitivity to learning-rate and weight-decay choices, and best controlled performance.

  • Optimizability describes optimization characteristics such as sensitivity to augmentation and optimizer choice, depth-related degradation, and modest mid-size-dataset performance.
  • Training length stability: Training-length stability is the gap between accuracy at 50, 100, or 200 epochs and approximate asymptotic accuracy at 400 epochs.The gap serves as a measure of convergence speed.
  • Optimizer stability: Optimizer stability is the accuracy gap between AdamW and SGD.SGD is practically attractive because it uses fewer hyperparameters and 50% less optimizer-state memory.
  • Hyperparameter stability: Hyperparameter stability evaluates performance across learning-rate and weight-decay values because both can substantially affect optimization results.
  • Peak performance: Peak performance is defined as 400-epoch performance using the best-performing optimizer and parsimoniously tuned learning rate and weight decay.

5 Stability Experiments

Replacing ViT’s patchify stem with a convolutional stem improves convergence and reduces sensitivity to optimizer and hyperparameter choices across model complexities.

  • Experimental Setup: The stability experiments use ImageNet-1k, controlled training settings, and model comparisons spanning ViT stems, model sizes, epoch schedules, and optimizers.Training uses a single half-period cosine decay with 5-epoch warm-up, batch size 2048, and shared regularization and augmentation settings.
  • 5.1 Training Length Stability: ViTC converges faster than ViTP and closes much of the convergence-rate gap to CNNs across model complexities.At 50 epochs, ViTP-1GF has a 10% error delta, while ViTC-1GF reduces this to about 6%.
  • 5.2 Optimizer Stability: ViTP suffers up to a 10% error drop with SGD, whereas ViTC reduces the SGD–AdamW gap to less than 0.2% in larger models and longer schedules.RegNetY shows only a ∼0.1-0.2% difference between the optimizers, within noise.
  • 5.2 Optimizer Stability: Under SGD, ∼60% of ViTC-18GF models fall within 4% ∆top-1 error of the best result, compared with less than 20% of ViTP-18GF models.Most ViTP-18GF runs do not converge.
  • 5.2 Optimizer Stability: ViTC significantly improves learning-rate and weight-decay stability over ViTP for both AdamW and SGD.The stability evaluation samples 64 learning-rate and weight-decay pairs around each model’s optimum and compares ∆top-1 error distributions.

6 Peak Performance

Under fair comparisons, the convolutional-stem ViTC is more competitive than the patchify-stem ViTP and achieves stronger peak performance across training speeds and dataset scales.

  • ImageNet-1k: ViTP underperforms RegNetY across the ImageNet-1k model-complexity spectrum and also underperforms ResNets, while ViTC outperforms CNNs in the middle-complexity range.These comparisons use ImageNet-1k training only.
  • ImageNet-21k Pretraining: With ImageNet-21k pretraining, ViTC strictly outperforms ViTP and RegNetY, whereas ViTP still does not outperform a state-of-the-art CNN.The comparison is plotted against ImageNet-1k validation top-1 error and epoch training time.
  • Evaluation: The controlled comparison evaluates ImageNet-1k validation top-1 error against epoch training time across model-complexity sweeps.Models share 400 epochs, EMA, and regularization and augmentation methods, while optimizer settings are selected by model family.
  • Main Finding: The convolutional stem improves peak performance across model-complexity and dataset-scale settings, while ViTP struggles to compete with state-of-the-art CNNs.The paper reports improvements only when using both large-scale pretraining data and the proposed convolutional stem.

7 Conclusion

The paper links ViT optimization challenges to its large-stride, large-kernel patchify stem and shows that a simple convolutional replacement improves optimization and accuracy across model and dataset scales.

  • Replacing ViT’s patchify stem with a simple convolutional stem makes training converge faster and work well with either AdamW or SGD.
  • The convolutional stem improves learning-rate and weight-decay stability while raising ImageNet top-1 performance by approximately 1–2%.
  • These improvements remain consistent from 1GF to 36GF models and from ImageNet-1k to ImageNet-21k.
  • The results support adding a small amount of convolutional inductive bias during ViT’s early visual processing.

Appendix A: Stem Design Ablation Experiments

Stem ablations indicate that larger patchify layers degrade ViT error and optimizer stability, while normalization and nonlinearity choices have comparatively small effects.

  • Stem design: Increasing the patchify-layer pixel size from p = 2 to p = 16 reliably worsens error and optimizer stability.
  • Stem design: The standard convolutional stem outperforms four alternative stems that each retain one patchify layer, despite comparable overall flops.
  • Stem normalization and non-linearity: Adding BN and ReLU after the default patchify stem slightly worsens best top-1 error without changing learning-rate or weight-decay stability.
  • Stem normalization and non-linearity: Replacing BN with LN in the convolutional stem marginally degrades best top-1 error and stability, while the 100-epoch error gap remains small.

Appendix B: Deeper Model Ablation Experiments

In deeper ViTs, the convolutional stem improves stability despite contributing only about 2% of total flops, allowing varied optimizer settings without major performance loss.

  • Deeper models: Despite accounting for only ∼2% of total flops, the convolutional stem solidly improves deeper ViT stability over its patchify counterpart.
  • Deeper models: A range of learning-rate and weight-decay choices trains deeper ViTs without a large drop in top-1 performance or additional modifications.
  • Deeper models: The deeper 48-block patchify ViT has better learning-rate and weight-decay stability than the 4GF and 18GF comparison models over the sampled range.

Appendix C: Larger Model ImageNet-21k Experiments

The 72GF experiments extend the convolutional-stem comparison to larger models, where the stem still improves error but a new training instability appears under transferred hyperparameters.

  • Larger models: The 72GF convolutional-stem model uses the same stem design as the 18GF and 36GF models, without removing a transformer block because its added flops are marginal.
  • Larger models: Using 36GF hyperparameters for 72GF models reveals training-error spikes that may recover within an epoch or persist for several epochs.
  • Larger models: For 72GF models, the convolutional stem more often produces multi-epoch recovery spikes, while the patchify stem more often produces within-epoch spikes.
  • Larger models: Reducing weight decay and freezing an ImageNet-21k-pretrained convolutional stem makes 72GF ImageNet-21k training feasible.
  • Larger models: On ImageNet-1k fine-tuning, the 72GF patchify and convolutional models reach 14.2% and 13.6% top-1 error, respectively.
  • Runtime measurement: For the broader model comparisons, activations correlate more strongly with runtime than flops or parameters, with Pearson r values of 0.93, 0.75, and 0.71.

Appendix F: Regularization and Data Augmentation

The appendix compares a simplified ViT training setup with DeiT’s recipe and examines convergence, regularization, and evaluation robustness. The simplified setup converges faster and can slightly improve top-1 error, but EMA remains important for larger models.

  • Training setup: The simplified setup is designed to resemble modern CNN training while retaining competitive accuracy with DeiT.It was developed using baseline ViTP-4GF and ViTP-18GF models.
  • Convergence: Removing augmentations and using model EMA accelerates convergence while slightly improving reproduced DeiT top-1 error.Figure 10 compares 4GF and 18GF models across 50-, 100-, 200-, and 400-epoch schedules.
  • Regularization: A 5–10× higher weight decay permits disabling erasing, stochastic depth, and repeating augmentation without the failures reported under DeiT’s settings.The cited comparison reports failed training or significant accuracy drops under the original learning-rate and weight-decay choices.
  • Regularization: Model EMA can be removed for ViTP-4GF but remains essential for ViTP-18GF, especially with 400-epoch training.The larger model can train without EMA only when additional augmentation and regularization are retained.
  • Convergence: Faster convergence is observed for both ViTP and ViTC models under the simplified setup.The appendix attributes this improvement to removing repeating augmentation, which had been shown to slow convergence.

Appendix G: ImageNet-V2 Evaluation

The appendix evaluates selected models on ImageNet-V2, a test set collected using the original ImageNet procedure. It compares ImageNet validation performance with ImageNet-V2 top-1 error to assess whether model rankings remain consistent.

  • Evaluation setup: The evaluation uses 400-epoch or ImageNet-21k models from Table 2, selecting whichever is better for each model.These models are then evaluated on the ImageNet-V2 test set.
  • Evaluation setup: ImageNet-V2 is described as a new test set collected following the original ImageNet procedure.The appendix contrasts it with the original ImageNet validation set used in the main experiments.
  • Results: Figure 11 compares top-1 errors on the original validation set and ImageNet-V2 to examine whether model rankings are preserved.The passage reports that rankings are mostly preserved within the estimated noise level.
Loading 2106.14881v3…