Early Convolutions Help Transformers See Better

1 Introduction . Vision transformer ( ViT ) models [ 13 ] offer an alternative design paradigm to convolutional neural networks ( CNNs ) [ 24 ] . ViTs replace the inductive bias towards local processing inherent in convolutions with global processing performed by multi-headed self-attention [ 43 ] . The hope is that this design has the potential to improve performance on vision tasks , akin to the trends observed in natural language processing [ 11 ] . While investigating this conjecture , researchers face another unexpected difference between ViTs and CNNs : ViT models exhibit substandard optimizability . ViTs are sensitive to the choice of optimizer [ 41 ] ( AdamW [ 27 ] vs. SGD ) , to the selection of dataset specific learning hyperparameters [ 13 , 41 ] , to training schedule length , to network depth [ 42 ] , etc . These issues render former training recipes and intuitions ineffective and impede research . Convolutional neural networks , in contrast , are exceptionally easy and robust to optimize . Simple training recipes based on SGD , basic data augmentation , and standard hyperparameter values have been widely used for years [ 19 ] . Why does this difference exist between ViT and CNN models ? In this paper we hypothesize that the issues lies primarily in the early visual processing performed by ViT . ViT “ patchifies ” the input image into p×p non-overlapping patches to form the transformer encoder ’ s input set . This patchify stem is implemented as a stride-p p×p convolution , with p = 16 as a default value . This large-kernel plus large-stride convolution runs counter to the typical design 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . choices used in CNNs , where best-practices have converged to a small stack of stride-two 3×3 kernels as the network ’ s stem ( e.g. , [ 30 , 36 , 39 ] ) . To test this hypothesis , we minimally change the early visual processing of ViT by replacing its patchify stem with a standard convolutional stem consisting of only ∼5 convolutions , see Figure 1 . To compensate for the small addition in flops , we remove one transformer block to maintain parity in flops and runtime . We observe that even though the vast majority of the computation in the two ViT designs is identical , 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 . In extensive experiments we show that replacing the ViT patchify stem with a more standard convolutional stem ( i ) allows ViT to converge faster ( §5.1 ) , ( ii ) enables , for the first time , the use of either AdamW or SGD without a significant drop in accuracy ( §5.2 ) , ( iii ) brings ViT ’ s stability w.r.t . learning rate and weight decay closer to that of modern CNNs ( §5.3 ) , and ( iv ) yields improvements in ImageNet [ 10 ] top-1 error of ∼1-2 percentage points ( §6 ) . We consistently observe these improvements across a wide spectrum of model complexities ( from 1G flops to 36G flops ) and dataset scales ( ImageNet-1k to ImageNet-21k ) . These results show that injecting some convolutional inductive bias into ViTs can be beneficial under commonly studied settings . We did not observe evidence that the hard locality constraint in early layers hampers the representational capacity of the network , as might be feared [ 9 ] . In fact we observed the opposite , as ImageNet results improve even with larger-scale models and larger-scale data when using a convolution stem . Moreover , under carefully controlled comparisons , we find that ViTs are only able to surpass state-of-the-art CNNs when equipped with a convolutional stem ( §6 ) . We conjecture that restricting convolutions in ViT to early visual processing may be a crucial design choice that strikes a balance between ( hard ) inductive biases and the representation learning ability of transformer blocks . Evidence comes by comparison to the “ hybrid ViT ” presented in [ 13 ] , which uses 40 convolutional layers ( most of a ResNet-50 ) and shows no improvement over the default ViT . This perspective resonates with the findings of [ 9 ] , who observe that early transformer blocks prefer to learn more local attention patterns than later blocks . Finally we note that exploring the design of hybrid CNN/ViT models is not a goal of this work ; rather we demonstrate that simply using a minimal convolutional stem with ViT is sufficient to dramatically change its optimization behavior . In summary , the findings presented in this paper lead us to recommend using a standard , lightweight convolutional stem for ViT models in the analyzed dataset scale and model complexity spectrum as a more robust and higher performing architectural choice compared to the original ViT model design . 2 Related Work . Convolutional neural networks ( CNNs ) . The breakthrough performance of the AlexNet [ 23 ] CNN [ 15 , 24 ] on ImageNet classification [ 10 ] transformed the field of recognition , leading to the development of higher performing architectures , e.g. , [ 19 , 36 , 37 , 48 ] , and scalable training methods [ 16 , 21 ] . These architectures are now core components in object detection ( e.g. , [ 34 ] ) , instance segmentation ( e.g. , [ 18 ] ) , and semantic segmentation ( e.g. , [ 26 ] ) . CNNs are typically trained with stochastic gradient descent ( SGD ) and are widely considered to be easy to optimize . Self-attention in vision models . Transformers [ 43 ] are revolutionizing natural language processing by enabling scalable training . Transformers use multi-headed self-attention , which performs global information processing and is strictly more general than convolution [ 6 ] . Wang et al . [ 46 ] show that ( single-headed ) self-attention is a form of non-local means [ 2 ] and that integrating it into a ResNet [ 19 ] improves several tasks . Ramachandran et al . [ 32 ] explore this direction further with stand-alone self-attention networks for vision . They report difficulties in designing an attention-based network stem and present a bespoke solution that avoids convolutions . In contrast , we demonstrate the benefits of a convolutional stem . Zhao et al . [ 53 ] explore a broader set of self-attention operations with hard-coded locality constraints , more similar to standard CNNs . Vision transformer ( ViT ) . Dosovitskiy et al . [ 13 ] apply a transformer encoder to image classification with minimal vision-specific modifications . As the counterpart of input token embeddings , they partition the input image into , e.g. , 16×16 pixel , non-overlapping patches and linearly project them to the encoder ’ s input dimension . They report lackluster results when training on ImageNet-1k , but demonstrate state-of-the-art transfer learning when using large-scale pretraining data . ViTs are sensitive to many details of the training recipe , e.g. , they benefit greatly from AdamW [ 27 ] compared to SGD and require careful learning rate and weight decay selection . ViTs are generally considered to be difficult to optimize compared to CNNs ( e.g. , see [ 13 , 41 , 42 ] ) . Further evidence of challenges comes from Chen et al . [ 4 ] who report ViT optimization instability in self-supervised learning ( unlike with CNNs ) , and find that freezing the patchify stem at its random initialization improves stability . ViT improvements . ViTs are gaining rapid interest in part because they may offer a novel direction away from CNNs . Touvron et al . [ 41 ] show that with more regularization and stronger data augmentation ViT models achieve competitive accuracy on ImageNet-1k alone ( cf . [ 13 ] ) . Subsequently , works concurrent with our own explore numerous other ViT improvements . Dominant themes include multi-scale networks [ 14 , 17 , 25 , 45 , 50 ] , increasing depth [ 42 ] , and locality priors [ 5 , 9 , 17 , 47 , 49 ] . In [ 9 ] , d ’ Ascoli et al . modify multi-head self-attention with a convolutional bias at initialization and show that this prior improves sample efficiency and ImageNet accuracy . Resonating with our work , [ 5 , 17 , 47 , 49 ] present models with convolutional stems , but do not analyze optimizability ( our focus ) . Discussion . Unlike the concurrent work on locality priors in ViT , our focus is studying optimizability under minimal ViT modifications in order to derive crisp conclusions . Our perspective brings several novel observations : by adding only ∼5 convolutions to the stem , ViT can be optimized well with either AdamW or SGD ( cf . all prior works use AdamW to avoid large drops in accuracy [ 41 ] ) , it becomes less sensitive to the specific choice of learning rate and weight decay , and training converges faster . We also observe a consistent improvement in ImageNet top-1 accuracy across a wide spectrum of model complexities ( 1G flops to 36G flops ) and dataset scales ( ImageNet-1k to ImageNet-21k ) . These results suggest that a ( hard ) convolutional bias early in the network does not compromise representational capacity , as conjectured in [ 9 ] , and is beneficial within the scope of this study . 3 Vision Transformer Architectures . Next , we review vision transformers [ 13 ] and describe the convolutional stems used in our work . The vision transformer ( ViT ) . ViT first partitions an input image into non-overlapping p×p patches and linearly projects each patch to a d-dimensional feature vector using a learned weight matrix . A patch size of p = 16 and an image size of 224×224 are typical . The resulting patch embeddings ( plus positional embeddings and a learned classification token embedding ) are processed by a standard transformer encoder [ 43 , 44 ] followed by a classification head . Using common network nomenclature , we refer to the portion of ViT before the transformer blocks as the network ’ s stem . ViT ’ s stem is a specific case of convolution ( stride-p , p×p kernel ) , but we will refer to it as the patchify stem and reserve the terminology of convolutional stem for stems with a more conventional CNN design with multiple layers of overlapping convolutions ( i.e. , with stride smaller than the kernel size ) . ViTP models . Prior work proposes ViT models of various sizes , such as ViT-Tiny , ViT-Small , ViT-Base , etc . [ 13 , 41 ] . To facilitate comparisons with CNNs , which are typically standardized to 1 gigaflop ( GF ) , 2GF , 4GF , 8GF , etc. , we modify the original ViT models to obtain models at about these complexities . Details are given in Table 1 ( left ) . For easier comparison with CNNs of similar flops , and to avoid subjective size names , we refer the models by their flops , e.g. , ViTP -4GF in place of ViT-Small . We use the P subscript to indicate that these models use the original patchify stem . Convolutional stem design . We adopt a typical minimalist convolutional stem design by stacking 3×3 convolutions [ 36 ] , followed by a single 1×1 convolution at the end to match the d-dimensional input of the transformer encoder . These stems quickly downsample a 224×224 input image using overlapping strided convolutions to 14×14 , matching the number of inputs created by the standard patchify stem . We follow a simple design pattern : all 3×3 convolutions either have stride 2 and double the number of output channels or stride 1 and keep the number of output channels constant . We enforce that the stem accounts for approximately the computation of one transformer block of the corresponding model so that we can easily control for flops by removing one transformer block when using the convolutional stem instead of the patchify stem . Our stem design was chosen to be purposefully simple and we emphasize that it was not designed to maximize model accuracy . ViTC models . To form a ViT model with a convolutional stem , we simply replace the patchify stem with its counterpart convolutional stem and remove one transformer block to compensate for the convolutional stem ’ s extra flops ( see Figure 1 ) . We refer to the modified ViT with a convolutional stem as ViTC . Configurations for ViTC at various complexities are given in Table 1 ( right ) ; corresponding ViTP and ViTC models match closely on all complexity metrics including flops and runtime . Convolutional stem details . Our convolutional stem designs use four , four , and six 3×3 convolutions for the 1GF , 4GF , and 18GF models , respectively . The output channels are [ 24 , 48 , 96 , 192 ] , [ 48 , 96 , 192 , 384 ] , and [ 64 , 128 , 128 , 256 , 256 , 512 ] , respectively . All 3×3 convolutions are followed by batch norm ( BN ) [ 21 ] and then ReLU [ 29 ] , while the final 1×1 convolution is not , to be consistent with the original patchify stem . Eventually , matching stem flops to transformer block flops results in an unreasonably large stem , thus ViTC-36GF uses the same stem as ViTC-18GF . Convolutions in ViT . Dosovitskiy et al . [ 13 ] also introduced a “ hybrid ViT ” architecture that blends a modified ResNet [ 19 ] ( BiT-ResNet [ 22 ] ) with a transformer encoder . In their hybrid model , the patchify stem is replaced by a partial BiT-ResNet-50 that terminates at the output of the conv4 stage or the output of an extended conv3 stage . These image embeddings replace the standard patchify stem embeddings . This partial BiT-ResNet-50 stem is deep , with 40 convolutional layers . In this work , we explore lightweight convolutional stems that consist of only 5 to 7 convolutions in total , instead of the 40 used by the hybrid ViT . Moreover , we emphasize that the goal of our work is not to explore the hybrid ViT design space , but rather to study the optimizability effects of simply replacing the patchify stem with a minimal convolutional stem that follows standard CNN design practices .

Summary: The authors perform an extensive study investigating the stem of ViT models and replacing it with one that consists of a stack of downsampling 3x3 convolutions. The study reveals that this simple change makes ViT training more robust, faster to converge, but also scales better (cf. i21k pre-training).