Bootstrapping Semantic Segmentation with Regional Contrast

1 INTRODUCTION . Semantic segmentation is an essential part of applications such as scene understanding and autonomous driving , whose goal is to assign a semantic label to each pixel in an image . Significant progress has been achieved by use of large datasets with high quality human annotations . However , labelling images with pixel-level accuracy is time consuming and expensive ; for example , labelling a single image in CityScapes can take more than 90 minutes ( Cordts et al. , 2016 ) . When deploying semantic segmentation models in practical applications where only limited labelled data are available , high quality ground-truth annotation is a significant bottleneck . To reduce the need for labelled data , there is a recent surge of interest in leveraging unlabelled data for semi-supervised learning . Previous methods include improving segmentation models via adversarial learning ( Hung et al. , 2019 ; Mittal et al. , 2019 ) and self-training ( Zou et al. , 2019 ; 2018 ; Zhu et al. , 2020 ) . Others focus on designing advanced data augmentation strategies to generate pseudo image-annotation pairs from unlabelled images ( Olsson et al. , 2021 ; French et al. , 2020 ) . In both semi-supervised and supervised learning , a segmentation model often predicts smooth label maps , because neighbouring pixels are usually of the same class , and rarer high-frequency regions are typically only found in object boundaries . This learning bias produces blurry contours and regularly mis-labels rare objects . After carefully examining the label predictions , we further observe that wrongly labelled pixels are typically confused with very few other classes ; e.g . a pixel labelled as rider has a much higher chance of being wrongly classified as person , compared to train or bus . By understanding this class structure , learning can be actively focused on the challenging pixels to improve overall segmentation quality . Here we propose ReCo , a contrastive learning framework designed at a regional level . Specifically , ReCo is a new loss function which helps semantic segmentation not only to learn from local context ( neighbouring pixels ) , but also from global semantic class relationships across the entire dataset . ReCo performs contrastive learning on a pixel-level dense representation , as visualised in Fig . 1 . For each semantic class in a mini-batch , ReCo samples a set of pixel-level representations ( queries ) , and encourages them to be close to the class mean representation ( positive keys ) , and simultaneously pushes them away from representations sampled from other classes ( negative keys ) . For pixel-level contrastive learning with high-resolution images , it is impractical to sample all pixels . In ReCo , we actively sample a sparse set of queries and keys , consisting of less than 5 % of all available pixels . We sample negative keys from a learned distribution based on the relative distance between the mean representation of each negative key and the query class . This distribution can be interpreted as a pairwise semantic class relationship , dynamically updated during training . We sample queries for those having a low prediction confidence . Active sampling helps ReCo to rapidly focus on the most confusing pixels for each semantic class , and requires minimal additional memory . ReCo enables a high-accuracy segmentation model to be trained with very few human annotations . We evaluate ReCo in a semi-supervised setting , with two different modes : i ) Partial Dataset Full Labels — a sparse subset of training images , where each image has full ground-truth labels , and the remaining images are unlabelled ; ii ) Partial Labels Full Dataset — all images have some labels , but covering only a sparse subset of pixels within each image . In both settings , we show that ReCo can consistently improve performance across all methods and datasets . 2 RELATED WORK . Semantic Segmentation One recent direction is in designing more effective deep convolutional neural networks . Fully convolutional networks ( FCNs ) ( Long et al. , 2015 ) are the foundation of modern segmentation network design . They were later improved with dilated/atrous convolutions with larger receptive fields , capturing more long range information ( Chen et al. , 2017 ; 2018 ) . Alternative approaches include encoder-decoder architectures ( Ronneberger et al. , 2015 ; Kirillov et al. , 2019 ) , sometimes using skip connections ( Ronneberger et al. , 2015 ) to refine filtered details . A parallel direction is to improve optimisation strategies , by designing loss functions that better respect class imbalance ( Lin et al. , 2017 ) or using rendering strategy to refine uncertain pixels from high-frequency regions improving the label quality ( Kirillov et al. , 2020 ) . ReCo is built upon this line of research , to improve segmentation by providing additional supervision on hard pixels . Semi-supervised Classification and Segmentation The goal of semi-supervised learning is to improve model performance by taking advantage of a large amount of unlabelled data during training . Here consistency regularisation and entropy minimisation are two common strategies . The intuition is that the network ’ s output should be invariant to data perturbation and geometric transformation . Based on these strategies , many semi-supervised methods have been developed for image classification ( Sohn et al. , 2020 ; Tarvainen & Valpola , 2017 ; Berthelot et al. , 2019 ; Kuo et al. , 2020 ) . However , for segmentation , generating effective pseudo-labels and well-designed data augmentation are non-trivial . Some solutions improved the quality of pseudo-labelling , using adversarial learning ( Hung et al. , 2019 ; Mittal et al. , 2019 ) or enforcing consistency from different augmented images ( French et al. , 2020 ; Olsson et al. , 2021 ) . In this work , we show that we can improve the performance of current semi-supervised segmentation methods by jointly training with a suitable auxiliary task . Contrastive Learning Contrastive learning learns a similarity function to bring views of the same data closer in representation space , whilst pushing views of different data apart . Most recent contrastive frameworks learn similarity scores based on global representations of the views , parameterising data with a single vector ( He et al. , 2020 ; Chen et al. , 2020 ; Khosla et al. , 2020 ) . Dense representations , on the other hand , rely on pixel-level representations and naturally provide additional supervision , capturing fine-grained pixel correspondence . Contrastive pre-training based on dense representations has recently been explored , and shows better performance in dense prediction tasks , such as object detection and keypoint detection ( Wang et al. , 2021b ; O. Pinheiro et al. , 2020 ) . Contrastive Learning for Semantic Segmentation Contrastive learning has been recently studied to improve semantic segmentation , with a number of different design strategies . Zhang et al . ( 2021 ) and Zhao et al . ( 2021 ) both perform contrastive learning via pre-training , based on the generated auxiliary labels and ground-truth labels respectively , but at the cost of huge memory consump- tion . In contrast , ours performs contrastive learning whilst requiring much less memory , via active sampling . In concurrent work , ( Wang et al. , 2021a ; Alonso et al. , 2021 ) also perform contrastive learning with active sampling . However , whilst both these methods are applied to a stored feature bank , ours focuses on sampling features on-the-fly . Active sampling in Alonso et al . ( 2021 ) is further based on learnable , class-specific attention modules , whilst ours only samples features based on relation graphs and prediction confidence , without introducing any additional computation overhead , which results in a simpler and much more memory-efficient implementation . 3 RECO – REGIONAL CONTRAST . 3.1 PIXEL-LEVEL CONTRASTIVE LEARNING . Let ( X , Y ) be a training dataset with training images x ∈ X and their corresponding C-class pixellevel segmentation labels y ∈ Y , where y can be either provided in the original dataset , or generated automatically as pseudo-labels . A segmentation network f is then optimised to learn a mapping fθ : X 7→ Y , parameterised by network parameters θ . This segmentation network f can be decomposed into two parts : an encoder network : φ : X 7→ Z , and a decoder classification head ψc : Z 7→ Y . To perform pixel-level contrastive learning , we additionally attach a decoder representation head ψr on top of the encoder network φ , parallel to the classification head , mapping the encoded feature into a higher m-dimensional dense representation with the same spatial resolution as the input image : ψr : Z 7→ R , R ∈ Rm . This representation head is only applied during training to guide the classifier using the ReCo loss as an auxiliary task , and is removed during inference . A pixel-level contrastive loss is a function which encourages queries rq to be similar to the positive key r+k , and dissimilar to the negative keys r − k . All queries and keys are sampled from the decoder representation head : rq , r + , − k ∈ R. In ReCo , we use a pixel-level contrastive loss across all available semantic classes in each mini-batch , with the distance between keys and queries measured by their normalised dot product . The general formation of the ReCo loss Lreco is then defined as : Lreco = ∑ c∈C ∑ rq∼Rcq − log exp ( rq · rc , +k /τ ) exp ( rq · rc , +k /τ ) + ∑ r−k ∼R c k exp ( rq · r−k /τ ) , ( 1 ) for which C is a set containing all available classes in the current mini-batch , τ is the temperature control of the softness of the distribution , Rcq represents a query set containing all representations whose labels belong to class c , Rck represents a negative key set containing all representations whose labels do not belong to class c , and rc , +k represents the positive key which is the mean representation of class c. Suppose P is a set containing all pixel coordinates with the same resolution as R , these queries and keys are then defined as : Rcq = ⋃ [ u , v ] ∈P 1 ( y [ u , v ] = c ) r [ u , v ] , Rck = ⋃ [ u , v ] ∈P 1 ( y [ u , v ] 6= c ) r [ u , v ] , rc , +k = 1 |Rcq| ∑ rq∈Rcq rq . ( 2 ) 3.2 ACTIVE HARD SAMPLING ON QUERIES AND KEYS . Contrastive learning on all pixels in high-resolution images would be computationally expensive . Here , we introduce active hard sampling strategies to optimise only a sparse set of queries and keys . Active Key Sampling When classifying a pixel , a semantic network might be uncertain only over a very small number of candidates , among all available classes . The uncertainty from these candidates typically comes from a close spatial ( e.g . rider and bicycle ) or semantic ( e.g . horse and cow ) relationship . To reduce this uncertainty , we propose to sample negative keys non-uniformly , based on the relative distance between each negative key class and the query class . This involves building a pair-wise class relationship graph G , with G ∈ R|C|×|C| , computed and dynamically updated for each mini-batch . This pair-wise relationship is measured by the normalised dot product between the mean representation from a pair of two classes and is defined as : G [ p , q ] = ( rp , +k · r q , + k ) , ∀p , q ∈ C , and p 6= q . ( 3 ) We further apply SoftMax to normalise these pair-wise relationships among all negative classes j for each query class c , which produces a probabilistic distribution : exp ( G [ c , i ] ) / ∑ j∈C , j 6=c exp ( G [ c , j ] ) . We sample negative keys for each class i based on this distribution , to learn the corresponding query class c. This procedure allocates more samples to hard , confusing classes chosen specifically for each query class , helping the segmentation network to learn a more accurate decision boundary . Active Query Sampling Due to the natural class imbalance in semantic segmentation , it is easy to over-fit on common classes , such as the road and building classes in the CityScapes dataset , or the background class in the Pascal VOC dataset . These common classes contribute to the majority of pixel space in training images , and so randomly sampling queries will under-sample rare classes and provide minimal supervision to these classes . Therefore , we instead sample hard queries — for those whose corresponding pixel prediction confidence is below a defined threshold . Accordingly , ReCo ’ s loss would then guide the segmentation network by providing appropriate supervision on these less certain pixels . The easy and hard queries are defined as follows , and visualised in Fig . 2 , Rc , easyq = ⋃ rq∈Rcq 1 ( ŷq > δs ) rq , Rc , hardq = ⋃ rq∈Rcq 1 ( ŷq ≤ δs ) rq , ( 4 ) where ŷq is the predicted confidence of label c after the SoftMax operation corresponding to the same pixel location as rq , and δs is the user-defined confidence threshold .

Summary: This work presents a new semi-supervised learning framework based on a pixel-level contrastive learning with active query/key sampling. For a given query feature with a class c, the positive key sample is generated as a mean vector of all samples with the class c, while the negative key sample is sampled by considering the pair-wise relationship between classes in (3) such that more samples are from hard cases. The active query is adaptively sampled by considering rare queries in order to overcome the class imbalance in semantic segmentation as in (4). The proposed contrastive loss, call ReCo, is applied to semi-supervised semantic segmentation task with two cases: full labels and partial labels. Extensive results demonstrate the outstanding performance when using ReCo loss within existing semi-supervised learning methods.