Contrastive learning of global and local features for medical image segmentation with limited annotations

Krishna Chaitanya; Ertunc Erdil; Neerav Karani; Ender Konukoglu

Source-linked AI summary

Contrastive learning of global and local features for medical image segmentation with limited annotations

Krishna Chaitanya, Ertunc Erdil, Neerav Karani, Ender Konukoglu

arXiv:2006.10511v2Jun 18, 2020cs.CVcs.LGeess.IVstat.ML

TL;DR

Limited medical-image annotations make supervised segmentation difficult, motivating pre-training from unlabeled data. The paper extends contrastive learning with volumetric domain cues and local representations, and reports substantial improvements in limited-annotation MRI segmentation, including performance within 8% of benchmark performance using two labeled volumes with augmentation.

Problem
Large annotated datasets are difficult to obtain for medical image segmentation, limiting supervised deep learning when annotations are scarce.
Method
The paper extends contrastive pre-training with domain-specific similarity across volumetric images and a local contrastive loss for representations useful in dense prediction.
Results
The proposed global and local strategies substantially outperform competing methods across three MRI datasets and limited training-volume settings.
Takeaways & Limitations
Combining the proposed pre-training with augmentation and semi-supervised methods substantially narrows the performance gap with the benchmark.
Takeaways & Limitations
Medical datasets may contain far fewer unlabeled images than natural-image datasets, making very large pre-training batch sizes impractical.

Abstract

from arXiv · show

A key requirement for the success of supervised deep learning is a large labeled dataset - a condition that is difficult to meet in medical image analysis. Self-supervised learning (SSL) can help in this regard by providing a strategy to pre-train a neural network with unlabeled data, followed by fine-tuning for a downstream task with limited annotations. Contrastive learning, a particular variant of SSL, is a powerful technique for learning image-level representations. In this work, we propose strategies for extending the contrastive learning framework for segmentation of volumetric medical images in the semi-supervised setting with limited annotations, by leveraging domain-specific and problem-specific cues. Specifically, we propose (1) novel contrasting strategies that leverage structural similarity across volumetric medical images (domain-specific cue) and (2) a local version of the contrastive loss to learn distinctive representations of local regions that are useful for per-pixel segmentation (problem-specific cue). We carry out an extensive evaluation on three Magnetic Resonance Imaging (MRI) datasets. In the limited annotation setting, the proposed method yields substantial improvements compared to other self-supervision and semi-supervised learning techniques. When combined with a simple data augmentation technique, the proposed method reaches within 8% of benchmark performance using only two labeled MRI volumes for training, corresponding to only 4% (for ACDC) of the training data used to train the benchmark. The code is made public at https://github.com/krishnabits001/domain_specific_cl.

1 Introduction

The paper motivates contrastive pre-training for volumetric medical image segmentation when annotations are limited. It targets two gaps: insufficiently distinctive local representations and contrasting strategies that do not exploit medical-image structure.

Large labeled datasets enable state-of-the-art medical image segmentation, but assembling them is challenging.
Contrastive learning trains representations by making transformed views similar and different images dissimilar using unlabeled data.
Existing contrastive methods largely emphasize global representations rather than distinctive local representations useful for per-pixel segmentation.
The paper proposes domain-specific contrasting strategies for volumetric MRI and CT data, plus a local contrastive loss for local regions.
Across three MRI datasets, combining global and local strategies substantially improves performance over pre-training, self-supervised, and semi-supervised alternatives.

2 Related works

Prior work uses unlabeled data through pretext tasks, contrastive learning, semi-supervised learning, and augmentation. This paper adapts contrastive learning to encoder-decoder segmentation by emphasizing local representations and medical-image structure.

Pretext-task methods learn representations from labels generated freely from unlabeled images, such as orientation prediction, inpainting, and context restoration.
Contrastive learning enforces similar representations for similar pairs and dissimilar representations for dissimilar pairs, often using transformed views as similar examples.
The proposed approach differs by targeting encoder-decoder architectures for pixel-wise prediction rather than encoder architectures for image-wise tasks.
Its local representations are intended to distinguish different image areas for the decoder, complementing image-wide representations.
Other limited-annotation approaches include semi-supervised self-training, adversarial training, and data augmentation.

3 Methods

The method extends global contrastive pre-training for volumetric medical-image segmentation with domain-specific structural cues and a local contrastive loss for distinctive pixel-level representations.

Global contrastive loss: Global contrastive learning encourages transformed views of the same image to have similar representations and views of different images to have dissimilar representations.The loss is optimized over similar image pairs and dissimilar images formed from the remaining batch samples.
Domain-specific global strategies: Aligned volumetric images are divided into corresponding partitions, which are treated as capturing similar anatomical areas across different volumes.The same grouping is used across volumes under the assumed rough alignment.
Domain-specific global strategies: The proposed global strategies construct batches from sampled volumes and use partition correspondence to control which images form similar and dissimilar pairs.GD− excludes corresponding partitions from negatives, whereas GD additionally makes corresponding partitions across volumes positive pairs.
Local contrastive loss: The local contrastive loss trains decoder blocks to make different local regions dissimilar while preserving similarity of corresponding regions across intensity transformations.This complements the encoder’s global representations for pixel-level prediction tasks.
Local contrastive loss: Local loss computation divides transformed feature maps into local regions, pairing corresponding regions as positives and treating other regions in both maps as negatives.The feature maps are produced by the encoder, decoder blocks, and a shallow projection network.

