[Paper Review] Early Convolutions Help Transformers See Better

Tete Xiao, Mannat Singh+4 more|arXiv (Cornell University)|Jun 28, 2021

Advanced Neural Network ApplicationsComputer Science56 references352 citations

TL;DR

Replacing ViT's patchify stem with a lightweight convolutional stem (ViT C) markedly improves optimization stability, enables SGD or AdamW without accuracy loss, and gains 1–2% top-1 on ImageNet across model sizes and data scales.

ABSTRACT

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.

Motivation & Objective

• Motivate and understand why ViT optimizability lags behind CNNs.
• Investigate whether early visual processing (patchify stem) causes optimization issues in ViT.
• Evaluate a minimal convolutional stem as a drop-in replacement and compare training dynamics with the original ViT stem.
• Assess convergence speed, optimizer compatibility, hyperparameter stability, and peak performance across diverse model sizes and data scales.

Proposed method

• Replace the ViT patchify stem (stride-p, p×p) with a lightweight convolutional stem of ~5 layers.
• Remove one transformer block to keep FLOPs and runtime parity.
• Train ViT P (original patchify stem) and ViT C (convolutional stem) across multiple sizes (1G–36G FLOPs) and ImageNet scales (1k–21k).
• Evaluate optimization behavior using training length stability, optimizer stability (AdamW vs SGD), and hyperparameter stability (lr, wd) via error distribution functions (EDFs).
• Compare peak performance under controlled training settings with EMA weights and standard augmentation recipes.

Experimental results

Research questions

• RQ1 Does replacing the patchify stem with a convolutional stem affect ViT optimization stability and convergence speed?
• RQ2 Can ViT models with a convolutional stem be effectively trained with SGD as well as AdamW without significant accuracy loss?
• RQ3 How does the convolutional stem impact learning rate and weight decay stability compared to the patchify stem?
• RQ4 What is the impact on peak ImageNet top-1 performance across different model complexities and dataset scales?
• RQ5 Is the convolutional stem sufficient to make ViTs competitive with CNNs under standard training recipes?

Key findings

• ViT C converges faster than ViT P across model complexities.
• ViT C enables training with either AdamW or SGD without significant accuracy loss, unlike ViT P which struggles with SGD.
• ViT C shows greater stability to learning rate and weight decay, reducing hyperparameter sensitivity to near CNN-like levels.
• Across 1G–36G FLOPs and ImageNet 1k–21k, ViT C yields ~1–2 percentage points lower top-1 error than ViT P on ImageNet-1k.
• Under fair peak-performance comparisons, ViT C outperforms ViT P and competitive CNNs in the middle-complexity range; with ImageNet-21k pretraining, ViT C surpasses all.