Dynamic Group Transformer: A General Vision Transformer Backbone with Dynamic Group Attention

1 Introduction

Recently, Transformer has shown a great potential for various vision tasks [Dosovitskiy et al., 2020; Touvron et al., 2020; Liu et al., 2021b; Dong et al., 2021]. The pioneer Vision Transformer [Dosovitskiy et al., 2020] (ViT) stacked multiple Transformer blocks to process non-overlapping image patch (i.e., visual token) sequences for image classification. However, the global self-attention in Transformer makes each query attend to all keys, which has the quadratic complexity to sequence length, and results in expensive computation costs and memory usage, especially for high-resolution images.

∗Equal contribution

† Interns at the Institute of Deep Learning, Baidu Research

‡Corresponding author

To improve the efficiency of the global self-attention, the state-of-the-art methods [Liu et al., 2021b; Huang et al., 2021; Dong et al., 2021; Fang et al., 2021] focused on how to divide the global image into multiple local regions (or windows). Each query only attends to a few keys within the manually designed local regions (windows). For example, Swin Transformer [Liu et al., 2021b] computed the self-attention within each local window and employed a shifted mechanism to make cross-window connections (Figure 1 (b)). Different from Swin Transformer, CSwin [Dong et al., 2021] proposed cross-shaped window self-attention for computing attention in the horizontal and vertical stripes in parallel that form a cross-shaped window (Figure 1 (c)). Shuffle Transformer [Huang et al., 2021] presented a spatial shuffle operation to make information flow across windows (Figure 1 (d)). MSG-Transformer [Fang et al., 2021] proposed to compute attention in local regular windows, and used additional MSG tokens to get connections between them (Figure 1 (e)). These window-based methods achieved excellent performances and were superior to the CNN counterparts, however, they rely on hand-crafted window partition mechanisms. Furthermore, these partition methods are data-agnostic and ignore the input content, as a result, it is likely that one query maybe attends to irrelevant keys/values.

To address the issues mentioned above, a good idea is to dynamically select relevant keys/values for each query. However, it leads to unreasonably high memory usage and computation complexity. We propose dynamic group attention (DG-Attention), which dynamically divides all queries into multiple groups and selects the most relevant keys/values for each group. Specifically, the input visual tokens (or feature vectors) are divided adaptively into multiple groups according to their similarity to all cluster centroids. Therefore, such a partition mechanism is adaptive to input images. Then, we use the cluster centroid of each group to select the most relevant keys/values subset from the whole keys/values set, and the self-attention is conducted within each group. It enables our model to focus on relevant keys/values without any spatial constraint. And also, through the dynamic grouping, our DG-Attention does not cause a high memory usage or large computation cost. For example, as shown in Figure 1 (f), the red point (query) can attend to its relevant region denoted with red solid line boundaries, and the blue point can attend to the regions with blue solid line boundaries. Benefiting from the data-dependent and flexible group mechanism, our DGAttention shows superiority to the other window-based selfattention illustrated in Figure 1.

Based on the proposed DG-Attention, we design a general vision transformer backbone for image classification, named Dynamic Group Transformer (DGT). We scale our approach up to get a family of models, including DGT-T (24M), DGTS (52M), and DGT-B (90M). They achieve significantly a better performance than previous methods. Our DGT-T can achieve Top-1 classification accuracy of 83.8% on ImageNet1k, 50.2% mIoU on ADE20K for semantic segmentation, 47.7% box mAP for object detection, and 43.0% mask mAP on COCO for instance segmentation, outperforming the stateof-the-art methods. Furthermore, our largest variant DGTB is also superior to the previous methods, achieving 85.0% Top-1 accuracy on ImageNet-1K, 51.2% mIoU on ADE20K, 49.1% box mAP, and 44.1% mask mAP on COCO dataset.

2 RelatedWork

This section briefly reviews related works, including improving efficiency and enhancing inductive bias for Vision Transformer.

