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: The paper introduces a new spin on the ViT architecture. Similar to other recent developments, it proposes a pyramidal style architecture where subsequent layers process the input image at decreasing resolutions. This allows the approach to be significantly more memory and compute efficient compared to the original ViT architecture, while still obtaining solid results on several tasks and benchmarks. At the core of the method is a regional-to-local attention approach, where very coarse regional tokens are processed by a normal multi-head attention yielding global attention and each coarse regional patch is additionally represented by a set of smaller local patches, which are also processed together with the corresponding regional token to get finer-grained local attention. Each layer in the network follows this principle and between different sections of the network both the regional and local tokens are downsampled with strided conv layers.