PaperlayerView exact revision on arXiv

Source-linked AI summary

Rethinking Semi-Supervised Medical Image Segmentation: A Variance-Reduction Perspective

Chenyu You, Weicheng Dai, Yifei Min, Fenglin Liu, David A. Clifton, S Kevin Zhou, Lawrence Hamilton Staib, James S Duncan

arXiv:2302.01735v5cs.CVcs.AIcs.LGeess.IV

TL;DR

Extremely limited labels and long-tailed medical data make robust segmentation challenging, while sampling in contrastive learning can introduce variance. ARCO uses stratified group sampling with variance-reduction estimation; the paper reports improved variance reduction and consistent performance improvements across eight benchmarks.

  • Problem

    Medical segmentation needs robustness with extremely limited labels in long-tailed data, where sampling can produce large feature variations and collapse concerns.

    Source: §1 · pp. 1–2

  • Method

    ARCO is a semi-supervised contrastive learning framework that uses stratified group sampling and variance-reduction estimation to select informative pixels.

    Source: §1 · p. 2

  • Results

    ARCO consistently improves overall segmentation quality across five medical and three semantic datasets and across label ratios.

    Source: §1 · pp. 2–3

  • Takeaways & Limitations

    The analysis presents Stratified Group and Stratified-Antithetic Group sampling as practical means to improve variance reduction and label efficiency.

    Source: §1 · p. 2, §5 · p. 10

Abstract

from arXiv · show

For medical image segmentation, contrastive learning is the dominant practice to improve the quality of visual representations by contrasting semantically similar and dissimilar pairs of samples. This is enabled by the observation that without accessing ground truth labels, negative examples with truly dissimilar anatomical features, if sampled, can significantly improve the performance. In reality, however, these samples may come from similar anatomical regions and the models may struggle to distinguish the minority tail-class samples, making the tail classes more prone to misclassification, both of which typically lead to model collapse. In this paper, we propose ARCO, a semi-supervised contrastive learning (CL) framework with stratified group theory for medical image segmentation. In particular, we first propose building ARCO through the concept of variance-reduced estimation and show that certain variance-reduction techniques are particularly beneficial in pixel/voxel-level segmentation tasks with extremely limited labels. Furthermore, we theoretically prove these sampling techniques are universal in variance reduction. Finally, we experimentally validate our approaches on eight benchmarks, i.e., five 2D/3D medical and three semantic segmentation datasets, with different label settings, and our methods consistently outperform state-of-the-art semi-supervised methods. Additionally, we augment the CL frameworks with these sampling techniques and demonstrate significant gains over previous methods. We believe our work is an important step towards semi-supervised medical image segmentation by quantifying the limitation of current self-supervision objectives for accomplishing such challenging safety-critical tasks.

1 Introduction

The introduction frames reliable medical segmentation as requiring robustness and label efficiency under extremely limited labels and long-tailed data. It presents ARCO’s stratified sampling and variance-reduction approach and reports theoretical and empirical benefits.

1 Introduction

  • Medical segmentation models are expected to remain robust and converge quickly despite extremely limited labels in long-tailed data.The introduction notes that supervised approaches rely on large labeled datasets, requiring annotation and clinical expertise.Source: §1 · p. 1
  • Random pixel sampling in non-smooth feature spaces can cause large feature variation, motivating informative-pixel sampling as a route to variance reduction.The introduction connects robustness concerns to constant-feature or lower-dimensional collapse and proposes sampling informative pixels in semi-supervised contrastive learning.Source: §1 · p. 2
  • ARCO partitions images into equal-sized class grids and samples semantically close pixels within each grid, with minimal additional memory.The framework proposes Stratified Group and Stratified-Antithetic Group sampling to select informative pixels.Source: §1 · p. 2
  • ARCO consistently improves segmentation quality across all label ratios and five medical and three semantic datasets.The introduction also reports competitive improvements across all eight benchmarks for models trained with different labeled ratios.Source: §1 · pp. 2–3
  • Theoretical analysis reports improved variance reduction and label efficiency for ARCO’s gradient estimator.The introduction also describes an optimization guarantee for the variance-reduction analysis.Source: §1 · pp. 2–3
  • ARCO applies variance-reduction techniques to semi-supervised contrastive learning for medical image segmentation, targeting model robustness and label efficiency.The paper presents this as a way to address variance introduced by common sampling practice.Source: §1 · p. 2

