Source-linked AI summary
Uncertainty-aware Self-ensembling Model for Semi-supervised 3D Left Atrium Segmentation
Lequan Yu, Shujun Wang, Xiaomeng Li, Chi-Wing Fu, Pheng-Ann Heng
TL;DR
Because reliable annotation of 3D MR images is expensive, the paper develops an uncertainty-aware teacher–student framework for semi-supervised left-atrium segmentation. The method combines supervised and perturbation-consistency losses while filtering teacher targets by uncertainty. It improves segmentation using unlabeled data and outperforms state-of-the-art semi-supervised methods.
Problem
Reliable 3D MR segmentation annotations are expensive and time-consuming, whereas unlabeled data is generally abundant.
Method
UA-MT trains a student from teacher targets using supervised segmentation and perturbation-consistency losses, with uncertainty estimation filtering unreliable targets.
Results
UA-MT outperforms the state-of-the-art semi-supervised methods, while unlabeled data improves segmentation over supervised-only V-Net across labeled-data settings.
Takeaways & Limitations
The framework demonstrates the potential of uncertainty-aware self-ensembling for challenging semi-supervised medical image segmentation.
Abstract
from arXiv · showhide
Training deep convolutional neural networks usually requires a large amount of labeled data. However, it is expensive and time-consuming to annotate data for medical image segmentation tasks. In this paper, we present a novel uncertainty-aware semi-supervised framework for left atrium segmentation from 3D MR images. Our framework can effectively leverage the unlabeled data by encouraging consistent predictions of the same input under different perturbations. Concretely, the framework consists of a student model and a teacher model, and the student model learns from the teacher model by minimizing a segmentation loss and a consistency loss with respect to the targets of the teacher model. We design a novel uncertainty-aware scheme to enable the student model to gradually learn from the meaningful and reliable targets by exploiting the uncertainty information. Experiments show that our method achieves high performance gains by incorporating the unlabeled data. Our method outperforms the state-of-the-art semi-supervised methods, demonstrating the potential of our framework for the challenging semi-supervised problems.
1 Introduction
The paper addresses left-atrium segmentation when expert annotations are costly by leveraging abundant unlabeled 3D MR data. It introduces an uncertainty-aware teacher–student framework that promotes perturbation-consistent predictions while filtering unreliable targets.
- Expert annotation of 3D MR images is expensive and tedious, motivating semi-supervised left-atrium segmentation with limited labeled and abundant unlabeled data.
- Prior semi-supervised approaches used self-training, adversarial learning, and self-ensembling to exploit unlabeled medical images.
- UA-MT trains a student from a teacher using supervised segmentation loss on labeled data and consistency loss on teacher targets for all inputs.
- The teacher estimates prediction uncertainty with Monte Carlo sampling, allowing the student to learn progressively from meaningful and reliable targets.
2 Method
UA-MT combines supervised and consistency training with teacher-generated targets whose uncertainty determines which predictions guide the student. The teacher uses self-ensembling and Monte Carlo Dropout, while an uncertainty-aware voxel-level loss emphasizes reliable targets during training.
- 2.1 Semi-supervised Segmentation: UA-MT trains a student with supervised loss on labeled data and uncertainty-guided consistency loss using teacher targets.The teacher generates targets for both labeled and unlabeled inputs, while uncertainty guides the student toward reliable predictions.
- 2.1 Semi-supervised Segmentation: The teacher’s weights are updated as an exponential moving average of the student’s weights, ensembling information across training steps.The EMA decay controls the teacher’s updating rate.
- Uncertainty Estimation: Monte Carlo Dropout estimates voxel-level uncertainty from T stochastic teacher passes with random dropout and Gaussian input noise.Predictive entropy is used to approximate uncertainty, and the uncertainty of a volume is represented voxel by voxel.
- Uncertainty-Aware Consistency Loss: The uncertainty-aware consistency loss applies voxel-level MSE only to predictions whose uncertainty falls below a threshold.An indicator function selects the most certain teacher targets for student learning.
- 2.3 Technique Details: Training progressively moves from relatively certain to more uncertain cases by filtering less data as training continues.The implementation uses V-Net, dropout layers for uncertainty estimation, EMA decay α = 0.99, and a Gaussian warm-up for the consistency weight.
3 Experiments and Results
Experiments on 3D GE-MRI left-atrium segmentation show that UA-MT leverages limited labels and abundant unlabeled scans effectively, outperforming supervised and competing semi-supervised methods.
- Dataset and setup: The Atrial Segmentation Challenge dataset provides 100 isotropic GE-MRI scans, split into 80 training and 20 evaluation scans.Training scans were cropped around the heart and normalized to zero mean and unit variance.
- Dataset and setup: The main experiment uses 16 labeled and 64 unlabeled training scans to compare Bayesian V-Net with UA-MT.Evaluation uses Dice, Jaccard, average surface distance, and 95% Hausdorff Distance.
- Quantitative results: 2.85% Dice and 4.15% Jaccard improvements over Bayesian V-Net demonstrate the benefit of incorporating unlabeled data.Bayesian V-Net reaches 86.03% average Dice and 76.06% Jaccard with labeled data alone.
- Ablation and reference comparisons: UA-MT approaches fully supervised performance and its gains mainly arise from unlabeled data, because UA-MT-UN performs very similarly while applying consistency loss only to unlabeled data.The fully supervised V-Net trained with all 80 scans serves as the upper-line reference.
- Comparison with competing methods: UA-MT achieves the best performance among compared state-of-the-art semi-supervised methods, although its ASD is comparable with ASDNet.The comparison uses the same Bayesian V-Net backbone for competing methods.
- Robustness and qualitative analysis: UA-MT consistently improves supervised-only Bayesian V-Net with both 10% and 30% labeled scans, while visualizations show higher overlap, fewer false positives, and high uncertainty near boundaries.The labeled subsets contain 8 and 24 scans, respectively.
4 Conclusion
The paper concludes that UA-MT is an uncertainty-aware semi-supervised method for 3D left-atrium segmentation that exploits unlabeled data through perturbation consistency and uncertainty-guided targets.
- Conclusion: UA-MT encourages consistent segmentations under different perturbations to use unlabeled 3D MR images.The conclusion identifies this as the method’s central semi-supervised mechanism.
- Conclusion: Model uncertainty is used to improve teacher-target quality, and comparisons with other semi-supervised methods support the method’s effectiveness.Future work includes evaluating alternative uncertainty-estimation methods and applying the framework to other medical-image segmentation problems.