MixFormer: Mixing Features across Windows and Dimensions

Although performing self-attention inside non-overlapped windows brings computational efficiency²²2It has linear computational complexity concerning image size, as shown in Table 1., it results in a limited receptive field due to no cross-window connections being extracted. Several methods resort to shift [34], expand [45, 57], shuffle [26], or convolution [26, 61] to model connections across windows. Considering that convolution layers are designed to model local relations, we choose the efficient alternative (depth-wise convolution) as a promising way to connect windows.

Attention then moves to adopt a proper way to combine local-window self-attention and depth-wise convolution. Previous methods [34, 45, 57, 26, 61] fill the goal by stacking these two operators successively. However, capturing intra-window and cross-window relations in successive steps make these two types of relations less interweaved.

In this paper, we propose a parallel design that enlarges the receptive fields by simultaneously modeling intra-window and cross-window relations. As illustrated in Figure 1, local-window self-attention and depth-wise convolution lie in two parallel paths. In detail, they use different window sizes. A $7\times 7$ window is adopted in local-window self-attention, following previous works [45, 23, 64, 34]. While in depth-wise convolution, a smaller kernel size $3\times 3$ is applied considering the efficiency³³3The results in Table 8 show that $3\times 3$ is a good choice to achieve balance in accuracy and efficiency.. Moreover, as their FLOPs are different, we adjust the number of channels according to the FLOPs proportion in Table 1. Then, their outputs are normalized by different normalization layers [27, 1] and merged by concatenation. The merged feature is sent to the successive Feed-Forward Network (FFN) to mix the learned relations across channels, generating the final output feature.

The parallel design benefits two-folds: First, combining local-window self-attention with depth-wise convolution across branches models connections across windows, addressing the limited receptive fields issue. Second, parallel design models intra-window and cross-window relations simultaneously, providing opportunities for feature interweaving across branches and achieving better feature representation learning.

In general, sharing weights limits the modeling ability in the shared dimension. A common way to solve the dilemma is to generate data-dependent weights as done in dynamic networks [24, 52, 4, 30]. Local-window self-attention computes weights on the fly on the spatial dimension while sharing weights across channels, resulting in the weak modeling ability issue on the channel dimension. We focus on this issue in this subsection.

To enhance the modeling capacity of local-window self-attention on the channel dimension, we try to generate channel-wise dynamic weights [24]. Given that depth-wise convolution shares weights on the spatial dimension while focusing on the channel. It can provide complementary clues for local-window self-attention and vice versa. Thus, we propose bi-directional interactions (in Figure 1 and Figure 2) to enhance modeling ability in the channel and spatial dimension for local-window self-attention and depth-wise convolution respectively. The bi-directional interactions consist of the channel and spatial interaction among the parallel branches. The information in the depth-wise convolution branch flows to the other branch through the channel interaction, which strengthens the modeling ability in the channel dimension. Meanwhile, the spatial interaction enables spatial relations to flow from the local-window self-attention branch to the other. As a result, the proposed bi-directional interactions provide complementary clues for each other. Next, we present the designs of the channel and spatial interactions in detail.

For the channel interaction, we follow the design of the SE layer [24], as shown in Figure 2. The channel interaction contains one global average pooling layer, followed by two successive $1\times 1$ convolution layers with normalization (BN [27]) and activation (GELU [20]) between them. At last, we use sigmoid to generate attention in the channel dimension. Although our channel interaction shares the same design with the SE layer [24], they differ in two aspects: (1) The input of the attention module is different. The input of our channel interaction comes from another parallel branch, while the SE layer is performed in the same branch. (2) We only apply the channel interaction to the value in the local-window self-attention instead of applying it to the module’s output as the SE layer does.

For the spatial interaction, we also adopt a simple design, which consists of two $1\times 1$ convolution layers with followed BN [27] and GELU [20]. The detailed design is presented in Figure 2. These two layers reduce the number of channels to one. At last, a sigmoid layer is used to generate the spatial attention map. Same as we did in the channel interaction, the spatial attention is generated by another branch, where the local-window self-attention module is applied. It has a larger kernel size ($7\times 7$) than the depth-wise $3\times 3$ convolution and focuses on the spatial dimension, which provides strong spatial clues for the depth-wise convolution branch.

