SegFormer: Simple and Efficient Design for Semantic Segmentation with Transformers

1 Introduction Semantic segmentation is a fundamental task in computer vision and enables many downstream applications . It is related to image classification since it produces per-pixel category prediction instead of image-level prediction . This relationship is pointed out and systematically studied in a seminal work [ 1 ] , where the authors used fully convolutional networks ( FCNs ) for semantic segmentation tasks . Since then , FCN has inspired many follow-up works and has become a predominant design choice for dense prediction . Since there is a strong relation between classification and semantic segmentation , many stateof-the-art semantic segmentation frameworks are variants of popular architectures for image classification on ImageNet . Therefore , designing backbone architectures has remained an active area ∗Corresponding authors : Zhiding Yu and Ping Luo 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . in semantic segmentation . Indeed , starting from early methods using VGGs [ 1 , 2 ] , to the latest methods with significantly deeper and more powerful backbones [ 3 ] , the evolution of backbones has dramatically pushed the performance boundary of semantic segmentation . Besides backbone architectures , another line of work formulates semantic segmentation as a structured prediction problem , and focuses on designing modules and operators , which can effectively capture contextual information . A representative example in this area is dilated convolution [ 4 , 5 ] , which increases the receptive field by “ inflating ” the kernel with holes . Witnessing the great success in natural language processing ( NLP ) , there has been a recent surge of interest to introduce Transformers to vision tasks . Dosovitskiy et al . [ 6 ] proposed vision Transformer ( ViT ) for image classification . Following the Transformer design in NLP , the authors split an image into multiple linearly embedded patches and feed them into a standard Transformer with positional embeddings ( PE ) , leading to an impressive performance on ImageNet . In semantic segmentation , Zheng et al . [ 7 ] proposed SETR to demonstrate the feasibility of using Transformers in this task . SETR adopts ViT as a backbone and incorporates several CNN decoders to enlarge feature resolution . Despite the good performance , ViT has two important limitations : 1 ) ViT outputs single-scale lowresolution features instead of multi-scale ones , and 2 ) it has very high computational cost on large images . To address these limitations , Wang et al . [ 8 ] proposed a pyramid vision Transformer ( PVT ) , a natural extension of ViT with pyramid structures for dense prediction . PVT shows considerable improvements over the ResNet counterpart on object detection and semantic segmentation . However , together with other emerging methods such as Swin Transformer [ 9 ] and Twins [ 10 ] , these methods mainly consider the design of the Transformer encoder , neglecting the contribution of the decoder for further improvements . This paper introduces SegFormer , a cutting-edge Transformer framework for semantic segmentation that jointly considers efficiency , accuracy , and robustness . In contrast to previous methods , our framework redesigns both the encoder and the decoder . The key novelties of our approach are : • A novel positional-encoding-free and hierarchical Transformer encoder . • A lightweight All-MLP decoder design that yields a powerful representation without complex and computationally demanding modules . • As shown in Figure 1 , SegFormer sets new a state-of-the-art in terms of efficiency , accuracy and robustness in three publicly available semantic segmentation datasets . First , the proposed encoder avoids interpolating positional codes when performing inference on images with resolutions different from the training one . As a result , our encoder can easily adapt to arbitrary test resolutions without impacting the performance . In addition , the hierarchical part enables the encoder to generate both high-resolution fine features and low-resolution coarse features , this is in contrast to ViT that can only produce single low-resolution feature maps with fixed resolutions . Second , we propose a lightweight MLP decoder where the key idea is to take advantage of the Transformer-induced features where the attentions of lower layers tend to stay local , whereas the ones of the highest layers are highly non-local . By aggregating the information from different layers , the MLP decoder combines both local and global attention . As a result , we obtain a simple and straightforward decoder that renders powerful representations . We demonstrate the advantages of SegFormer in terms of model size , run-time , and accuracy on three publicly available datasets : ADE20K , Cityscapes , and COCO-Stuff . On Citysapces , our lightweight model , SegFormer-B0 , without accelerated implementations such as TensorRT , yields 71.9 % mIoU at 48 FPS , which , compared to ICNet [ 11 ] , represents a relative improvement of 60 % and 4.2 % in latency and performance , respectively . Our largest model , SegFormer-B5 , yields 84.0 % mIoU , which represents a relative 1.8 % mIoU improvement while being 5 × faster than SETR [ 7 ] . On ADE20K , this model sets a new state-of-the-art of 51.8 % mIoU while being 4 × smaller than SETR . Moreover , our approach is significantly more robust to common corruptions and perturbations than existing methods , therefore being suitable for safety-critical applications . Code will be publicly available . 2 Related Work . Semantic Segmentation . Semantic segmentation can be seen as an extension of image classification from image level to pixel level . In the deep learning era [ 12–14 ] , FCN [ 1 ] is the fundamental work of semantic segmentation , which is a fully convolution network that performs pixel-to-pixel classification in an end-to-end manner . After that , researchers focused on improving FCN from different aspects such as : enlarging the receptive field [ 15–17 , 5 , 2 , 4 , 18 ] ; refining the contextual information [ 19– 27 ] ; introducing boundary information [ 28–35 ] ; designing various attention modules [ 36–44 ] ; or using AutoML technologies [ 45–49 ] . These methods significantly improve semantic segmentation performance at the expense of introducing many empirical modules , making the resulting framework computationally demanding and complicated . More recent methods have proved the effectiveness of Transformer-based architectures for semantic segmentation [ 7 , 44 ] . However , these methods are still computationally demanding . Transformer backbones . ViT [ 6 ] is the first work to prove that a pure Transformer can achieve state-of-the-art performance in image classification . ViT treats each image as a sequence of tokens and then feeds them to multiple Transformer layers to make the classification . Subsequently , DeiT [ 50 ] further explores a data-efficient training strategy and a distillation approach for ViT . More recent methods such as T2T ViT [ 51 ] , CPVT [ 52 ] , TNT [ 53 ] , CrossViT [ 54 ] and LocalViT [ 55 ] introduce tailored changes to ViT to further improve image classification performance . Beyond classification , PVT [ 8 ] is the first work to introduce a pyramid structure in Transformer , demonstrating the potential of a pure Transformer backbone compared to CNN counterparts in dense prediction tasks . After that , methods such as Swin [ 9 ] , CvT [ 56 ] , CoaT [ 57 ] , LeViT [ 58 ] and Twins [ 10 ] enhance the local continuity of features and remove fixed size position embedding to improve the performance of Transformers in dense prediction tasks . Transformers for specific tasks . DETR [ 50 ] is the first to use Transformers for end-to-end object detection framework without non-maximum suppression ( NMS ) . Other works have also used Transformers in tasks such as tracking [ 59 , 60 ] , super-resolution [ 61 ] , re-id [ 62 ] , colorization [ 63 ] , retrieval [ 64 ] and multi-modal learning [ 65 , 66 ] . For semantic segmentation , SETR [ 7 ] adopts ViT [ 6 ] as a backbone to extract features , achieving impressive performance . However , these Transformer-based methods have very low efficiency and , thus , difficult to deploy in real-time applications . 3 Method . As depicted in Figure 2 , SegFormer consists of two main modules : ( 1 ) a hierarchical Transformer encoder ; and ( 2 ) a lightweight All-MLP decoder to predict the final mask . Given an image with size H ×W × 3 , we first divide it into patches of size 4 × 4 . Unlike ViT which uses 16 × 16 , using fine-grained patches favors semantic segmentation . Second , we use these patches as input to the hierarchical Transformer encoder to get multi-level features with resolution { 1/4 , 1/8 , 1/16 , 1/32 } of the original image . We then pass these multi-level features to the All-MLP decoder to predict the segmentation mask with a H4 × W 4 ×Ncls resolution , where Ncls is the number of categories . In the remainder of this section , we first detail the proposed encoder and decoder designs and then summarize the main differences of our approach compared to SETR . 3.1 Hierarchical Transformer Encoder . We design a series of Mix Transformer encoders ( MiT ) , MiT-B0 to MiT-B5 , with the same architecture but different sizes . On top of the hierarchical architecture and efficient self-attention module in PVT [ 8 ] , we further propose several novel features including overlapped patch merging and positional-encoding-free design which will be shown to greatly benefit the segmentation tasks . Hierarchical Feature Representation . Unlike ViT [ 6 ] , our encoder generates multi-level multi-scale features given an input image . These features provide both high-resolution coarse features and lowresolution fine-grained features that boost the performance of semantic segmentation . Specifically , given an input image with size H×W ×3 , we perform patch merging to obtain a hierarchical feature map Fi with a resolution of H2i+1 × W 2i+1 × Ci , where i ∈ { 1 , 2 , 3 , 4 } , and Ci+1 is larger than Ci . Efficient Self-Attention . A major bottleneck of the above hierarchical feature representation is the quadratic self-attention complexity with long sequence inputs from higher resolution features . Recall that in the original multi-head self-attention , each of the heads Q , K , V have the same dimensions N × C , where N = H ×W is the length of the sequence , the self-attention is estimated as : Attention ( Q , K , V ) = Softmax ( QKT√ dhead ) V. ( 1 ) We instead adopt the sequence reduction process introduced in [ 8 ] . This process uses a reduction ratio R to reduce the length of the sequence of as follows : K̂ = Reshape ( N R , C ·R ) ( K ) K = Linear ( C ·R , C ) ( K̂ ) , ( 2 ) where K is the sequence to be reduced , Reshape ( NR , C ·R ) ( K ) refers to reshape K to the one with shape of NR × ( C · R ) , and Linear ( Cin , Cout ) ( · ) refers to a linear layer taking a Cin-dimensional tensor as input and generating a Cout-dimensional tensor as output . Therefore , the new K has dimensions NR × C. As a result , the complexity of the self-attention mechanism is reduced from O ( N2 ) to O ( N 2 R ) . In our experiments , we set R to [ 64 , 16 , 4 , 1 ] from stage-1 to stage-4 . Overlapped Patch Merging . Given an image patch , the patch merging process used in ViT , unifies a N × N × 3 patch into a 1 × 1 × C vector . This can easily be extended to unify a 2 × 2 × Ci feature path into a 1× 1×Ci+1 vector to obtain hierarchical feature maps . Using this , we can shrink our hierarchical features from F1 ( H4 × W 4 × C1 ) to F2 ( H 8 × W 8 × C2 ) , and then iterate for any other feature map in the hierarchy . This process was initially designed to combine non-overlapping image or feature patches . Therefore , it fails to preserve the local continuity around those patches . Instead , we use an overlapping patch merging process . To this end , we define K , S , and P , where K is the patch size , S is the stride between two adjacent patches , and P is the padding size . In our experiments , we set K = 7 , S = 4 , P = 3 , and K = 3 , S = 2 , P = 1 to perform overlapping patch merging to produces features with the same size as the non-overlapping process . Similar to the original patch embedding in ViT [ 6 ] , this operation can be implemented by “ nn.Conv2D ” in PyTorch . Positional-Encoding-Free Design . The resolution of the PE in ViT is fixed . One thus needs to interpolate the PE when the test resolution differs from training . This leads to the drop of accuracy , which is undesirable since the resolution mismatch is common in semantic segmentation . We instead introduce Mix-FFN where we consider the effect of zero padding to the leak location information [ 67 ] by directly using a 3 × 3 Conv in the feed-forward network ( FFN ) . Mix-FFN is formulated as : xout = MLP ( GELU ( Conv3×3 ( MLP ( xin ) ) ) ) + xin , ( 3 ) where xin is the feature from the self-attention module . Mix-FFN mixes a 3 × 3 convolution and an MLP into each FFN . In our experiments , we will show that a 3 × 3 convolution is sufficient to provide positional information for Transformers . In particular , we use depth-wise convolutions for reducing the number of parameters and improving efficiency . It should be mentioned that CPVT [ 52 ] also alleviates this issue by using a 3 × 3 Conv to generate conditional PE at different resolutions and then add it to the feature map . Our work conceptually goes one step further as we argue that adding PE to feature map is not necessary in semantic segmentation . Another recent work CvT [ 56 ] introduced 3× 3 Convs to model the spatial relationship among tokens . Despite the converging design , our work differs in both motivation and application as we aim to totally remove PEs to handle the training/testing resolution mismatch issue in semantic segmentation . Our intuition started from [ 67 ] whereas the same intuition was not discussed in CvT .

Summary: This work proposed an efficient Transformer-based semantic segmentation model called SegFormer. In essence, it consists of a hierarchical Transformer encoder and a lightweight MLP decoder. Compared to ViT, the transformer encoder employs overlapped patch merging, self-attention with reduced (key, value) for efficient computation, and a depthwise convolution inserted between two MLP layers in FFN to replace position embedding. The authors also conducted an analysis based on effective receptive field and revealed that transformer's encoder with large receptive field enables relative lightweight decoder design. Extensive experiments demonstrate that SegFormer achieves strong speed-accuracy trade-off and better robustness.