PaperlayerView exact revision on arXiv

Source-linked AI summary

Contrastive Learning for Label-Efficient Semantic Segmentation

Xiangyun Zhao, Raviteja Vemulapalli, Philip Mansfield, Boqing Gong, Bradley Green, Lior Shapira, Ying Wu

arXiv:2012.06985v4cs.CVcs.AIcs.LG

TL;DR

Semantic segmentation requires costly dense annotations, and cross-entropy-trained CNNs overfit when labeled data are limited. The paper introduces supervised pixel-wise contrastive pretraining followed by cross-entropy fine-tuning, improving class feature structure and segmentation performance without additional data. Across Cityscapes and PASCAL VOC 2012, it reports large gains and often matches or outperforms ImageNet pretraining.

  • Problem

    Limited pixel-level annotations make semantic segmentation difficult because cross-entropy-trained CNNs can overfit small labeled datasets.

  • Method

    Pretrain the feature extractor with pixel-wise label-based contrastive loss, then fine-tune the entire network with cross-entropy loss.

  • Results

    23.6 mIoU points improve performance on PASCAL VOC 2012 with 2118 labeled images, from 39.1 under cross-entropy to 62.7 with the proposed strategy.

  • Takeaways & Limitations

    In many limited-label settings, contrastive pretraining matches or outperforms ImageNet pretraining without using additional data.

  • Takeaways & Limitations

    The authors plan to explore combining the contrastive loss with consistency-based losses for semi-supervised learning.

Abstract

from arXiv · show

Collecting labeled data for the task of semantic segmentation is expensive and time-consuming, as it requires dense pixel-level annotations. While recent Convolutional Neural Network (CNN) based semantic segmentation approaches have achieved impressive results by using large amounts of labeled training data, their performance drops significantly as the amount of labeled data decreases. This happens because deep CNNs trained with the de facto cross-entropy loss can easily overfit to small amounts of labeled data. To address this issue, we propose a simple and effective contrastive learning-based training strategy in which we first pretrain the network using a pixel-wise, label-based contrastive loss, and then fine-tune it using the cross-entropy loss. This approach increases intra-class compactness and inter-class separability, thereby resulting in a better pixel classifier. We demonstrate the effectiveness of the proposed training strategy using the Cityscapes and PASCAL VOC 2012 segmentation datasets. Our results show that pretraining with the proposed contrastive loss results in large performance gains (more than 20% absolute improvement in some settings) when the amount of labeled data is limited. In many settings, the proposed contrastive pretraining strategy, which does not use any additional data, is able to match or outperform the widely-used ImageNet pretraining strategy that uses more than a million additional labeled images.

1. Introduction

The paper targets semantic segmentation with limited pixel-level annotations, where cross-entropy-trained CNNs overfit and performance declines. It proposes supervised pixel-wise contrastive pretraining followed by cross-entropy fine-tuning, yielding compact, separated features and strong results without additional data.

  • Motivation: Dense pixel-level annotation is costly, motivating segmentation with limited labeled images and no ImageNet data.Cityscapes labeling averages 90 minutes per image, while ImageNet has licensing and domain limitations.
  • Problem: Cross-entropy-trained segmentation models lose performance with fewer labeled images because the loss does not explicitly promote compact classes or large inter-class margins.The loss focuses on creating class decision boundaries, allowing CNNs to overfit small datasets.
  • Method: The method pretrains the feature extractor with supervised pixel-wise label-based contrastive loss, then fine-tunes the full network with cross-entropy.Same-class pixel features are encouraged to be close, while different-class features are encouraged to be far apart.
  • Results: 23.6 mIoU points separate the proposed strategy from cross-entropy training on 2118 PASCAL VOC 2012 images, with values of 62.7 versus 39.1.The corresponding feature distributions show more compact and separated class support regions under the proposed strategy.
  • Contribution: The loss models relationships among same- and different-class pixels in feature space regardless of image location, unlike local region-based losses.The paper evaluates three variants of the pixel-wise label-based contrastive loss.
  • Results: The proposed strategy produces large gains on Cityscapes and PASCAL VOC 2012 with limited labels and uses no additional data.In most settings, it outperforms ImageNet pretraining, which uses more than a million additional labeled images.

2. Related works

