Embedding-Free Transformer with Inference Spatial Reduction for Efficient Semantic Segmentation

The importance of modeling the global context has been demonstrated by the self-attention mechanism in the transformer. Beyond the general attention method, various attention methods have been studied. [58, 59] proposed the spatial reduction attention mechanism, which reduces the key-value resolution for efficiency. [63] leveraged the pyramid pooling to reduce the key-value in multi-scale resolution. Based on the spatial reduction attention structure, [23, 75, 76] exploited the convolutional layer in the attention. The window-based attention method [40, 39] considered the local window regions for efficiency. [13] proposed the local window attention with global attention. The convolution-based attention [62, 65, 23, 15] used the convolutional operation to consider local context with global context. The channel reduction attention method [31] reduced the query and key channels. However, all these self-attention methods are based on the query, key and value embeddings. Different from these methods, we propose the efficient Embedding-Free Attention module by focusing on that the global non-linearity is important in the attention mechanism.

Since ViT [20] achieved the great performance in the image classification task, the transformer-based architectures have also been studied on the semantic segmentation, one of the most fundamental vision tasks. SETR [78] was the first semantic segmentation model to adopt the transformer architecture as a backbone with convolutional decoder. Beyond introducing the effective encoder structures, recent method [65] proposed the efficient encoder-decoder structures for the semantic segmentation. SegFormer [65] introduced a mix transformer encoder and a purely MLP-based decoder. FeedFormer [50] introduced a cross attention-based decoder to refer the low-level feature information of the transformer encoder. VWFormer [66] used the transformer encoder and exploited the window-based attention for considering the multi-scale representation in the decoder. We introduce the efficient Encoder-Decoder Attention TransFormer model for the semantic segmentation to effectively capture the global context at both the encoder and the decoder.

In the all-attention decoder, the output features $\textbf{F}_{i}$ of each encoder stage $i\in\{2,3,4\}$ are first fed into the EFT blocks in each decoder stage $j\in\{3,2,1\}$, where $j$ denotes the index of the decoder stages. Then, the features $\widehat{\textbf{F}}_{j}\in\mathbb{R}^{H_{j}\times W_{j}\times C_{j}}$ of each decoder stage are up-sampled to $H_{2}\times W_{2}$ resolution using the bilinear interpolation. These up-sampled features $\textbf{U}_{j}\in\mathbb{R}^{H_{2}\times W_{2}\times C_{j}}$ are then concatenated and passed to linear layers for fusion. Finally, the final prediction mask is projected into the number of classes $C_{cls}$ mask by another linear layer. This process is formulated as:

where $\textbf{M}\in\mathbb{R}^{H_{2}\times W_{2}\times C_{cls}}$ is the final prediction mask.

Different from previous SRA, our inference spatial reduction (ISR) method reduces the key-value spatial resolution at the inference phase. Our method achieves the computational efficiency by changing the hyperparameter associated with the ‘reduction ratio $R$’ of the average pooling in the EFA module. Our ISR can be used in the self-attention structures because the self-attention has a special structure where reducing the resolution of key and value does not affect the shape of the input and output features. Due to this structure, the reduction ratio can be adjusted during inference without affecting the resolution of the input and output features.

However, reducing the key and value resolution largely at training has the advantage of computational efficiency, but leads to the performance degradation because the query cannot consider enough information from the key and value. To address this issue, our ISR alleviates the trade-off gap between the computational cost and the accuracy by reducing the resolution of the key and value at inference. In this part, we describe that our ISR is applied to the our EDAFormer, which is the optimized architecture for applying our ISR effectively.

As shown in Fig. 1, our EDAformer uses the proposed transformer blocks in both encoder-decoder structures. Each pooling-based SRA used in each encoder stage and decoder stage has a corresponding reduction ratio setting that reduces the key and value resolution. At training as illustrated in Fig. 3, the reduction ratios $t_{E}^{i}$ of each encoder stage are set to [8, 4, 2, 1], which are the default setting of other previous models [58, 65, 59] using SRA. The reduction ratios $t_{D}^{j}$ of the decoder stage that takes each encoder features are set to [1, 2, 4], which are equal to the reduction ratios of the corresponding encoder stage. $t_{E}$ and $t_{D}$ denote the reduction ratio of the encoder and decoder at training, respectively. The computational complexity of the previous attention is as follows:

where $\mathrm{\Omega}$ and $\mathtt{SRA}$ denote the computational complexity and the spatial reduction attention. $h$, $w$ and $c$ represent the height, the width and the channel of the features, respectively. $r$ is the reduction ratio at training phase.

