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: Contrastive learning requires positives and negatives, and their definitions have an impact on the transfer of contrastive representations to supervised tasks. This work contributes a new definition of positives for the purpose of semantic segmentation, where the positives come not only *within* an image but *across* images. The intra-image (within) and inter-image (across) positives are derived by clustering spatially-local features across a dataset. These initial clustering features can be drawn from image-wise contrastive learning, transfer from another supervised dataset, or an oracle that is trained on some or all of the target dataset. The clustering is a straightforward but effective application of k-means; the insight is in choosing the number of desired clusters. In this work an over-clustering is desired, so that clusters correspond to "auxiliary" labels that are not necessarily the same as the supervised task labels, and more fine-grained for contrasting. This difference in labeling separates auxiliary labels from pseudo-labels, which are predictions of the supervised task labels, and underlines that auxiliary labeling can be totally unsupervised unlike pseudo-labeling. A simple metric comparing auxiliary labels to ground truth segmentations is proposed ("proxy mIoU") as a surrogate of full training as a means to cross-validate auxiliary labeling design choices.  Experiments compare the proposed method with existing contrastive learning methods (DenseContrast, MoCo, PixPro) and one recent semi-supervised semantic segmentation method (PseudoSeg, ICLR'21) on the standard datasets of Cityscapes and ADE20K plus the older, smaller dataset of NYUDv2. The downstream semantic segmentation results for all methods are measured with the DeepLabV3 architecture with a ResNet-50 backbone for standardized fair comparison. The proposed method improves on the state-of-the-art, as well as the common standard of supervised ImageNet pre-training, on all three datasets vs. both the contrastive and semi-supervised methods. The margin of improvement is largest with the least supervision, but even the fully-supervised case improves by 1-2 points. Ablations examine different choices of representation for auxiliary labeling and a variety of design choices and hyperparameters for contrastive learning.