Source-linked AI summary

Momentum Contrastive Voxel-wise Representation Learning for Semi-supervised Volumetric Medical Image Segmentation

Chenyu You, Ruihan Zhao, Lawrence Staib, James S. Duncan

arXiv:2105.07059v4cs.CVcs.LGeess.IV

TL;DR

Limited expert annotation and possible representational collapse complicate contrastive learning for volumetric medical segmentation. CVRL combines voxel-wise and dimensional contrastive objectives to learn low- and high-level representations; the authors report state-of-the-art performance on two benchmarks with accurate boundaries under limited annotations.

  • Problem

    Contrastive learning for medical image segmentation must learn useful representations without expert supervision, while avoiding collapse and accounting for dissimilar anatomical structures within an image.

    Source: §1 · p. 2
  • Method

    CVRL combines voxel-wise and dimensional contrastive learning to capture low- and high-level cues across batch and feature dimensions for 3D segmentation.

    Source: §1.2 · p. 4, Dimensional Contrastive Objective · p. 5, §3 · p. 8
  • Results

    CVRL is reported to achieve state-of-the-art performance on two benchmark datasets and generate more accurate boundaries with very limited annotations.

    Source: §1 · p. 3, §3 · p. 9
  • Takeaways & Limitations

    The authors report that dimensional contrastive learning retains the hardness-aware property of standard contrastive learning.

    Source: §1 · p. 3, §3 · p. 9

    The problem setting assumes N labeled examples and M unlabeled examples, with N ≪ M.

    Source: §1.1 · p. 3

Abstract

from arXiv · show

Contrastive learning (CL) aims to learn useful representation without relying on expert annotations in the context of medical image segmentation. Existing approaches mainly contrast a single positive vector (i.e., an augmentation of the same image) against a set of negatives within the entire remainder of the batch by simply mapping all input features into the same constant vector. Despite the impressive empirical performance, those methods have the following shortcomings: (1) it remains a formidable challenge to prevent the collapsing problems to trivial solutions; and (2) we argue that not all voxels within the same image are equally positive since there exist the dissimilar anatomical structures with the same image. In this work, we present a novel Contrastive Voxel-wise Representation Learning (CVRL) method to effectively learn low-level and high-level features by capturing 3D spatial context and rich anatomical information along both the feature and the batch dimensions. Specifically, we first introduce a novel CL strategy to ensure feature diversity promotion among the 3D representation dimensions. We train the framework through bi-level contrastive optimization (i.e., low-level and high-level) on 3D images. Experiments on two benchmark datasets and different labeled settings demonstrate the superiority of our proposed framework. More importantly, we also prove that our method inherits the benefit of hardness-aware property from the standard CL approaches.

1 Introduction

CVRL addresses limited 3D annotations and possible dimensional collapse by combining voxel-wise and dimensional contrastive learning for semi-supervised volumetric segmentation. The paper reports state-of-the-art results on two benchmarks and high-quality global shapes.

1 Introduction

  • Manual 3D annotations are costly and expertise-intensive, while medical imaging datasets are relatively small and contrastive methods may collapse along feature dimensions.The introduction also notes that predefined augmentations may ignore anatomical feasibility.Source: §1 · pp. 1–2
  • CVRL is reported to outperform state-of-the-art methods on two benchmark datasets and produce object segmentations with high-quality global shapes.The method is designed for limited-annotation clinical scenarios.Source: §1 · p. 3

1.1 Overview

  • CVRL combines supervised and unsupervised objectives to learn voxel-wise representations from 3D unlabeled data in a setting with far fewer labeled than unlabeled examples.The framework regularizes the embedding space and uses training-voxel geometric and spatial context.Source: §1.1 · p. 3

1.2 Unsupervised Contrastive Learning

  • Voxel-wise contrast pulls corresponding locations across two augmentations of an image together and pushes other locations or inputs apart.The objective uses InfoNCE, with teacher features as keys and student features as queries.Source: §1.2 · p. 4

Dimensional Contrastive Objective

  • Dimensional contrast pairs each feature channel with its teacher counterpart to promote diverse channel information; the authors report that it retains standard contrastive learning’s hardness-aware property.The method applies the contrastive objectives to low- and high-level features.Source: §1 · p. 3, Dimensional Contrastive Objective · p. 5

2 Experiments

Experiments evaluate CVRL on two benchmark datasets using 10% and 20% labeled data. CVRL outperforms other end-to-end semi-supervised methods on LA, shows gains on pancreas, and improves over its baseline as components are added.

2 Experiments

  • The experiments use LA and pancreas datasets under 10% and 20% labeled ratios, with V-Net as the network backbone for two networks.LA has 100 scans, with 80 for training and 20 for evaluation; pancreas has 82 scans, with 62 for training and 20 for evaluation.Source: §2 · p. 6
  • Visual comparisons report that CVRL produces more accurate segmentation borders and shapes than other methods.The authors describe improvements in the 10% annotation setting as difficult to achieve.Source: Figure 3 · p. 7, §2 · p. 8
  • Segmentation performance is evaluated with Dice, Jaccard, 95HD, and ASD.All evaluated methods are trained for 10,000 iterations with batch size 6.Source: §2 · p. 7
  • At 20% labels, CVRL outperforms other end-to-end semi-supervised methods in Dice (90.45%), Jaccard (83.02%), ASD (1.81), and 95HD (6.56); at 10%, it improves Dice (88.56%) and Jaccard (78.89%).The 20% setting uses 16 labeled and 64 unlabeled scans; the 10% setting uses 8 labeled and 72 unlabeled scans.Source: §2 · p. 7
  • On pancreas, CVRL significantly outperforms state-of-the-art methods, with relative Dice improvement of 3.25–5.21%.The paper reports the detailed pancreas evaluation in an appendix table.Source: §2 · p. 8
  • Adding Lhigh, Llow, and Lcon to the MT baseline improves Dice by 5.47% and Jaccard by 7.14%, while reducing ASD by 2.72 and 95HD by 11.43.The ablation uses the LA dataset with 10% labeled data.Source: §2 · p. 8, Table 2 · p. 8

3 Conclusion

CVRL learns voxel-wise representations for volumetric medical image segmentation using low- and high-level cues. The paper reports state-of-the-art performance and more accurate boundaries with very limited annotations.

3 Conclusion

  • CVRL is a semi-supervised contrastive representation distillation framework that uses low- and high-level cues to learn voxel-wise representations for volumetric medical image segmentation.Source: §3 · p. 8
  • Voxel-wise and dimensional contrastive learning promote diversity and exploit complex relations among training voxels; the paper identifies hardness-awareness as key to dimensional contrastive learning.Source: §3 · p. 9
  • Experiments report state-of-the-art performance and more accurate boundaries with very limited annotations.Source: §3 · p. 9
Loading 2105.07059v4…