Per-Pixel Classification is Not All You Need for Semantic Segmentation

1 Introduction . The goal of semantic segmentation is to partition an image into regions with different semantic categories . Starting from Fully Convolutional Networks ( FCNs ) work of Long et al . [ 28 ] , most deep learning-based semantic segmentation approaches formulate semantic segmentation as per-pixel classification ( Figure 1 left ) , applying a classification loss to each output pixel [ 8 , 46 ] . Per-pixel predictions in this formulation naturally partition an image into regions of different classes . Mask classification is an alternative paradigm that disentangles the image partitioning and classification aspects of segmentation . Instead of classifying each pixel , mask classification-based methods predict a set of binary masks , each associated with a single class prediction ( Figure 1 right ) . The more flexible mask classification dominates the field of instance-level segmentation . Both Mask R-CNN [ 19 ] and DETR [ 3 ] yield a single class prediction per segment for instance and panoptic segmentation . In contrast , per-pixel classification assumes a static number of outputs and can not return a variable number of predicted regions/segments , which is required for instance-level tasks . Our key observation : mask classification is sufficiently general to solve both semantic- and instancelevel segmentation tasks . In fact , before FCN [ 28 ] , the best performing semantic segmentation methods like O2P [ 4 ] and SDS [ 18 ] used a mask classification formulation . Given this perspective , a natural question emerges : can a single mask classification model simplify the landscape of effective approaches to semantic- and instance-level segmentation tasks ? And can such a mask classification model outperform existing per-pixel classification methods for semantic segmentation ? To address both questions we propose a simple MaskFormer approach that seamlessly converts any existing per-pixel classification model into a mask classification . Using the set prediction mechanism proposed in DETR [ 3 ] , MaskFormer employs a Transformer decoder [ 37 ] to compute a set of pairs , ∗Work partly done during an internship at Facebook AI Research . 1Project page : https : //bowenc0221.github.io/maskformer 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . each consisting of a class prediction and a mask embedding vector . The mask embedding vector is used to get the binary mask prediction via a dot product with the per-pixel embedding obtained from an underlying fully-convolutional network . The new model solves both semantic- and instance-level segmentation tasks in a unified manner : no changes to the model , losses , and training procedure are required . Specifically , for semantic and panoptic segmentation tasks alike , MaskFormer is supervised with the same per-pixel binary mask loss and a single classification loss per mask . Finally , we design a simple inference strategy to blend MaskFormer outputs into a task-dependent prediction format . We evaluate MaskFormer on five semantic segmentation datasets with various numbers of categories : Cityscapes [ 13 ] ( 19 classes ) , Mapillary Vistas [ 31 ] ( 65 classes ) , ADE20K [ 49 ] ( 150 classes ) , COCOStuff-10K [ 2 ] ( 171 classes ) , and ADE20K-Full [ 49 ] ( 847 classes ) . While MaskFormer performs on par with per-pixel classification models for Cityscapes , which has a few diverse classes , the new model demonstrates superior performance for datasets with larger vocabulary . We hypothesize that a single class prediction per mask models fine-grained recognition better than per-pixel class predictions . MaskFormer achieves the new state-of-the-art on ADE20K ( 55.6 mIoU ) with Swin-Transformer [ 27 ] backbone , outperforming a per-pixel classification model [ 27 ] with the same backbone by 2.1 mIoU , while being more efficient ( 10 % reduction in parameters and 40 % reduction in FLOPs ) . Finally , we study MaskFormer ’ s ability to solve instance-level tasks using two panoptic segmentation datasets : COCO [ 26 , 22 ] and ADE20K [ 49 ] . MaskFormer outperforms a more complex DETR model [ 3 ] with the same backbone and the same post-processing . Moreover , MaskFormer achieves the new state-of-the-art on COCO ( 52.7 PQ ) , outperforming prior state-of-the-art [ 38 ] by 1.6 PQ . Our experiments highlight MaskFormer ’ s ability to unify instance- and semantic-level segmentation . 2 Related Works . Both per-pixel classification and mask classification have been extensively studied for semantic segmentation . In early work , Konishi and Yuille [ 23 ] apply per-pixel Bayesian classifiers based on local image statistics . Then , inspired by early works on non-semantic groupings [ 11 , 33 ] , mask classification-based methods became popular demonstrating the best performance in PASCAL VOC challenges [ 16 ] . Methods like O2P [ 4 ] and CFM [ 14 ] have achieved state-of-the-art results by classifying mask proposals [ 5 , 36 , 1 ] . In 2015 , FCN [ 28 ] extended the idea of per-pixel classification to deep nets , significantly outperforming all prior methods on mIoU ( a per-pixel evaluation metric which particularly suits the per-pixel classification formulation of segmentation ) . Per-pixel classification became the dominant way for deep-net-based semantic segmentation since the seminal work of Fully Convolutional Networks ( FCNs ) [ 28 ] . Modern semantic segmentation models focus on aggregating long-range context in the final feature map : ASPP [ 6 , 7 ] uses atrous convolutions with different atrous rates ; PPM [ 46 ] uses pooling operators with different kernel sizes ; DANet [ 17 ] , OCNet [ 45 ] , and CCNet [ 21 ] use different variants of non-local blocks [ 39 ] . Recently , SETR [ 47 ] and Segmenter [ 34 ] replace traditional convolutional backbones with Vision Transformers ( ViT ) [ 15 ] that capture long-range context starting from the very first layer . However , these concurrent Transformer-based [ 37 ] semantic segmentation approaches still use a per-pixel classification formulation . Note , that our MaskFormer module can convert any per-pixel classification model to the mask classification setting , allowing seamless adoption of advances in per-pixel classification . Mask classification is commonly used for instance-level segmentation tasks [ 18 , 22 ] . These tasks require a dynamic number of predictions , making application of per-pixel classification challenging as it assumes a static number of outputs . Omnipresent Mask R-CNN [ 19 ] uses a global classifier to classify mask proposals for instance segmentation . DETR [ 3 ] further incorporates a Transformer [ 37 ] design to handle thing and stuff segmentation simultaneously for panoptic segmentation [ 22 ] . However , these mask classification methods require predictions of bounding boxes , which may limit their usage in semantic segmentation . The recently proposed Max-DeepLab [ 38 ] removes the dependence on box predictions for panoptic segmentation with conditional convolutions [ 35 , 40 ] . However , in addition to the main mask classification losses it requires multiple auxiliary losses ( i.e. , instance discrimination loss , mask-ID cross entropy loss , and the standard per-pixel classification loss ) . 3 From Per-Pixel to Mask Classification . In this section , we first describe how semantic segmentation can be formulated as either a per-pixel classification or a mask classification problem . Then , we introduce our instantiation of the mask classification model with the help of a Transformer decoder [ 37 ] . Finally , we describe simple inference strategies to transform mask classification outputs into task-dependent prediction formats . 3.1 Per-pixel classification formulation . For per-pixel classification , a segmentation model aims to predict the probability distribution over all possibleK categories for every pixel of anH×W image : y = { pi|pi ∈ ∆K } H·Wi=1 . Here ∆K is theKdimensional probability simplex . Training a per-pixel classification model is straight-forward : given ground truth category labels ygt = { ygti |y gt i ∈ { 1 , . . . , K } } H·Wi=1 for every pixel , a per-pixel cross- entropy ( negative log-likelihood ) loss is usually applied , i.e. , Lpixel-cls ( y , ygt ) = ∑H·W i=1 − log pi ( y gt i ) . 3.2 Mask classification formulation . Mask classification splits the segmentation task into 1 ) partitioning/grouping the image intoN regions ( N does not need to equal K ) , represented with binary masks { mi|mi ∈ [ 0 , 1 ] H×W } Ni=1 ; and 2 ) associating each region as a whole with some distribution over K categories . To jointly group and classify a segment , i.e. , to perform mask classification , we define the desired output z as a set of N probability-mask pairs , i.e. , z = { ( pi , mi ) } Ni=1 . In contrast to per-pixel class probability prediction , for mask classification the probability distribution pi ∈ ∆K+1 contains an auxiliary “ no object ” label ( ∅ ) in addition to the K category labels . The ∅ label is predicted for masks that do not correspond to any of the K categories . Note , mask classification allows multiple mask predictions with the same associated class , making it applicable to both semantic- and instance-level segmentation tasks . To train a mask classification model , a matching σ between the set of predictions z and the set of N gt ground truth segments zgt = { ( cgti , m gt i ) |c gt i ∈ { 1 , . . . , K } , m gt i ∈ { 0 , 1 } H×W } N gt i=1 is required . 2 Here cgti is the ground truth class of the i th ground truth segment . Since the size of prediction set |z| = N and ground truth set |zgt| = N gt generally differ , we assume N ≥ N gt and pad the set of ground truth labels with “ no object ” tokens ∅ to allow one-to-one matching . For semantic segmentation , a trivial fixed matching is possible if the number of predictionsN matches the number of category labels K. In this case , the ith prediction is matched to a ground truth region with class label i and to ∅ if a region with class label i is not present in the ground truth . In our experiments , we found that a bipartite matching-based assignment demonstrates better results than the fixed matching . Unlike DETR [ 3 ] that uses bounding boxes to compute the assignment costs between prediction zi and ground truth z gt j for the matching problem , we directly use class and mask predictions , i.e. , −pi ( cgtj ) + Lmask ( mi , m gt j ) , where Lmask is a binary mask loss . To train model parameters , given a matching , the main mask classification loss Lmask-cls is composed of a cross-entropy classification loss and a binary mask loss Lmask for each predicted segment : Lmask-cls ( z , zgt ) = ∑N j=1 [ − log pσ ( j ) ( cgtj ) + 1cgtj 6=∅Lmask ( mσ ( j ) , m gt j ) ] . ( 1 ) 2Different mask classification methods utilize various matching rules . For instance , Mask R-CNN [ 19 ] uses a heuristic procedure based on anchor boxes and DETR [ 3 ] optimizes a bipartite matching between z and zgt . Note , that most existing mask classification models use auxiliary losses ( e.g. , a bounding box loss [ 19 , 3 ] or an instance discrimination loss [ 38 ] ) in addition to Lmask-cls . In the next section we present a simple mask classification model that allows end-to-end training with Lmask-cls alone . 3.3 MaskFormer . We now introduce MaskFormer , the new mask classification model , which computes N probabilitymask pairs z = { ( pi , mi ) } Ni=1 . The model contains three modules ( see Fig . 2 ) : 1 ) a pixel-level module that extracts per-pixel embeddings used to generate binary mask predictions ; 2 ) a transformer module , where a stack of Transformer decoder layers [ 37 ] computes N per-segment embeddings ; and 3 ) a segmentation module , which generates predictions { ( pi , mi ) } Ni=1 from these embeddings . During inference , discussed in Sec . 3.4 , pi and mi are assembled into the final prediction . Pixel-level module takes an image of size H ×W as input . A backbone generates a ( typically ) low-resolution image feature map F ∈ RCF×HS ×WS , where CF is the number of channels and S is the stride of the feature map ( CF depends on the specific backbone and we use S = 32 in this work ) . Then , a pixel decoder gradually upsamples the features to generate per-pixel embeddings Epixel ∈ RCE×H×W , where CE is the embedding dimension . Note , that any per-pixel classificationbased segmentation model fits the pixel-level module design including recent Transformer-based models [ 34 , 47 , 27 ] . MaskFormer seamlessly converts such a model to mask classification . Transformer module uses the standard Transformer decoder [ 37 ] to compute from image features F and N learnable positional embeddings ( i.e. , queries ) its output , i.e. , N per-segment embeddings Q ∈ RCQ×N of dimension CQ that encode global information about each segment MaskFormer predicts . Similarly to [ 3 ] , the decoder yields all predictions in parallel . Segmentation module applies a linear classifier , followed by a softmax activation , on top of the per-segment embeddings Q to yield class probability predictions { pi ∈ ∆K+1 } Ni=1 for each segment . Note , that the classifier predicts an additional “ no object ” category ( ∅ ) in case the embedding does not correspond to any region . For mask prediction , a Multi-Layer Perceptron ( MLP ) with 2 hidden layers converts the per-segment embeddingsQ to N mask embeddings Emask ∈ RCE×N of dimension CE . Finally , we obtain each binary mask prediction mi ∈ [ 0 , 1 ] H×W via a dot product between the ith mask embedding and per-pixel embeddings Epixel computed by the pixel-level module . The dot product is followed by a sigmoid activation , i.e. , mi [ h , w ] = sigmoid ( Emask [ : , i ] T · Epixel [ : , h , w ] ) . Note , we empirically find it is beneficial to not enforce mask predictions to be mutually exclusive to each other by using a softmax activation . During training , the Lmask-cls loss combines a cross entropy classification loss and a binary mask loss Lmask for each predicted segment . For simplicity we use the same Lmask as DETR [ 3 ] , i.e. , a linear combination of a focal loss [ 25 ] and a dice loss [ 30 ] multiplied by hyper-parameters λfocal and λdice respectively .

Summary: This paper proposes a transformer-based model, called MaskFormer for semantic segmentation task. MaskFormer treats semantic segmentation as a mask classification problem, and predicts a set of binary masks with the corresponding class labels.  The proposed approach can be easily extended for panoptic segmentaion task. Experimental evaluations for semantic segmentation and panoptic segmentation tasks demonstrate the effectiveness of proposed method.