Thanks to the above two designs, we mitigate two core issues in local-window self-attention. We integrate them to build a new transformer block, Mixing Block, upon the standard window attention block. As shown in Figure 1, the Mixing Block consists of two efficient operations in a parallel design, bi-directional interactions (Figure 2), and an FFN (Feed-Forward Networks) [46]. It can be formulated as follow:

	$\displaystyle\hat{X}^{l+1}$	$\displaystyle=\mathrm{MIX}(\mathrm{LN}(X^{l}),\mathrm{W\text{-}MSA},\mathrm{CONV})+X^{l},$		(1)
	$\displaystyle X^{l+1}$	$\displaystyle=\mathrm{FFN}(\mathrm{LN}(\hat{X}^{l+1}))+\hat{X}^{l+1}$		(2)

Where MIX represents a function that achieves feature mixing between the W-MSA (Window-based Multi-Head Self-Attention) branch and the CONV (Depth-wise Convolution) branch. The MIX function first projects the input feature to parallel branches by two linear projection layers and two norm layers. Then it mixes the features by following the steps shown in Figure 1 and Figure 2. For FFN, we keep it simple and follow previous works [44, 34], which is an MLP that consists of two linear layers with one GELU [20] between them. Moreover, we also try to add depth-wise convolution as done in PVTv2 [48] and HRFormer [61], which does not give many improvements over the MLP design (Table 9). Thus, to keep the block simple, we use MLP in FFN.

Based on the obtained block, we design an efficient and general-purpose vision transformer, MixFormer, with pyramid feature maps. There are four stages with downsampling rates of $\{4,8,16,32\}$ respectively. MixFormer is a hybrid vision transformer, which uses convolution layers in both stem layers and downsampling layers. Besides, we introduce a projection layer in the tail of the stages. The projection layer increases the feature’s channels to $1280$ with a linear layer followed by an activation layer, aiming to preserve more details in the channel before the classification head. It gives a higher performance in classification, especially works well with smaller models. Same design can be found in previous efficient networks, such as MobileNets [22, 41] and EfficeintNets [43]. The sketch of our MixFormer is given in Figure 3.

We stack the blocks in each stage manually and format several models in different sizes, whose computational complexities ranges from $0.7$G (B1) to $3.6$G (B4). The number of blocks in different stages is set by following a recipe: putting more blocks in the last two stages, which is roughly verified in Table 10. As shown in Table 2, we present the detailed settings of the models.

We first verify our method by classification on ImageNet-1K [9]. To make a fair comparison with previous works [44, 49, 34], we train all models for $300$ epochs with an image size of $224\times 224$ and report Top-1 validation accuracy. We apply an AdamW optimizer using a cosine decay schedule. By following the rule that smaller models need less regularization, we adjust the training settings gently when training models in different sizes. Details are in Appendix C.

Table 3 compares our MixFormer with efficient ConvNets [43, 40] and various Vision Transformers [44, 49, 53, 34, 31, 57, 26]. MixFormer performs on par with EfficientNet [43] and outperforms RegNet [40] by significant margins under various computational budgets (from B1 to B4). We note that it is nontrivial to achieve such results for vision transformer-based models, especially on small models (FLOPs $<1.0$G ). Previous works such as DeiT [44] and PVT [49] show dramatic performance drops when reducing model complexities ($-7.7\%$ from DeiT-S to DeiT-T and $-4.7\%$ from PVT-S to PVT-T). Compared with Swin Transformer [34] and its variants [31, 57, 26], MixFormer shows better performance with less computational costs. In detail, MixFormer-B4 achieves $83.0\%$ Top-1 accuracy with only $3.6$G FLOPs. It outperforms Swin-T [34] by $1.7\%$ while saving $20\%$ computational costs and gives comparable results with Swin-S [34] but being $2.4\times$ efficient. The competitive advantage of MixFormer maintains when it comes to LG-Transformer [31], Focal Transformer [57] and Shuffle Transformer [26]. Moreover, our MixFormer also scales well to smaller and larger models. More results are provided in Appendix A.

