Source-linked AI summary

Semi-supervised Contrastive Learning for Label-efficient Medical Image Segmentation

Xinrong Hu, Dewen Zeng, Xiaowei Xu, Yiyu Shi

arXiv:2109.07407v2cs.CV

TL;DR

Medical image segmentation depends on dense pixel-wise labels, which are difficult and time-consuming to annotate. This paper incorporates limited labels into contrastive pre-training through supervised local contrast alongside self-supervised global contrast. On two public datasets, the combined approach substantially improves segmentation performance over state-of-the-art methods.

  • Problem

    Medical segmentation requires abundant pixel-wise labels that are difficult to collect, while prior contrastive segmentation pre-training does not use labels.

    Source: §1 · pp. 1–2
  • Method

    The framework combines self-supervised global contrast with supervised local contrast that uses limited pixel-wise annotations during pre-training.

    Source: §1 · p. 2, §2.2 · p. 5
  • Results

    On two public medical-image datasets with partial labels, combining supervised local and global contrast substantially improves segmentation performance over the state of the art.

    Source: §4 · p. 8
  • Takeaways & Limitations

    The image-level representations from unsupervised global contrast are complementary to pixel-level representations from supervised local contrast.

    Source: §3 · p. 8

    In self-supervised local contrast, spatial transformations can mismatch corresponding pixel locations, and sampled points may all be background, limiting local-feature learning.

    Source: §2.2 · pp. 4–5

Abstract

from arXiv · show

The success of deep learning methods in medical image segmentation tasks heavily depends on a large amount of labeled data to supervise the training. On the other hand, the annotation of biomedical images requires domain knowledge and can be laborious. Recently, contrastive learning has demonstrated great potential in learning latent representation of images even without any label. Existing works have explored its application to biomedical image segmentation where only a small portion of data is labeled, through a pre-training phase based on self-supervised contrastive learning without using any labels followed by a supervised fine-tuning phase on the labeled portion of data only. In this paper, we establish that by including the limited label in formation in the pre-training phase, it is possible to boost the performance of contrastive learning. We propose a supervised local contrastive loss that leverages limited pixel-wise annotation to force pixels with the same label to gather around in the embedding space. Such loss needs pixel-wise computation which can be expensive for large images, and we further propose two strategies, downsampling and block division, to address the issue. We evaluate our methods on two public biomedical image datasets of different modalities. With different amounts of labeled data, our methods consistently outperform the state-of-the-art contrast-based methods and other semi-supervised learning techniques.

1 Introduction

Medical image segmentation relies on abundant pixel-wise labels, but acquiring them requires domain knowledge and time. The paper uses available labels during contrastive pre-training and reports higher Dice than state-of-the-art methods across different labeled-data amounts.

1 Introduction

  • Dense pixel-wise labels are needed to guide supervised segmentation, but medical-image annotation requires domain knowledge and is time-consuming.Source: §1 · pp. 1–2
  • Contrastive learning is presented as a candidate for label-efficient learning because it extracts downstream-useful features from unlabeled data.The introduction also notes that semi-supervised methods can use both labeled and unlabeled data.Source: §1 · p. 2
  • Prior contrastive segmentation methods pre-train without labels and then fine-tune with supervision on the labeled data, leaving labels unused during pre-training.Source: §1 · p. 2
  • Across different labeled-data amounts, the proposed framework produces higher Dice than state-of-the-art methods.It combines self-supervised global contrast with supervised local contrast and uses downsampling and block division to reduce the latter's computational complexity.Source: §1 · p. 2
  • The framework adds supervised local contrast to self-supervised global contrast, using available labels during pre-training.The supervised local loss is intended to draw same-class features together and distinguish features from different classes.Source: §1 · p. 2

2 Methods

The framework combines unlabeled global contrastive pre-training with label-guided local contrastive pre-training, then uses the resulting model to initialize fine-tuning. Downsampling and block division reduce the computation required by the supervised local loss.

2.1 Overview of the Semi-supervised Framework

  • Unlabeled global contrast trains the encoder on image-level features, then limited labeled data trains the encoder and decoder on pixel-level features before fine-tuning.The two-step pre-training produces an initialization for subsequent fine-tuning on labeled data.Source: §2.1 · pp. 3–4

2.2 Supervised Local Contrastive learning

  • Supervised local contrast treats pre-convolution pixel features as embeddings, drawing same-class features together and separating different classes while excluding background points from the loss anchors.Background features remain in the negative set; positive and negative sets are defined by whether annotations match the anchor pixel's class.Source: §2.2 · pp. 4–5

2.3 Downsampling and Block Division

  • The supervised loss scales as O(h4); fixed-stride downsampling (stride 4) and block-wise averaging (16×16 blocks) reduce computation, as illustrated by the feature-selection strategies.Downsampling reduces loss complexity by s4; block division has complexity O(h2h′2), and all-background blocks are skipped.Source: §2.3 · pp. 5–6, Figure 2 · p. 5

3 Experiments and Results

Experiments compare contrastive and semi-supervised methods on Hippocampus MRI and MMWHS cardiac CT datasets. Combining global contrast with supervised local contrast improves segmentation over the state of the art.

3 Experiments and Results

  • Experiments use 260 Hippocampus MRI volumes and 20 MMWHS cardiac CT volumes, independently split four times into folds; each split allocates training, validation, and test data at ratios of 3:1:1 and 2:1:1, respectively.Source: §3 · p. 6
  • Comparisons include global contrast, global+local(self), Mixup, TCSM, and random initialization alongside the proposed stride and block variants.Source: §3 · p. 7
  • Supervised local-only variants achieve results comparable to existing contrastive methods despite using labeled training data only, while those methods use all volumes.The local-only variants are local(stride) and local(block).Source: §3 · p. 7
  • All contrastive learning methods achieve higher Dice scores than random initialization.Source: §3 · p. 7
  • Combining global contrast with supervised local contrast substantially improves segmentation over the state of the art; global+local(block) achieves the highest Dice on MMWHS.The comparison includes subsequent segmentation Dice scores on two datasets.Source: Table 1 · p. 7, §3 · p. 8
  • At 40% labeled data on MMWHS, the proposed segmentation predictions perform best among the methods shown.The figure compares random, global+local(self), global+local(stride), and global+local(block) predictions.Source: §3 · p. 8, Figure 3 · p. 8

4 Conclusion

The paper proposes supervised local contrast to learn pixel-wise representations and reduce its computational cost, and reports improved segmentation from combining it with global contrast. Its experiments use two public medical image datasets with partial labels.

4 Conclusion

  • Combining supervised local contrast with global contrast substantially improves segmentation over the state of the art on two public datasets with partial labels.Source: §4 · p. 8
  • The proposed supervised local contrastive loss learns pixel-wise representations, with downsampling and block division proposed to reduce its computational complexity.Source: §4 · p. 8
  • Fig. 4 visualizes MMWHS embedding features after t-SNE, with colors representing classes and features taken before the last point-wise convolution.Source: Figure 4 · p. 9
Loading 2109.07407v2…