Early Convolutions Help Transformers See Better

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.

arXiv commentNeurIPS 2021

References in corpus (14)

• Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift
• Faster R-CNN: Towards Real-Time Object Detection with Region Proposal Networks
• Decoupled Weight Decay Regularization
• Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour
• ConViT: Improving Vision Transformers with Soft Convolutional Inductive Biases
• AutoAugment: Learning Augmentation Policies from Data
• Swin Transformer: Hierarchical Vision Transformer using Shifted Windows
• Pyramid Vision Transformer: A Versatile Backbone for Dense Prediction without Convolutions
• Augment your batch: better training with larger batches
• MultiGrain: a unified image embedding for classes and instances
• Tokens-to-Token ViT: Training Vision Transformers from Scratch on ImageNet
• FBNetV3: Joint Architecture-Recipe Search using Predictor Pretraining
• An Empirical Study of Training Self-Supervised Vision Transformers
• Incorporating Convolution Designs into Visual Transformers

Cited by in corpus (20)

• Improving Chest X-Ray Report Generation by Leveraging Warm Starting
• MobileViT: Light-weight, General-purpose, and Mobile-friendly Vision Transformer
• 3DSAM-adapter: Holistic adaptation of SAM from 2D to 3D for promptable tumor segmentation
• SwinIR: Image Restoration Using Swin Transformer
• On the Role of ViT and CNN in Semantic Communications: Analysis and Prototype Validation
• A Survey on State-of-the-art Deep Learning Applications and Challenges
• PSLT: A Light-weight Vision Transformer with Ladder Self-Attention and Progressive Shift
• Prototype Learning Guided Hybrid Network for Breast Tumor Segmentation in DCE-MRI
• ClST: A Convolutional Transformer Framework for Automatic Modulation Recognition by Knowledge Distillation
• STB-VMM: Swin Transformer Based Video Motion Magnification
• Homogeneous Tokenizer Matters: Homogeneous Visual Tokenizer for Remote Sensing Image Understanding
• On the Integration of Self-Attention and Convolution
• Little Strokes Fell Great Oaks: Boosting the Hierarchical Features for Multi-exposure Image Fusion
• Parameterization of Cross-Token Relations with Relative Positional Encoding for Vision MLP
• Hire-MLP: Vision MLP via Hierarchical Rearrangement
• Transformer Compressed Sensing via Global Image Tokens
• UniNet: Unified Architecture Search with Convolution, Transformer, and MLP
• Semi-Supervised Vision Transformers
• Ripple Attention for Visual Perception with Sub-quadratic Complexity
• SwinSF: Image Reconstruction from Spatial-Temporal Spike Streams