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 paper introduces the SegFormer model, a Transformer-based model for semantic segmentation (i.e. dense pixel classification) in images. The model uses prior architectural innovations such as a pyramidal structure (progressive downsampling) and a factorized version of self-attention for computational efficiency. The main novelties over prior work on Transformer-based semantic segmentation models and image classification models are: 1) the use of convolutional filters to propagate position information from the image boundaries (as opposed to using explicit positional encodings) and 2) skip-connections between individual layers directly to the output layer, using spatial upsampling. Experimental evaluation demonstrates improved mIoU scores on default semantic image classification benchmarks compared to a baseline SETR (Segmentation Transformer) model, while achieving computational benefits thanks to the adoption of a more efficient Transformer architecture.