PaperlayerView exact revision on arXiv

Source-linked AI summary

SimCVD: Simple Contrastive Voxel-Wise Representation Distillation for Semi-Supervised Medical Image Segmentation

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

arXiv:2108.06227v4cs.CVcs.AIcs.LG

TL;DR

Medical image segmentation is constrained by scarce manual annotations, and existing semi-supervised methods can lack robustness and explicit geometric and semantic modeling. SimCVD addresses these gaps with dropout-based contrastive learning of signed distance maps and pair-wise structural distillation. On the LA dataset, it reaches 90.85% Dice with 20% labels and 89.03% with 10% labels, improving on previous best results.

  • Problem

    Scarce manual annotations and limited robustness and geometric modeling in semi-supervised segmentation motivate better voxel-wise representation learning.

  • Method

    SimCVD uses two dropout views to contrast signed distance map representations and distills pair-wise structural similarities within a mean-teacher framework.

  • Results

    90.85% average Dice with 20% labels and 89.03% with 10% labels on the LA dataset, improving previous best results.

  • Takeaways & Limitations

    SimCVD achieves new state-of-the-art results on two benchmarks in an extreme few-annotation setting while using an end-to-end framework.

Abstract

from arXiv · show

Automated segmentation in medical image analysis is a challenging task that requires a large amount of manually labeled data. However, most existing learning-based approaches usually suffer from limited manually annotated medical data, which poses a major practical problem for accurate and robust medical image segmentation. In addition, most existing semi-supervised approaches are usually not robust compared with the supervised counterparts, and also lack explicit modeling of geometric structure and semantic information, both of which limit the segmentation accuracy. In this work, we present SimCVD, a simple contrastive distillation framework that significantly advances state-of-the-art voxel-wise representation learning. We first describe an unsupervised training strategy, which takes two views of an input volume and predicts their signed distance maps of object boundaries in a contrastive objective, with only two independent dropout as mask. This simple approach works surprisingly well, performing on the same level as previous fully supervised methods with much less labeled data. We hypothesize that dropout can be viewed as a minimal form of data augmentation and makes the network robust to representation collapse. Then, we propose to perform structural distillation by distilling pair-wise similarities. We evaluate SimCVD on two popular datasets: the Left Atrial Segmentation Challenge (LA) and the NIH pancreas CT dataset. The results on the LA dataset demonstrate that, in two types of labeled ratios (i.e., 20% and 10%), SimCVD achieves an average Dice score of 90.85% and 89.03% respectively, a 0.91% and 2.22% improvement compared to previous best results. Our method can be trained in an end-to-end fashion, showing the promise of utilizing SimCVD as a general framework for downstream tasks, such as medical image synthesis, enhancement, and registration.

I. INTRODUCTION

SimCVD addresses scarce annotation, robustness, and geometric-information challenges in semi-supervised medical image segmentation through contrastive and structured distillation. It combines boundary-aware representations, signed distance maps, pair-wise structural knowledge, and dropout-based views.

  • Limited labeled data creates practical challenges for accurate medical image segmentation despite strong deep-learning performance.
  • Prior semi-supervised methods reduce labeling costs but can remain less robust than fully supervised models and lose geometric information such as object boundaries.
  • SimCVD learns voxel-wise representations from unlabeled data using two dropout views and contrastive learning over signed distance map representations.
  • The framework jointly predicts segmentation maps and signed distance maps, using distance maps as global shape constraints on labeled data.
  • SimCVD further distills pair-wise structural knowledge from a mean-teacher model while avoiding co-adapting and over-fitting.

II. RELATED WORK

The related work motivates semi-supervised segmentation as a way to exploit abundant unlabeled data when annotations are limited. SimCVD is positioned as an end-to-end voxel-wise contrastive distillation framework for boundary-aware and structured representation learning.

  • Semi-supervised medical segmentation methods use limited labeled data with abundant unlabeled data to improve accuracy and reduce labeling cost.
  • Existing approaches include adversarial learning, knowledge distillation, self-supervised learning, and contrastive learning.
  • The framework matches relational similarity among voxel pairs from mean-teacher encoded feature maps to improve voxel-wise representations.
  • SimCVD targets extremely few annotations by learning boundary-aware voxel-wise representations through contrastive distillation.
  • Its base network jointly performs classification and signed-distance-map regression, followed by structured distillation.

