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: The paper propose to inject inductive bias (IB) for vision tasks to boost the performance of vision transformers (ViT). The main contributions of the paper are two-fold: a reduction cell that aggregate the features from different scales by controlling the dilation of a  few of parallel convolution branches and a parallel convolution branch with the multi-head self-attention (MHSA) block. The results on ImageNet is promising compared with the selected baselines but I found it is hard to get a conclusive result on downstream classification tasks.