Prior segmentation work reduces annotation demands through self-supervision, weaker labels, region-based losses, metric learning, or alternative feature modeling. This paper differs by using supervised contrastive learning in a two-stage training scheme and evaluating it in the limited-label setting.

  • Training strategies: Joint training of cross-entropy and contrastive losses differs from the proposed two-stage strategy, which experiments report as more effective.The cited joint-training work also does not demonstrate contrastive learning in the limited-labeled-data setting.
  • Annotation-efficient segmentation: Semi-supervised and weakly supervised methods reduce pixel-level annotation needs using unlabeled images, bounding boxes, image-level labels, scribbles, or points.The proposed strategy does not use additional data or annotations and is described as complementary to these approaches.
  • Relationship modeling: Region-based losses model pixel relationships in label space, while the proposed contrastive loss models them in feature space.The cited region-based approaches focus on relationships within local neighborhoods; the proposed loss is not location-restricted.
  • Metric learning: Prior metric-learning methods often restrict relationships to local neighborhoods or object instances, whereas the proposed loss contrasts similar pixels with many dissimilar pixels irrespective of location.This distinguishes the proposed relationship structure from independent pairwise similarity and dissimilarity losses.
  • Alternative feature modeling: A mixture-of-vMF approach trains feature extractors and uses spherical K-Means followed by nearest-neighbor label retrieval during inference.Its inference procedure differs from the proposed pixel-wise contrastive pretraining and softmax fine-tuning pipeline.

3. Proposed approach

The proposed approach pretrains segmentation features with pixel-wise label-based contrastive losses, then fine-tunes the full network with cross-entropy. Its variants organize pixel representations around class labels using within-image or cross-image relationships.

  • The method extends supervised contrastive learning to pixel-level semantic segmentation with three label-based contrastive losses.
  • Within-image loss: Within-image contrastive loss clusters pixels according to their labels, using original and sometimes distorted image features.Distortions generate the paired image with probability p = 0.8, and the loss is averaged across images in a minibatch.
  • Cross-image loss: Cross-image contrastive loss adds same-class positive pixels from another image while excluding its negatives as easier comparisons.The loss pairs each image with another random minibatch image and averages across image pairs.
  • Training strategy: The training pipeline pretrains the CNN feature extractor with a projection head, discards that head, and fine-tunes the entire network with a softmax classifier.The projection-head features are used for contrastive loss; fine-tuning uses pixel-wise cross-entropy.
  • Training strategy: Softmax fine-tuning rearranges within-image feature clusters so they fall on the correct side of the pixel-classification decision boundary.

4. Experiments

Experiments on Cityscapes and PASCAL VOC 2012 evaluate contrastive pretraining under limited-label settings and compare variants, training choices, and datasets. Contrastive pretraining consistently improves performance, with especially large gains on PASCAL VOC 2012.

  • Performance gains: Performance gains exceed 4 points on Cityscapes below 600 labeled images and reach about 30 points on PASCAL VOC 2012 across label counts.The results are reported on validation splits; Cityscapes uses 2,975 training images and PASCAL VOC 2012 uses 10,582 training images with 21 classes.
  • Implementation and evaluation: Longer training caused overfitting when the number of labeled images was low.This observation is reported as an experiment footnote.
  • Performance gains: 1059 labeled PASCAL VOC 2012 images outperform cross-entropy training with 5295 images, reducing labeling requirements by 2× while improving performance.This comparison is reported as a validation result for the proposed contrastive pretraining strategy.
  • Contrastive-loss variants: Cross-image contrastive loss outperforms within-image loss by 1.8 points on Cityscapes with 343 labeled images.Across most settings, the cross-image variant exceeds the within-image variant by at least 0.8 points.
  • Implementation and evaluation: The batch contrastive-loss variant has quadratic memory complexity in the number of pixels, so experiments randomly sample 10K minibatch pixels to avoid GPU memory issues.The batch variant treats all pixels in a minibatch as one bag for loss computation.
  • Dataset differences: PASCAL VOC 2012 shows larger improvements than Cityscapes, possibly because its background category contains diverse visual content from many object classes.The paper presents this explanation as a conjecture and reports an experiment that removes the background category.

4.5. Comparison with ImageNet pretraining

