PaperlayerView exact revision on arXiv

Source-linked AI summary

Local contrastive loss with pseudo-label based self-training for semi-supervised medical image segmentation

Krishna Chaitanya, Ertunc Erdil, Neerav Karani, Ender Konukoglu

arXiv:2112.09645v1cs.CVcs.AIcs.LGstat.ML

TL;DR

Medical image segmentation needs accurate pixel-level representations but is limited by the cost of expert labels, while existing local contrastive methods lack reliable semantic correspondences. This paper introduces a pseudo-label-driven local contrastive loss trained jointly with supervised segmentation on limited annotations, and reports high performance across three public MRI datasets with gains over compared methods.

  • Problem

    Expert annotation requirements make high-accuracy medical image segmentation time-consuming and expensive, motivating methods that use unlabeled data with limited labels.

  • Method

    The method jointly trains a local pixel-wise contrastive loss on labeled and pseudo-labeled images with segmentation loss applied only to the limited labeled set.

  • Results

    Across three MRI datasets, the method obtains high segmentation performance and higher gains than compared state-of-the-art semi-supervised and concurrent contrastive methods.

  • Takeaways & Limitations

    The approach produces more compact class clusters and greater inter-class separability than a self-training model.

Abstract

from arXiv · show

Supervised deep learning-based methods yield accurate results for medical image segmentation. However, they require large labeled datasets for this, and obtaining them is a laborious task that requires clinical expertise. Semi/self-supervised learning-based approaches address this limitation by exploiting unlabeled data along with limited annotated data. Recent self-supervised learning methods use contrastive loss to learn good global level representations from unlabeled images and achieve high performance in classification tasks on popular natural image datasets like ImageNet. In pixel-level prediction tasks such as segmentation, it is crucial to also learn good local level representations along with global representations to achieve better accuracy. However, the impact of the existing local contrastive loss-based methods remains limited for learning good local representations because similar and dissimilar local regions are defined based on random augmentations and spatial proximity; not based on the semantic label of local regions due to lack of large-scale expert annotations in the semi/self-supervised setting. In this paper, we propose a local contrastive loss to learn good pixel level features useful for segmentation by exploiting semantic label information obtained from pseudo-labels of unlabeled images alongside limited annotated images. In particular, we define the proposed loss to encourage similar representations for the pixels that have the same pseudo-label/ label while being dissimilar to the representation of pixels with different pseudo-label/label in the dataset. We perform pseudo-label based self-training and train the network by jointly optimizing the proposed contrastive loss on both labeled and unlabeled sets and segmentation loss on only the limited labeled set. We evaluated on three public cardiac and prostate datasets, and obtain high segmentation performance.

I. INTRODUCTION

Medical image segmentation benefits from supervised deep learning but is constrained by the expense and labor of expert annotations. Semi-supervised and self-supervised methods use unlabeled data, yet local contrastive learning remains limited when local regions lack consistent semantic labels.

  • I. INTRODUCTION: Supervised segmentation achieves high accuracy but requires large datasets annotated by clinical experts.Obtaining such annotations is time-consuming and expensive.
  • I. INTRODUCTION: Semi-supervised and self-supervised methods reduce reliance on labeled examples by leveraging unlabeled images alongside limited annotations.Reported strategies include pseudo-label self-training, entropy minimization, consistency regularization, and data augmentation.
  • I. INTRODUCTION: Contrastive learning seeks similar representations for positive pairs and dissimilar representations for negative pairs, making pair definition central to representation quality.Earlier approaches commonly define positives through differently transformed views of an image.
  • A. Motivation: Local representations are crucial for segmentation, but existing local contrastive methods struggle to define semantically similar and dissimilar regions without labels.Prior methods use spatial proximity, augmentations, or surrogate masks rather than reliable semantic labels.
  • A. Motivation: Super-pixel segmentation can assign different label identifiers to anatomically similar regions across images, hindering cross-image similarity and dissimilarity constraints.It also may produce labels that do not resemble the ground truth semantics.

B. Contributions

