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 discusses per-pixel classification and mask classification for semantic segmentation. It shows that mask classification is sufficiently general to solve semantic and instance segmentation in a unified manner. Thorough experiments are conducted on several benchmarks for both semantic and panoptic segmentation.