[Paper Review] Mine yOur owN Anatomy: Revisiting Medical Image Segmentation with Extremely Limited Labels

Chenyu You, Weicheng Dai+7 more|arXiv (Cornell University)|Sep 27, 2022

Radiomics and Machine Learning in Medical ImagingMedicine22 citations

TL;DR

MONA is a semi-supervised 2D medical image segmentation framework that learns diverse yet consistent anatomical features under extremely limited labels, achieving state-of-the-art results on multiple benchmarks.

ABSTRACT

Recent studies on contrastive learning have achieved remarkable performance solely by leveraging few labels in the context of medical image segmentation. Existing methods mainly focus on instance discrimination and invariant mapping. However, they face three common pitfalls: (1) tailness: medical image data usually follows an implicit long-tail class distribution. Blindly leveraging all pixels in training hence can lead to the data imbalance issues, and cause deteriorated performance; (2) consistency: it remains unclear whether a segmentation model has learned meaningful and yet consistent anatomical features due to the intra-class variations between different anatomical features; and (3) diversity: the intra-slice correlations within the entire dataset have received significantly less attention. This motivates us to seek a principled approach for strategically making use of the dataset itself to discover similar yet distinct samples from different anatomical views. In this paper, we introduce a novel semi-supervised 2D medical image segmentation framework termed Mine yOur owN Anatomy (MONA), and make three contributions. First, prior work argues that every pixel equally matters to the model training; we observe empirically that this alone is unlikely to define meaningful anatomical features, mainly due to lacking the supervision signal. We show two simple solutions towards learning invariances - through the use of stronger data augmentations and nearest neighbors. Second, we construct a set of objectives that encourage the model to be capable of decomposing medical images into a collection of anatomical features in an unsupervised manner. Lastly, we both empirically and theoretically, demonstrate the efficacy of our MONA on three benchmark datasets with different labeled settings, achieving new state-of-the-art under different labeled semi-supervised settings.

Motivation & Objective

• Address the challenge of medical image segmentation with very few labels and abundant unlabeled data.
• Investigate whether contrastive learning can learn meaningful, consistent anatomical features under long-tail class distributions.
• Develop a framework that promotes tail-aware sampling, anatomical consistency, and view diversity to improve segmentation.

Proposed method

• Introduce MONA with two stages: relational semi-supervised pre-training and anatomical contrastive reconstruction fine-tuning.
• Use pixel-level anchors and class-averaged positives with a memory bank to handle long-tail class distributions (tailness).
• Incorporate equivariance (consistency) loss to enforce robustness to image transformations.
• Employ a nearest-neighbors loss in output space to promote semantic diversity without excessive memory use.
• Combine global/local instance discrimination with pseudo-labels and supervised segmentation loss during pre-training.
• Provide theoretical and empirical analysis showing label efficiency and complementarity with other CL frameworks.

Figure 1: Examples of three benchmarks ( i.e. , ACDC, LiTS, MMWHS) with long-tail class distributions.

Experimental results

Research questions

• RQ1 Can contrastive learning be effectively adapted for medical image segmentation under extremely limited labels?
• RQ2 How can long-tail class distributions, intra-class variation, and inter-class diversity be addressed in a pixel-level, anatomy-aware framework?
• RQ3 Do tailness, consistency, and diversity jointly improve segmentation performance across datasets and label regimes?

Key findings

• MONA outperforms state-of-the-art SSL methods across ACDC, LiTS, and MMWHS benchmarks under 1%, 5%, and 10% label regimes.
• On ACDC with 1% labels, MONA achieves 82.6% DSC and 2.03 mm ASD, a substantial improvement over prior methods.
• On LiTS with 1% labels, MONA achieves 64.1% DSC and 20.9 mm ASD, and maintains superiority across other label ratios.
• MONA consistently yields sharper anatomical boundaries and better boundary adherence in qualitative visualizations.
• Ablation studies show that combining tailness, consistency, and diversity yields the best performance, with augmentation strategies further enhancing robustness.
• The authors provide theoretical justification for MONA’s effectiveness in label-efficient learning and demonstrate compatibility with multiple contrastive-learning frameworks.

Figure 2: Overview of the MONA framework including two stages: (1) relational semi-supervised pre-training, (2) our proposed anatomical contrastive reconstruction fine-tuning. Note that U and L denote unlabeled and labeled data.