2 Related work

Related work spans medical segmentation architectures and optimization, semi-supervised learning, and contrastive representation learning. The paper positions variance-reduced estimation as a way to address instability and class imbalance in dense contrastive segmentation.

2 Related work

  • Medical segmentation work commonly uses FCNs or UNets, while architecture and optimization strategies address representation quality, imbalance, and uncertain pixels.The paper instead focuses on a model-agnostic, label-efficiency framework using massive unlabeled data and extremely limited labels.Source: §2 · p. 3
  • Semi-supervised medical segmentation includes consistency regularization and self-training, but the cited approaches are described as failing to prevent collapse due to class imbalance.The paper focuses on extremely limited labels and long-tailed anatomical classes, with additional supervision on hard pixels.Source: §2 · p. 3
  • Contrastive learning pulls similar representations together and pushes dissimilar ones apart, while global representations may be insufficient for dense prediction and dense methods may face class imbalance.Recent medical segmentation work mitigates imbalance but is described as lacking stability; this paper focuses on variance-reduced estimation and reports a theoretical robustness guarantee.Source: §2 · pp. 3–4
  • The Figure 1 pipeline combines relational semi-supervised pre-training with anatomical contrastive reconstruction fine-tuning.Unlabeled-data branches use EMA-teacher outputs for instance discrimination during pre-training and representation maps and pseudo-labels during fine-tuning.Source: Figure 1 · p. 4

3 Methodology

This section motivates ARCO as a response to unreliable pixel-level contrastive sampling and develops stratified sampling as a variance-reduction approach. It presents theoretical and optimization arguments for the sampling methods.

3.1 Preliminaries and setup

  • MONA learns representations by bringing semantically close neighbors together and pushing other non-neighbors apart, providing the contrastive-learning context for ARCO.The section describes ARCO as a simplification of the MONA pipeline without additional complex augmentation strategies.Source: §3.1 · pp. 4–5

3 Methodology

  • ARCO uses Stratified Group (SG) and Stratified-Antithetic Group (SAG) sampling to select informative pixels within a two-phase contrastive training workflow.The phases are relational semi-supervised pre-training and anatomical contrastive fine-tuning.Source: §3 · p. 4, §3.1 · p. 5, §3.2 · p. 5

3.2 Motivation and Challenges

  • Naïve pixel sampling can have high variance, miss semantically similar pixels, and leave existing sampling rules poorly understood in memory-constrained clinical settings.ARCO instead partitions pixels into same-size grids and samples semantically close pixels within each grid; SAG adds reflection sampling.Source: §3.2 · p. 5, Figure 2 · p. 5

3.3 Stratified Group Sampling

  • With proportional group sizes, SG is unbiased and has variance no greater than naïve sampling, while SAG's variance is at most twice SG's.SG has strictly lower variance when the pixel groups do not all share the same mean over h(x; p), a condition described as almost-sure in medical images.Source: §3.3 · pp. 6–7
  • SG's variance-reduced gradient estimates support faster convergence in theory, and experiments report faster loss decay with smaller error bars than other methods.For general non-convex loss functions, the stated convergence result is to some local minimum; the convergence rate depends on gradient-estimate standard deviation.Source: §3.3 · p. 7

4 Experiments

Across medical and semantic segmentation benchmarks, ARCO improves segmentation performance with limited labels. Ablations support the roles of its losses, augmentations, and sampling strategies, with SG showing faster and more stable convergence.

