CoAtFormer: Vision Transformer with Composite Attention

1 Introduction

In the past years, Convolution Neural Networks (CNNs) have become a defacto choice for a wide variety of computer vision tasks [Simonyan and Zisserman, 2014; He et al., 2015; Ren et al., 2015; He et al., 2017a] since AlexNet [Krizhevsky et al., 2012]. However, convolution operations merely capture local dependencies of pixels, which neglect the dependency modeling between distant pixels to some extent [Wang et al., 2017]. Recently, self-attention based Transformers [Vaswani et al., 2017] have achieved excellent performance in Natural Language Processing (NLP) benchmarks and became the dominant architecture for various applications. Meanwhile, numerous researchers attempt to introduce the Transformer-based architectures into vision domains, and attain promising performance in various tasks such as image classification [Dosovitskiy et al., 2020; Touvron et al., 2020], object detection [?; Zhu et al., 2020], and semantic segmentation [Zheng et al., 2020]. Vision Transformer (ViT) [Dosovitskiy et al., 2020], which utilises a sequence of embedded image patches as input to stacked standard Transformer encoders, is the first fully Transformer-based architecture that demonstrate comparable performance to CNNs. While ViTs present strong capabilities to model the long-range dependencies, their computational complexity grows quadratically with the image size, limiting its application to high resolution images [Yang et al., 2021] in which modeling multi-scale global context is essential for accurate representation modeling.

To reduce the computational cost, some works [Liu et al., 2021; Dong et al., 2021; Tu et al., 2022; Ding et al., 2022], most notably Swin Transformer [Liu et al., 2021], have limited attention region in a spatial local window. However, the limited receptive field of local attention challenges the capability of global self-attention to model global contextual information. Subsequent works such as MaxViT [Tu et al., 2022] attempted to mitigate this issue by developing highly intricate self-attention modules with increased model size. How to efficiently integrate various long-range interactions to balance the model complexity and generalizability under a computation budget still remains challenging.

In this work, we develop a novel Transformer block, called composite attention block, that capably serves as a fundamental architecture component which can model global channel and spatial interactions in a single module. In contrast to vanilla self-attention, composite attention exhibits greater efficiency and flexibility, specifically adapting seamlessly to varying input lengths with linear complexity. Compared to window/local attention, composite attention performs stronger perception capacity by simultaneously capturing global spatial and channel dependencies. Moreover, with only linear complexity, composite attention can serve as a basic stand-alone attention block in any layer of a neural network, even in high-resolution stages.

It is worthwhile to note that our proposed composite attention contains two parallel branches to capture various global context. In composite attention block, one branch exploits a dynamic channels module to model global channel dependencies, while the other branch employs an efficient spatial attention module to capture long-range spatial interactions. Concretely, We replace the self-attention, as originally introduced by Vaswani et al. [Dosovitskiy et al., 2020], with a efficient spatial attention. Instead of capturing the pairwise interactions between keys, queries, we model the interactions between a learnable global token embedding and queries only, computing the global context-aware spatial attention weights with linear complexity. Additionally, we use a simple graph neural network as the core component of dynamic channel attention to encourage the communication across channels.

Based on the composite attention mechanism, we further propose a simple but effective vision Transformer architecture named "CoAtFormer" by hierarchically stacking repeated blocks composed of composite attention and convolutions. This architecture exhibits significantly stronger modeling power while limiting computation cost. As a general-purpose vision backbone, the CoAtFormer demonstrates state-of-the-art (SOTA) performance for a broad range of visual tasks including image classification, object detection and segmentation. Specifically, for image classification using ImageNet-1K dataset, CoAtFormer with 24M, 37M, 73M parameters achieve new SOTA performance of 84.4%, 85.3%,

85.9% Top-1 accuracy and without using any extra data. As Figure 1 shows, CoAtFormer consistently outperforms both GCViT [Hatamizadeh et al., 2023], MaxViT [Tu et al., 2022] and ConvNeXt [Liu et al., 2022] models by a significant margin. Furthermore, for object detection and instance segmentation using MS COCO dataset, our model achieves a box mAP of 54.9 for object and a mask mAP of 47.5, surpassing recent state-of-the-art MaxViT counterpart by +1.5 and +1.8 respectively.