4 Experiments and Results

Experiments on three MRI datasets evaluate the proposed pre-training under limited annotation, showing gains from domain-specific global contrasts, local contrastive loss, and combining pre-training with augmentation or semi-supervised methods.

Experimental setup: The evaluation uses three publicly available MRI datasets and assesses fine-tuned UNet models with 1, 2, or 8 labeled training volumes.Pre-training uses unlabeled images, while segmentation performance is measured on a held-out test set using mean Dice scores over six runs.
Global contrasting strategies: Domain-specific global contrasting strategies GD− and GD substantially improve performance over random contrasting across all datasets and training-set sizes.The authors attribute these gains to exploiting slice correspondence across volumes, which provides richer similarity cues than random augmentations.
Local contrastive loss: Local contrastive pre-training of the decoder provides an additional performance boost over encoder-only pre-training with random decoder initialization.The reported experiments use three pre-trained decoder blocks, and LR performs better than LD in 6 of 9 settings.
Limitations: Correspondence between local regions across different volumes is not guaranteed, so forcing similarity between assumed corresponding regions may adversely affect learning.This limitation motivates caution when applying correspondence-based local contrasting strategies.
Comparison with other methods: The complete proposed method performs substantially better than competing pre-training methods across all datasets and labeled-volume counts, especially with 1 or 2 volumes.The comparison includes pretext-task pre-training, data augmentation, and adversarial training, using a common architecture across methods.
Combined methods: Applying data augmentation and semi-supervised methods after proposed pre-training yields further performance gains and helps close the gap with the fully supervised benchmark.The authors describe the benefits of the proposed initialization as complementary to those obtained from other methods.

5 Conclusion

The paper extends contrastive pre-training for medical image segmentation by learning local representations and using domain-specific contrasting strategies. Across three MRI datasets, both improvements produce substantial gains in limited-annotation settings, with benefits complementary to augmentation and semi-supervised learning.

Conclusion: The method addresses the need for accurate medical image segmentation with fewer annotated training examples.The paper targets the persistent challenge of obtaining large annotated datasets.
Conclusion: The proposed pre-training adds a local contrastive loss for dense prediction and a problem-specific strategy based on naturally occurring clusters in the data.These components are designed to learn local representations and select similar or dissimilar image pairs for contrastive loss computation.
Conclusion: Experiments on three MRI datasets show substantial performance gains from both proposed improvements under limited annotation.The conclusion reports the gains across the evaluated MRI datasets and limited-annotation settings.
Conclusion: The proposed initialization provides benefits orthogonal to those of data augmentation and semi-supervised learning.Combining the approaches is presented as a route toward reducing the need for large annotated training sets.

6 Broader Impact

Large expert-labeled datasets are difficult and costly to assemble, motivating less data-hungry segmentation methods. The proposed self-supervised strategy uses unlabeled images to support high performance with few annotations, while healthcare deployment still requires extensive validation and attention to domain shift, uncertainty, and interpretability.

Radiologist scarcity and costly annotation create a bottleneck for deploying supervised segmentation systems in clinical settings.The paper frames reducing annotation requirements as crucial for achieving high performance with few annotations.
The method pre-trains on unlabeled images to learn global and local representations before fine-tuning segmentation with limited annotations.It incorporates domain-specific contrasting strategies and a local contrastive loss for segmentation-relevant features.
Extensive validation remains essential because medical imaging systems may face domain shift, uncertainty-quantification requirements, and interpretability requirements.These issues can arise when acquisition settings or imaging modalities vary.

7 Network architecture and training details

The model uses a UNet-based encoder–decoder architecture with separate projection networks and staged global and local contrastive pre-training. Training uses fixed optimization settings, dataset-specific splits, and fine-tuning with one, two, or eight labeled volumes.

The network is a UNet-based encoder–decoder with six convolutional encoder blocks and separate heads for global and local contrastive learning.The global head g1 uses dense layers of dimensions 3200 and 128; local pre-training adds l decoder blocks and g2.
Global pre-training appends g1 to the encoder, after which g1 is discarded before local pre-training adds decoder blocks and g2.The staged design successively pre-trains the encoder and partial decoder using global and local losses.
13 local regions are selected from each feature map when computing the local contrastive loss.
Each training stage uses Adam for 10,000 iterations with batch size 40 and learning rate 10^-3; τ=0.1 is selected for all experiments.The validation set supports model selection during fine-tuning.
Images are bias-corrected with N4, while pre-training and testing volumes are split to roughly 50% of each dataset.Fine-tuning uses one, two, or eight training volumes and a fixed validation set of two volumes.
On a Titan X GPU, approximate training times are 2 hours for global pre-training, 4 hours for local pre-training, and 2 hours for fine-tuning.Pre-training convergence was reported as consistently stable.