4.1 Main Results

  • Across five medical benchmarks and 1%, 5%, and 10% label settings, ARCO methods outperform SSL comparators; ARCO-SG's average Dice gains over MONA on ACDC are 2.9%, 1.8%, and 1.7%.The evaluations also cover three semantic segmentation datasets.Source: §4.1 · pp. 7–8, Table 1 · p. 8
  • At 1% labels, ARCO-SG improves average Dice over MONA by 2.2% on LA, 3.3% on LiTS, and 5.1% on seven-class MMWHS.The reported average ASD reductions are 6.64 on LA, 4.7 on LiTS, and 2.26 on MMWHS.Source: §4.1 · p. 9
  • ARCO's predictions show clearer anatomical boundaries; on LiTS, ARCO-SG and ARCO-SAG retain sharp boundaries around ambiguous tail-class samples such as small lesions.On MMWHS, the approaches also produce smoother, more shape-consistent segmentations.Source: Figure 4 · p. 8, §4.1 · p. 9

4.2 Ablation Studies

  • Ablations on ACDC with 1% labels find that Lcontrast and Lnn boost performance, and SG and SAG provide greater robustness gains than naïve sampling.Each tested loss function also yields performance gains.Source: §4.2 · p. 9
  • On ACDC with a 10% label ratio, SG sampling shows faster convergence and greater training stability than NS and SAG; SAG slightly outperforms NS.Figure 5 presents trajectories over three independent trials, with solid lines and shaded areas denoting mean and variance.Source: Figure 5 · p. 10, §4.2 · p. 10
  • Each tested augmentation improves performance; using all three defines the comparison's upper bound, while using none defines its lower bound.The ablations are on ACDC with 1% labeled data.Source: §4.2 · pp. 9–10

5 Conclusion and Discussion of Broader Impact

The conclusion presents ARCO's stratified sampling methods as practical approaches to variance reduction and improved robustness and label efficiency. It also identifies fairness and privacy as future research challenges.

5 Conclusion and Discussion of Broader Impact

  • ARCO combines stratified group sampling solutions with variance-reduction estimation to improve robustness and label efficiency in medical image segmentation.The conclusion describes the sampling methods as practical means of variance reduction and presents the framework as an adaptable training toolkit.Source: §5 · p. 10
  • The authors identify fairness and privacy challenges in medical image analysis as future research directions.Source: §5 · p. 10

A Variance-Reduction Perspective

  • The appendix includes additional theoretical analysis and studies of generalization across label ratios and frameworks.The proof excerpt states that sampling from each group is independent and introduces a variance expression for the estimate.Source: A Variance-Reduction Perspective · p. 16, §A.1 · pp. 16–17

A.1 Proof of Lemma 3.1

  • The proof excerpt describes SAG as sampling half of each group's pixels and selecting the remainder as reflections of the first half.It considers the variance of a sampled pixel and its reflection.Source: §A.1 · p. 16

A.2 Proof of Theorem 3.2

The proof establishes unbiasedness for the stratified-group estimate through proportional sampling, then relates gradient-estimate variance to convergence behavior under stated assumptions.

A.2 Proof of Theorem 3.2

  • Proportional sampling makes every population pixel equally likely to be selected, establishing that the bHSG estimate is unbiased.The proof uses equal selection probability across pixels to establish the result.Source: §A.2 · p. 17
  • Naïve sampling draws pixels uniformly from all pixels, and its variance is analyzed by conditioning first on the group and then on the pixel.The proof applies the law of total variance to this two-step sampling process.Source: §A.2 · p. 17

