MaiT: integrating spatial locality into image transformers with attention masks 

1 INTRODUCTION . Convolutional neural networks ( CNNs ) ( Krizhevsky et al. , 2012 ; He et al. , 2016 ; Tan & Le , 2019 ) have been the de facto model for computer vision ( CV ) tasks , which are inherently equipped with inductive biases such as translation equivariance and locality . Recently , vision transformers ( Dosovitskiy et al. , 2021 ; Touvron et al. , 2021a ) are gaining momentum translating the success from transformer-based models in natural language processing ( NLP ) tasks ( Vaswani et al. , 2017 ; Devlin et al. , 2019 ) . Self-attention heads excel at capturing the long-range dependencies in sequences but struggle at focusing on local information . Unlike CNNs which have gone through many iterations of optimization , vision transformers are not very efficient and still require huge computing power and a large number of Flops . Naturally , combining the benefits of both CNNs and vision transformers is promising to further boost performances of CV models . The question remains how to effectively integrate inductive bias such as spatial locality into transformers . One direction is to utilize convolutional blocks to extract spatial information by adapting either the patch-token embedding layer , self-attention module or feed-forward layers , to form CNN-transformer hybrid structures as in Li et al . ( 2021b ) ; Srinivas et al . ( 2021 ) ; Graham et al . ( 2021 ) ; Wu et al . ( 2021a ) ; Yuan et al . ( 2021a ) ; Guo et al . ( 2021 ) . However , forcefully inserting convolutional operations into transformers may potentially constrain the learning capacity of transformers . To capture the spatial information without significantly changing the transformer model architecture , Chu et al . ( 2021b ) introduce extra positional encoding . Han et al . ( 2021 ) ; Chen et al . ( 2021 ) fuse local and global representations using multiple transformer blocks or branches to simultaneously process images at different scales such as pixel-level , small-patch or large patch . Yuan et al . ( 2021b ) apply a layerwise tokens-to-tokens transformation to capture local structure . These approaches usually come with the cost of extra parameters and model complexity , thus potentially lowering the inference speed . d ’ Ascoli et al . ( 2021 ) ; Zhou et al . ( 2021 ) improve the self-attention for better representation with gated positional self-attention and learnable transformation matrix respectively . Hu et al . ( 2019 ) , Ramachandran et al . ( 2019 ) adapt the self-attention module and improve the performance of CNN models . Liu et al . ( 2021 ) captures locality within shifted windows in a hierarchical structure . Though it could save computation in some cases , the added model complexity may lower its potentials . Different from prior works , we try to incorporate spatial locality without changing its architecture or adding extra parameters/FLOPs . We propose attention masks to guide the attention heads to focus on local information . Masked attention heads extract local dependencies more efficiently by allowing information aggregation only from the closest neighbors . This liberates other unmasked heads to learn global information more effectively . We name the modified model , Masked attention image Transformer ( MaiT ) , which is built on top of DeiT ( Touvron et al. , 2021a ) . MaiT gathers both local and global information at the same time from different heads . Moreover , the regularization effects from attention masks facilitate the training of deep transformers by guiding the attention map learning and promoting diversity across transformer layers . We proved that less is more in this specific case with attention masks . MaiT achieves up to 2.5 % higher top-1 accuracy on ImageNet ( Deng et al. , 2009 ) with the same model architecture as DeiT ( Touvron et al. , 2021a ) . Additionally , deep MaiT outperforms CaiT ( Touvron et al. , 2021b ) by up to 1.7 % in top-1 accuracy with fewer parameters and simpler structure . In summary , we make three major contributions in this work : 1 ) . we propose attention masks to encode the spatial locality into self-attention heads without structural change or extra parameter/computation , while improving model efficiency . 2 ) . We present a quick empirical search strategy for the masking scheme exploration , along with an automatic end-to-end search alternative . 3 ) . We also reveal the importance of the locality across layers and the performance impact of attention masks on deep vision transformers . Note though MaiT is demonstrated with DeiT/CaiT , the attention mask is applicable to other vision transformers as well . 2 RELATED WORK . Spatial locality is an integral part of the convolutional operation with weight filters attending to local regions of input feature maps . Vision transformer ( ViT ) ( Dosovitskiy et al. , 2021 ) is the first pure transformer-based model on vision tasks , but it requires a large private labeled dataset JFT300M ( Sun et al. , 2017 ) to achieve competitive performances . Data-efficient image transformer ( DeiT ) ( Touvron et al. , 2021a ) improves upon ViT models by introducing stronger data augmentation , regularization , and knowledge distillation . Class-attention in image transformer ( CaiT ) ( Touvron et al. , 2021b ) extends DeiT by increasing the number of transformer layers . To overcome the difficulties of training deeper transformers , CaiT introduces LayerScale and class-attention layers , which increase the parameters and model complexity . Tokens-to-Token vision transformer ( T2T ) ( Yuan et al. , 2021b ) proposes an image transformation by recursive token aggregation to capture local structure . Stand-alone self-attention ( Ramachandran et al. , 2019 ) applies local self-attention layer to replace spatial convolution and outperform original ResNet models . Even though sharing value and key spatially is parameter efficient in this approach , content-based information is lost . Transformer-iN-Transformer ( TNT ) ( Han et al. , 2021 ) models both patch-level and pixel-level representations and applies outer and inner transformer blocks to extract global and local information respectively . ConViT ( d ’ Ascoli et al. , 2021 ) proposes the gated positional self-attention to incorporate soft convolutional biases . CrossViT ( Chen et al. , 2021 ) proposes a dual-branch transformer architecture for multi-scale feature extraction . For pixel-level prediction tasks such as semantic segmentation , object detection , Pyramid Vision Transformer ( PVT ) ( Wang et al. , 2021a ) introduces a progressive shrinking pyramid and spatialreduction attention with fine-grained image patches . DETR ( Carion et al. , 2020 ) adapts transformers for object detection tasks . Swin Transformer ( Liu et al. , 2021 ) applies hierarchical transformer with shifted windows of varying sizes . Twins ( Chu et al. , 2021a ) deploys interleaved locally-grouped self-attention and global sub-sample attention layers to improve performances . To optimize transformer and save computation , Wu et al . ( 2021b ) uses centroid attention to extract and compress input information , Jaegle et al . ( 2021 ) iteratively distill inputs into latent space with attention bottlenecks , Wang et al . ( 2021b ) dynamically adjusts the number of tokens with multiple cascading transformers , and Wu et al . ( 2020 ) introduced semantic token to replace pixel-based transformers to save computation . There are hybrid architectures fusing convolutional and transformer blocks , such as LocalViT ( Li et al. , 2021b ) , BoTNet ( Srinivas et al. , 2021 ) , LeViT ( Graham et al. , 2021 ) , BossNet ( Li et al. , 2021a ) , CvT ( Wu et al. , 2021a ) , CoaT ( Xu et al. , 2021 ) , CMT ( Guo et al. , 2021 ) for higher accuracy and faster inference . Unlike prior literature , our work explores the intrinsic capability of pure transformer block on incorporating spatial locality and the impact of the locality along the depth direction . This work is also inspired by the emerging graph attention network ( Veličković et al. , 2018 ) , borrowing the concept of message passing and information aggregation from nearest neighbors . Note attention masks have been used in NLP tasks as a sparsification method to reduce the computation complexity , as well as to capture local information , as in Guo et al . ( 2019 ) ; Child et al . ( 2019 ) ; Beltagy et al . ( 2020 ) ; Ainslie et al . ( 2020 ) . 3 VISION TRANSFORMER PRELIMINARIES . The transformer architecture introduced by Vaswani ( Vaswani et al. , 2017 ) inspired many model variants with remarkable success in NLP tasks . ViT ( Dosovitskiy et al. , 2021 ) extends pure transformer-based architecture into CV applications . Instead of pixel-level processing , ViT splits the original images into a sequence of patches as inputs and transforms them into patch tokens , for better computation efficiency . In general , ViT consists of 3 fundamental modules : embedding layer , multi-head self-attention , and feed-forward network . To process images in transformer , the original RGB images ( 224×224 ) is flattened into a sequence of N ( 14×14 ) patches . Each patch has a fixed size ( typical 16x16 pixels ) . Patches are then transformed into patch embedding with hidden dimensions ( D ) of 192 , 384 , and 768 for tiny , small , and base models respectively in ViT/DeiT . In addition to patch tokens , the embedding layer also integrates positional information , classification and knowledge distillation through the positional token , class token and distillation token , respectively . Positional token is added into the patch embedding with a trainable positional embedding . However , this positional embedding is added only in the embedding layers . The spatial information is largely lost in the transformer layers since all-to-all attention is invariant to the order of the patches . The class token is another trainable vector ( 1×D ) , concatenated to the patch tokens ( total N+1 ) . It is used to collect information from the patch tokens to make output predictions , while also spreading information among patch tokens during training . Distillation token is sometimes added for knowledge transfer from teacher models , such as a CNN model . When training the distilled version of the model , a distillation token is further concatenated to the patch token along with the class token ( total N+2 ) . Multi-head self-attention ( MHA ) module has multiple parallel attention heads , where each of them comprises three main components : Key ( K ) , Query ( Q ) , and Value ( V ) , Key and Query are trained and multiplied to estimate how much weights on each corresponding token in Value for output : Attention ( K , Q , V ) = Softmax ( QKT√ d ) V ( 1 ) Where softmax is applied to each row of the input product matrix ( QKT ) and √ d provides appropriate normalization . Multiple attention heads in MHA attend to different parts of the input simultaneously . Considering H heads in MHA layer , the hidden dimension D is split equally across all heads ( D = H × d ) . Feed-forward network ( FFN ) follows after the MHA module , containing two linear transformation layers separated by Gelu activation . The hidden dimension expands by 4x after the first linear layer from D to 4D , and is reduced back to D in the second linear layer . Both MHA and FFN use skip-connections with layer normalization as the residual operation . 4 MASKED ATTENTION HEAD . Spatial locality plays a crucial role in computer vision tasks . CNN models capture it using the sliding filter of shared weights , typically with a receptive field of 3×3 , 5×5 , or 7×7 . In contrast to CNN models , the locality is not explicitly introduced in the transformer structure . With attention masks , we can explicitly insert locality into self-attention modules without introducing any extra parameter or computation . The key idea is to apply a mask on the all-to-all attention products ( i.e . QKT ) to reinforce the weights from the closest neighbors by aggregating information only from tokens selected by the mask . Figure 1 illustrates an example of our proposed masking scheme . The orange box shows a 3 × 3 mask , where only the direct neighboring patches are selected . Specifically , Patch 16 only gathers information from the closest neighbors of Patch 1 , 2 , 3 , 15 , 17 , 29 , 30 , and 31 , and ignores the rest of patches . This is different from the typical all-to-all attention module , where Patch 16 attends to all 0-195 patches . We can easily expand the depth of the attention mask beyond the closest neighbors , to second-level neighbors ( green box in Figure 1 ) or more . Note that the class token ( and distillation token ) still attends to all the patches to collect and spread information during forward and backward passes . Since each attention product selected by the mask is calculated by Q and K , the masked attention head also retains the content-based locality information . The attention mask is added before Softmax , regulating the distribution of attention maps to focus more on the closest neighbors : Masked Attention ( K , Q , V ) = Softmax ( M QKT√ d ) V ( 2 ) where M ∈ { 0 , 1 } ( N+1 ) × ( N+1 ) is a binary attention mask , encoding the spatial locality into the attention head by passing through only the weights from close neighbors and setting the rest to zero . Note it ’ s important to add the mask before Softmax because it allows the model to learn the importance of the locality flexibly . More precisely , unselected patches appear as e0 = 1 in the numerator of the softmax operation . Thus , if the attention product result of the closest neighbors is meaningfully larger than zero ( i.e . M QKT 0 ) , it suggests that local information dominates . However , if those results are negative or close to zero , it implies that local information is insignificant and global information is more important . Therefore , inserting masks before the softmax operation allows models to enforce locality or bypass it .

Summary: This paper attempts to improve the classic vision transformers, or specifically the DeiT model by introducing the Masked Attention Head. Instead of focusing on aggregating global information in the original self-attention heads, this paper introduce local information via the proposed Masked Attention Head into self-attention. Experiments on ImageNet show that the proposed approach does matter for lifting the model performance of the original DeiT model.