We validate the effectiveness of MixFormer on downstream tasks. We train Mask R-CNN [18] on the COCO2017 train split and evaluate the models on the val split. Two training schedules ($1\times$ and $3\times$) are adopted to show a consistent comparison with previous methods [19, 44, 34, 31, 57]. For the $1\times$ schedule, we train for $12$ epochs with a single size (resizing the shorter side to $800$ while keeping its longer side no more than $1333$) [18]. While in $3\times$ schedule (36 epochs), we use multi-scale training by randomly resizing the shorter side to the range of $[480,800]$ (See Appendix C for more details). Expect for Mask R-CNN [18], we also provide comparisons with previous works based on a stronger model, Cascade Mask R-CNN [3, 18], where a $3\times$ schedule is conducted.

Table 4 shows that MixFormer consistently outperforms other competitors [19, 49, 53, 34, 31, 57] under various model sizes with Mask R-CNN [18]. In particular, MixFormer-B4 achieves +$\boldsymbol{2.9}$(+$\boldsymbol{1.6}$) higher box mAP and +$\boldsymbol{2.1}$(+$\boldsymbol{1.4}$) higher mask mAP than the Swin-T [34] baseline with $1\times$ ($3\times$) schedule. Moreover, MixFormer keeps its efficiency in detection and instance segmentation, enabling higher performance with less computational costs than other networks [19, 34]. It is a surprise that our MixFormer-B1 (only with 0.7G) performs strongly with Mask R-CNN ($1\times$), which exceeds ResNet-50 (with 4.1G) [19] by 2.3 box mAP and 2.9 mask mAP. The results suggest that implications for designing high-performance small models on detection are highlighted in MixFormer.

We also evaluate MixFormer with Cascade Mask R-CNN [18, 3], which is a stronger variant of Mask R-CNN [18]. MixFormer-B4 provides robust improvements over Swin-T [34] (Table 5) regardless of different detectors, as it shows similar gains (+1.1/1.2 box/mask mAP v.s. +1.6/1.4 box/mask mAP) with the ones on Mask R-CNN ($3\times$) (Table 4).

Our experiments are conducted on ADE20K [66] using UperNet [54]. For training recipes, we mainly follow the settings in [34]. We report mIoU of our models in single scale testing (ss) and multi-scale testing (ms). Details are provided in Appendix C.

In Table 6, MixFormer-B4 consistently achieves better mIoU performance than previous networks. It seems that the connections across windows and dimensions in the Mixing Block provide more benefits on semantic segmentation as the gains are larger than the ones on detection tasks (Table 4,Table 5) with the same backbones. In particular, MixFormer-B4 outperforms Swin-T [34] by $\boldsymbol{2.2}$ mIoU.

Moreover, other variants of MixFormer (from B1 to B3) also achieve higher performance while being more efficient than previous networks. Notably, MixFormer-B3 obtains $45.5$ mIoU (comparable with Swin-T [34] but less FLOPs), which achieves on par results with OCRNet [60] with HRNet-W48 [47] ($45.7$ mIoU). Note that HRNet [47] is carefully designed to aggregate the features in different stages, while MixFormer simply constructs pyramid feature maps, indicating the strong potential for further improvements on dense prediction tasks.

We provide ablations with respect to our designs on MixFormer-B1. We report all variations of different designs on ImageNet-1K [9] classification, COCO [33] detection and segmentation, and ADE20K [66] semantic segmentation. To make quick evaluations, we only train MixFormer-B1 for $200$ epochs on ImageNet-1K [9]. Then, the pre-trained models are adopted by Mask R-CNN [18] ($1\times$) on MS COCO [33] and UperNet [54] ($160k$) on ADE20K [66]. Note that, the differences in pre-train models provide slightly different results with the ones in Table 4 and Table 6.