8 Illustration of slice correspondence in medical volumetric images

Figure 2 samples 2D slices from four partitions across three volumes. Although corresponding slices vary in shape and intensity, slices from the same partition retain shared global cardiac-anatomy information.

Each row presents four 2D slices from four partitions of one selected volume, with partition number indicated by s.The illustration covers three different volumes.
Table 3 is labeled as the dataset split rather than an illustration of slice correspondence.
Slices from the same partition across different volumes vary in shape and intensity but contain the same global cardiac-anatomy information.

9 Ablation studies

The ablations examine batch size and the number of partitions used during global contrastive pre-training. On ACDC, larger batches did not improve performance, while increasing partitions degraded Dice scores, partly because misaligned volumes create unrealistic pairings.

Ablation studies: The ablation study uses the ACDC dataset and the GD global contrasting strategy for encoder pre-training.
6.1.1 Batch Size: Batch size is evaluated at 40, 250, and 450 with four partitions per volume during encoder pre-training.
6.1.1 Batch Size: Higher batch sizes did not improve medical-image results, and performance deteriorated at batch size 450.The authors therefore report that large batches may not be required for high pre-training performance on the evaluated datasets.
6.1.1 Batch Size: Further evaluation on more datasets is required before drawing a conclusive statement about batch-size effects in medical datasets.
6.1.1 Batch Size: Medical datasets contain fewer unlabeled images than natural-image datasets, making batch sizes such as 2048 or 4096 potentially impractical.The ACDC dataset contains about 100 volumes and roughly 1000 2D images.
6.1.2 Number of partitions per 3D volume: The partition ablation evaluates S=3, 4, and 6 at batch size 40, and Dice scores degrade as S increases.With approximately 10 images per ACDC volume, larger S can force clusters corresponding to individual images and produce unrealistic pairings.

9.2 Local Contrastive Loss

The ablation evaluates decoder depth and local-region size for local contrastive pre-training on ACDC. Three decoder blocks and 3 × 3 regions generally provide the strongest results, while deeper decoder pre-training performs less well.

Decoder depth: Decoder-block counts l = 1, 2, 3 yielded higher gains than l = 4, 5 in the ACDC ablation.The authors hypothesize that pre-training the entire decoder gives local regions a smaller receptive field and insufficient information for useful representations.
Local-region size: A 3 × 3 local region produced around 2% higher performance than 1 × 1 for d3 on ACDC.The authors associate the improvement with the larger region containing more information through a higher receptive field.
Experimental setup: The experiment varies decoder blocks l = {1, 2, 3, 4, 5} and local-region sizes K × K of 1 × 1 and 3 × 3 before fine-tuning with |Xtr| = 1, 2.Table 6 reports mean Dice on the ACDC test set after pre-training the encoder and selected decoder blocks.
Cross-dataset check: Across the remaining datasets, 3 × 3 local regions worked better in most settings, as observed for ACDC.The additional ablation used dl = 3 with GD and LD sampling strategies.

9.3 Combination of Local and Global Contrastive Losses

This section compares combinations of local and global contrastive losses and joint versus stage-wise pre-training. The reported ACDC results favor stage-wise training over joint optimization of encoder and decoder losses.

Loss combinations: The study combines local losses LR and LD with global losses GR and GD for decoder lengths dl, l = 2, 3, 4.These combinations are evaluated across datasets and training-set sizes |Xtr| = 1, 2, 8.
Joint versus stage-wise training: Joint pre-training trains the encoder and decoder in one step, unlike the two-step stage-wise procedure.Table 9 reports the ACDC comparison between these pre-training strategies.
Joint versus stage-wise training: Stage-wise training outperformed joint training on ACDC, with DSC values of 0.725 for |Xtr| = 1 and 0.789 for |Xtr| = 2.Joint training updates encoder weights using Lnet = Lg + λl ∗ Ll, whereas stage-wise training uses Lg for encoder updates.

10 Experiments with Natural Image Datasets

The proposed local contrastive pre-training is also evaluated for natural-image segmentation on Cityscapes under limited annotation. The experiment compares no pre-training, global-only pre-training, and combined global-local pre-training.

Experimental design: Cityscapes evaluation compares combined local-global pre-training with no pre-training and global-only pre-training for segmentation.The decoder uses local contrastive loss LR, while the encoder uses global contrastive loss with random strategy GR.
Data split: The Cityscapes test set Xts contains 705 images, while 2770 images comprise Xpre for unlabeled pre-training and limited-label fine-tuning.Labeled training and validation subsets are selected from Xpre, and Xts is reserved for final evaluation.
Implementation constraints: Images were down-scaled from 1024,2048 to 256,512 because memory constraints made a batch size of around 40 difficult at original dimensions.Downsampling caused some small objects to vanish or become negligible, so they were assigned to the background class.
Augmentation: Training used random cropping followed by random color jitter for augmentation.The color-jitter operations varied brightness, contrast, saturation, and hue.
Evaluation metric: Table 10 reports mean Dice over Xts for selected Cityscapes labels under random initialization, global-only pre-training, and the proposed pre-training.The table is the final test-set comparison for the natural-image experiment.