A.3 Further Details on Convergence

  • Proposition 1 gives the average-gradient-norm convergence rate a fast 1/T term and a slower term that depends on gradient-estimate variance.Under Assumptions A.1 and A.2, the variance-dependent term is dominant; the authors state variance reduction can therefore accelerate convergence and improve stability.Source: §A.3 · p. 18
  • The authors describe SG and SAG as universal variance-reduction techniques for pixel- or voxel-level sampling under mild assumptions A.1 and A.2.They connect the result to 2D/3D image sampling and report faster loss decay and a narrower error bar for SG in Figure 5.Source: §A.3 · p. 18

B Datasets

The experiments use five medical segmentation datasets: three 2D benchmarks and two 3D benchmarks, with label ratios defined by patient counts.

B Datasets

  • Experiments are conducted on five representative datasets in semi-supervised medical image segmentation: 2D benchmarks ACDC, LiTS, and MMWHS, and 3D benchmarks LA and in-house MP-MRI.Source: B Datasets · p. 18
  • ACDC contains 200 3D cardiac cine MRI scans and three classes: left ventricle, myocardium, and right ventricle.The split is 120 training, 40 validation, and 40 test scans; label ratios are based on patients.Source: B Datasets · pp. 18–19
  • LiTS contains 131 contrast-enhanced abdominal CT volumes labeled for liver and tumor, with 100 scans for training and 31 for testing.Its 1%, 5%, and 10% label ratios are patient-based.Source: B Datasets · p. 19
  • MMWHS contains 20 cardiac MRI scans with seven anatomical classes, split into 15 training and five test scans.Its label ratios are patient-based.Source: B Datasets · p. 19
  • The LA benchmark has 100 gadolinium-enhanced MRI scans for one left-atrium class, while in-house MP-MRI has 160 scans for one liver class.LA uses 80 training and 20 test scans for its fixed 5% and 10% splits; MP-MRI uses 100 training, 40 validation, and 20 test scans.Source: B Datasets · p. 19

C Implementation Details

The implementation evaluates ARCO against prior 2D and 3D methods and uses a two-stage pipeline combining relational pretraining with anatomical contrastive fine-tuning. The reported visual and table scopes include LiTS results across three label ratios.

C Implementation Details

  • Experiments use 2D UNet or 3D VNet backbones, an FPN representation head with a separate 512-dimension output layer, and SGD with momentum 0.9, weight decay 10−4, and initial learning rate 0.01.Pretraining uses 100 epochs (batch size 6), while fine-tuning uses 200 epochs (batch size 8). For 2D, augmentations are random rotation, random cropping, and horizontal flipping; for 3D they are random rotation and random flipping, and 3D results are evaluated with DSC and ASD.Source: C Implementation Details · pp. 19–20

D 2D/3D Methods in Comparison

  • The comparison uses 2D UNet and 3D VNet backbones and includes fully supervised, limited-supervised, and SSL methods.The compared SSL methods include GCL, SS-Net, ACTION, and MONA, which use contrastive objectives.Source: D 2D/3D Methods in Comparison · p. 20

E Framework Overview

  • On LiTS, ARCO-SG and ARCO-SAG produce more accurate predictions on anatomical regions and boundaries than the other SSL methods.The accompanying quantitative-comparison table covers DSC and ASD across 1%, 5%, and 10% label ratios, but the supplied caption gives no cell values.Source: Figure 6 · p. 21, Table 5 · p. 21
  • Fine-tuning uses pixel-level contrastive loss to emphasize tail-class hard pixels and a nearest-neighbor loss to exploit inter-instance relationships.The total loss also includes unsupervised cross-entropy and supervised segmentation losses; the nearest-neighbor search uses a FIFO memory bank.Source: E Framework Overview · p. 22
  • ARCO uses relational semi-supervised pretraining followed by anatomical contrastive fine-tuning, with the student trained using an EMA teacher on unlabeled data.Pretraining combines supervised segmentation on labeled data with instance discrimination on augmented and mined unlabeled views.Source: E Framework Overview · pp. 20, 22, F Model Architecture · p. 22

