MUSTER: A Multi-scale Transformer-based Decoder for Semantic Segmentation

TL;DR

This work tackles the overlooked role of decoders in transformer-based semantic segmentation by introducing MUSTER, a multi-scale transformer-based decoder that symmetrically couples with hierarchical encoders via Multi-head Skip Attention (MSKA) and a novel Fuse&Upsample mechanism. A lighter variant, Light-MUSTER, further reduces computation by redesigning attention and masking while preserving performance. On ADE20K and Cityscapes, MUSTER achieves competitive mIoU with significantly reduced FLOPs (up to ~61.3% fewer operations) and demonstrates strong robustness in weakly supervised settings. The results underscore the decoder’s central role in dense pixel-level prediction and offer a flexible, efficient framework for plug-and-play integration with a variety of backbone architectures.

Abstract

In recent works on semantic segmentation, there has been a significant focus on designing and integrating transformer-based encoders. However, less attention has been given to transformer-based decoders. We emphasize that the decoder stage is equally vital as the encoder in achieving superior segmentation performance. It disentangles and refines high-level cues, enabling precise object boundary delineation at the pixel level. In this paper, we introduce a novel transformer-based decoder called MUSTER, which seamlessly integrates with hierarchical encoders and consistently delivers high-quality segmentation results, regardless of the encoder architecture. Furthermore, we present a variant of MUSTER that reduces FLOPS while maintaining performance. MUSTER incorporates carefully designed multi-head skip attention (MSKA) units and introduces innovative upsampling operations. The MSKA units enable the fusion of multi-scale features from the encoder and decoder, facilitating comprehensive information integration. The upsampling operation leverages encoder features to enhance object localization and surpasses traditional upsampling methods, improving mIoU (mean Intersection over Union) by 0.4% to 3.2%. On the challenging ADE20K dataset, our best model achieves a single-scale mIoU of 50.23 and a multi-scale mIoU of 51.88, which is on-par with the current state-of-the-art model. Remarkably, we achieve this while significantly reducing the number of FLOPs by 61.3%. Our source code and models are publicly available at: https://github.com/shiwt03/MUSTER.

Figures (14)

• Figure 1: mIoU vs FLOPs on ADE20K dataset. Our MUSTER achieves 8.9% higher mIoU than original Swin Transformer with same FLOPs, and cuts 53.5% FLOPs while remaining almost the same mIoU on ADE20K.
• Figure 2: Feature maps of 16 heads of the second stage in MUSTER, where red and orange areas are the focus of a single head.
• Figure 3: The left part is a pyramid backbone, which is not limited to the Swin Transformer. The right part is our proposed MUSTER, which is a multi-scale transformer-based decoder.
• Figure 4: Multi-head Skip Attention Architecture. The shape of $Q, K, V$ should be $M^2 \times \frac{C}{N_{heads}}$, where $N_{head}$ means count head in Multi-head Skip Attention.
• Figure 5: Fuse&Upsample