Twins: Revisiting the Design of Spatial Attention in Vision Transformers

1 Introduction . Recently , Vision Transformers [ 1–3 ] have received increasing research interest . Compared to the widely-used convolutional neural networks ( CNNs ) in visual perception , Vision Transformers enjoy great flexibility in modeling long-range dependencies in vision tasks , introduce less inductive bias , and can naturally process multi-modality input data including images , videos , texts , speech signals , and point clouds . Thus , they have been considered to be a strong alternative to CNNs . It is expected that vision transformers are likely to replace CNNs and serve as the most basic component in the next-generation visual perception systems . One of the prominent problems when applying transformers to vision tasks is the heavy computational complexity incurred by the spatial self-attention operation in transformers , which grows quadratically in the number of pixels of the input image . A workaround is the locally-grouped self-attention ( or self-attention in non-overlapped windows as in the recent Swin Transformer [ 4 ] ) , where the input is spatially grouped into non-overlapped windows and the standard self-attention is computed only within each sub-window . Although it can significantly reduce the complexity , it lacks the connections between different windows and thus results in a limited receptive field . As pointed out by many previous works [ 5–7 ] , a sufficiently large receptive field is crucial to the performance , particularly for dense prediction tasks such as image segmentation and object detection . Swin [ 4 ] proposes a shifted window operation to tackle the issue , where the boundaries of these local windows are gradually moved as the network proceeds . Despite being effective , the shifted windows may have uneven sizes . The uneven windows result in difficulties when the models are deployed with ONNX or TensorRT , ∗Corresponding author . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . which prefers the windows of equal sizes . Another solution is proposed in PVT [ 8 ] . Unlike the standard self-attention operation , where each query computes the attention weights with all the input tokens , in PVT , each query only computes the attention with a sub-sampled version of the input tokens . Although its computational complexity in theory is still quadratic , it is already manageable in practice . From a unified perspective , the core in the aforementioned vision transformers is how the spatial attention is designed . Thus , in this work , we revisit the design of the spatial attention in vision transformers . Our first finding is that the global sub-sampled attention in PVT is highly effective , and with the applicable positional encodings [ 9 ] , its performance can be on par or even better than state-of-the-art vision transformers ( e.g. , Swin ) . This results in our first proposed architecture , termed Twins-PCPVT . On top of that , we further propose a carefully-designed yet simple spatial attention mechanism , making our architectures more efficient than PVT . Our attention mechanism is inspired by the widely-used separable depthwise convolutions and thus we name it spatially separable self-attention ( SSSA ) . Our proposed SSSA is composed of two types of attention operations— ( i ) locally-grouped self-attention ( LSA ) , and ( ii ) global sub-sampled attention ( GSA ) , where LSA captures the fine-grained and short-distance information and GSA deals with the long-distance and global information . This leads to the second proposed vision transformer architecture , termed Twins-SVT . It is worth noting that both attention operations in the architecture are efficient and easy-to-implement with matrix multiplications in a few lines of code . Thus , all of our architectures here have great applicability and can be easily deployed . We benchmark our proposed architectures on a number of visual tasks , ranging from image-level classification to pixel-level semantic/instance segmentation and object detection . Extensive experiments show that both of our proposed architectures perform favorably against other state-of-the-art vision transformers with similar or even reduced computational complexity . 2 Related Work . Convolutional neural networks . Characterized by local connectivity , weight sharing , shiftinvariance and pooling , CNNs have been the de facto standard model for computer vision tasks . The top-performing models [ 10–13 ] in image classification also serve as the strong backbones for downstream detection and segmentation tasks . Vision Transformers . Transformer was firstly proposed by [ 14 ] for machine translation tasks , and since then they have become the state-of-the-art models for NLP tasks , overtaking the sequence-tosequence approach built on LSTM . Its core component is multi-head self-attention which models the relationship between input tokens and shows great flexibility . In 2020 , Transformer was introduced to computer vision for image and video processing [ 1–3 , 9 , 15– 17 , 17–32 ] . In the image classification task , ViT [ 1 ] and DeiT [ 2 ] divide the images into patch embedding sequences and feed them into the standard transformers . Although vision transformers have been proved compelling in image classification compared with CNNs , a challenge remains when it is applied to dense prediction tasks such as object detection and segmentation . These tasks often require feature pyramids for better processing objects of different scales , and take as inputs the highresolution images , which significantly increase the computational complexity of the self-attention operations . Recently , Pyramid Vision Transformer ( PVT ) [ 8 ] is proposed and can output the feature pyramid [ 33 ] as in CNNs . PVT has demonstrated good performance in a number of dense prediction tasks . The recent Swin Transformer [ 4 ] introduces non-overlapping window partitions and restricts self-attention within each local window , resulting in linear computational complexity in the number of input tokens . To interchange information among different local areas , its window partitions are particularly designed to shift between two adjacent self-attention layers . The semantic segmentation framework OCNet [ 34 ] shares some similarities with us and they also interleave the local and global attention . Here , we demonstrate this is a general design paradigm in vision transformer backbones rather than merely an incremental module in semantic segmentation . Grouped and Separable Convolutions . Grouped convolutions are originally proposed in AlexNet [ 35 ] for distributed computing . They were proved both efficient and effective in speeding up the networks . As an extreme case , depthwise convolutions [ 12 , 36 ] use the number of groups that is equal to the input or output channels , which is followed by point-wise convolutions to aggregate the information across different channels . Here , the proposed spatially separable self-attention shares some similarities with them . Positional Encodings . Most vision transformers use absolute/relative positional encodings , depending on downstream tasks , which are based on sinusoidal functions [ 14 ] or learnable [ 1 , 2 ] . In CPVT [ 9 ] , the authors propose the conditional positional encodings , which are dynamically conditioned on the inputs and show better performance than the absolute and relative ones . 3 Our Method : Twins . We present two simple yet powerful spatial designs for vision transformers . The first method is built upon PVT [ 8 ] and CPVT [ 9 ] , which only uses the global attention . The architecture is thus termed Twins-PCPVT . The second one , termed Twins-SVT , is based on the proposed SSSA which interleaves local and global attention . 3.1 Twins-PCPVT . PVT [ 8 ] introduces the pyramid multi-stage design to better tackle dense prediction tasks such as object detection and semantic segmentation . It inherits the absolute positional encoding designed in ViT [ 1 ] and DeiT [ 2 ] . All layers utilize the global attention mechanism and rely on spatial reduction to cut down the computation cost of processing the whole sequence . It is surprising to see that the recently-proposed Swin transformer [ 4 ] , which is based on shifted local windows , can perform considerably better than PVT , even on dense prediction tasks where a sufficiently large receptive field is even more crucial to good performance . In this work , we surprisingly found that the less favored performance of PVT is mainly due to the absolute positional encodings employed in PVT [ 8 ] . As shown in CPVT [ 9 ] , the absolute positional encoding encounter difficulties in processing the inputs with varying sizes ( which are common in dense prediction tasks ) . Moreover , this positional encoding also breaks the translation invariance . On the contrary , Swin transformer makes use of the relative positional encodings , which bypasses the above issues . Here , we demonstrate that this is the main cause why Swin outperforms PVT , and we show that if the appropriate positional encodings are used , PVT can actually achieve on par or even better performance than the Swin transformer . Here , we use the conditional position encoding ( CPE ) proposed in CPVT [ 9 ] to replace the absolute PE in PVT . CPE is conditioned on the inputs and can naturally avoid the above issues of the absolute encodings . The position encoding generator ( PEG ) [ 9 ] , which generates the CPE , is placed after the first encoder block of each stage . We use the simplest form of PEG , i.e. , a 2D depth-wise convolution without batch normalization . For image-level classification , following CPVT , we remove the class token and use global average pooling ( GAP ) at the end of the stage [ 9 ] . For other vision tasks , we follow the design of PVT . Twins-PCPVT inherits the advantages of both PVT and CPVT , which makes it easy to be implemented efficiently . Our extensive experimental results show that this simple design can match the performance of the recent state-of-the-art Swin transformer . We have also attempted to replace the relative PE with CPE in Swin , which however does not result in noticeable performance gains , as shown in our experiments . We conjecture that this maybe due to the use of shifted windows in Swin , which might not work well with CPE . Architecture settings We report the detailed settings of Twins-PCPVT in Table 2 ( in supplementary ) , which are similar to PVT [ 8 ] . Therefore , Twins-PCPVT has similar FLOPs and number of parameters to [ 8 ] . 3.2 Twins-SVT . Vision transformers suffer severely from the heavy computational complexity in dense prediction tasks due to high-resolution inputs . Given an input of H ×W resolution , the complexity of selfattention with dimension d is O ( H2W 2d ) . Here , we propose the spatially separable self-attention ( SSSA ) to alleviate this challenge . SSSA is composed of locally-grouped self-attention ( LSA ) and global sub-sampled attention ( GSA ) . Locally-grouped self-attention ( LSA ) . Motivated by the group design in depthwise convolutions for efficient inference , we first equally divide the 2D feature maps into sub-windows , making self-attention communications only happen within each sub-window . This design also resonates with the multi-head design in selfattention , where the communications only occur within the channels of the same head . To be specific , the feature maps are divided into m × n sub-windows . Without loss of generality , we assume H % m = 0 and W % n = 0 . Each group contains HWmn elements , and thus the computation cost of the self-attention in this window is O ( H 2W 2 m2n2 d ) , and the total cost is O ( H2W 2 mn d ) . If we let k1 = Hm and k2 = W n , the cost can be computed as O ( k1k2HWd ) , which is significantly more efficient when k1 H and k2 W and grows linearly with HW if k1 and k2 are fixed . Although the locally-grouped self-attention mechanism is computation friendly , the image is divided into non-overlapping sub-windows . Thus , we need a mechanism to communicate between different sub-windows , as in Swin . Otherwise , the information would be limited to be processed locally , which makes the receptive field small and significantly degrades the performance as shown in our experiments . This resembles the fact that we can not replace all standard convolutions by depth-wise convolutions in CNNs . Global sub-sampled attention ( GSA ) . A simple solution is to add extra standard global selfattention layers after each local attention block , which can enable cross-group information exchange . However , this approach would come with the computation complexity of O ( H2W 2d ) . Here , we use a single representative to summarize the important information for each of m × n sub-windows and the representative is used to communicate with other sub-windows ( serving as the key in self-attention ) , which can dramatically reduce the cost to O ( mnHWd ) = O ( H 2W 2d k1k2 ) . This is essentially equivalent to using the sub-sampled feature maps as the key in attention operations , and thus we term it global sub-sampled attention ( GSA ) . If we alternatively use the aforementioned LSA and GSA like separable convolutions ( depth-wise + point-wise ) . The total computation cost is O ( H 2W 2d k1k2 + k1k2HWd ) . We have H 2W 2d k1k2 + k1k2HWd ≥ 2HWd √ HW . The minimum is obtained when k1 · k2 = √ HW . We note that H = W = 224 is popular in classification . Without loss of generality , we use square sub-windows , i.e. , k1 = k2 . Therefore , k1 = k2 = 15 is close to the global minimum for H = W = 224 . However , our network is designed to include several stages with variable resolutions . Stage 1 has feature maps of 56 × 56 , the minimum is obtained when k1 = k2 = √ 56 ≈ 7 . Theoretically , we can calibrate optimal k1 and k2 for each of the stages . For simplicity , we use k1 = k2 = 7 everywhere . As for stages with lower resolutions , we control the summarizing window-size of GSA to avoid too small amount of generated keys . Specifically , we use the size of 4 , 2 and 1 for the last three stages respectively . As for the sub-sampling function , we investigate several options including average pooling , depthwise strided convolutions , and regular strided convolutions . Empirical results show that regular strided convolutions perform best here . Formally , our spatially separable self-attention ( SSSA ) can be written as ẑlij = LSA ( LayerNorm ( zl−1ij ) ) + zl−1ij , zlij = FFN ( LayerNorm ( ẑlij ) ) + ẑlij , ẑl+1 = GSA ( LayerNorm ( zl ) ) + zl , zl+1 = FFN ( LayerNorm ( ẑl+1 ) ) + ẑl+1 , i ∈ { 1 , 2 , .... , m } , j ∈ { 1 , 2 , .... , n } ( 1 ) where LSA means locally-grouped self-attention within a sub-window ; GSA is the global sub-sampled attention by interacting with the representative keys ( generated by the sub-sampling functions ) from each sub-window ẑij ∈ Rk1×k2×C . Both LSA and GSA have multiple heads as in the standard self-attention.The PyTorch code of LSA is given in Algorithm 1 ( in supplementary ) . Again , we use the PEG of CPVT [ 9 ] to encode position information and process variable-length inputs on the fly . It is inserted after the first block in each stage . Model variants . The detailed configure of Twins-SVT is shown in Table 3 ( in supplementary ) . We try our best to use the similar settings as in Swin [ 4 ] to make sure that the good performance is due to the new design paradigm . Comparison with PVT . PVT entirely utilizes global attentions as DeiT does while our method makes use of spatial separable-like design with LSA and GSA , which is more efficient . Comparison with Swin . Swin utilizes the alternation of local window based attention where the window partitions in successive layers are shifted . This is used to introduce communication among different patches and to increase the receptive field . However , this procedure is relatively complicated and may not be optimized for speed on devices such as mobile devices . Swin Transformer depends on torch.roll ( ) to perform cyclic shift and its reverse on features . This operation is memory unfriendly and rarely supported by popular inference frameworks such as NVIDIA TensorRT , Google TensorflowLite , and Snapdragon Neural Processing Engine SDK ( SNPE ) , etc . This hinders the deployment of Swin either on the server-side or on end devices in a production environment . In contrast , Twins models don ’ t require such an operation and only involve matrix multiplications that are already optimized well in modern deep learning frameworks . Therefore , it can further benefit from the optimization in a production environment . For example , we converted Twins-SVT-S from PyTorch to TensorRT , and its throughput is boosted by 1.7× . Moreover , our local-global design can better exploit the global context , which is known to play an important role in many vision tasks . Finally , one may note that the network configures ( e.g. , such as depths , hidden dimensions , number of heads , and the expansion ratio of MLP ) of our two variants are sightly different . This is intended because we want to make fair comparisons to the two recent well-known transformers PVT and Swin . PVT prefers a slimmer and deeper design while Swin is wider and shallower . This difference makes PVT have slower training than Swin . Twins-PCPVT is designed to compare with PVT and shows that a proper positional encoding design can greatly boost the performance and make it on par with recent state-of-the-art models like Swin . On the other hand , Twins-SVT demonstrates the potential of a new paradigm as to spatially separable self-attention is highly competitive to recent transformers .

Summary: This paper proposes to combine CPE with PVT. In addition, local attention is combined with a subsampled global attention layer so that the model can capture both local details and global relations efficiently. The design is verified on image classification, detection and segmentation.