The main contributions of our work are summarized as follows:

• We propose a strong vision backbone, CoAtFormer, that can capture various global interactions throughout every stage of the network.
• We develop a simple but effective composite attention block composed of dynamic channel and efficient spatial attention, enjoining global context in linear complexity.
• Large amounts of performance analysis shows that CoAtFormer outperforms previous SOTA transformers on ImageNet dataset, with a significantly lesser number of parameters. In addition, CoAtFormer also has superior performance on other downstream tasks.

2 Related work

CNNs. Since AlexNet [Krizhevsky et al., 2012], convolutional neural networks (CNNs) have shown remarkable success in various vision applications [Chen et al., 2017; Tan and Le, 2019; Howard et al., 2017; Sandler et al., 2018; Simonyan and Zisserman, 2014; He et al., 2015]. VGGNet [Simonyan and Zisserman, 2014] and InceptionNets [Szegedy et al., 2014; Szegedy et al., 2015] show that a deep neural network consisted of convolutional layers and pooling layers can attain adequate performance in image recognition. ResNet [He et al., 2015] show stronger generalization ability by introducing skip connections every two layers to the base architecture. ConvNeXt [Liu et al., 2022] has re-introduced core designs of vision transformers and demonstrate a pure CNN can achieve performance comparable to vision transformers on various vision tasks.

Vision transformers. Since Transformers [Vaswani et al., 2017; Devlin et al., 2019] achieve tremendous successes in wide natural language processing (NLP) tasks, many efforts [Dosovitskiy et al., 2020; Liu et al., 2021; Dong et al., 2021; Tu et al., 2022; Ding et al., 2022; Hatamizadeh et al., 2023; Yang et al., 2023] have been devoted to developing stronger Transformer based architectures for various vision applications. The pioneering work ViT [Dosovitskiy et al., 2020] directly applies the transformer encoder architecture to a sequence of image patches. However, ViT requires large datasets such as JFT300M [Sun et al., 2017] for training. DeiT [Touvron et al., 2020] utilizes a new training paradigm to enable training of high-performance ViT architecture with fewer data. PVT [Wang et al., 2021] leverages the pyramid structure to generate multi-scale feature maps for generl pixel-level dense prediction tasks. MaxViT [Tu et al., 2022] introduces multi-axis attention to capture both local and global context. DaViT [Ding et al., 2022] presents the dual attention mechanism, which contains spatial self-attention and channel self-attention. DualViT [Yao et al., 2023] incorporates a critical semantic pathway to obtain global semantics with reduced order of complexity. SMT [Lin et al., 2023] proposes Multi-Head Mixed Convolution (MHMC) module and Scale-Aware Aggregation (SAA) module to enhance the convolutional modulation.

Channel-wise Attention. Channel attention mechanisms have been widely adopted in CNNs [Hu et al., 2017; Woo et al., 2018; Qin et al., 2020]. One of the most successful methods is SENet [Hu et al., 2017], which learns channel attention. It first squeezes the feature maps with global average pooling and captures the cross-channel relationships using two fully connected layers. FcaNet [Qin et al., 2020] learns many valuable frequency components by compressing channels using the discrete cosine transform (DCT). Some vision transformers architectures apply channel-wise attention to reduce the computational costs. XCiT [El-Nouby et al., 2021] proposes cross-covariance attention (XCA) to compute channel attention maps. DaViT [Ding et al., 2022] introduces channel group attention (CGA) to perform image-level interactions within each group.

Our work presents dynamic channel attention to capture global channel interactions in Transformers. We demonstrate its power when combined with efficient spatial attention, forming our composite attention mechanism.

3 Method

3.1 Overall Architecture

We propose Composite Attention Vision Transformers (CoAtFormer), a general, efficient, yet effective Transformer backbone capturing both local and global denpendices. The overall architecture of CoAtFormer is illustrated in Figure 2. For an input image with size of H ×W × 3, we leverage the Convolutional Token Embedding layer (two 3×3 convolution layer with a stride 2) to obtain H4 × W4 ×C1 feature maps. Following the design in modern CNNs (e.g., ResNet [He et al., 2015]), the whole network has four stages to generate feature maps of different scales which are important for dense prediction tasks. To produce the hierarchical representation, A downsampling layer (3 × 3 convolution, stride 2) is used between two consecutive stages to reduce the number of tokens and increase the channel dimension. In each stage, several CoAtFormer blocks are stacked sequentially for feature transformation while maintaining the number of tokens. The CoAtFormer block is able to capture local-global representations.

