[Paper Review] Translation Consistent Semi-supervised Segmentation for 3D Medical Images

Yuyuan Liu, Yu Tian+5 more|arXiv (Cornell University)|Mar 28, 2022

Advanced Neural Network Applications7 citations

TL;DR

This paper proposes TraCoCo, a novel semi-supervised 3D medical image segmentation method that enforces consistency across translated views of unlabelled data to reduce reliance on spatial context and improve foreground object generalization. It introduces a Confident Regional Cross-Entropy (CRC) loss for robust, convergent training and extends 3D CutMix augmentation, achieving state-of-the-art performance on LA, Pancreas, and BRaTS19 datasets with minimal labelled data.

ABSTRACT

3D medical image segmentation methods have been successful, but their dependence on large amounts of voxel-level annotated data is a disadvantage that needs to be addressed given the high cost to obtain such annotation. Semi-supervised learning (SSL) solve this issue by training models with a large unlabelled and a small labelled dataset. The most successful SSL approaches are based on consistency learning that minimises the distance between model responses obtained from perturbed views of the unlabelled data. These perturbations usually keep the spatial input context between views fairly consistent, which may cause the model to learn segmentation patterns from the spatial input contexts instead of the segmented objects. In this paper, we introduce the Translation Consistent Co-training (TraCoCo) which is a consistency learning SSL method that perturbs the input data views by varying their spatial input context, allowing the model to learn segmentation patterns from visual objects. Furthermore, we propose the replacement of the commonly used mean squared error (MSE) semi-supervised loss by a new Cross-model confident Binary Cross entropy (CBC) loss, which improves training convergence and keeps the robustness to co-training pseudo-labelling mistakes. We also extend CutMix augmentation to 3D SSL to further improve generalisation. Our TraCoCo shows state-of-the-art results for the Left Atrium (LA) and Brain Tumor Segmentation (BRaTS19) datasets with different backbones. Our code is available at https://github.com/yyliu01/TraCoCo.

Motivation & Objective

• To address the limitation of existing semi-supervised learning (SSL) methods in 3D medical imaging that rely too heavily on spatial input context rather than foreground object features.
• To reduce confirmation bias and model convergence to local minima by introducing a co-training framework with spatially translated views.
• To improve training convergence and robustness to pseudo-labeling errors through a new loss function that focuses on confident positive and negative predictions.
• To enhance generalization by extending 3D CutMix augmentation to semi-supervised 3D segmentation.
• To achieve state-of-the-art performance on major 3D medical segmentation benchmarks with limited annotated data.

Proposed method

• TraCoCo employs a co-training strategy where two independently initialized models are trained with consistency enforced across spatially translated views of the same unlabelled 3D volume.
• Translation consistency is enforced by computing segmentation agreement only within the intersection region of the translated views, reducing dependence on background patterns.
• A novel Confident Regional Cross-Entropy (CRC) loss is introduced, which focuses training on the most confident predictions (positive and negative) from the student model to improve convergence and robustness.
• The CRC loss is applied only to regions predicted with high confidence, minimizing the impact of noisy or incorrect pseudo-labels.
• 3D CutMix augmentation is extended to 3D SSL by mixing patches from different volumes during training, improving generalization and data efficiency.
• Connected component post-processing is applied to filter out isolated false positive regions smaller than 1/1500 of the input volume, further improving segmentation accuracy.

Experimental results

Research questions

• RQ1 Can enforcing consistency across spatially translated views of unlabelled 3D medical images reduce model dependence on background patterns and improve generalization to foreground objects?
• RQ2 Does a confidence-based loss function that focuses on high-confidence pseudo-labels improve training convergence and robustness compared to standard MSE-based losses?
• RQ3 Can 3D CutMix augmentation improve generalization in semi-supervised 3D medical image segmentation with limited annotated data?
• RQ4 How does the proposed TraCoCo framework compare to state-of-the-art methods in terms of segmentation accuracy across diverse 3D medical imaging tasks?
• RQ5 What is the impact of post-processing techniques such as connected component filtering on the final segmentation performance of the model?

Key findings

• TraCoCo achieves state-of-the-art performance on the Left Atrium (LA), Pancreas-CT (Pancreas), and Brain Tumor Segmentation (BRaTS19) datasets across all evaluation metrics, including Dice, mIoU, and 95HD.
• The Confident Regional Cross-Entropy (CRC) loss outperforms MSE loss by producing tighter confidence clustering for foreground and background voxels, leading to better segmentation accuracy and faster convergence.
• On the Pancreas-CT dataset, CRC loss results in a 1.5% improvement in Dice score and a 2.1% improvement in mIoU compared to MSE loss.
• Connected component post-processing with a threshold of 1/1500 of the input volume reduces false positives and improves all metrics, with the best results achieved when combined with CRC loss.
• The addition of 3D CutMix augmentation further improves generalization, particularly in low-data regimes, as shown by consistent performance gains across all datasets.
• The method runs in 1.386 seconds per volume on average (1.412 with post-processing), with only 0.098 seconds additional overhead for post-processing on BraTS19, demonstrating practical efficiency.