The paper proposes an end-to-end semi-supervised framework that uses pseudo-labels to define pixel-level contrastive relationships while restricting segmentation loss to labeled data. It distinguishes the approach from concurrent methods and evaluates it against established self-supervised and semi-supervised strategies.

  • B. Contributions: The proposed approach jointly trains segmentation and local pixel-level contrastive objectives using limited labeled data and pseudo-labels for unlabeled images.Pseudo-labels are periodically re-estimated during training.
  • B. Contributions: Pseudo-labels are used only for contrastive loss computation, unlike concurrent approaches that also use them in segmentation or cross-entropy losses.The framework uses one network without external memory or entropy minimization.
  • B. Contributions: The local loss encourages intra-class similarity and inter-class separability across the dataset, whereas prior unsupervised local methods enforce these relationships only within an image.This difference addresses inconsistent semantic identifiers from unsupervised segmentation methods.
  • B. Contributions: The method is evaluated on three public MRI datasets against state-of-the-art self-supervised and semi-supervised methods, and the results improve the best-performing baselines.The paper also reports a detailed ablation study of framework components and consistency-based pseudo-label selection.
  • Related Work: Contrastive-learning and pretext-task methods learn representations from unlabeled data, while supervised local contrastive methods rely on many labeled examples.The proposed work targets the semi-supervised setting with limited annotations.
  • Related Work: Unsupervised local contrastive methods use augmented views, domain cues, or surrogate masks such as saliency maps and super-pixels.Semi-supervised methods differ in how they combine pseudo-labels with contrastive and segmentation objectives.
  • Related Work: Consistency regularization enforces similar outputs for perturbed versions of an image, while other semi-supervised methods use adversarial training, entropy minimization, or hybrid objectives.These approaches complement self-training and contrastive strategies in the broader literature.

III. METHODS

The method jointly uses limited labeled data and unlabeled data with pseudo-labels to learn pixel-level representations for semi-supervised segmentation. It combines supervised segmentation loss on labeled images with local contrastive loss over the combined dataset, encouraging intra-class similarity and inter-class separation.

  • The method targets discriminative pixel-level representations with intra-class affinity and inter-class separability for semantic segmentation.
  • It defines labeled data as X_L and unlabeled data as X_U with pseudo-labels, forming the full dataset X = X_L ∪ X_U.The setting uses limited labeled sets, including |X_L| = 1, 2, or 8.
  • An initial supervised optimization trains the encoder-decoder and segmentation branch on X_L to obtain initial pseudo-labels for unlabeled images.The initial segmentation loss uses the limited labeled set and initializes the network parameters randomly.
  • A second optimization jointly trains the encoder-decoder, segmentation branch, and contrastive branch using segmentation loss on X_L and local contrastive loss on X.Unlabeled pseudo-labels are estimated initially and updated every P iterations; the contrastive loss has weighting parameter λ_cont.
  • The local contrastive loss matches pixel representations from the same class across images while separating them from representations of other classes.For computational efficiency, the method matches pixels to class mean representations and subsamples pixel coordinates to fit GPU memory.
  • The loss can match representations within the same image or across different images, with cosine similarity and temperature scaling used in its formulation.These are described as intra-image and inter-image representation matching implementations.

A. Datasets

The study uses three public MRI datasets covering cardiac, prostate, and multi-structure segmentation tasks.

  • The ACDC dataset contains 100 short-axis MR-cine T1 3D cardiac volumes with annotations for three structures.
  • The Prostate dataset contains 48 T2-weighted MRI 3D prostate volumes with expert annotations.

B. Pre-processing

All datasets undergo bias correction, percentile-based normalization, resampling, and cropping or zero padding to fixed image dimensions.

  • Images are bias-corrected using N4 implemented in the ITK toolkit.
  • Each 3D volume is min-max normalized using its 1st and 99th intensity percentiles.
  • 2D images and labels are resampled to fixed in-plane resolutions using bilinear and nearest-neighbor interpolation, respectively.
  • Images are cropped or zero-padded to fixed dimensions of 192 × 192 for ACDC and Prostate, and 160 × 160 for MMWHS.

C. Network Architecture

The network uses a shared U-Net encoder-decoder followed by separate segmentation and contrastive-learning branches.

  • The architecture is U-Net-based, with a common encoder-decoder network denoted cθ.
  • Two smaller decoder branches separately support segmentation through gξ and contrastive learning through hφ.
  • The encoder has six convolutional blocks, while the decoder has five convolutional blocks with upsampling.
  • All layers except the final layers of both branches use batch normalization and ReLU activation.

D. Experimental Setup

Training begins with supervised segmentation on limited labels, then alternates pseudo-label estimation with joint optimization using labeled and unlabeled data.

  • The first 5000 iterations optimize segmentation loss using only the labeled set and update the shared and segmentation-specific parameters.
  • Pseudo-labels for the unlabeled set are estimated using the network trained in the initial supervised stage.
  • The second optimization updates encoder-decoder, segmentation, and contrastive-specific parameters through joint training.
  • Random intensity transformations are used for contrastive-loss inputs, while labeled segmentation batches receive geometric and intensity augmentations.
  • The default contrastive sampling uses three positive pixel representations per class and image, with temperature τ = 0.1.
  • Experiments use unlabeled and test splits across ACDC, Prostate, and MMWHS, with labeled sets of 1, 2, or 8 3D volumes.
  • Performance is measured by mean Dice similarity coefficient over foreground structures across six runs with randomly sampled labeled and validation sets.