Under these reduction ratio settings, we train our EDAFormer to get pretrained weights. After that, at inference phase, it is possible to optionally adjust the inference computational reduction by selecting the reduction ratios at the discretion of the user. As shown in Fig. 3, $r_{E}^{i}$ and $r_{D}^{j}$ denote the reduction ratio of the encoder and decoder at inference, respectively. They are formulated as:

where $a_{E}^{i}$ and $a_{D}^{j}$ denote the additional reduction ratio of the encoder and decoder at inference, respectively. After applying our ISR, the computational complexity is as follows:

where $\mathtt{ISR}$ is the inference spatial reduction and $a$ is the additional reduction ratio at inference. Therefore, one of the advantage of our ISR is that it is simple to obtain the computational reduction on the pretrained model without additional training. Our ISR reduces the performance degradation compared with reducing by $r^{2}a^{2}$ at training. Empirically, the optimal setting is [16,8,2,1]-[2,4,8] in the encoder-decoder, which has the best reduction ratio of the performance degradation to the computational cost reduction.

To verify the effectiveness of considering the globality at the decoder, we compared the different operations at the Embedding-Free Attention (EFA) position of the EFT block in Table 3 (a). The applied operations are the local context operation (i.e., DW Conv, Conv) and the global context operation (i.e., w/ embedding attention, w/o embedding attention). Our w/o embedding structure improved 1.6% and 2.4% mIoU compared to the depth-wise convolution and the standard convolution, respectively. These results show that capturing the global context in the decoder is important for the mIoU performance improvement. While w/ embedding method outperformed the local context operation by capturing global context, our EFA further improved mIoU by 0.8% with the lightweight model parameter and FLOPs. This indicates that our EFA module better models the global context.

Our decoder, a {3-2-1} structure, is the hierarchical structure with six EFT blocks that assigns more attention blocks to high-level semantic features. In Table 3 (b), we verified the effectiveness of our decoder structure compared with three cases. The case of {2-2-2} structure assigned two EFT blocks equally to all decoder stages. The cases of {1-2-3}, {1-4-1} and our {3-2-1} allocated more EFT blocks to the decoder stage-3, 2 and 1, respectively. As a result, our {3-2-1} structure assigning more attention to higher level features shows better performance of 0.8%, 1.7%, 1.8% mIoU compared to {2-2-2},{1-2-3}, and {1-4-1}, respectively. These results indicate that allocating the additional attention layers to the higher level features, which contain richer semantic information, is more effective for semantic segmentation performance.

In Table 4, we verified the effectiveness of our Inference Spatial Reduction (ISR) method in the proposed EDAFormer-T and EDAFormer-B, and empirically found the optimal reduction ratio. At training, our EDAFormer was trained with the base setting of [8,4,2,1]-[1,2,4]. At inference, We experimented on applying our ISR to only decoder (i.e.[8,4,2,1]-[2,4,8]), part of the encoder-decoder (i.e.[16,8,2,1]-[2,4,8]), and all of the encoder-decoder (i.e.[16,8,4,2]-[2,4,8]). The setting of [16,8,2,1]-[2,4,8] showed the optimal performance for improving the computational efficiency compared to the accuracy degradation. Compared to EDAFormer-T with the base setting, EDAFormer-T with the optimal setting reduced the computation by 16.1%, 37.4% and 16.1% on ADE20K, Cityscapes and COCO-Stuff, respectively. The performance dropped by only 0.2% mIoU on ADE20K and did not drop on COCO-Stuff and Cityscapes. Furthermore, EDAFormer-B reduced the computation by 8.1% with only 0.1% mIoU degradation on ADE20K and COCO-Stuff, and reduced the computation by 25.3% without performance degradation on Cityscapes. These results indicate that our ISR method is simple, yet significantly reduces the computational cost with little performance degradation. In addition, our method showed the impressive effectiveness by only adjusting the reduction ratio at the inference without fine-tuning. Our ISR is effective without the fine-tuning, but we trained the models with 40K iterations for fine-tuning to further compensate for performance degradation at higher reduction ratio of [16,8,4,2]-[2,4,8]. As a result, EDAFormer-T showed a 0.2% drop in mIoU on ADE20K, and 0.1% drops in mIoU on Cityscapes and COCO-Stuff. EDAFormer-B showed 0.3% and 0.2% drops in mIoU on ADE20K and COCO-Stuff, and no drop in mIoU on Cityscapes.

In Table 5 (a), we compared our w/ ISR with w/o ISR, which used the same reduction ratio of [16,8,2,1]-[2,4,8] at both training and inference. Our EDAFormer with our ISR was trained with the reduction ratio of [8,4,2,1]-[1,2,4] and adjusted the ratio to [16,8,2,1]-[2,4,8] at inference. Despite the same computation at inference phase, the result with our ISR showed better mIoU than the case w/o ISR, with both 0.5% improvements for our EDAFormer-T and EDAFormer-B, respectively. Therefore, our model w/ ISR, which considers enough information of the key and value during training, can achieve better performance than the model that cannot consider enough information by reducing the resolution of the key and value during training.