Improving Efficiency for Vision Transformer. There are two main categories of methods to reduce the computation demand for Vision Transformer. (1) Pruning Token. It aims to remove redundant tokens and reduce the number of tokens entering into attention modules to save the computation demand. For example, DynamicViT [Rao et al., 2021] pruned tokens in each layer with Gumbel softmax. IA-Red [Pan et al., 2021] used reinforcement Learning to achieve a similar effect. Such methods achieved good performances on image classification, but they are not friendly enough for downstream dense prediction tasks. (2) Designing efficient attention mechanisms. Such methods mainly explored how to make each query attend to partial keys/values for reducing the computational cost. PVT [Wang et al., 2021a] directly downsampled the keys and values in each block. Swin transformer [Liu et al., 2021b] divided all queries/keys/values into multiple windows and computed the self-attention within each local window. Similarly, CSwin transformer [Dong et al., 2021] expanded the window into a cross-shaped window. MSG-Transformer [Fang et al., 2021] used additional MSG tokens to make connections between windows. Different from these methods that employed pre-designed, handcrafted window partition mechanisms, our method dynamically divides all queries into multiple groups and selects the most relevant keys/values for each group.

Enhancing Inductive Bias for Vision Transformer. Vision transformers have shown successes in various computer vision tasks, due to their ability to model long-range dependencies within an image. However, recent works also showed that inductive bias could be incorporated for vision transformers. CPE [Chu et al., 2021b] used convolution layers to generate the conditional position encoding. CVT [Wu et al., 2021] employed convolution layers to generate the queries, keys and values. CMT [Guo et al., 2021] also incorporated the convolution layers into the FFN.

3 Method

In this section, we first introduce our Dynamic Group attention (DG-Attention). Then, we present the composition of the Dynamic Group Transformer Block. Finally, we describe the overall architecture and variant configurations of our Dynamic Group Transformer (DGT) backbone.

3.1 Dynamic Group Attention

To make each query attend to relevant keys/values, we propose a Dynamic Group Attention (DG-Attention). It dynamically divides all queries into multiple groups and selects the most relevant keys/values for each group to compute the selfattention. As shown in Figure 2(c), given an input feature map X ∈ RH×W×C (C is the channel number, H and W denotes the height and width, respectively), we first get query embeddings {XiQ}Li=1 , key embeddings {XiK }Li=1 , and value embeddings {XiV }Li=1 , where L = H ×W. For simplicity, we assume there is only one head in the DGAttention. It's easy to expand to the multi-head situation where each head has its queries, keys, values, and cluster centroids. Then, we use k-means clustering algorithm to dynamically divide all queries into G different query groups (clusters) XQ = {XQj |XQj ∈ RNj×C }Gj=1 , where j is the group index and Nj is the number of queries in the j th group. Meanwhile, we use top-k operations to find the k most relevant keys and values for each query group, which are denoted as XK = {XKj |XKj ∈ Rk×C }Gj=1 and XV = {XVj |XVj ∈ Rk×C }Gj=1 , respectively.

Specifically, for the j th query group, we compute the dot product between its cluster centroid ej and all keys {XiK }Li=1 , and then select the k most relevant elements according to the dot product sorting, which can be formulated as follow:

Then, we update the cluster centroid as below:

where τ is a hyper-parameter to control the update speed. We empirically set τ to 0.1 × lr, where lr is the learning rate.

Computation Complexity Analysis

We analyze the computation complexity of our DG-Attention and the global self-attention to further reveal the efficiency of our method. Here, we only consider the process of computing the attention maps and weighted sums of values for clarity. Given input features of size L ×C, the global self-attention has a computational complexity of

In DG-Attention, each query only attends to k keys, so the basic computation complexity of DG-Attention is 2kLC. Besides, grouping queries and selecting the most significant k keys require an additional computation cost of 2LGC + kGlogL. Therefore, the total computation complexity of our DG-Attention is

The ratio of the computation complexity of our DGAttention and Global self-attention is:

where Top-k is the function that returns the indices of top k values, and idj is an index vector. Then, the self-attention is conducted within each group:

where SA denotes the Self-Attention, Yj is the updated output of the j th query group. Finally, {Yj}Gj=1 are scattered into the output Y ∈ RL×C according to their original spatial position indexes.

As each group has a different number of queries, this algorithm cannot be implemented using the general matrix multiplication. We implement this algorithm using CUDA, and the detail can be found in the supplementary material.

To make the training stable, we update the cluster centroids with exponential moving average after each iteration. Specifically, for the j th cluster centroid, we compute the current cluster centroid as follow:

where L is larger than G and k. For high-resolution inputs, the ratio ΩDG-Attention ΩGlobal << 1. Typically, for the ImageNet classification task, k is set to 98 for the first three stages, while the corresponding L is 3136, 784, and 196. Thus, the ratio is 0.05, 0.19, and 0.75 for DGT-T. Besides, k is independent of the shapes of the parameters in our models, so we can adjust k to balance the performance and computation efficiency.

3.2 Dynamic Group Transformer Block

The Dynamic Group Transformer block is shown in Figure 2(b). It first employs the widely-used conditional position embeddings (CPE) [Chu et al., 2021b] to generate the positional information. Then, DG-Attention is applied to model spatial relevant dependencies flexibly and dynamically. Last, IRFFN (Inverted Residual Feed-forward Network) [Guo et al., 2021] further is employed to capture local dependencies. The forward process of the lth block can be formulated as follows:

where LN(·) denotes Layer Normalization, Xl and Xl-1 are the output of the lth and l-1th block, respectively.

3.3 Overall Architecture and Variants

Our Dynamic Group Transformer (DGT) consists of a convolutional stem, four hierarchical stages, and a classifier head, as shown in Figure 2 (a). The stem is designed to extract local dependency, similar to [Guo et al., 2021], which consists of one 3×3 convolution layer with stride = 2 and two 3×3 convolution layers with stride = 1. After the stem, each stage contains a patch merging layer and multiple transformer blocks. The first three stages use the DGT block, and the last stage applies global self-attention (GSA) block, which is achieved by replacing the DG-Attention with global self-attention in DGT block. We decrease the number of tokens and double the channel dimension by using a 3×3 convolutional layer with stride = 2 before each stage to produce a hierarchical representation. The final classifier head consists of two linear layers.

Finally, we design three different variants, including DGTTiny (DGT-T), DGT-Small (DGT-S), and DGT-Base (DGTB), whose detailed configurations are shown in Table 1. For all variants, the number of blocks in each stage is fixed with [1,2,17,2]. In each DGT block, the expand ratios of IRFFN are set to 4, the number of groups G is 48. The number of selected keys/values k is 98 for image classification on ImageNet [Deng et al., 2009]. The main differences among all variants are the channel dimension and the number of heads in DGT blocks. Besides, to train the model stably, we apply post LayerNorm and cosine attention [Liu et al., 2021a] in DGT-S and DGT-B.

4 Experiments

We first compare our Dynamic Group Transformer (DGT) with the state-of-the-art backbones on ImageNet-1K [Deng et al., 2009] for image classification. To further verify the effectiveness and generalization of our backbone, we perform experiments on ADE20K [Zhou et al., 2017] for semantic segmentation, COCO [Lin et al., 2014] for object detection and instance segmentation. Finally, we analyze the key design of our Dynamic Group Transformer.

4.1 Image Classification on ImageNet-1K

We conduct experiments on ImageNet-1K [Deng et al., 2009] dataset, which has 1.28M images in the training set and 50K images in the validation set. Detailed settings are described in the supplementary material.

Results

Table 2 compares the performance of our DGT with the stateof-the-art CNN models and vision transformer backbones on ImageNet-1K validation set. We can see that our DGT variants outperform the state-of-the-art CNN models and vision transformer models when using similar FLOPs. DGT-T achieves 83.8% top-1 accuracy with 4.3G FLOPs and outperforms the CNN models Reg-4G and Efficient B4 by 3.8% and 0.9%, respectively. Meanwhile, our DGT outperforms the advanced Transformer-based backbones, and is +1.6% and +0.8% higher than Swin and CSwin Transformer, respectively, for all variants under the similar model size and FLOPs. For example, our DGT-T can surpass PVT-S, SwinT, and CSwin-S by 4.0%, 2.5%, and 1.1%, respectively. Our DGT-B can outperform Swin-B and CSwin-B by 1.7% and 0.8%, respectively. These results demonstrate the effectiveness and efficiency of our approach.

4.2 Semantic Segmentation on ADE20K

To demonstrate the superiority of our Dynamic Group Transformer for semantic segmentation. We conduct experiments on ADE20K with the widely-used UperNet [Xiao et al., 2018] framework for fair comparisons to other backbones. Detailed implementation can be found in the supplementary material.

Results