3.2 CoAtFormer Block

The proposed CoAtFormer block contains a conv encoder and a composite attention block, as illustrated in Figure 3. We will detail these parts in the following.

Conv Encoder. The baseline MaxViT [Tu et al., 2022] employs MBConv block [Sandler et al., 2018] as a local token mixer. Although MBConv block has been widely used in efficient models [Tu et al., 2022; Yang et al., 2022a], replacing them with "modernize" convolution block [Liu et al., 2022] does not increase the computational cost. Further, it improves the performance and generalization without increasing the parameters. Using conv encoder before attention provides an additional benefit, as depth-wise convolutions can be considered as conditional position encoding (CPE) [Chu et al., 2021], eliminating the need for explicit positional encoding layers in our model. Specifically, a 3 × 3 depth-wise convolution (DWConv) is first applied to capture the local spatial interactions between pixels. Then, the derived features are fed into two point-wise convolutions with GELU activation. Finally, we introduce a residual connection [He et al., 2015] to enable information to flow across the network. Formally, given an input tensor X ∈ RC×H×W (H, W, C are its height, width, and channels), the conv encoder is represented as follows:

where BN, GELU, Conv1×1, and DWConv3×3 denote Batch Normalization [Ioffe and Szegedy, 2015], Gaussian error Linear Unit [Hendrycks and Gimpel, 2016], 1 × 1 point-wise convolution, and 3 × 3 depth-wise convolution, respectively.

Composite Attention Block. The detailed architecture of the composite attention block is presented in Figure 3. The composite attention block aims to learn enriched local-global features. It begins with local convolutional layers to extract local representations, followed by the composite attention mechanism. As the core element in a composite attention block, the composite attention contains two parallel branches. The two branches share the same input, but focus on relationships of different global context, which can be complementary to each other. In each branch, we first extract global channel or spatial dependencies using Dynamic Channel Attention (DCA) Module or Efficient Spatial Attention (ESA) Module, respectively, then the outputs of two branches can be merged using concatenation. We then describe our dynamic channel attention module and efficient spatial attention module.

Dynamic Channel AttentionModule. In Figure 3, the top branch in composite attention aims to model cross-channel interactions. For an input feature tensor X ∈ RC×H×W , we first squeeze global spatial information into a channel descriptor. This is achieved by using T-Pool layer to reduce the spatial dimension of the input tensor to three by concatenating the average pooled, max pooled and std pooled features across that dimension. Mathematically, the channel responses T ∈ RC×3 are calculated by:

These channel features can be viewed as a set of unordered nodes which are denoted as V = {v1, v2, ...vC}. We then can build a graph G = (V, E), where E stands for all the edges. Each edge reflects the relation weight between two nodes. Based on this graph, we employ a double-layer graph convolution network [Kipf and Welling, 2016] to propagate information between nodes, the information diffusion process can be expressed as:

where W1 ∈ R3×d1 and W2 ∈ Rd1×d2 are state update matrix to be learned, h denotes a non-linear activation, which is GELU [Hendrycks and Gimpel, 2016] in our work.

A ∈ RC×C is a correlation matrix for propagating information, which contains the relation weight between nodes. In our experiments, A is randomly initialized and trained in an end-to-end manner along with the whole model. After graphlevel processing, we aggregate node features F ∈ RC×d2 by taking average of node feature map on the feature dimension then exploit a sigmoid activation to generate channel-wise weights:

Consequently, the output feature maps of DCA module can be formulated as:

where ∗ denotes element-wise multiplication.

Compared to the recent SOTA channel group attention (CGA) [Ding et al., 2022], our proposed DCA enables a more comprehensive communication through a double-layer graph convolution network by treating each channel as a node in the graph. We will verify the effectiveness of DCA in subsequent sections.