To verify the effectiveness of our embedding-free structure for ISR. We experiment the ablated model that w/ embedding attention is adopt to our EFA position in all-attention decoder. We also compared with the ablated model (i.e., w/ embedding) by applying our ISR to the decoder stages in Table 5 (b). The w/ embedding structure showed the gradual performance degradation as the reduction ratio increased, and the reduction ratio of [8,4,2,1]-[4,8,12] showed the performance decrease of 0.8% mIoU. However, our structure showed no performance degradation up to the reduction ratio of [8,4,2,1]-[3,6,9], and only a 0.1% drop in mIoU at the reduction ratio of [8,4,2,1]-[4,8,12]. This indicate that our w/o embedding structure is effective with proposed ISR method.

In Table 7 (a), We experimented to compare which method is better in terms of the mIoU and inference speed (FPS) for the key-value spatial reduction. The bipartite matching-based pooling had no mIoU degradation even though it was applied to every encoder-decoder stage. However, the bipartite matching can reduce maximum 50% of tokens, which corresponds to a reduction ratio of $r=1.4$ ($\approx\sqrt{2}$). This is because it divides the tokens into two sets and merges them. In addition, this method has the additional latency caused by the matching algorithm. Therefore, the bipartite matching showed similar FPS compared to w/o ISR even though they reduce the computation of the attention. The max pooling showed a drop of 0.3% mIoU, and the overlapped pooling was slightly slower than the average pooling. Therefore, we adopted the average pooling method to reduce the tokens, which is a simple operation for general purposes and is most effective in terms of performance with inference speed.

In Table 7 (b), we represented the inference speed (FPS) comparisons of various reduction ratios. We measured the inference speed by using a single RTX 3090 GPU without any additional accelerating techniques. Compared to base setting, applying our ISR shows 29.4% and 47.1% FPS improvements in the reduction ratios of [16,8,2,1]-[2,4,8] and [16,8,4,2]-[2,4,8], respectively. The inference speed became faster as the computational cost was reduced by increasing the reduction ratio. These results indicate that the the computational reduction by our ISR leads to the improvement of the actual inference speed.

Our ISR can be universally applied not only to our EDAFormer, but also to other transformer-based models by using the additional spatial reduction at the inference. To verify generalizability of our ISR, we applied ours to various models in Table 7. The transformer-based backbones are trained with our decoder for the semantic segmentation task. For the convolutional self-attention models (i.e., CvT [62], MViT [72] and LVT [67]), our ISR significantly reduced computation by 34.8$\sim$41.6% with 0.1$\sim$0.3% performance degradation. Our method also showed the effective computational reduction with less performance degradation for window-based attention models (i.e., Swin [40] and DaViT [18]), spatial reduction attention-based models (i.e., PVTv2 [59] and MiT [65]) and segmentation models (i.e., SegFormer [65] and FeedFormer [50]). The result for FeedFormer using the cross-attention decoder showed that our method is also effective in the cross-attention mechanism. These results indicate that our ISR framework can be effectively extended to various transformer-based architecture using different attention methods, and our EDAFormer is especially the optimized architecture for applying our ISR effectively.

Fig. 4 visualized the features and prediction maps of the EDAFormer-B decoder stage-2 before and after applying the ISR. Firstly, we visualized the attention score maps representing the similarity score between the query and key. When ISR was applied, the resolution of the attention score map was reduced because the resolution of the key was reduced. Compared to the similarity scores without applying the ISR, the similarity scores between the query and key applying the ISR were well maintained. In other words, the attention regions before and after applying ISR were similar, even though we reduce the key tokens rather than the attention score map. Therefore, this means that applying our ISR can maintain the semantic similarity scores in the global regions.

Secondly, we compared the output features after operating between the attention score map and values. Surprisingly, the output features before and after applying ISR showed almost the same results. Therefore, these results indicate that the information obtained from the self-attention operation is maintained even though the spatial reduction is applied to the key and value in inference. Thirdly, when comparing the prediction maps, the results before and after applying the ISR are almost same. This means that the effect of ISR can be applied not only to the decoder stage-2, but also to the whole EDAFormer network.

In Fig. 5, we visualized our segmentation predictions on ADE20K, Cityscapes and COCO-Stuff, compared with the embedding-based transformer model (i.e. SegFormer [65]). Our EDAFormer better predicted the finer details near object boundaries. Our model also better segmented the large regions (e.g., road, roof and truck) than SegFormer. Furthermore, our model predicted the objects of the same category (e.g., sofa) that were far apart more precisely than SegFormer. This indicates that our embedding-free attention structure can capture enough global spatial information.