  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.
