Looking Beyond Single Images for Contrastive Semantic Segmentation Learning

1 Introduction . Training semantic segmentation models typically requires a large amount of densely annotated images . Producing such dense annotations for an adequate number of images at sufficient quality is an extremely laborious and costly task [ 17 ] . Self-supervised and weakly-supervised learning techniques offer the possibility to significantly reduce the manual labeling effort that is required to train high-performance machine learning models . Contrastive learning [ 14 , 22 ] has recently emerged as a promising technique for self-supervised representation learning that can leverage large amounts of unlabeled data for augmenting and extending existing labeled datasets without costly manual annotation . However , recent works on contrastive learning primarily focused on image/instance-level learning . Few works have studied contrastive learning in dense image labeling tasks , such as semantic segmentation [ 50 , 51 ] , as dense pixel-level segmentation poses unique challenges . At its core , self-supervised contrastive learning is rooted in the availability of training samples that can be grouped automatically into positive pairs of similar concepts , and negative pairs of dissimilar concepts . The notion of similarity thereby depends strongly on the task at hand . The appropriate selection of such pairs is of paramount importance for the performance of the resulting models . While it is often straightforward to generate negative pairs , automatically generating positive pairs can present a formidable challenge that requires intimate knowledge about the data and task at hand . Image augmentations are both simple and effective to generate positive pairs for image/instancelevel learning [ 14 , 15 ] but are too coarse-grained to learn representations at a pixel level . ∗Equal advising . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . Recent works [ 51,56 ] adapt the idea of using image augmentations to generate positive pairs for more fine-grained pixel-level representation learning . While surprisingly effective , image augmentations have only a limited capability of covering the full visual spectrum of how two semantically related pixels or patches can look . Mid and high-level information such as context , texture , large visual shifts inside a semantic class ( for example , the class `` car '' may comprise a wide range of different shapes and colors ) , and even projective distortions due to the imaging process , may be extremely relevant but can not be simulated realistically with augmentations alone . As a consequence , methods that exclusively rely on augmentations are limited in the contrastive information they can provide to downstream tasks . A different line of work thus exploits ground truth semantic labels to establish correspondences between images to generate more diverse and informative positive pairings [ 50 , 52 ] . However , the need for labeled data impedes scaling these approaches to the massive amounts of data required to fully reap the benefits of contrastive learning [ 14 , 22 ] . In this work , we explore the generation of auxiliary labels to establish fine-grained , pixel-level correspondences across images to select contrastive pairs . We distinguish auxiliary labels from the more commonly employed pseudo-labels [ 66 ] , as auxiliary labels need not represent the same classes as the ground truth labels . This concept is illustrated in Figure 1 . To this end , we leverage feature embeddings that we generate with existing feature extractors . The feature extractors have either been trained using coarse image-level supervision , on a small set of pixel-level labeled data , or using image-level , self-supervised contrastive learning approaches [ 16 , 22 ] . We cluster feature embeddings and assign auxiliary labels to individual pixels based on cluster membership . We also introduce a simple , yet effective , voting strategy to spatially regularize the auxiliary label maps . The auxiliary labels tend to cluster into semantically and perceptually meaningful groups which can be used to establish cross-image correspondences . This allows us to generate a diverse and realistic set of positive pairs that can be used for fine-grained contrastive visual representation learning . An example that illustrates our approach is shown in Figure 2 . We additionally introduce a proxy metric that allows us to quickly assess the quality of any auxiliary labeling without the need to train a model . Finally , we present a comprehensive experimental evaluation of various factors and hyper-parameters that influence the performance of contrastive training in semantic segmentation . Our experiments across three datasets and two network architectures show that contrastive learning with auxiliary labels consistently boosts semantic segmentation performance when compared to standard ImageNet pre-training and outperforms existing approaches for contrastive and semisupervised semantic segmentation . 2 Related Work . Approaches to self-supervised learning via a contrastive learning objective have recently demonstrated impressive performance compared to fully supervised methods in a wide variety of tasks [ 14–16 , 22 , 24 ] . In self-supervised learning , the contrastive objective is used as a pretext task to learn general visual representations , which can be well transferred to downstream tasks via fine-tuning . The intuition of the contrastive objective is straightforward : similar visual concepts should be mapped to nearby representations in latent space , whereas dissimilar concepts should be mapped further apart . A crucial ingredient is the definition of similarity and hence the selection of similar and dissimilar samples during training . Seminal work adopted instance discrimination [ 14 , 22 ] . That is , positive samples are generated from the same image via different transformations ( views ) . All other images available for training represent potential negative samples . Recent work investigated different strategies to generate more informative positive and negative pairs , including mutual view selection [ 46 , 47 ] , hard mining [ 30 , 42 , 50 ] , adversarial perturbations [ 25 , 26 , 29 , 32 ] , and using large mini-batches or a momentum memory bank [ 22 , 37 , 53 ] . Another line of work exploited ground truth labels to find positive and negative pairs for contrastive learning [ 6 , 31 , 49 , 50 , 52 , 61 ] . While this approach may improve performance in fully supervised settings , it can not be applied to unlabeled supplementary data . Contrastive learning for segmentation . Recent works adapted contrastive learning to dense prediction tasks such as semantic segmentation [ 10 , 51 , 56 , 65 ] . DenseCL applied the contrastive objective to the pixel level [ 51 ] . ReSim learned spatial features from sliding windows [ 54 ] . Xie et al . generated views by smoothing over nearby feature representations to aggregate local context [ 56 ] . Chen et al . exploited foreground/backward masks [ 11 ] while Van Gansbeke et al . leveraged object masks from pre-trained saliency detectors to group samples [ 48 ] . All of the above-mentioned approaches generate positive views from the same image . This limits the diversity of positive samples severely , as instances of the same category in different images can not be paired . In contrast to existing works , our auxiliary-labeling strategy allows us to establish positive correspondences across images and consequently to pair more diverse examples of the same concept . Pseudo-labels in semi-supervised semantic segmentation . Pseudo-labels are a surrogate ground truth for unlabeled data . Many existing works [ 2 , 12 , 35 , 63 , 63 , 65 ] train a teacher network on labeled data in order to produce surrogate annotations on the unlabeled data available . A student network is then trained on both the ground truth and the pseudo-labels . The surrogate annotations are usually far from perfect , forcing the student to learn from noisy ( and sometimes just misleading ) data , which may negatively affect its performance [ 3 , 45 ] . An underlying problem in this setup is the need for compatibility between pseudo labels and ground truth labels . This forces a ( possibly wrong ) decision on the teacher network , while potentially more fine-grained knowledge present in the feature activations used to make the decision is completely ignored [ 39 ] . The auxiliary labels we construct from feature extractors exploit exactly this type of knowledge . By clustering feature representations , we obtain a more fine-grained set of labels that does not need to match the ground truth label set . Besides providing a more comprehensive perspective on the unlabeled data , auxiliary labels also allow us to leverage a wider range of feature extractors than just a ( teacher ) method built for the same task . Unsupervised clustering for representations learning . Several works proposed unsupervised clustering objectives for deep representation learning , including approaches to image or instance-level representation learning [ 5 , 21 , 27 , 28 , 55 , 57 , 58 , 58 , 64 ] as well as image segmentation [ 28 , 33 , 38 ] . A line work clusters features and uses the cluster assignments as pseudo-labels for self-training [ 4 , 7–9 ] . While we generate auxiliary labels also by clustering features , we leverage powerful existing feature extractors instead of a self-trained network . Furthermore , instead of learning image-level visual representations , we address the challenges of dense prediction tasks , such as semantic segmentation . 3 Method . The key to contrastive learning is the definition of positive and negative pairs , i.e. , which samples should share a representation , and which samples should be mapped to different representations . For annotated datasets , the decision on how to construct positive and negative pairs is clear : samples with matching labels define the positives while samples with different labels define the negatives [ 31,50,52 ] . However , the need for annotations severely limits the data available for learning . Missing any information about correspondences in the dataset , unsupervised approaches take a conservative stance and implicitly assume all randomly paired samples to mismatch . Matching samples are synthesized by sampling a single image and transforming it in different ways . The resulting views represent positive correspondences for training [ 14–16 , 22 , 24 ] . An obvious problem with this approach is that correct pairs ( e.g. , the same car in different images ) will be labeled as negative . At the same time , invariances are only learned w.r.t . the simple transformations applied . To mitigate this issue , we generate auxiliary labels that allow establishing positive pairs across the whole training set akin to supervised methods . However , our auxiliary labels are generated using existing feature extractors , which alleviates the need for dense annotations and unlocks training on large-scale unlabeled data . Figure 3 provides an overview of our method . Auxiliary label generation . Our approach for generating auxiliary labels consists of three steps : From a set of unlabeled images , we obtain visual features with an existing feature extractor . We cluster the features across the whole dataset and assign individual pixels in the input images to their corresponding cluster center . In a final step , we refine the resulting label maps through a super-pixel voting algorithm . A schematic overview of our approach is shown in Figure 4 . In principle , any feature extraction method can be used for generating auxiliary labels and we study several choices . First , we employ a ResNet-50 [ 23 ] , trained through self-supervised contrastive learning [ 22 ] . It represents a feature extractor that has seen no manual annotations . Second , we use the same ResNet-50 architecture but trained for image classification on the ImageNet dataset to study the influence of coarse , image-level labels for training the feature extractor . Finally , we explore several instances of a DeepLabV3 [ 13 ] semantic segmentation network that has been trained for semantic segmentation on various datasets . When using ResNet , we extract features before the final adaptive average pooling layer . This yields a representation at 1/32 of the input resolution . We extract multi-scale features by rescaling the input image to various resolutions and passing it to the feature extractor . We concatenate the feature maps from multiple scales , upsample the representation to 1/8 of the input image size through bilinear interpolation , and normalize the resulting feature vectors using L2 normalization . For DeepLabV3 feature extractors , we extract features before the penultimate layer ( counting the final linear layer as the last layer ) and follow the same multi-scale feature extraction procedure as before . We apply our feature extraction approach to all images in an unlabeled dataset and generate a set of auxiliary classes using K-means clustering [ 40 ] . Since the total number of features in an unlabeled dataset can be extremely large ( on the order of tens or hundreds of millions of feature vectors ) , we perform the clustering step on features from 5 000 images that we sample uniformly at random . We annotate the remaining images in the dataset by assigning pixels to the closest auxiliary class in terms of the L2 distance between the pixel feature vector and the cluster center . To remove outliers and false auxiliary labels , we drop the 5 % pixels with the largest distance to the nearest cluster center . An example result of our labeling strategy is shown in Figure 2 . The resolution and spatial consistency of auxiliary labels depends on the feature extractor employed . Some may produce only noisy or low-resolution feature maps . We thus refine the auxiliary label maps with a simple superpixel voting strategy . We compute a superpixel over-segmentation [ 1 ] and assign the most common label within a superpixel to all pixels of the superpixel . This aligns label map boundaries to image edges , removes outliers , and alleviates errors from low-resolution features . Contrastive training . We use the auxiliary labels to generate positive and negative pairs both within an image as well as across images , and leverage them for pre-training semantic segmentation models . We train two standard segmentation architectures in our experiments [ 13 , 60 ] . The final classification layer of both models is replaced with a linear layer that outputs a 128-dimensional feature map at 1/8 of the input resolution . For each training iteration , we randomly sample 10 000 feature vectors as anchor vectors v ∈ R128 and a different set of 40 000 feature vectors as contrastive feature candidates ( v+ ∈ R128 or v− ∈ R128 ) . All anchors are selected from the current mini-batch , whereas we select 10 000 contrastive features from the current mini-batch and the remaining 30 000 contrastive features from a memory bank [ 22 ] . For every anchor v , we construct a set of associated positive pairs P ( v ) = { ( v , v+ ) } from all contrastive candidates where the auxiliary label of a candidate matches the anchor . We similarly create the set of negative pairs N ( v ) = { ( v , v− ) } from all contrast feature vectors where the auxiliary labels differ . Our training loss is inspired by the supervised contrastive loss of Khosla et al . [ 31 ] . We define the contrastive distance for a single pair as L ( v , v+ ) = − log exp ( v · v +/τ ) exp ( v · v+/τ ) + ∑ ( v , v− ) ∈N ( v ) exp ( v · v−/τ ) , ( 1 ) where we set the temperature τ = 0.07 in our experiments . To compute the loss over all pairs , we introduce a confidence weight that is based on the expected reliability of a positive pair . To this end , we partition the set of positive pairs into three subsets : the set P0 consists of pairs for which both the anchor and the contrast sample belong to the same superpixel . We assign particularly high confidence to these pairs due to the inherent high spatial regularity of superpixels and feature extractors . We assign medium weight to the set P1 that consists of pairs sampled from the same image . Finally , the set P2 consists of samples where the anchor and contrastive features are from different images . Compared with samples in P0 and P1 , these cross-image samples are more likely to be wrongly labeled as positive pairs . We thus assign the lowest weight to these samples . Our final loss is given by L = ∑2 i=0 λi ∑ ( v , v+ ) ∈Pi L ( v , v + ) ∑2 i=0 λi|Pi| , ( 2 ) where λ0,1,2 = { 10 , 4 , 1 } for all our experiments . Auxiliary label quality . Our label generation strategy supports various technical choices that affect the efficacy of the auxiliary label maps . For example , different feature extractors or a different number of auxiliary classes in the clustering step will produce auxiliary labelings that will lead to different performance for a given downstream task . We introduce a novel metric that can be quickly computed and can serve as a guiding metric to choose a specific auxiliary-labeling strategy when a representative set of labeled images is available . Let S = { Sj } Kj=1 be the set of auxiliary segments , where Sj is defined as the set of pixels with the auxiliary label j ( computed over the complete validation set ) . Similarly , let G = { Gi } Mi=1 be the set of true segments , where Gi is the set of pixels that correspond to the ground truth class i . We define the Proxy mean Intersection-over-Union ( PmIoU ) as PmIoU ( S , G ) = 1 max { K , M } ∑ Sj max Gi { |Sj ∩Gi| |Sj ∪Gi| } . ( 3 ) The metric assigns each auxiliary segment Sj to the ground truth segment Gi that yields the highest Intersection-over-Union ( IoU ) among all possible assignments , as we do not know the correct correspondence between auxiliary labels and ground truth labels ( additionally , in most cases a one-toone correspondence will not exist ) . Intuitively , PmIoU reflects the highest mIoU that is achievable by optimally matching the auxiliary labels with the ground truth classes . In the ideal case , where the auxiliary labels exactly coincide with the ground truth labels , the PmIoU equals one . On the other hand , if the auxiliary labels represent a strong over-segmentation or are crossing spatial boundaries between different ground truth segments , the PmIoU will be small . Figure 5 shows PmIoU together with the mIoU after fine-tuning a model that was pre-trained with the corresponding auxiliary labels . We observe that PmIoU correlates well with mIoU across auxiliary labeling strategies . This indicates that PmIoU is an excellent and , most importantly , computationally cheap metric to evaluate different auxiliary-labeling strategies . Practical considerations . In practice , we find that a large amount of positive pairs improves convergence speed , while a sufficiently large amount of negative pairs are required to train models that perform well on the downstream task . We thus implement a memory-efficient version of our loss that doesn ’ t explicitly store the inner products between all pairs . This allows us to scale the total number of pairs per mini-batch substantially . With 10 000 anchor feature vectors and 40 000 contrast feature vectors , we can , depending on the composition of the mini-batch and auxiliary labels , see on average 10 million positive pairs in a single mini-batch . We find that with this very large sample size , 20 epochs suffice to train strong representations for segmentation models . For comparison , MoCo [ 22 ] uses only 256 positive pairs per mini-batch and consequently requires more than 200 training epochs . Supervised semantic segmentation models commonly employ an auxiliary loss on mid-level features to improve performance [ 13 , 60 ] . We find that an auxiliary contrastive loss on mid-level features can likewise improve the representational power of our contrastively trained models . State-of-the-art semantic segmentation models output features at a lower resolution than the input image ( usually 1/8 of the input size ) . We train with the contrastive loss both at the low output resolution that is defined by the model , as well as at the original resolution of the input image ( by upsampling the feature maps ) . We give equal weight to both losses during training . In addition to large mini-batches , momentum contrast [ 22 ] can provide candidate pairs with high diversity and is commonly employed for effective contrastive learning . We also adopt this technique .

Summary: In this paper, the authors propose a a method to learn pixel-level representations based on contrastive learning. The method rely on pre-trained feature extractor (trained unsupervised, on image-level and pixel-level supervision) to generate ‘auxiliary labels’ . The auxiliary labels, computed through K-means, are used to find positive pairs at a pixel-level. The authors show that the proposed approach achieve improved results in semantic segmentation accuracy when compared to ImageNet pretraining and other contrastive approaches.