[Paper Review] SegFormer: Simple and Efficient Design for Semantic Segmentation with Transformers

Enze Xie, Wenhai Wang+4 more|arXiv (Cornell University)|May 31, 2021

Advanced Image and Video Retrieval Techniques3,230 citations

TL;DR

SegFormer presents a hierarchical, positional-encoding-free Transformer encoder plus a lightweight All-MLP decoder, achieving strong accuracy with high efficiency across ADE20K, Cityscapes, and COCO-Stuff.

ABSTRACT

We present SegFormer, a simple, efficient yet powerful semantic segmentation framework which unifies Transformers with lightweight multilayer perception (MLP) decoders. SegFormer has two appealing features: 1) SegFormer comprises a novel hierarchically structured Transformer encoder which outputs multiscale features. It does not need positional encoding, thereby avoiding the interpolation of positional codes which leads to decreased performance when the testing resolution differs from training. 2) SegFormer avoids complex decoders. The proposed MLP decoder aggregates information from different layers, and thus combining both local attention and global attention to render powerful representations. We show that this simple and lightweight design is the key to efficient segmentation on Transformers. We scale our approach up to obtain a series of models from SegFormer-B0 to SegFormer-B5, reaching significantly better performance and efficiency than previous counterparts. For example, SegFormer-B4 achieves 50.3% mIoU on ADE20K with 64M parameters, being 5x smaller and 2.2% better than the previous best method. Our best model, SegFormer-B5, achieves 84.0% mIoU on Cityscapes validation set and shows excellent zero-shot robustness on Cityscapes-C. Code will be released at: github.com/NVlabs/SegFormer.

Motivation & Objective

• Motivate a semantic segmentation framework that blends Transformer encoders with lightweight decoders for efficiency and robustness.
• Develop a hierarchical, positional-encoding-free Transformer encoder that outputs multi-scale features suitable for dense prediction.
• Design a compact All-MLP decoder to fuse multi-level features without heavy backbones or complex modules.
• Demonstrate state-of-the-art performance and robustness on ADE20K, Cityscapes, and COCO-Stuff with scalable MiT encoders.

Proposed method

• Introduce MiT (Mix Transformer) encoders with hierarchical feature maps at 1/4, 1/8, 1/16, and 1/32 resolutions.
• Employ overlapping patch merging to build multi-scale features without losing spatial locality.
• Adopt efficient self-attention with sequence reduction to lower complexity from O(N^2) to O(N^2/R).
• Replace standard ViT-style FFN with Mix-FFN, combining a 3x3 depthwise convolution with an MLP to inject location information without fixed positional embeddings.
• Use a lightweight All-MLP decoder that upsamples and fuses multi-level features via simple Linear layers and MLPs to produce the segmentation mask.

Experimental results

Research questions

• RQ1 Can a hierarchical, positional-encoding-free Transformer encoder produce high-resolution, multi-scale features suitable for semantic segmentation?
• RQ2 Does a lightweight All-MLP decoder suffice to fuse multi-level Transformer features for accurate per-pixel predictions?
• RQ3 How do SegFormer variants scale in accuracy, parameters, FLOPs, and speed across standard segmentation benchmarks?
• RQ4 Is the proposed Mix-FFN and overlapping patch merging robust to test-resolution changes and different datasets?

Key findings

• SegFormer-B0 achieves strong real-time performance with 3.8M parameters and 8.4G FLOPs on ADE20K, outperforming real-time counterparts on multiple metrics.
• SegFormer-B5 reaches 84.0% mIoU on Cityscapes validation with significantly smaller and faster operation than prior best methods.
• On ADE20K, SegFormer-B4 achieves 50.3% mIoU with 64M parameters, surpassing the previous best while being ~5x smaller.
• SegFormer sets a new state-of-the-art on ADE20K (51.8% mIoU) with a much smaller model than SETR, and achieves 83.8-84.0% on Cityscapes with efficient networks.
• SegFormer demonstrates strong robustness to natural corruptions (Cityscapes-C), outperforming prior methods by large margins in several corruption categories.
• Across COCO-Stuff, SegFormer-B5 reaches 46.7% mIoU with 84.7M parameters, outpacing comparable methods like SETR by ~0.9% while being substantially smaller.