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: The paper proposes two new transformer architectures to address image classification and object detection and segmentation tasks, focusing on the attention mechanisms in each. The first method, Twins-PCPVT, combines the pyramidal subsampling of attention in PVT with conditional positional encoding of CPVT, the second, Twins-SVT, proposes to combine local grouping to manage computational complexity in attention with global subsampling to improve receptive field size, drawing analogies to separable filters.  Both methods are empirically analyzed on relevant datasets (Imagenet1k, MSCOCO, ADE20k) in comparison to state of the art methods (PVT, CPVT, Swin) and show improvements in accuracy and run-time speed.