Learning Debiased and Disentangled Representations for Semantic Segmentation

1 Introduction . Semantic segmentation is a pixelwise classification task that is applicable to a wide range of practical problems including autonomous driving , medical image diagnosis , scene understanding , and many others . Ever since deep neural networks have been adopted for semantic segmentation [ 1–5 ] , its accuracy has increased gradually with the introduction of stronger network architectures , improved training schemes , and large-scale datasets . Despite such phenomenal achievement , semantic segmentation approaches still suffer from the chronic limitations caused by class imbalance and stereotyped scene context in datasets . The class imbalance issues are prevalent in semantic segmentation , where certain classes occupy larger area and/or appear more frequently than others . For example , the “ road ” class in the Cityscapes dataset [ 6 ] accounts for over 36 % of pixels of all training images while the “ motorcycle ” class does only 0.1 % of pixels . Furthermore , the scene context is not sufficiently diversified in training datasets . For instance , the objects in the “ motorcycle ” class is always observed in the vicinity of the “ road ” class , and the spatial layout of the two classes are almost identical in all images . Consequently , the models often learn substantially correlated representations and fail to accurately delineate the boundaries in the absence of frequently co-located objects . These issues cause the model to learn suboptimal feature representations , especially for classes that are under-represented . These feature representations are often entangled with one another , causing confusion between classes , and are also susceptible to encode adverse biases as well , which may be detrimental to the model ’ s overall performance . On the other hand , for classes with abundant labels , the model is able to leverage the large variations in the data and learn more robust feature 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . representations that are less dependent on the presence of other classes . As such , the model exhibits strong performance on classes with sufficient labels but weaker performances on rare classes . We propose DropClass , a simple yet effective training scheme for semantic segmentation that aims to alleviate the vulnerability of dataset bias and feature entanglement . Given an input example , we first extract a class-specific feature map for each class using Grad-CAM [ 7 ] . Then , we drop the feature map corresponding to a randomly sampled class and aggregate the remaining feature maps to generate the predictions . Finally , we employ our proposed loss functions to facilitate training . This procedure is somewhat similar to Dropout [ 8 ] , but is specifically designed with a different goal in mind ; while Dropout aims to reduce co-adaptations among features , DropClass aims to disentangle the class representations and reduce the underlying biases that arise from inter-class relationships . Our training scheme requires no additional data or annotations . More importantly , since DropClass is model agnostic , it can be plugged into any existing network architecture straightforwardly . Our contributions can be summarized as follows : • Our work takes the first step to tackling dataset biases in the semantic segmentation task . We describe how dataset biases in semantic segmentation are different from those in classification , and devise a benchmark to better analyze the biases found in segmentation models . • We propose DropClass , a model-agnostic training scheme that results in learning more debiased and disentangled feature representations for semantic segmentation . • Our experimental results show that the model trained with DropClass exhibits stronger performance compared to the baseline across multiple network architectures and datasets . Furthermore , we observe significant performance gains especially on under-represented classes , and provide analysis that further validate the effectiveness of our training scheme . 2 Proposed Method . We first define the semantic segmentation task and its relevant notations . For a dataset with a set of classes , C , we sample an input image x ∈ Rh×w×3 and its corresponding label y ∈ Rh×w×|C| . The feature extractor g ( · ) takes x as input and produces an intermediate feature representation A ∈ Rh1×w1×k , where k denotes an arbitrary channel size of the feature map that depends on the choice of network architecture . The feature map A is then fed into the classifier , h ( · ) , to produce ŷ ∈ Rh×w×|C| , which serves as the final output of the model before the softmax operation , i.e. , the logits . Altogether , we denote the entire model as a composition of the feature extractor and classifier : ŷ = h ( A ) = h ( g ( x ) ) . Note that throughout the paper , we use subscripts to refer to specific spatial and channel locations on tensors . 2.1 Motivation . Dataset bias issues are introduced in a plethora of literature dealing with classification problems [ 9– 13 ] , where the primary concerns are class imbalance in the dataset and spurious correlation between attributes in images . Although the dense prediction tasks including semantic segmentation may suffer from additional challenges related to dataset bias , their potential limitations have still been hardly discussed . Hence , this work first points out a critical drawback exposed to semantic segmentation algorithms caused by the stereotyped co-occurrence of multiple classes as well as the issues inherent from the classification tasks . We argue that , besides the class imbalance and attribute correlation issues , inter-class relationships with spatial co-occurrence or feature similarity could be the culprit of dataset biases in semantic segmentation models . Consider the class relationship between the “ car ” and “ road ” classes in road scenes segmentation . The two classes have a strong spatial correlation since the vast majority of “ car ” samples will be spatially located above the “ road ” class . Models can easily exploit this spatial correlation in the process of maximizing the pixel-level accuracy . However , this could have adverse effects when the model encounters an image where “ cars ” do not appear above the “ road ” class . Moreover , the representations of a class are often entangled with one another since the features of co-occurring classes in an image are mixed in the convolution layers , and this tendency is aggravated with the semantic similarity between classes , e.g. , “ bicycle↔ motorcycle ” and “ person↔ rider ” . In both cases , under-represented , rare classes typically become the victims of the feature entanglement . In this sense , dataset bias and feature entanglement are not mutually exclusive , but dataset bias is manifested as feature entanglement in segmentation tasks . The fundamental goal of our work is to develop a training scheme that allows the model to learn debiased and disentangled representations for each class , such that the model can make robust predictions for the under-represented classes without depending on the feature representations of other classes . In essence , this is achieved by removing the class-specific features of a random class at each iteration and providing a proper loss as described in Section 2.3 . Our training process not only breaks up inter-class relationships through a stochastic training scheme , but also exposes the classifier to a diverse set of class-specific feature combinations given the same input . 2.2 DropClass for Debiasing and Disentangling . The feature representations of multiple classes should not be entangled with one another in order to facilitate less confusion and stronger discriminability . To achieve this , we remove the feature dependencies among classes by selectively discarding the information related to a randomly chosen class at each training iteration . The difficulty is that neural networks are black-box in nature and produce highly entangled feature representations . To alleviate this issue , we leverage Grad-CAM [ 7 ] , a well-known gradient-based method of highlighting important class-specific information from entangled feature representations . Grad-CAM uses the backpropagated gradients to produce a coarse localization map that highlights important regions in the image with respect to a certain label , and is widely used to visualize classdiscriminative features in deep neural networks . In similar fashion , we employ the gradients of feature Ai , j , k with respect to the class logit ŷc as the degree of importance that Ai , j , k has on the prediction for class c. Since semantic segmentation is a pixel-level classification task , we average the gradients ofAi , j , k with respect to every pixel location that corresponds to class c. This is used as the importance score of Ai , j , k for class c , and is denoted by Sci , j , k : Sci , j , k = 1 h× w h∑ u=1 w∑ v=1 ∂ŷu , v , c ∂Ai , j , k . ( 1 ) The importance score for a certain class can be calculated for all locations ( i , j , k ) of the feature map A to obtain an importance score map Sc , with the same dimensions . Then , we perform element-wise multiplication of the feature and importance score maps to obtain class-specific feature representations , Ac , which is given by Ac = ReLU ( A Sc ) , ( 2 ) where denotes the Hadamard product . The ReLU activation is used to preserve all features of A that have a positive influence on ŷc and filter out all negative ones . With the highlighted feature representation Ac , ∀c ∈ C , we are now able to discard information related to a certain class from the output of the feature extractor . In every iteration of the training process , a class z is randomly sampled as z ∼ U ( 1 , |C| ) , ( 3 ) where U ( 1 , |C| ) denotes a discrete uniform distribution across elements of C. Through the stochastic drop of the class z , we deactivate Az and obtain a different combination of feature representations at each iteration , which is given by A′drop = 1 |C| |C|∑ c=1 1 [ c 6= z ] ·Ac , ( 4 ) where 1 [ · ] represents an indicator function . After passing A′drop through a convolutional layer to compensate missing information , we obtain a new feature representation Adrop = conv ( A′drop ) that contains information for all classes except class z . This feature is passed through the classifier to generate a new logit , ŷdrop . Since Adrop may still contain the information about the dropped class z , our approach attempts to suppress the information regarding class z in the other class-specific feature maps as well , using the objective function as defined in Section 2.3 . Overall , our DropClass procedure is illustrated in Figure 1 . 2.3 Objective Function . We train the model with multiple loss functions , each with a specific goal in mind . We first evaluate the cross-entropy loss : LCE = − 1 h× w h∑ u=1 w∑ v=1 |C|∑ c=1 yu , v , c · log ( p̂u , v , c ) , ( 5 ) where p̂ = softmax ( ŷ ) applied over the channel dimension . To account for the dropped class z , we design a modified cross-entropy loss that masks the pixel-wise loss at all locations , where class z is the ground-truth : LCE_drop = − 1 h× w h∑ u=1 w∑ v=1 |C|∑ c=1 yu , v , c · log ( p̃u , v , c ) · 1 [ yu , v , z 6= 1 ] , ( 6 ) where p̃ = softmax ( ŷdrop ) applied over the channel dimension . Then , LCE for the model ’ s original output and LCE_drop for the DropClass output are combined by a weighted sum : Lseg = ( 1− λ ) · LCE + λ · LCE_drop , ( 7 ) where λ is a scaling constant that balances the two loss terms . The scaling constant λ plays an important role in ensuring the stability and effectiveness of training . The process outlined in Section 2.2 is based on the assumption that the gradient of Ai , j , k with respect to ŷc accurately quantifies the importance of Ai , j , k on the prediction of class c. In the initial stages of training , however , ŷ is likely to be largely inaccurate , and thus , the gradient may not accurately reflect the relationship between the feature and class prediction . Hence , relying on this gradient to perform DropClass from the early stages of training may lead to an unstable training process . To alleviate this issue , we train the model primarily with LCE at the early stages of training and linearly decrease its weight as training progresses . Conversely , we gradually increase the weight for LCE_drop , thereby facilitating learning of debiased and disentangled representations towards the end of training . This is achieved by linearly increasing the value of λ from 0 to 1 across the duration of training . To further stabilize training , we linearly increase the probability of dropping out any class from 0 to 1 as well . Finally , we add an additional objective term Lsup that helps suppress the output probability of the dropped class z : Lsup = 1 h× w h∑ u=1 w∑ v=1 p̃u , v , z . ( 8 ) Even after removing information about class z , some information could still persist in the dropped feature Adrop , which would allow the model to predict a score for class z , albeit with low confidence . Hence , by suppressing the probability for class z directly , the model is encouraged to learn Algorithm 1 Training scheme for DropClass . Require : Dataset D = { ( x ( i ) , y ( i ) } , set of classes C , feature extractor g ( · ) and classifier h ( · ) . 1 : for i = 1 , 2 , . . . do 2 : Compute A = g ( x ( i ) ) , ŷ = h ( A ) . 3 : for c = 1 , 2 , . . . do 4 : Compute Sci , j , k = 1 h×w ∑h u=1 ∑w v=1 ∂ŷu , v , c ∂Ai , j , k . . Eq . ( 1 ) 5 : Compute Ac = ReLU ( A Sc ) . . Eq . ( 2 ) 6 : end for 7 : Sample a class z to drop and deactivate Az . . Eq . ( 3 ) 8 : Gather all Ac for ∀c ∈ C \ { z } and feed them into the classifier h ( · ) . . Eq . ( 4 ) 9 : Compute loss function and update parameters. . Eq . ( 9 ) 10 : end for class-specific representations , i.e. , purge any redundant information related to class z from feature representations of all other classes c ∈ C \ { z } . Note that this loss is only calculated when a class is dropped . Altogether , the final objective function proposed for DropClass is : Ltotal = Lseg + αLsup , ( 9 ) where α is a hyper-parameter that weighs the relative importance of Lsup . A summary of our proposed training scheme can be found in Algorithm 1 .

Summary: Recent semantic segmentation methods often learn biased representation from the dataset with highly co-located objects (e.g., motorcycle and road) as well as the class imbalance. This paper proposes a training scheme of dropping the class-specific representations in the highly-entangled features, mitigating the bias induced by the dropped class. The proposed DropClass is also model-agnostic and demonstrates the improved segmentation results in Cityscapes and Pascal VOC datasets.