B. Task Formulation

SimCVD is formulated for a dataset with far fewer labeled than unlabeled 3D images. Teacher and student networks process perturbed views of each volume to produce corresponding prediction maps.

  • Training data contain N labeled and M unlabeled 3D images, with N ≪ M.
  • Each labeled volume includes a binary voxelwise ground-truth label and a signed distance map measuring each voxel’s distance to the object boundary.
  • SimCVD uses a mean-teacher network and a student network, with teacher parameters updated through an exponential moving average of student parameters.
  • The two networks receive perturbed versions of the same volume, including affine transformations and random crops.
  • The teacher and student produce corresponding confidence-score probability maps from their respective augmented inputs.

C. Base Architecture

The base architecture uses V-Net with encoder-decoder networks and jointly predicts segmentation probabilities and signed distance maps. The regression branch supplies geometric structure and continuity-related constraints.

  • The V-Net backbone comprises encoder-decoder networks whose outputs include probability maps in [0, 1] and signed distance maps in [−1, 1].
  • Multi-task learning combines voxelwise classification with signed distance map regression.
  • The classification branch predicts a probability map, while the regression branch predicts the object’s signed distance map.
  • The regression branch uses a hyperbolic tangent function and is intended to encode geometric structure and implicitly enforce continuity and smoothness.
  • Training on labeled data uses segmentation loss and mean squared error, with the signed distance map loss imposed as a geometric constraint.

D. Boundary-aware Contrastive Distillation

SimCVD combines boundary-aware contrastive distillation with pair-wise structural distillation to learn robust voxel-wise representations from unlabeled medical images. Its objectives enforce geometric consistency, spatial relationships, and prediction consistency.

  • Motivation: SimCVD addresses weak geometric modeling and limited robustness by distilling boundary-aware and structural representations.Prior methods often distill only in prediction space, while SimCVD targets geometric information and generalization.
  • Boundary-aware contrastive learning: SimCVD creates boundary-aware features by adding each input volume to its teacher and student signed distance maps.This fusion combines distance information with image intensity information before projection.
  • Boundary-aware contrastive learning: Two independent dropout masks define positive views, while slices from different locations or inputs provide negative samples for the InfoNCE objective.The same slice across the two views is treated as positive; other slices are negatives.
  • Pair-wise distillation: Pair-wise distillation transfers spatial relationships between encoder features using cosine similarities across voxel pairs.The local structural objective complements global boundary-aware representations for dense segmentation.
  • Overall objective: The unsupervised objective combines boundary-aware contrastive, pair-wise distillation, and consistency losses for unlabeled data.These terms are combined with the supervised loss in the overall training objective, with λ, β, and γ controlling their weights.

A. Dataset and Pre-processing

The evaluation uses the LA MR dataset, consisting of 100 annotated 3D scans, with 80 scans for training and 20 for evaluation.

  • LA dataset: The LA dataset contains 100 3D gadolinium-enhanced MR scans with expert annotations.The scans have isotropic resolution of 0.625 × 0.625 × 0.625 mm3.
  • Data split: The experimental split uses 80 scans for training and 20 scans for evaluation.The study follows the experimental setting and preprocessing methods of prior work.

B. Implementation Details

All evaluated methods use a common PyTorch training setup with fixed optimization, augmentation, and evaluation procedures.

  • Training setup: All methods are trained for 6000 iterations on an NVIDIA 1080Ti GPU with batch size 4.The implementation uses PyTorch and SGD with momentum-related settings described in the study.
  • Data augmentation: Data augmentation consists of random rotation, flipping, and cropping.These are the standard augmentation techniques used for all evaluated methods.
  • Evaluation: No post-processing step is used for fairness.The training schedule includes parameter ramp-up during optimization.
  • Evaluation: Evaluation uses Dice, Jaccard, 95HD, and ASD with sliding-window inference.The stride is 18×18×4 for LA and 16×16×16 for pancreas.

