RegionViT: Regional-to-Local Attention for Vision Transformers

1 INTRODUCTION . Transformers ( Vaswani et al. , 2017 ) based on self-attention come naturally with the ability to learn long-range dependencies in sequential data . As important as it is to language modeling ( Devlin et al. , 2019 ) , such ability is also highly desired for many vision tasks where contextual modeling plays a significant role . For this reason , self-attention and transformers have been receiving an increasing attention in the vision community ( Bello , 2021 ; Srinivas et al. , 2021 ; Zhao et al. , 2020 ; Ramachandran et al. , 2019a ; Bello et al. , 2019 ; Hu et al. , 2019 ; Ramachandran et al. , 2019b ; Wang et al. , 2018 ) . Especially , the recent Vision Transformers ( ViT ) ( Dosovitskiy et al. , 2021 ) demonstrates comparable image classification results against the firmly established and prevalent CNNs in computer vision ( He et al. , 2016 ; Tan & Le , 2019 ; Brock et al. , 2021 ) , albeit relying on a huge amount of training data . It has since then led to an explosion of interest in further investigating its potential for a wide variety of vision applications ( Wu et al. , 2021 ; Wang et al. , 2021 ; Heo et al. , 2021 ; Zhang et al. , 2021a ; Li et al. , 2021 ; Graham et al. , 2021 ; Liu et al. , 2021 ; Chu et al. , 2021a ; Yan et al. , 2021 ; Chen et al. , 2021b ) . The ViT inherits the entire architecture from the vanilla transformer ( Vaswani et al. , 2017 ) , which is designed for natural language processing tasks , and some of those designs thus may not meet the needs of vision tasks . For example , the transformer has an isotropic network structure with a fixed number of tokens and unchanged embedding size , which loses the capability to model the context with different scales and allocates computations at different scales . As opposed to this , a majority of CNNs adopt a popular pyramid architecture to compute multi-scale features efficiently . Recent vision transformers such as PVT ( Wang et al. , 2021 ) and PiT ( Heo et al. , 2021 ) also follow a similar pyramid structure as CNNs , showing improvement on both computation and memory efficiency as well as on model accuracy . Another critical bottleneck of the transformer is that the self-attention module has a quadratic cost in memory and computation with regard to the sequence length ( i.e. , the number of tokens ) . This issue is even worse in ViT as images are 2-D , suggesting a quadratic relationship between the number of tokens and image resolution . As a result , ViT indicates a quadruple complexity w.r.t image resolution . The highly compute- and memory-intensive self-attention makes it challenging to train vision transformer models at fine-grained patch sizes . It also significantly undermines the applications of these models to tasks such as object detection and semantic segmentation , which benefit from or require fine feature information computed from high-resolution images . To address the aforementioned computational limitations of vision transformers , in this work , we develop a memory-friendly and efficient self-attention method for transformer models to reach their promising potential for vision applications . We propose a novel coarse-to-fine mechanism to compute self-attention in a hierarchical way . Specifically , our approach first divides the input image into a group of nonoverlapping patches of large size ( e.g. , 28×28 ) , on which regional tokens are computed via linear projection . Similarly , local tokens are created for each region using a smaller patch size ( e.g. , 4×4 ) . We then use a standard transformer to process regional and local tokens separately . To enable communication between the two types of tokens , we first perform self-attention on regional tokens ( regional attention ) and then jointly attend to the local tokens of each region including their associated regional token ( local attention ) . By doing so , regional tokens pass global contextual information to local tokens efficiently while being able to effectively learn from local tokens themselves . For clarity , we represent this two-stage attention mechanism as Regional-to-Local Attention , or R2L attention for short ( see Figure 1 for an illustration ) . Since both regional and local attention involve much fewer tokens , our R2L attention requires substantially less memory than regular global self-attention used in vision transformers . For example , in our default setting , the memory saving using R2L attention can be up to as much as 73 % . We demonstrate the effectiveness of our approach on image classification and several downstream vision tasks including object detection and action recognition . To summarize , our key contributions in this work are as follows : 1 . We propose a new vision transformer ( RegionViT ) based on regional-to-local attention to learn both local and global features . Our proposed regional-to-local attention alleviates the overhead of standard global attention ( too many tokens ) and the weakness of pure local attention ( no interaction between regions ) used in existing vision transformers . 2 . Our regional-to-local attention reduces the memory complexity significantly as compared to standard self-attention , leading to a savings in memory complexity by about O ( N/M2 ) , where M is the window size of a region and N is the total number of tokens . This effectively allows us to train a more deep network for better performance while with comparable complexity . 3 . Our models outperform or on par with several concurrent works on vision transformer that exploit pyramid structure for image classification . Experiments also demonstrate that our models work well on several downstream classification tasks , including object detection and action recognition . 2 RELATED WORK . CNNs with Attention . Attention has been widely used for enhancing the features in CNNs , e.g. , SENet ( Hu et al. , 2018 ) and CBAM ( Woo et al. , 2018 ) . Various methods have been proposed that combine self-attention with CNNs ( Bello , 2021 ; Srinivas et al. , 2021 ; Zhao et al. , 2020 ; Ramachandran et al. , 2019a ; Bello et al. , 2019 ; Hu et al. , 2019 ; Ramachandran et al. , 2019b ; Wang et al. , 2018 ) . E.g. , SAN ( Zhao et al. , 2020 ) and SASA ( Ramachandran et al. , 2019a ) attempts to replace all convolutional layers with local self-attention network . While these works show promising results , their complexity is relatively high due to the self-attention . On the other hand , BoTNet ( Srinivas et al. , 2021 ) replaces few convolutional layers with a slightly-modified self-attention to balance computation and accuracy . LambdaNet work ( Bello , 2021 ) uses the approximated self-attention to reduce the overhead of selfattention and make the network efficient . Few works use attention approach to define the regional of interest for the fine-grained visual recognition ( Zheng et al. , 2020 ; Ding et al. , 2021 ; Wharton et al. , 2021 ) . In contrast , our model utilizes regional-to-local attention to alleviate the workload of global self-attention by local attention while still keeping the global information via regional attention . Vision Transformer . ViT ( Dosovitskiy et al. , 2021 ) has recently achieved comparable results to CNNs when using a huge amount of data , e.g. , JFT300M ( Sun et al. , 2017 ) . DeiT ( Touvron et al. , 2020 ; Wightman , 2019 ) subsequently proposes an efficient training scheme that allows vision transformer to work on par with CNNs while training only on ImageNet1K ( Deng et al. , 2009 ) . Despite promising performance of ViT , the architecture is directly borrowed from natural language processing , which might not be suitable for vision applications . Motivated by this , two main line of research works have been developed to improve ViT . One is to enhance different components of the vision transformer ( Yuan et al. , 2021 ; Han et al. , 2021 ; Touvron et al. , 2021 ; Jiang et al. , 2021 ) while still using isotropic structure ( i.e. , fixed token numbers and channel dimension ) like ViT , e.g. , while T2T-ViT ( Yuan et al. , 2021 ) introduces a Tokens-to-Token ( T2T ) transformation to encode local structure for each token , CrossViT propose a dual-path architecture , each with different scales , to learn multi-scale features ( Chen et al. , 2021a ) . CaiT ( Touvron et al. , 2021 ) proposes a layer-scalar to train a deeper network for better performance and LV-ViT ( Jiang et al. , 2021 ) modified how the model is trained when CutMix ( Yun et al. , 2019 ) augmentation is applied on ViT . Another parallel thread for improving vision transformer is in incorporating CNN-like pyramid structure into ViT ( Wu et al. , 2021 ; Wang et al. , 2021 ; Heo et al. , 2021 ; Zhang et al. , 2021a ; Li et al. , 2021 ; Graham et al. , 2021 ; Liu et al. , 2021 ; Chu et al. , 2021a ; Yan et al. , 2021 ; Chen et al. , 2021b ) . PVT ( Wang et al. , 2021 ) and PiT ( Heo et al. , 2021 ) introduce the pyramid structure into ViT , which makes them more suitable for objection detection as it can provide multi-scale features . LocalViT ( Li et al. , 2021 ) and ConT ( Yan et al. , 2021 ) mix convolutions with self-attention to encode locality information . Swin ( Liu et al. , 2021 ) , ViL ( Zhang et al. , 2021a ) and Twins ( Chu et al. , 2021a ) limited the self-attention into a local region and then propose different methods to allow the interaction among each local region . Our model also utilizes pyramid structure and limits the self-attention in a local region ; while we propose the regional-to-local attention to exchange the information among each region efficiently . Difference from PVT , Swin , ViL and Twins . Table 1 summarizes the difference between our proposed approach , RegionViT with the closely related works . PVT ( Wang et al. , 2021 ) uses pyramid structure but the scope of self-attention is still global . While ViL ( Zhang et al. , 2021a ) limits self-attention into a local region , they adopt overlapping windows as convolutions to process all the tokens . On the other hand , Swin ( Liu et al. , 2021 ) , Twins ( Chu et al. , 2021a ) and ours adopt non-overlapped windows . While Swin shifts the position of windows alternatively in the consecutive transformer encoder to allow the interaction between regions , Twins subsamples local tokens as the global tokens , and then use the global tokens as the keys for self-attention to achieve the interaction between regions . By contrast , our proposed method utilizes a extra set of regional tokens with regional-to-local attention to perform self-attention on regional tokens only for the global information and then each regional token is sent to the associated local tokens to pass the global information . One major difference from others is that the local self-attention in our approach includes one extra regional token to exchange the global information with local tokens . 3 METHOD . Our method is built on top of a vision transformer , so we first present a brief background of ViT and then describe our method ( RegionViT ) on regional-to-local attention for vision transformers . Vision Transformer . As opposed to CNN-based approaches for image classification , ViT is a purely attention-based counterpart , borrowed from NLP . It consists of stacked transformer encoders , each of which contains a multihead self-attention ( MSA ) and a feed-forward network ( FFN ) with layer normalization and residual shortcut . To classify an image , ViT first splits it into patches of fixed size , e.g . 16×16 , and then transforms them into tokens by linear projection . A class token is additionally prepended to the patch tokens to form the input sequence . Before the token are fed into transformer encoders , a learnable absolute positional embedding is added to each token to learn the position information . At the end of the network , the class token is used as the final feature representation for classification . Mathematically , ViT can be expressed as x0 = [ xcls||xpatch ] + xpos yk = xk−1 + MSA ( LN ( xk−1 ) ) , xk = yk + FFN ( LN ( yk ) ) , ( 1 ) where xcls ∈ R1×C and xpatch ∈ RN×C are the class token and patch tokens respectively and xpos ∈ R ( 1+N ) ×C is the position embedding , and [ || ] denotes the tensor concatenation . k , N and C are the layer index , the number of patch tokens and dimension of the embedding , respectively . While the vanilla ViT is ideally capable of learning global interaction among all the patch tokens , the memory complexity of self-attention becomes high when there are many tokens as the complexity is quadratically linear to the length of the input sequence . Moreover , the use of isotropic structure limits the capability of extending the vanilla ViT model to many vision applications that require high-resolution details , e.g. , object detection . To address these issues , we propose regional-to-local attention for vision transformer , RegionViT for short . An Overview of Our Proposed Approach ( RegionViT ) . Figure 2 illustrates the architecture of the proposed vision transformer , which consists of two tokenization processes that convert an image into regional ( upper path ) and local tokens ( lower path ) . Each tokenization is a convolution with different patch sizes , e.g. , in Figure 2 , the patch size of regional tokens is 282 while 42 is used for local tokens with dimensions projected to C , which means that one regional token covers 72 local tokens based on the spatial locality , leading to the window size of a local region to 72 . At stage 1 , two sets of tokens are passed through the proposed regional-to-local transformer encoders . However , for the later stages , to balance the computational load and to have feature maps at different resolutions , we deploy a downsampling process to halve the spatial resolution while doubling the channel dimension like CNN on both regional and local tokens before going to the next stage . Finally , at the end of the network , we average the remaining regional tokens as the final embedding for the classification while the detection uses all local tokens at each stage since it provides more fine-grained location information . By having the pyramid structure , the ViT can generate multi-scale features and hence it could be easily extended to more vision applications , e.g. , object detection , rather than image classification only . We explain the main components of the regional-to-local transformer encoder in the next section . Regional-to-Local Transformer Encoder . Figure 3 shows the proposed regional-to-local ( R2L ) transformer encoder which includes R2L attention and feed-forward network ( FFN ) . Specifically , within R2L attention , the regional self-attention ( RSA ) first involves all regional tokens to learn the global information efficiently as the number of regional tokens is few and then local self-attention ( LSA ) takes the local tokens with the associated regional token to learn local feature while the involved regional token could provide global information at the same time . Both RSA and LSA are multihead self-attention ( MSA ) but with different input tokens . Finally , the FFN is applied to enhance the features . We also add layer normalization ( LN ) and residual shortcuts as in standard transformer encoders . Mathematically , given the regional and local tokens , xd−1r and x d−1 l as the inputs at layer d , the R2L transformer encoder can be expressed as : ydr = x d−1 r + RSA ( LN ( x d−1 r ) ) , y d i , j = [ y d ri , j || { x d−1 li , j , m , n } m , n∈M ] zdi , j = y d i , j + LSA ( LN ( y d i , j ) ) , x d i , j = z d i , j + FFN ( LN ( z d i , j ) ) ( 2 ) where i , j are the spatial index with respect to regional tokens whilem , n are the index of local token in the window size M2 . Note that the input of LSA , ydi , j , includes one regional token and corresponding local tokens , and thus , the information between local and regional tokens are exchanged ; on the other hand , the outputs , xdr and x d l , can be extracted from x d i , j like the top-right equation .

Summary: This paper introduces local inductive bias into vanilla vision transformer by adopting the pyramid structure and employing regional-to-local attention. The regional-to-local (R2L) attention is the main contribution of the paper. Compared to the vanilla self-attention, R2L has two steps: 1) self-attention on regional tokens; 2) self-attention on local tokens and their corresponding regional token.