Focal Attention for Long-Range Interactions in Vision Transformers

Recently , Vision Transformer and its variants have shown great promise on various computer vision tasks . The ability of capturing local and global visual dependencies through self-attention is the key to its success . However , this also brings challenges due to quadratic computational overhead , especially for the high-resolution vision tasks ( e.g. , object detection ) . Many recent works have attempted to reduce the cost and improve model performance by applying either coarse-grained global attention or fine-grained local attention . However , both approaches cripple the modeling power of the original self-attention mechanism of multi-layer Transformers , leading to sub-optimal solutions . In this paper , we present focal attention , a new attention mechanism that incorporates both fine-grained local and coarse-grained global interactions . In this new mechanism , each token attends its closest surrounding tokens at fine granularity and the tokens far away at coarse granularity , and thus can capture both short- and long-range visual dependencies efficiently and effectively . With focal attention , we build a new variant of Vision Transformer models , called Focal Transformers , which achieve superior performance over the state-of-theart ( SoTA ) Vision Transformers on a range of public image classification and object detection benchmarks . In particular , our Focal Transformer models with a moderate size of 51.1M and a large size of 89.8M achieve 83.6 % and 84.0 % Top-1 accuracy , respectively , on ImageNet classification at 224 × 224 . When employed as the backbones , Focal Transformers achieve consistent and substantial improvements over the current SoTA Swin Transformers [ 43 ] across 6 different object detection methods . Our largest Focal Transformer yields 58.7/59.0 box mAPs and 50.9/51.3 mask mAPs on COCO mini-val/test-dev , and 55.4 mIoU on ADE20K for semantic segmentation , creating new SoTA on three of the most challenging computer vision tasks . Our code is available at : https : //github . com/microsoft/Focal-Transformer . 1 Introduction . Nowadays , Transformer [ 57 ] has become a prevalent model architecture in natural language processing ( NLP ) [ 20 , 6 ] . In the light of its success in NLP , there is an increasing effort on adapting it to computer vision ( CV ) [ 47 , 50 ] . Since its promise firstly demonstrated in Vision Transformer ( ViT ) [ 21 ] , we have witnessed a flourish of full-Transformer models for image classification [ 55 , 60 , 64 , 43 , 76 , 56 ] , object detection [ 8 , 85 , 79 , 18 ] and semantic segmentation [ 58 , 62 ] . Beyond these static image tasks , it has also been applied on various temporal understanding tasks , such as action recognition [ 40 , 78 , 10 ] , object tracking [ 13 , 59 ] , scene flow estimation [ 38 ] . The self-attention mechanism is arguably the key component that differentiates Transformers from the widely used convolutional neural networks ( CNNs ) [ 37 ] in computer vision . At each Transformer layer , self-attention enables global content-dependent interactions among different image regions for 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . modeling short- and long-range dependencies , respectively . Through the visualization of full selfattention results1 , we indeed observe that self-attention learns to attend local surroundings ( like CNNs ) and the global contexts at the same time , as illustrated in Fig . 1 ( Left ) . Nevertheless , when dealing with high-resolution vision tasks such as object detection or segmentation , an efficient implementation of a global and fine-grained self-attention becomes non-trivial due to the quadratic computational cost with respect to the number of tokens in feature maps . Recent works have alternatively exploited either a coarse-grained global self-attention [ 60 , 64 ] or a fine-grained local self-attention [ 43 , 76 , 56 ] , for the sake of reducing the computational cost . However , both approaches cripple the power of the original full self-attention i.e. , the ability to simultaneously capture local and global visual dependencies . In this paper , we present a new attention mechanism to capture both short- and long-range interactions in Transformer layers for high-resolution input images . Considering that the visual dependencies between the nearby ( local ) regions are usually much stronger than the dependencies between the regions that are far away , we perform the fine-grained attention only in local regions while the coarse-grained attention globally . As depicted in Fig . 1 ( Right ) , a query token in the feature map attends its closest local surroundings at the finest granularity as itself . However , when it goes to the regions far away , it attends to summarized tokens to capture coarse-grained visual dependencies . We call this new mechanism focal attention , as each token attends the others in a focal manner . We will show in this study that focal attention allows to effectively model visual dependencies among all regions covering the whole high-resolution feature maps while introducing much less number of tokens in the computation than that in the standard self-attention mechanism . Equipped with focal attention , a series of Focal Transformers are developed and validated via a comprehensive empirical study across three core vision tasks , including image classification , object detection and segmentation . Results show that Focal Transformers consistently outperform the SoTA Vision Transformers across various settings ( i.e. , in model sizes and complexities ) . Notably , the small Focal Transformer with 51.1M parameters achieves 83.6 % top-1 accuracy on ImageNet-1K , and the base model with 89.8M parameters obtains 84.0 % top-1 accuracy . In the fine-tuning experiments for object detection , Focal Transformers consistently outperform the SoTA Swin Transformers [ 43 ] across six popular object detection methods . Our largest Focal Transformer model achieves 59.0 box mAP and 51.3 mask mAP on COCO test-dev for object detection and instance segmentation , respectively , and 55.4 mIoU on ADE20K for semantic segmentation . These results demonstrate that focal attention is highly effective in modeling the global interactions in Vision Transformers . 2 Related work . Vision Transformers . Vision Transformer ( ViT ) is first introduced in [ 21 ] . It applies a standard Transformer , originally developed for NLP [ 57 ] , to encode an image by analogously splitting the image into a sequence of visual tokens . It has demonstrated superior performance to CNNs such as ResNet [ 33 ] on multiple image classification benchmarks , when trained with sufficient data [ 21 ] and carefully designed data augmentation and regularization methods [ 55 ] . The results thus inspire researchers to explore the applications of ViT on various vision tasks beyond image classification , such as self-supervised learning [ 14 , 9 , 39 ] , object detection [ 8 , 85 , 79 , 18 ] and semantic segmentation [ 58 , 62 , 81 ] . There are also increasing number of studies for improving ViT via data-efficient 1DeiT-Tiny model , checkpoint downloaded from https : //github.com/facebookresearch/deit . training [ 55 ] , improved patch embedding/encoding [ 16 , 71 , 31 ] , integrating convolutional projections into transformers [ 64 , 70 ] , and using multi-scale architectures and efficient self-attention mechanisms for high-resolution vision tasks [ 60 , 64 , 43 , 76 , 15 ] . Recent surveys include [ 36 , 30 , 36 ] . This paper focuses on improving the self-attention mechanism of ViT for encoding high-resolution images . Efficient global and local self-attention . In many real-world tasks , Transformers need to cope with a large number of input tokens , such as long documents in NLP and high-resolution images in computer vision ( CV ) . Recently , many efficient self-attention mechanisms have been proposed to deal with the quadratic computational and memory cost incurred by the standard self-attention mechanism . On one hand , a number of works in both NLP and CV resort to coarse-grained global self-attention ( i.e. , attending the down-sampled or summarized tokens ) to capture the long-range interactions [ 49 , 46 , 60 , 64 , 31 , 23 ] . Although this approach improves the model efficiency , it loses the detailed context information surrounding the query tokens . On the other hand , to make the computational cost manageable , various local fine-grained attention mechanism ( i.e. , attending neighboring tokens within a pre-set window size ) are used for both NLP [ 3 , 74 , 1 ] and CV [ 56 , 43 , 76 ] . In this paper , we argue that both global and local attentions are important for model performance . This is also validated by some recent studies that aim to improve CNNs by incorporating ways of modeling global attentions [ 35 , 63 , 61 , 68 , 2 , 7 , 51 ] . The standard self-attention mechanism used by ViT can indeed learned both types of attentions , as shown in Fig . 1 ( Left ) . But it often incurs a prohibitively high cost for high-resolution images . To the best of our knowledge , the proposed focal attention provides the first mechanism to incorporate local and global attention in a single Transformer layer 2 . It can capture both short- and long-range interactions as standard self-attention but in a much more efficient and effective way , especially for high-resolution images . 3 Method . 3.1 Model architecture . To accommodate high-resolution dense prediction tasks , we employ a multi-scale model architecture as in [ 60 , 76 , 43 ] . As shown in Fig . 2 , an image I ∈ RH×W×3 is first partitioned into patches of size 4× 4 , resulting in H4 × W 4 visual tokens with dimension 4× 4× 3 . Then , we use a patch embedding layer , consisting of a convolutional layer with filter size and stride both equal to 4 , to project these patches into hidden features with dimension d. We then pass this spatial feature map to the four stages of Focal Transformer blocks . In each stage i ∈ { 1 , 2 , 3 , 4 } , the Focal Transformer block consists of Ni Focal Transformer layers . After each stage , we use a patch embedding layer to reduce the spatial size of feature map by factor 2 and increase the feature dimension by 2 . For image classification tasks , we take the average of the output from the last stage and send it to a classification layer . For object detection , the feature maps from the last 3 or all 4 stages are fed to a particular object detector head , 2A similar focal mechanism has been used in CNNs for NLP [ 27 ] . depending on the specific detection method we choose to use . The model capacity can be customized by varying the input feature dimension d and the number of Focal Transformer layers . Standard self-attention can capture both short- and long-range interactions at fine-grain , but suffers from high computational cost when it performs attention on high-resolution feature maps as noted in [ 76 ] . Take stage 1 in Fig . 2 as an example . For a feature map of size H4 × W 4 × d , the complexity of self-attention is O ( ( H4 × W 4 ) 2d ) , resulting in an explosion of time and memory cost , considering that min ( H , W ) could be 800 or even larger for object detection . In the next section , we describe how we address this issue with the proposed focal attention mechanism . 3.2 Token-wise focal attention Focal attention is proposed to make the Transformer layers suitable for encoding high-resolution input images . Instead of attending all tokens at fine-grain , we attend the fine-grain tokens only locally , but the summarized ones ( i.e. , the coarse-grained tokens generated by sub-window pooling , which is illustrated in Fig . 4 and will be described later ) globally . As such , focal attention can cover the same amount of image regions as standard self-attention but with much less cost . In Fig . 3 , we show the size of the receptive field for standard self-attention and our focal attention as a function of the number of attended tokens . For a given query position , by reducing the granularity of its surroundings based on their distance to the query , focal attention can have significantly larger receptive fields at the same cost measured by the number of visual tokens , compared to the standard self-attention mechanism . Theoretically , the focal attention mechanism enables global interaction with much less time and memory cost , because it attends a much smaller number of surrounding ( summarized ) tokens . In practice , however , extracting the surrounding tokens for each query position could incur high time cost since we need to duplicate the extraction of each token for all queries that the token surrounds . This issue had been extensively discussed in [ 56 , 76 , 43 ] and a common solution is to partition the input feature map into windows . Thus , in our Focal Transformers , we resort to performing focal attention at the window level . We elaborate the window-wise focal attention in the following .

Summary: This paper proposes a focal version of self-attention module for vision transformers. The query is kept at high resolution, but the number of keys is reduced depending on the relative position difference. This proposed attention is implemented within multi-scale windows. It shows good performance on top of recent transformer advances, on ImageNet and COCO.