F Model Architecture

  • The model outputs both a segmentation map and a representation map from each 2D or 3D medical image.The representation head is used during training and removed for inference.Source: F Model Architecture · p. 22

G More Experiment Results on LiTS

On LiTS, the authors report better segmentation by ARCO-SG and ARCO-SAG than by the other methods. Visual comparisons on MMWHS and LA also favor these variants for anatomical detail and boundary sharpness.

G More Experiment Results on LiTS

  • ARCO-SG and ARCO-SAG provide better LiTS segmentation results than all other methods.The passage points to both qualitative and quantitative results, but supplies no metric values.Source: G More Experiment Results on LiTS · p. 22
  • The paper uses a simplified MONA setup without additional complicated augmentation strategies because pixel-level contrastive sampling is its focus.Source: G More Experiment Results on LiTS · p. 22
  • The representation head is used only during training and removed during inference.Source: G More Experiment Results on LiTS · p. 22
  • On MMWHS, ARCO-SG and ARCO-SAG produce more accurate predictions on anatomical regions and boundaries than other SSL methods.The figure covers 1%, 5%, and 10% label ratios.Source: Figure 7 · p. 23
  • On LA, ARCO-SG and ARCO-SAG yield more accurate predictions and sharper boundaries than other SSL methods.The figure covers 1%, 5%, and 10% label ratios.Source: Figure 8 · p. 23

H More Experiment Results on MMWHS

The MMWHS experiments compare the baselines and ARCO methods at 1%, 5%, and 10% label ratios, with results reported in Table 6 and Figure 7.

H More Experiment Results on MMWHS

  • The MMWHS evaluation covers 1%, 5%, and 10% label ratios and reports baseline and ARCO results in Table 6 and Figure 7.The authors state that these experiments demonstrate the effectiveness of their models.Source: H More Experiment Results on MMWHS · p. 23

I More 3D Experiment Results on LA [84]

The LA experiments compare the baselines and ARCO methods at 1%, 5%, and 10% label ratios, with results reported in Table 9 and Figure 8.

I More 3D Experiment Results on LA [84]

  • The LA evaluation covers 1%, 5%, and 10% label ratios and reports baseline and ARCO results in Table 9 and Figure 8.The authors state that these experiments demonstrate their models’ robustness.Source: I More 3D Experiment Results on LA [84] · p. 23

J More 3D Experiment Results on MP-MRI

The experiments report ARCO results on MP-MRI and semantic segmentation benchmarks, compare sampling across contrastive-learning frameworks, and visualize loss landscapes.

J More 3D Experiment Results on MP-MRI

  • On MP-MRI, the authors compare baselines and ARCO at 1%, 5%, and 10% label ratios, reporting results in Table 9 and Figure 9.They describe these results as showing the superiority of their models.Source: J More 3D Experiment Results on MP-MRI · p. 23

K More Experiment Results on Semantic Segmentation

  • Across Pascal VOC, Cityscapes, and SUN RGB-D, ARCO-SG and ARCO-SAG improve on ReCo in all semi-supervised settings, with up to 2.4–7.8% relative gains for SG under the fewest-label settings.Visual comparisons report more accurate, sharper boundaries than ReCo across the listed label settings.Source: K More Experiment Results on Semantic… · p. 24

L Generalization across Label Ratios and Frameworks

  • SG and SAG improve segmentation performance across the tested contrastive-learning frameworks, with larger benefits at lower label ratios.The comparison covers MoCov2, kNN-MoCo, SimCLR, BYOL, ISD, and VICReg on ACDC.Source: L Generalization across Label Ratios… · pp. 24–25

M Final Checkpoint Loss Landscapes

N Ablation on Different Training Settings

The supplied ablation evidence examines hyperparameter selection, sampling across contrastive-learning frameworks, and qualitative segmentation results at different label settings.

Ablation on Different Training Settings

Related papers

Loading 2302.01735v5…