V. RESULTS

SimCVD improves semi-supervised segmentation across the LA and pancreas datasets under limited annotation, with gains reported over previous methods and strong visual quality.

  • LA results: 90.85% Dice and 83.80% Jaccard are achieved with 20% labeled LA data, improving previous best results of 89.94% and 81.82%.The Dice result is comparable to the fully supervised baseline.
  • LA results: 89.03% Dice is achieved with 10% labeled LA data, improving the previous best result of 87.49%.The corresponding reported Jaccard, ASD, and 95HD values are 80.34%, 2.59, and 8.34.
  • Ablation study: Ablation experiments compare removing SDM, adaptive max pooling, contrastive loss, pair-wise distillation, and SDM loss against the baseline and full SimCVD.The supplied ablation caption identifies the compared components but does not report their numerical outcomes.
  • Pancreas results: Up to 6.72% absolute Dice improvement is obtained on the pancreas CT dataset over previous methods.The model consistently outperforms all previous methods in the reported pancreas experiments.

VI. ABLATION STUDY

The ablation studies examine boundary-aware distillation, loss components, dropout, augmentation, and pooling choices in SimCVD. Results support signed-distance-map supervision, contrastive and pair-wise objectives, dropout masks, limited augmentation, and adaptive mean pooling.

  • Ablation on Model Component: Removing signed distance maps reduced Dice, Jaccard, ASD, and 95HD by 0.79%, 1.27%, 1.6, and 3.09 absolute differences, respectively.The comparison uses the 10% labeled setting and supports boundary-aware representations as a useful prior for segmentation accuracy.
  • Ablation on Model Component: SimCVD outperformed the best previous method by +0.71% Dice, +1.08% Jaccard, +1.20 ASD, and +6.35 95HD absolute differences.These results are reported for the 10% labeled setting.
  • Ablation on Loss Formulation: Removing Lcontrast decreased Dice by 3.90%, while removing Lpd decreased Dice and Jaccard by 0.92% and 1.45%, respectively.The ablations support complementary roles for contrastive learning and pair-wise distillation.
  • How to Interpret Dropout?: Dropout with p = 0.1 outperformed no dropout by 1.34% Dice, 2.11% Jaccard, 1.71 ASD, and 2.69 95HD absolute differences.The authors interpret dropout as a minimal form of data augmentation producing two views through different masks.
  • Effect of Augmentation Techniques: Adding local shuffle pixel, non-linear transformation, inpainting, or out-painting did not outperform the basic dropout mask and could hurt prediction performance.The study suggests that additional augmentation may introduce noise associated with representation collapse.
  • Effect of Pooling Size: Larger pooling sizes consistently improved performance, but increasing the size to 256 provided no further improvement; the implementation used pooling size 128.Adaptive mean pooling improved results relative to adaptive max pooling.

VII. CONCLUSION

SimCVD is a contrastive distillation framework for voxel-wise representation learning in medical segmentation. It uses signed distance-map views, boundary-aware representations, pair-wise structural distillation, and dropout, achieving state-of-the-art results on two benchmarks in an extreme few-annotation setting while remaining scoped for future multi-class extension.

  • VII. CONCLUSION: SimCVD predicts signed distance maps of object boundaries from two views of an input volume using two different dropout masks in an unsupervised contrastive objective.The framework is presented as a contrastive distillation approach for voxel-wise representation learning.
  • VII. CONCLUSION: SimCVD combines boundary-aware representations, dropout as minimal data augmentation, and pair-wise similarity distillation.The conclusion identifies these components as important to the approach and its performance improvements.
  • VII. CONCLUSION: SimCVD obtained new state-of-the-art results on two benchmarks in an extreme few-annotation setting.The conclusion reports this as the overall experimental outcome.
  • VII. CONCLUSION: The authors plan to extend the unsupervised training framework to multi-class medical image segmentation tasks.This identifies a future scope rather than a demonstrated result.
Loading 2108.06227v4…