Table 7 provides the comparison of the ways (successive design or parallel design) to combine local-window self-attention and depth-wise convolution. Our parallel design consistently outperforms the successive design across various vision tasks, which verifies the hypothesis that the parallel design enables better feature representation learning in Section 1. The models below use parallel design by default.

Table 7 shows the results of the proposed interactions. According to the results, we see that both channel and spatial interactions outperform the model without interactions across all different vision tasks. Combining two interactions promotes better performance, resulting in consistent improvements by $0.3\%$ Top-1 accuracy on ImageNet-1K, $0.9/0.7$ box/mask mAP on COCO, and $1.1$ mIoU on ADE20K. Given that we only use simple and light-weighted designs for bi-directional interactions, the gains are nontrivial, which indicates the effectiveness of providing complementary clues for local-window self-attention and depth-wise convolution.

Table 8 shows that the performance will drop significantly on various vision tasks ($-1.3$ Top-1 accuracy on ImageNet-1K, $-4.0/-3.0$ box/mask mAP on COCO, and $-3.3$ mIoU on ADE20K) if we reduce the window size of the depth-wise convolution from $3\times 3$ to $1\times 1$. This phenomenon means that it’s necessary for depth-wise convolution to use a window size (at least $3\times 3$) with the ability to connect across-window. Besides, when we increase the window size to $5\times 5$, no clear further gains are observed. Thus, we use a window size of $3\times 3$ regarding the efficiency.

We also investigate other designs in MixFormer, including applying shifted windows and inserting $3\times 3$ depth-wise convolution in FFN, which play significant roles in previous works [34, 61]. As presented in Table 9, shifted window fails to provide gains over MixFormer. We hypothesize that the depth-wise convolution builds connections among windows, removing the need for shift operation. Besides, although inserting $3\times 3$ depth-wise convolution in FFN can provide further gains, the room for improvements is limited with MixFormer. Thus, we use MLP in FFN by default.

Previous works usually put more blocks in the third stage and greatly increase the number of blocks in that stage when scaling models [19, 49, 34]. We show an alternative way that can achieve the goal. We roughly conduct experiments on the way of stacking blocks. In Table 10, we achieve slightly higher performance on various vision tasks under less computational complexities by putting more blocks in both the last two stages. We follow this recipe to build our MixFormer.

In Table 11, we conduct experiments on two more downstream tasks: keypoint detection and long-tail instance segmentation. Detailed experimental settings are provided in Appendix C.

COCO keypoint Detection: In Table 11, MixFormer-B4 outperforms baseline models [19, 34] by significant margins in all metrics. Moreover, MixFormer also shows clear advantages compared with HRFormer [61], which is specifically designed for dense prediction tasks.

LVIS $1.0$ Instance Segmentation: This task has $\sim 1000$ long-tailed distribution categories, which relies on the discriminative feature learned by the backbone. Results in Table 11 show that MixFormer outperforms the Swin-T [34] by $1.0$ AP${}^{\text{mask}}$, which demonstrates the robustness of the learned representation in MixFormer.

Summary: Considering the promising results given by MixFormer in previous tasks: object detection, instance segmentation, and semantic segmentation, MixFomer can serve as a general-purpose backbone and outperform its alternatives in 5 dense prediction tasks.

We apply our Mixing Block to typical ConvNets, ResNet50 [19] and MobileNetV2 [41]. Following [42], we replace all the blocks in the last stage with our Mixing Block in ConvNets. To make a fair comparison, we adjust the number of blocks to maintain the overall computational cost. Table 12 shows that the Mixing Block can provide gains on ConvNets [19, 41] as a alternative to ConvNet blocks. Specifically, Mixing Block brings $1.9\%$ and $1.6\%$ Top-1 accuracy on ImageNet-1K [9] over MobileNetV2 [41] and ResNet50 [19]. Moreover, we also provide the result of MobileNetV2 [41] with SE layer [24] and Non-Local [50] in Table 12. It gives inferiror performance than our mixing block.