The paper compares its label-based contrastive pretraining with ImageNet pretraining and examines alternative training and augmentation strategies. Contrastive pretraining generally matches or outperforms ImageNet pretraining without additional data, while some alternatives are less effective.

  • Practical considerations: ImageNet pretraining may be unsuitable for commercial products because the dataset is restricted to non-commercial research, while creating a proprietary substitute is expensive and time-consuming.The paper motivates contrastive pretraining as an approach that does not require this additional dataset.
  • Comparison with ImageNet pretraining: Contrastive-pretrained models match or outperform ImageNet-pretrained models in most cases, except below 600 Cityscapes labeled images.Contrastive pretraining uses no additional data, whereas ImageNet pretraining uses more than a million additional labeled images.
  • Training strategies: Joint training with contrastive and cross-entropy losses performs significantly worse than the proposed two-stage strategy.Joint training performs only slightly better than cross-entropy-only training.
  • Data distortions: Color distortions produce a small Cityscapes gain and no PASCAL VOC 2012 gain during contrastive pretraining.The results suggest that distortions effective for image recognition may not transfer directly to semantic segmentation.

4.7. Additional results

Additional experiments show that contrastive pretraining transfers to OCR and semi-supervised settings, producing strong gains without requiring additional labeled data.

  • Contrastive pretraining with OCR [61] approach: Contrastive pretraining can be combined with OCR and is applicable beyond the DeepLabV3+ baseline.The authors use DeepLabV3+ throughout but combine the strategy with OCR to demonstrate model compatibility.
  • Contrastive pretraining with OCR [61] approach: The OCR comparison is confounded because OCR performs worse than the DeepLabV3+ baseline when training data are scarce.The authors suggest that OCR has more learnable parameters and may be more prone to overfitting.
  • Semi-supervised setting: The proposed semi-supervised approach uses pseudo labels generated from predictions on unlabeled images.Pixels are assigned pseudo labels when the highest-scoring class exceeds a threshold.
  • Semi-supervised setting: 25-35 points: contrastive pretraining outperforms CCT by a large margin in the semi-supervised setting.Without contrastive pretraining, the pseudo-label approach performs similarly to CCT, linking the gap mainly to contrastive pretraining despite architectural differences.
  • Semi-supervised setting: The approach is competitive with ImageNet-pretrained CCT while using no additional data.The comparison supports the effectiveness of supervised contrastive pretraining relative to a method using ImageNet pretraining.

5. Conclusions and future work

The paper concludes that label-based contrastive pretraining improves limited-label semantic segmentation and can match or exceed supervised ImageNet pretraining without additional data.

  • Conclusions and future work: The proposed strategy first applies pixel-wise label-based contrastive loss, then fine-tunes the full network with cross-entropy loss.This increases intra-class compactness and inter-class separability, enabling a better pixel classifier.
  • Conclusions and future work: Across PASCAL VOC 2012 and Cityscapes, contrastive pretraining produces large gains with limited labeled data.The conclusion reports this pattern across both evaluated datasets.
  • Conclusions and future work: In many settings, contrastive pretraining matches or outperforms supervised ImageNet pretraining without additional data.The comparison is stated as a paper-level conclusion across the evaluated settings.
  • Conclusions and future work: Future work will combine the proposed contrastive loss with consistency-based losses for semi-supervised learning.The current semi-supervised experiments use pseudo labeling.

A. Performance on test splits

Additional test-split experiments show significant performance improvements from within-image contrastive pretraining on Cityscapes and PASCAL VOC 2012.

  • Performance on test splits: Significant performance improvements are obtained on the Cityscapes and PASCAL VOC 2012 test splits with within-image contrastive pretraining.These results are reported for the fully-supervised setting.

B. Performance gain in semi-supervised setting

In the semi-supervised setting, contrastive pretraining consistently improves performance across labeled and unlabeled-data configurations on Cityscapes and PASCAL VOC 2012.

  • Performance gain in semi-supervised setting: 2.8-7.4 points: Cityscapes mean IOU gains from contrastive pretraining in the semi-supervised setting.The gains occur across different amounts of labeled and unlabeled training data.
  • Performance gain in semi-supervised setting: Up to about 30 points: PASCAL VOC 2012 gains from contrastive pretraining in the semi-supervised setting.The improvement is reported across different amounts of labeled and unlabeled training data.
  • Performance gain in semi-supervised setting: 2×: labeling requirements are reduced while performance improves on PASCAL VOC 2012.This result is reported as similar to the fully-supervised setting.

C. Visual results

With 2118 labeled PASCAL VOC 2012 images, contrastive pretraining improves segmentation by reducing confusion between background and foreground classes, as well as among foreground classes.

  • Contrastive pretraining reduces confusion between background and various foreground classes in the segmentation results.
  • It also reduces confusion between different foreground classes in models trained on 2118 labeled PASCAL VOC 2012 images.
Loading 2012.06985v4…