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: This paper proposes a DropClass training scheme where class-specific feature maps are used to obtain debased and disentangled representations. This is achieved by randomly dropping out class-specific representation from the feature maps, implemented in an annealing manner, reinforced with a regularization term to further remove class-specific information from the remaining information. Experiments are performed on three segmentation networks on two distinct data sets. The presented method seems to achieves an improvement of performance in most classes, with the most noticeable margin of improvement in minority classes.