Efficient Spatial Attention Module. In Figure 3, the bottom branch of composite attention focuses on global spatial context. For an input feature tensor X ∈ RC×H×W , we reshape it to get the input token embedding X ∈ Rn×d , where n = H × W is the number of patches, d = C is the dimensions of the token embedding. The input token embedding X is linearly transformed into query Q, key K, and value V using three matrices Wq, Wk, and Wv, where Q,K,V ∈ Rn×d , Wq,Wk,Wv ∈ Rd×d . First, we compute the dot products of the query with a learnable global token vector g ∈ Rd , and then apply a softmax function to produce the global context scores cs ∈ Rn as:

The global context scores are able to capture the importance of each query element. Next, the global context scores are multiplied by the key and pooled, resulting in a single global context vector as follows:

Then, the global context vector is multiplied element-wise with the value to form the global spatial information and fed to a linear layer with weights Wo ∈ Rd×d to obtain the output tensor Yesa, it can be described as:

Finally, the output tensor Yesa is reshaped to the original input feature dimensions (C × H × W). Our proposed ESA module is comparatively cheap to compute compared to vanilla MHSA [Dosovitskiy et al., 2020] and has linear complexity with the token length.

After obtaining the outputs from these two branches, we concatenate them along the channel dimension and then project the result back to the original dimension:

where, Wm ∈ R2C×C is a weight matrix. We employ the FFN of vanilla ViT [Dosovitskiy et al., 2020], which consists of two linear layers and a GELU activation.

With the aforementioned these components, the CoAtFormer block can be formulated as:

where Yl, Ml and Zl denote the intermediate output features in the l-th block, respectively. BN denotes the batch normalization [Ioffe and Szegedy, 2015].

3.3 Architecture Variants

For a fair comparison with the other vision Transformer, we consider four different models with various number of parameters and computational complexity. Specifically, we introduce CoAtFormer-T(iny), CoAtFormerS(mall), CoAtFormer-B(ase) and CoAtFormer-L(arge) variants, which is corresponded to CoAtFormer-T, CoAtFormerS, CoAtFormer-B and CoAtFormer-L, respectively. In all these variants, the expansion ratio of each FFN is set as 3. The detail configurations of basic embedding channels C and number of blocks Ni are presented as following:

• CoAtFormer-T:C = 64, Ni = {1, 3, 12, 1}
• CoAtFormer-S:C = 80, Ni = {1, 3, 12, 1}
• CoAtFormer-B:C = 96, Ni = {1, 6, 18, 1}
• CoAtFormer-L:C = 128, Ni = {1, 10, 22, 1}

4 Experiments

To validate the efficacy of CoAtFormer as a general vision backbone, we conduct experiments on ImageNet-1K [Deng et al., 2009] classification, COCO [Lin et al., 2014] object detection and instance segmentation, and ADE20K [Zhou et al., 2017] semantic segmentation. We also perform comprehensive ablation studies to evaluate the effectiveness of each component of CoAtFormer.

4.1 Image Classification

Setup. For fair comparison, we follow the same training strategies as previous works [Touvron et al., 2020; Liu et al., 2021]. Specifically, we train all our models for 300 epochs with the input size of 224×224. We employ the AdamW optimizer with weight decay of 0.05. The default batch size and initial learning rate are set to 1024 and 0.001. Additionally, we explore the effectiveness of our models when pretrained on ImageNet-22K.

Results. In Table 1, we compare our proposed CoAtFormer with current state-of-the-art models. It shows that our CoAtFormer attains new state-of-the-art and consistently outperforms other models by large margins. Specifically, CoAtFormer-T achieves 84.4% Top-1 accuracy with only 4.6G FLOPs, surpassing Swin-T, DaViT-T and MaxViT-T by 3.2%, 1.6% and 0.8% respectively. In addition, when compared to small-sized and base-sized models, our CoAtFormer also achieves the best performance. Notably, CoAtFormer-S achieves a top-1 accuracy of 85.3% with only 37.5M parameters and 7.2GFLOPs of computation, outperforming many counterpart-Base models such as Swin-B [Liu et al., 2021], ConvNeXt-B [Liu et al., 2022], and MaxViT-B [Tu et al., 2022], which have over 80M parameters and 15GFLOPs of computation. For example, our CoAtFormer-S attains 2.0% and 0.7% higher accuracy than Swin-B and DaViT-B, respectively, using near half computations.