V. RESULTS AND DISCUSSION

The proposed method improves segmentation over baseline and compared semi-supervised methods by using pseudo-labels for contrastive representation learning. Both intra-image and inter-image pixel matching improve results, with inter-image matching performing better in most cases.

  • Semi-supervised methods substantially improve limited-label performance, with self-training and noisy student nearly closing the benchmark gap.With |X_L|=2, the gap is under 0.1 DSC for Prostate and MMWHS and under 0.2 DSC for ACDC.
  • The proposed method outperforms the baseline and compared semi-supervised methods.Using pseudo-labels only in contrastive loss may avoid erroneous pseudo-label information directly degrading segmentation-loss optimization.
  • Both intra-image and inter-image pixel representation matching improve performance over the compared methods.
  • Inter-image matching performs better than intra-image matching in 5 of 9 cases.For |X_L|=1, intra-image matching performs better on ACDC and MMWHS than inter-image matching.

C. Concurrent contrastive semi-supervised methods

The paper compares its joint pseudo-label contrastive approach with concurrent contrastive frameworks and contrastive-pretrained initialization. The proposed approach performs better, with larger improvements on ACDC than on Prostate and MMWHS.

  • Joint training can be more beneficial than disjoint contrastive pre-training and fine-tuning, while using pseudo-labels in segmentation loss can hinder gains in limited-label settings.
  • The proposed pseudo-label joint training provides better results than compared methods across all datasets and most training-label sizes.The comparison includes semi-supervised learning, data augmentation, and concurrent contrastive learning methods.
  • Compared with concurrent contrastive learning works, the proposed approach shows larger improvements on ACDC and relatively smaller improvements on Prostate and MMWHS.
  • Contrastive-pretrained initialization yields larger improvements for the proposed method over random initialization and self-training, especially on ACDC.The gains are smaller on Prostate and MMWHS.

E. Visualization of results of the proposed method

The proposed method produces more compact and separable pixel representations than self-training and continues improving through the second pseudo-labeling step. Its performance is generally stable across tested pseudo-labeling, sampling, and pseudo-label-quality settings, while very small contrastive-loss weights can hurt low-label performance.

  • Representation visualization: The proposed method forms more compact same-class clusters and better separates different cardiac-structure representations than self-training.The comparison concerns right ventricle, myocardium, and left ventricle pixel representations from ACDC unlabeled images.
  • Pseudo-labeling steps: Both methods improve mean DSC after pseudo-label updates, but the proposed method gains more from the first update across all runs.The experiment uses one labeled ACDC 3D volume, with pseudo-labels updated every 5000 iterations.
  • Pseudo-labeling steps: Self-training gains plateau after the second pseudo-labeling step, whereas the proposed method continues improving at that step before plateauing after the third.This behavior is reported across three runs on the ACDC setting with one labeled volume.
  • Hyper-parameter analysis: Very small contrastive-loss weights deteriorate performance when labeled examples are scarce, while the coefficient has no distinguishable effect with |X_L| = 8.The hyper-parameter analysis evaluates λ_cont on ACDC only because each experiment requires substantial computational resources.
  • Hyper-parameter analysis: Using highly confident pseudo-labels does not necessarily improve gains over using all predictions, and noisy labels can reduce inter-image matching performance with few labeled examples.At threshold 0, inter-image matching performs worse than intra-image matching for |X_L| = 1 or 2.

VI. CONCLUSION

The proposed semi-supervised framework combines pseudo-label-based pixel-wise contrastive learning with labeled-only segmentation loss to address limited annotations. Across three MRI datasets, it improves class representation structure and achieves strong gains in limited-annotation settings.

  • The method jointly trains pixel-wise contrastive loss on labeled and unlabeled images with segmentation loss applied only to labeled images.Iterative self-training improves pseudo-label quality during training.
  • The contrastive loss learns better intra-class compactness and inter-class separability than typical segmentation-loss-based self-training.
  • Evaluation on three MRI datasets shows higher performance gains than compared state-of-the-art semi-supervised and concurrent contrastive learning methods.
  • Unlabeled-image pretraining provides complementary network initialization that can be combined with the proposed method for higher gains.
  • Restricting segmentation loss to the labeled set is identified as an essential detail in the method’s performance gains.
Loading 2112.09645v1…