Source-linked AI summary
Semi-supervised Medical Image Segmentation through Dual-task Consistency
Xiangde Luo, Jieneng Chen, Tao Song, Yinan Chen, Guotai Wang, Shaoting Zhang
TL;DR
Medical-image segmentation needs costly expert annotations, while existing SSL methods mainly regularize predictions through network or data perturbations. This paper introduces task-level consistency between pixel-wise and level-set predictions using a differentiable transform, and reports state-of-the-art results on two 3D medical-image datasets. The method also requires less training time and computational cost than existing methods.
Problem
Medical-image annotations are expensive, and existing SSL methods mainly construct consistency through network- or data-level perturbations rather than explicit task-level regularization.
Method
A dual-task network predicts pixel-wise segmentation and a geometry-aware level-set representation, then applies a differentiable transform to enforce consistency between their maps on labeled and unlabeled data.
Results
The framework achieves state-of-the-art results on two 3D medical-image datasets and outperforms state-of-the-art semi-supervised segmentation frameworks on several clinical datasets.
Takeaways & Limitations
Task-level consistency provides a semi-supervised segmentation framework that captures geometric information while using unlabeled data without requiring multiple inference passes.
Abstract
from arXiv · showhide
Deep learning-based semi-supervised learning (SSL) algorithms have led to promising results in medical images segmentation and can alleviate doctors' expensive annotations by leveraging unlabeled data. However, most of the existing SSL algorithms in literature tend to regularize the model training by perturbing networks and/or data. Observing that multi/dual-task learning attends to various levels of information which have inherent prediction perturbation, we ask the question in this work: can we explicitly build task-level regularization rather than implicitly constructing networks- and/or data-level perturbation-and-transformation for SSL? To answer this question, we propose a novel dual-task-consistency semi-supervised framework for the first time. Concretely, we use a dual-task deep network that jointly predicts a pixel-wise segmentation map and a geometry-aware level set representation of the target. The level set representation is converted to an approximated segmentation map through a differentiable task transform layer. Simultaneously, we introduce a dual-task consistency regularization between the level set-derived segmentation maps and directly predicted segmentation maps for both labeled and unlabeled data. Extensive experiments on two public datasets show that our method can largely improve the performance by incorporating the unlabeled data. Meanwhile, our framework outperforms the state-of-the-art semi-supervised medical image segmentation methods. Code is available at: https://github.com/Luoxd1996/DTC
Introduction
The paper addresses costly medical-image annotations by introducing task-level consistency between pixel-wise segmentation and geometry-aware level-set predictions. Its dual-task framework uses differentiable transformation to exploit labeled and unlabeled data, improving fully supervised and semi-supervised performance.
- Motivation: Expensive expert annotations motivate semi-supervised segmentation methods that learn from limited labeled and larger unlabeled datasets.Accurate annotations are particularly costly for 3D volumetric medical images.
- Motivation: Existing SSL methods mainly enforce data-level consistency through perturbations, transformations, or adversarially matched prediction distributions.These approaches regularize predictions across altered inputs or across the unlabeled dataset.
- Method: The proposed method explicitly enforces task-level consistency after mapping different task outputs into a shared prediction space.This addresses the question of whether task-level regularization can replace implicit network- or data-level perturbation.
- Method: A dual-task network jointly predicts a pixel-wise classification map and a global level-set representation, with a differentiable transform converting the latter into a segmentation probability map.The level-set branch incorporates global shape and geometric information, while the transform enables consistency between the two task outputs.
- Results: The dual-task-consistency regularization outperforms separate and joint dual-task supervision in the fully supervised setting.This is reported as one of the paper’s principal contributions.
- Results: The framework outperforms state-of-the-art semi-supervised medical image segmentation methods on several clinical datasets while requiring less training time and computational cost.The authors also report direct applicability to semi-supervised medical image segmentation scenarios.
Related Works
Prior medical-image segmentation methods include handcrafted-feature models and deep-learning SSL approaches. The paper positions dual-task consistency as a task-level alternative to input perturbation and data-level regularization, while emphasizing its single-inference design.
- Traditional Methods: Traditional semi-supervised medical-image segmentation methods rely on handcrafted prior or clustering features whose effectiveness depends on representation capacity.Prior-based models require organ-specific prior information.
- Deep-Learning Methods: Deep-learning approaches include iterative pseudo-label refinement with Conditional Random Fields for unannotated images.The cited iterative framework predicts pseudo labels and refines them using a CRF.
- Consistency Regularization: Consistency regularization methods learn from unlabeled data by matching predictions across stochastic transformations, perturbations, or teacher-student passes.These methods focus on consistency under altered inputs rather than consistency between task branches.
- Task-Level Consistency: The proposed framework differs by enforcing consistency between pixel-level classification and level-set regression tasks within one network.A differentiable transform maps the level-set output into the pixel-wise prediction space.
- Task-Level Consistency: Unlike methods requiring multiple forward passes, the proposed approach needs only one inference pass to calculate its task-consistency loss.The paper contrasts this design with perturbation-based methods that perform two or more forward passes.
Methods
The framework uses dual-task predictions—a pixel-wise segmentation map and a level-set representation—and enforces consistency after differentiably mapping the level-set output into segmentation space. Supervised task losses apply to labeled data, while dual-task consistency uses both labeled and unlabeled data.
- Dual-task network: The model jointly predicts a pixel-wise classification map and a global-level level-set function to represent segmentation.The two branches capture pixel-level and geometric information.
- Task transform: A smooth Heaviside-based transform converts the level-set function into a differentiable pixel-wise probability map for consistency training.The approximate inverse transform is implemented as a sigmoid with its input multiplied by k.
- Consistency regularization: The dual-task consistency loss minimizes the difference between the direct segmentation prediction and the transformed level-set prediction.This explicitly regularizes the two task outputs in a shared prediction space.
- Supervised training: For labeled images, level-set ground truth is automatically generated from the segmentation mask using the task transform.The level-set branch is trained with an L2 loss against the transformed ground truth.
- Semi-supervised training: The segmentation and level-set supervised losses use labeled data, whereas the consistency loss uses both labeled and unlabeled data.A time-dependent Gaussian warming-up function controls the balance between supervised and consistency losses.
Experiments and Results
Experiments evaluate the framework on pancreas CT and left atrium MRI using predominantly 20% labeled data, task variants, and comparisons with semi-supervised baselines. The method improves data utilization, achieves strong segmentation performance, and reduces training cost relative to several alternatives.
- Datasets and Experimental Setup: Experiments use 100 left atrium MRI scans and 82 pancreas CT images, with results reported under 20% labeled and 80% unlabeled data.The left atrium dataset uses 80 scans for training and 20 for validation; the pancreas dataset uses 62 images for training and 20 for testing.
- Effectiveness of Dual-task-Consistency: The dual-task-consistency strategy improves dual-task VNet performance when trained with either 12 or 62 labeled pancreas scans.The comparison includes Seg, LSF, Seg + LSF, and Seg + LSF + DTC training strategies.
- Data Utilization Efficiency: The semi-supervised approach consistently outperforms supervised approaches across different labeled-data settings, with a narrower gap as more labels become available.With few labeled images, it can also obtain better segmentation results than the fully supervised method.
- Comparison with Other Semi-supervised Methods: On Pancreas CT, the framework outperforms the evaluated state-of-the-art semi-supervised methods on all evaluation metrics.The comparison includes DAN, Entropy Mini, CCT, MT, UA-MT, and SASSNet; the framework does not require multiple inference or iterative-update schemes.
- Comparison with Other Semi-supervised Methods: On Left Atrium MRI, the method achieves the best accuracy across all reported evaluation metrics, especially ASD and 95HD.Visualizations are described as having higher overlap with ground truth, fewer false positives, and more preserved details.
- Computational Cost: The framework requires less training time than MT, DAN, UA-MT, CCT, and SASSNet while using a simple network with fewer parameters.Compared with Entropy Mini and the fully supervised baseline, it achieves better accuracy with comparable computational cost.
Discussion and Conclusion
The paper presents a task-level consistency framework for semi-supervised medical image segmentation, combining pixel-wise classification with level set representations. It achieves state-of-the-art results on two 3D medical image datasets and can extend to additional tasks when differentiable transforms exist.
- The framework jointly predicts pixel-level classification maps and level set representations that capture global shape and geometric information.
- A task-transform layer enforces consistency between classification-map predictions and level set predictions during semi-supervised training.
- The method achieves state-of-the-art results on left atrial MR and pancreas CT datasets.
- The framework can incorporate additional tasks such as edge extraction and key-point estimation when differentiable transforms exist between tasks.