ViTAE: Vision Transformer Advanced by Exploring Intrinsic Inductive Bias

1 Introduction Transformers [ 79 , 17 , 40 , 14 , 46 , 61 ] have shown a domination trend in NLP studies owing to their strong ability in modeling long-range dependencies by the self-attention mechanism [ 67 , 81 , 51 ] . Such success and good properties of transformers has inspired following many works that apply them in various computer vision tasks [ 19 , 100 , 97 , 80 , 7 ] . Among them , ViT [ 19 ] is the pioneering pure transformer model that embeds images into a sequence of visual tokens and models the global dependencies among them with stacked transformer blocks . Although it achieves promising performance on image classification , it requires large-scale training data and a longer training schedule . One important reason is that ViT ∗Equal Contribution . Interns at JD Explore Academy . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . lacks intrinsic inductive bias ( IB ) in modeling local visual structures ( e.g. , edges and corners ) and dealing with objects at various scales like convolutions . Alternatively , ViT has to learn such IB implicitly from large-scale data . Unlike vision transformers , Convolution Neural Networks ( CNNs ) naturally equip with the intrinsic IBs of scale-invariance and locality and still serve as prevalent backbones in vision tasks [ 26 , 70 , 62 , 8 , 96 ] . The success of CNNs inspires us to explore intrinsic IBs in vision transformers . We start by analyzing the above two IBs of CNNs , i.e. , locality and scale-invariance . Convolution that computes local correlation among neighbor pixels is good at extracting local features such as edges and corners . Consequently , CNNs can provide plentiful low-level features at the shallow layers [ 94 ] , which are then aggregated into high-level features progressively by a bulk of sequential convolutions [ 32 , 68 , 71 ] . Moreover , CNNs have a hierarchy structure to extract multi-scale features at different layers [ 68 , 38 , 26 ] . Besides , intra-layer convolutions can also learn features at different scales by varying their kernel sizes and dilation rates [ 25 , 70 , 8 , 45 , 96 ] . Consequently , scale-invariant feature representation can be obtained via intra- or inter-layer feature fusion . Nevertheless , CNNs are not well suited to model long-range dependencies2 , which is the key advantage of transformers . An interesting question comes up : Can we improve vision transformers by leveraging the good properties of CNNs ? Recently , DeiT [ 76 ] explores the idea of distilling knowledge from CNNs to transformers to facilitate training and improve the performance . However , it requires an off-the-shelf CNN model as the teacher and consumes extra training cost . Different from DeiT , we explicitly introduce intrinsic IBs into vision transformers by re-designing the network structures in this paper . Current vision transformers always obtain tokens with singlescale context [ 19 , 93 , 80 , 86 , 47 , 69 , 77 ] and learn to adapt to objects at different scales from data . For example , T2T-ViT [ 93 ] improves ViT by delicately generating tokens in a soft split manner . Specifically , it uses a series of Tokens-to-Token transformation layers to aggregate single-scale neighboring contextual information and progressively structurizes the image to tokens . Motivated by the success of CNNs in dealing with scale variance , we explore a similar design in transformers , i.e. , intra-layer convolutions with different receptive fields [ 70 , 91 ] , to embed multi-scale context into tokens . Such a design allows tokens to carry useful features of objects at various scales , thereby naturally having the intrinsic scale-invariance IB and explicitly facilitating transformers to learn scale-invariant features more efficiently from data . On the other hand , low-level local features are fundamental elements to generate high-level discriminative features . Although transformers can also learn such features at shallow layers from data , they are not skilled as convolutions by design . Recently , [ 89 , 43 , 21 ] stack convolutions and attention layers sequentially and demonstrate that locality is a reasonable compensation of global dependency . However , this serial structure ignores the global context during locality modeling ( and vice versa ) . To avoid such a dilemma , we follow the “ divide-and-conquer ” idea and propose to model locality and long-range dependencies in parallel and then fuse the features to account for both . In this way , we empower transformers to learn local and long-range features within each block more effectively . Technically , we propose a new Vision Transformers Advanced by Exploring Intrinsic Inductive Bias ( ViTAE ) , which is a combination of two types of basic cells , i.e. , reduction cell ( RC ) and normal cell ( NC ) . RCs are used to downsample and embed the input images into tokens with rich multi-scale context while NCs aim to jointly model locality and global dependencies in the token sequence . Moreover , these two types of cells share a simple basic structure , i.e. , paralleled attention module and convolutional layers followed by a feed-forward network ( FFN ) . It is noteworthy that RC has an extra pyramid reduction module with atrous convolutions of different dilation rates to embed multi-scale context into tokens . Following the setting in [ 93 ] , we stack three reduction cells to reduce the spatial resolution by 1/16 and a series of NCs to learn discriminative features from data . ViTAE outperforms representative vision transformers in terms of data efficiency and training efficiency ( see Figure 1 ) , as well as classification accuracy and generalization on downstream tasks . Our contributions are threefold . First , we explore two types of intrinsic IB in transformers , i.e. , scale invariance and locality , and demonstrate the effectiveness of this idea in improving the feature learning ability of transformers . Second , we design a new transformer architecture named ViTAE based on two new reduction and normal cells to intrinsically incorporate the above two IBs . The proposed ViTAE embeds multi-scale context into tokens and learns both local and long-range features 2Despite projection in transformer can be viewed as 1 × 1 convolution [ 9 ] , the term of convolution here refers to those with larger kernels , e.g. , 3× 3 , which are widely used in typical CNNs to extract spatial features . effectively . Third , ViTAE outperforms representative vision transformers regarding classification accuracy , data efficiency , training efficiency , and generalization on downstream tasks . ViTAE achieves 75.3 % and 82.0 % top-1 accuracy on ImageNet with 4.8M and 23.6M parameters , respectively . 2 Related Work . 2.1 CNNs with intrinsic IB . CNNs have led to a series of breakthroughs in image classification [ 38 , 94 , 26 , 95 , 87 ] and downstream computer vision tasks . The convolution operations in CNNs extract local features from the neighbor pixels within the receptive field determined by the kernel size [ 42 ] . Following the intuition that local pixels are more likely to be correlated in images [ 41 ] , CNNs have the intrinsic IB in modeling locality . In addition to the locality , another critical topic in visual tasks is scale-invariance , where multi-scale features are needed to represent the objects at different scales effectively [ 49 , 90 ] . For example , to effectively learn features of large objects , a large receptive field is needed by either using large convolution kernels [ 90 , 91 ] or a series of convolution layers in deeper architectures [ 26 , 32 , 68 , 71 ] . To construct multi-scale feature representation , the classical idea is using image pyramid [ 8 , 1 , 55 , 4 , 39 , 16 ] , where features are hand-crafted or learned from a pyramid of images at different resolutions respectively [ 44 , 8 , 52 , 63 , 35 , 3 ] . Accordingly , features from the small scale image mainly encode the large objects while features from the large scale image respond more to small objects . In addition to the above inter-layer fusion way , another way is to aggregate multi-scale context by using multiple convolutions with different receptive fields within a single layer , i.e. , intra-layer fusion [ 96 , 71 , 70 , 70 , 72 ] . Either inter-layer fusion or intra-layer fusion empower CNNs an intrinsic IB in modeling scale-invariance . This paper introduces such an IB to vision transformers by following the intra-layer fusion idea and utilizing multiple convolutions with different dilation rates in the reduction cells to encode multi-scale context into each visual token . 2.2 Vision transformers with learned IB . ViT [ 19 ] is the pioneering work that applies a pure transformer to vision tasks and achieves promising results . However , since ViT lacks intrinsic inductive bias in modeling local visual structures , it indeed learns the IB from amounts of data implicitly . Following works along this direction are to simplify the model structures with fewer intrinsic IBs and directly learn them from large scale data [ 50 , 74 , 75 , 22 , 18 , 20 , 27 ] which have achieved promising results and been studied actively . Another direction is to leverage the intrinsic IB from CNNs to facilitate the training of vision transformers , e.g. , using less training data or shorter training schedules . For example , DeiT [ 76 ] proposes to distill knowledge from CNNs to transformers during training . However , it requires an off-the-shelf CNN model as a teacher , introducing extra computation cost during training . Recently , some works try to introduce the intrinsic IB of CNNs into vision transformers explicitly [ 23 , 58 , 21 , 43 , 15 , 89 , 83 , 92 , 6 , 47 , 11 ] . For example , [ 43 , 21 , 83 ] stack convolutions and attention layers sequentially , resulting in a serial structure and modeling the locality and global dependency accordingly . [ 80 , 28 ] design sequential stage-wise structures while [ 47 , 33 ] apply attention within local windows . However , these serial structure may ignore the global context during locality modeling ( and vice versa ) . [ 88 ] establishes connection across different scales at the cost of heavy computation . Instead , we follow the “ divide-and-conquer ” idea and propose to model locality and global dependencies simultaneously via a parallel structure within each transformer layer . Conformer [ 58 ] , the most relevant concurrent work to us , employs a unit to explore inter-block interactions between parallel convolution and transformer blocks . In contrast , in ViTAE , the convolution and attention modules are designed to be complementary to each other within the transformer block . In addition , Conformer is not designed to have inherent scale invariance IB . 3 Methodology . 3.1 Revisit vision transformer . We first give a brief review of vision transformer in this part . To adapt transformers to vision tasks , ViT [ 19 ] first splits an image x ∈ RH×W×C into tokens with a reduction ratio of p ( i.e. , xt ∈ R ( ( H×W ) /p 2 ) ×D ) , where H , W and C denote the height , width , and channel dimensions of the input image , D = Cp2 denotes the token dimension . Then , an extra class token is concatenated to the visual tokens before adding position embeddings in an element-wise manner . The resulting tokens are fed into the following transformer layers . Each transformer layer is composed of two parts , i.e. , a multi-head self-attention module ( MHSA ) and a feed forward network ( FFN ) . MHSA Multi-head self-attention extends single-head self-attention ( SHSA ) by using different projection matrices for each head . Specifically , the input tokens xt are first projected to queries ( Q ) , keys ( K ) and values ( V ) using projection matrices , i.e. , Q , K , V = xtWQ , xtQK , xtQV , where WQ/K/V ∈ RD×D denotes the projection matrix for query , key , and value , respectively . Then , the self-attention operation is calculated as : Attention ( Q , K , V ) = softmax ( QKT√ D ) V. ( 1 ) This SHSA module is repeated for h times to formulate the MHSA module , where h is the number of heads . The output features of the h heads are concatenated along the channel dimension and formulate the output of the MHSA module . FFN FFN is placed on top of the MHSA module and applied to each token identically and separately . It consists of two linear transformations with an activation function in between . Besides , a layer normalization [ 2 ] and a shortcut are added before and aside from the MHSA and FFN , respectively .

Summary: This paper proposed two types of basic cells to modify the vanilla vision transformer structure: reduction cell (RC) and normal cell (NC). RCs are used to down-sample and embed the input images into tokens and NCs are to jointly model locality and global dependencies in the token sequence. Specifically, they used paralleled attention module and convolutional layers followed by a feed-forward network (FFN) to construct these two cells. Experiments on ImageNet and downstream tasks verified the superiority of the proposed network design.