We further present the ImageNet-22K pre-training results in Table 1. CoAtFormer-L achieves an 88.7% top-1 accuracy, surpassing MOAT-3 [Yang et al., 2022a] by 0.5% while utilizing fewer parameters (157.5M vs. 190M) and lower FLOPs (105.8G vs. 141.2G). This highly highlights the remarkable scalability capabilities of CoAtFormer.

4.2 Object Detection and Instance Segmentation

Setup. We evaluate our models on object detection with COCO 2017 [Lin et al., 2014]. All the backbones are first pretrained using ImageNet-1K. The pretrained models are then plugged into two detectors, Mask R-CNN [He et al., 2017b] and Cascade Mask R-CNN [Cai and Vasconcelos, 2017] . We follow the finetuning strategy used in ConvNeXt [Liu et al., 2022] on the COCO training set.

Results. Table 2 reports the results of our models on MS COCO dataset. Using a Mask-RCNN detector, our CoAtFormer-T (49.1/44.8) backbone outperforms counterparts with ConvNeXt-T (46.2/41.7) by +2.9 and +3.1 and DaViT-T (47.4/42.9) by +1.7 and +1.9 in terms of box AP and mask AP, respectively. Using a Cascade Mask- RCNN detector, we also observe substantial gains across all model configurations. Furthermore, we observe a almost saturated mAP in Swin Transformer [Liu et al., 2021] and GCViT [Hatamizadeh et al., 2023] from small to base model, while the mAP of our model consistently improves with larger model size, demonstrating enhanced scalability.

4.3 Semantic Segmentation on ADE20k

Setup. We further benchmark our method for semantic segmentation using the ADE20K dataset. We use UperNet [Xiao et al., 2018] as the segmentation method and our CoAtFormer as the backbone. We follow the same training recipe proposed by [Liu et al., 2021]. Specifically, we train UperNet for 160k iterations with an input size of 512 × 512. We use the AdamW optimizer with a weight decay of 0.01 and set the batch size to 16.

Results. In Table 3, the results show that our CoAtFormer outperforms DaViT, FocalNet, and GCViT significantly under all configurations. Concretely, CoAtFormer-T (48.8), CoAtFormer-S (51.3) and CoAtFormer-B (52.1) backbones improve 1.8, 2.5 and 1.9 mIoU gains over counterpart models with GCViT [Hatamizadeh et al., 2023], respectively.

4.4 Ablation Study

Composite attention layout. We perform experiments on the arrangement of our composite attention mechanism. Three options with comparable computations are explored: (i) dynalic channel attention (DCA) first; (ii) efficient spatial attention (ESA) first; and (iii) parallel arrangement of both types of attention. The results are presented in Table 4. It can be observed that the three strategies yield similar performance, with a slight advantage for the 'parallel arrangement of both types of attention' approach.

Component analysis. Table 5 demonstrates the significance of composite attention block and conv encoder in our proposed architecture. Substituting composite attention block with convolution encoder leads to a decrease in accuracy by 2.9%, highlighting the effectiveness of composite attention block in our design. Moreover, removing conv encoder in all four stages of the network results in an additional accuracy reduction of 1.0%. In addition, We also measure the contributions of composite attention components (e.g., dynamic channel attention (DCA) and efficient spatial attention (ESA) ) in Table 5. Removing DCA module and ESA module decrease the accuracy by 1.3 and 1.5, respectively.

Different fusionmethod. We try two general operations as the alternatives to the fusion method in composite attention. One is element-wise addition, the other is weighted sum. The results are shown in Table 6. It is clear that concatenation and an project layer as the fusion module achieves the best model size and accuracy trade-off among the different fusion methods we evaluated.

5 Interpretability

In Figure 4, we can find that learned spatial attention distributions align the region of image semantics, and hence demonstrate the effictiveness of efficient spatial attention. Additionally, corresponding Grad-CAM maps present accurate object localization with most intricate details.

6 Conclusion

In this work, we propose a novel architecture named CoAtFormer, which can efficiently capture global channel and spatial contexts by utilizing dynamic channel attention and efficient spatial attention. The proposed CoAtFormer architectures take advantages of both CNNs and transformers to capture local and global information, improving the representation ability of the model. Extensive experiments on on ImageNet and other downstream vision applications demonstrate the effectiveness and superiority of the proposed architecture.