[Paper Review] Transformation Consistent Self-ensembling Model for Semi-supervised Medical Image Segmentation

Xiaomeng Li, Lequan Yu+4 more|arXiv (Cornell University)|Feb 28, 2019

Retinal Imaging and AnalysisMedicine66 references33 citations

TL;DR

The paper introduces TCSM_v2, a semi-supervised medical image segmentation method that enforces transformation-consistent predictions via a teacher-student EMA framework, evaluated on ISIC 2017, REFUGE, and LiTS datasets.

ABSTRACT

Deep convolutional neural networks have achieved remarkable progress on a variety of medical image computing tasks. A common problem when applying supervised deep learning methods to medical images is the lack of labeled data, which is very expensive and time-consuming to be collected. In this paper, we present a novel semi-supervised method for medical image segmentation, where the network is optimized by the weighted combination of a common supervised loss for labeled inputs only and a regularization loss for both labeled and unlabeled data. To utilize the unlabeled data, our method encourages the consistent predictions of the network-in-training for the same input under different regularizations. Aiming for the semi-supervised segmentation problem, we enhance the effect of regularization for pixel-level predictions by introducing a transformation, including rotation and flipping, consistent scheme in our self-ensembling model. With the aim of semi-supervised segmentation tasks, we introduce a transformation consistent strategy in our self-ensembling model to enhance the regularization effect for pixel-level predictions. We have extensively validated the proposed semi-supervised method on three typical yet challenging medical image segmentation tasks: (i) skin lesion segmentation from dermoscopy images on International Skin Imaging Collaboration (ISIC) 2017 dataset, (ii) optic disc segmentation from fundus images on Retinal Fundus Glaucoma Challenge (REFUGE) dataset, and (iii) liver segmentation from volumetric CT scans on Liver Tumor Segmentation Challenge (LiTS) dataset. Compared to the state-of-the-arts, our proposed method shows superior segmentation performance on challenging 2D/3D medical images, demonstrating the effectiveness of our semi-supervised method for medical image segmentation.

Motivation & Objective

• Motivate learning segmentation with limited labeled data in medical imaging.
• Propose a semi-supervised framework combining supervised loss on labeled data with a regularization loss on both labeled and unlabeled data.
• Introduce a transformation-consistent self-ensembling strategy for pixel-level predictions.
• Extend transformation regularization with scaling and a teacher model to improve targets and robustness.

Proposed method

• Use a weighted combination of supervised loss on labeled data and a regularization loss on both labeled and unlabeled data.
• Introduce transformation-consistent regularization that enforces agreement between predictions under input and output space transformations (rotation, flipping, scaling).
• Employ a teacher model that is the exponential-moving-average (EMA) of the student model to provide better targets for the consistency loss.
• Incorporate additional perturbations such as Gaussian noise and dropout to enrich regularization.
• Adapt the approach to both 2D and 3D segmentation networks (DenseUNet for 2D dermoscopy/refund datasets; 3D U-Net for LiTS).
• Train with a time-dependent ramp-up for the regularization weight and optimize the total loss as L + lambda(T) R.

Experimental results

Research questions

• RQ1 Can transformation-consistent self-ensembling improve semi-supervised segmentation performance with limited labels?
• RQ2 How do rotation, flipping, and scaling transformations affect pixel-level segmentation consistency between student and teacher predictions?
• RQ3 Does introducing a scaling transformation and EMA teacher improve regularization and accuracy across 2D and 3D medical imaging tasks?
• RQ4 How does TCSM_v2 perform on diverse modalities such as dermoscopy (skin lesions), fundus imagery (optic disc), and volumetric CT (liver)?

Key findings

• TCSM_v2 achieves superior performance over the supervised baseline and other semi-supervised methods on challenging 2D and 3D medical segmentation tasks.
• Ablation studies show transformation-consistent regularization, including rotation, scaling, and noise/dropout, contributes significantly to performance gains.
• In 2D dermoscopy experiments, the method improves multiple metrics (e.g., JA, DI, AC, SE, SP) compared to baselines and ablations.
• The approach demonstrates effectiveness across ISIC 2017 skin lesion segmentation, REFUGE optic disc segmentation, and LiTS liver segmentation datasets.
• Using a teacher model (EMA of the student) provides better targets for the consistency loss, enhancing regularization.