Table 3 shows the comparisons of UperNet [Xiao et al., 2018] with various advanced Transformer backbones on ADE20K validation set. We report both single-scale (SS) mIoU and multi-scale (MS) mIoU for a cleaner comparison. Our DGT variants outperforms the state-of-the-art methods consistently. Specifically, our DGT-T achieves 50.2% mIoU with single scale testing, outperforming the Swin-T and CrossFormer-S by 5.7% and 2.6%. Our DGT-S outperforms Swin-S and CrossFormer-B by 3.2% and 0.9% SS mIoU. Besides, our DGT-B achieves 51.2%/51.8% SS/MS mIoU, outperforming the Swin-B and CrossFormer-L by 3.1%/2.1% and 0.7%/0.4%. These results demonstrate the advantages of our Dynamic Group Transformer for semantic segmentation.

4.3 Object Detection and Instance Segmentation on COCO

We further evaluate our DGT backbone on COCO [Lin et al., 2014] dataset for object detection and instance segmentation. Following previous works [Liu et al., 2021b; Dong et al., 2021], we utilize Mask R-CNN [He et al., 2017] framework under 1x schedule. More details are provided in the supplementary material.

Results

The results on COCO dataset are shown in Table 4(a). All DGT variants outperform the state-of-the-art vision transformer backbones under similar FLOPs. Specifically, for object detection, our DGT-T, DGT-S, and DGT-B can achieve 47.7%, 48.4% and 49.1% box AP, which surpass Swin by 5.5%, 3.6%, and 2.2%, respectively. For instance segmentation, our DGT-T, DGT-S, and DGT-B are 3.9%, 2.6%, and 1.8% mask AP higher than the Swin. Besides, DGT-T outperforms CrossFormer-S by 2.3% bos AP on object detection and 1.6% mask AP on instance segmentation. DGT-S outperforms CrossFormer-B by 1.2% box AP on object detection and 0.8% mask AP on instance segmentation.

4.4 Ablation Studies

We conduct ablation studies for the key designs of our methods on the image classification task. All experiments are performed with the Tiny variant under the same training settings.

Effect of Hyper-parameters G and k

First, we validate the effect of the hyper-parameters G and k. G is the number of groups. A small G causes the queries in a group to be very different, so the selected keys cannot suit all queries. k determines how many keys each query attends to. There will be too much information loss if k is too small. The default setting of our model on ImageNet is G = 48 and k = 98. We compare the default setting with halving G from 48 to 24 and halving k from 98 to 49.

The results are shown in Table 4(b). Decreasing G or k leads to a poorer performance. Halving G and k decrease the top-1 accuracy by 0.2% and 0.1%, respectively. We can balance the performance and efficiency by adjusting G and k.

Table 4: Experiments on COCO and ablation studies.

Comparison with Related Transformer Blocks

To validate the design of our DGT block, which uses convolution layers to extract local dependency and uses the DGAttention to extract non-local dependency, we replace our DGT block with other blocks and compare their performance. We select three blocks: Swin block, CSwin block, and CMT block. Swin block and Cswin block use shifted windowbased self-attention and cross-shaped window self-attention, respectively. The results are shown in Figure 4(c). Our DGT block obviously outperforms Swin block, CSwin block and CMT block by 1.0%, 0.8% and 0.4%.

5 Visualization

We visualize the query groups and their corresponding selected keys with an example shown in Figure 3. One can find : 1) different queries prefer to attends to different keys according to their content. These three groups mainly contain the queries of the bird, branches, and background, and they also attend to the keys of the bird, branches, and background, respectively. 2) A query may prefer to attend to a long-range area rather than a short-range local region. These findings show the advantages of our DG-Attention. More visualization examples can be found in the supplementary material.

6 Conclusion

We have presented an effective dynamic attention mechanism named Dynamic Group Attention (DG-Attention), which dynamically divides input queries into multiple groups and selects relevant keys/values for each group. DG-Attention can model more relevant context dependencies than the previous pre-designed window-based local attention mechanism. Based on the proposed DG-Attention, we have developed a general Vision Transformer backbone, Dynamic Group Transformer (DGT), which can outperform the state-of-theart on ImageNet-1K for image classification. Furthermore, our DGT outperforms the existing Vision Transformer backbones on ADE20K for semantic segmentation, and COCO for object detection and instance segmentation.