[Paper Review] Double-Uncertainty Weighted Method for Semi-supervised Learning

Yixin Wang, Yao Zhang+5 more|arXiv (Cornell University)|Oct 19, 2020

Advanced Neural Network Applications13 references12 citations

TL;DR

This paper proposes a double-uncertainty weighted method for semi-supervised medical image segmentation using a teacher-student framework. By estimating both segmentation uncertainty and novel feature uncertainty via Monte Carlo dropout in a Bayesian deep learning setup, the method dynamically weights consistency loss to improve teacher model reliability and student learning, achieving state-of-the-art performance with minimal labeled data on two public datasets.

ABSTRACT

Though deep learning has achieved advanced performance recently, it remains a challenging task in the field of medical imaging, as obtaining reliable labeled training data is time-consuming and expensive. In this paper, we propose a double-uncertainty weighted method for semi-supervised segmentation based on the teacher-student model. The teacher model provides guidance for the student model by penalizing their inconsistent prediction on both labeled and unlabeled data. We train the teacher model using Bayesian deep learning to obtain double-uncertainty, i.e. segmentation uncertainty and feature uncertainty. It is the first to extend segmentation uncertainty estimation to feature uncertainty, which reveals the capability to capture information among channels. A learnable uncertainty consistency loss is designed for the unsupervised learning process in an interactive manner between prediction and uncertainty. With no ground-truth for supervision, it can still incentivize more accurate teacher's predictions and facilitate the model to reduce uncertain estimations. Furthermore, our proposed double-uncertainty serves as a weight on each inconsistency penalty to balance and harmonize supervised and unsupervised training processes. We validate the proposed feature uncertainty and loss function through qualitative and quantitative analyses. Experimental results show that our method outperforms the state-of-the-art uncertainty-based semi-supervised methods on two public medical datasets.

Motivation & Objective

• To address the challenge of limited labeled medical imaging data, which is expensive and time-consuming to obtain.
• To improve the reliability of teacher-student semi-supervised learning by incorporating uncertainty estimates into the training process.
• To propose and validate a novel feature uncertainty metric that captures uncertainty across convolutional kernel channels.
• To design a learnable uncertainty consistency loss that interacts with model predictions to refine teacher outputs without ground-truth supervision.
• To balance supervised and unsupervised training by using uncertainty as a dynamic weighting mechanism.

Proposed method

• The teacher model is trained using Monte Carlo dropout to approximate a Bayesian Neural Network, enabling estimation of both segmentation uncertainty and feature uncertainty.
• Feature uncertainty is computed across channels of a convolutional layer (specifically the 4th upsampling layer) to assess the reliability of feature extraction per channel.
• A learnable uncertainty consistency loss is introduced to iteratively refine the teacher’s predictions by minimizing inconsistency between student and teacher outputs, weighted by uncertainty estimates.
• The total loss combines supervised cross-entropy loss on labeled data and uncertainty-weighted consistency loss on both labeled and unlabeled data.
• Uncertainty estimates serve as adaptive weights in the consistency loss, reducing influence of high-uncertainty predictions during training.
• The method is evaluated on two public medical imaging benchmarks: 2018 Atrial Segmentation Challenge and MICCAI 2019 KiTS Challenge.

Experimental results

Research questions

• RQ1 Can feature uncertainty, a novel form of uncertainty across convolutional kernel channels, improve model interpretability and performance in semi-supervised segmentation?
• RQ2 Can a learnable uncertainty consistency loss that interacts with model predictions improve teacher model accuracy without ground-truth supervision?
• RQ3 How does uncertainty-based weighting of consistency loss affect the balance between supervised and unsupervised training in semi-supervised learning?
• RQ4 Does the proposed method outperform existing uncertainty-based semi-supervised approaches when labeled data is scarce?
• RQ5 To what extent does uncertainty estimation help in filtering out unreliable predictions and improving generalization in medical image segmentation?

Key findings

• On the KiTS 2019 kidney segmentation dataset, the method achieved a Dice score of 0.9206 with only 10 labeled scans, outperforming all baselines.
• With 32 labeled scans, the method reached a Dice score of 0.9365, demonstrating consistent improvement over state-of-the-art methods.
• The proposed uncertainty consistency loss reduced teacher model prediction uncertainty over training, as shown by decreasing mean segmentation uncertainty values in Figure 2(c).
• The feature uncertainty analysis revealed that uncertainty values across channels decreased over training, indicating improved feature extraction reliability.
• The method significantly improved performance even with very few labeled samples, such as 4 labeled scans, where it ranked first in Dice score.
• The uncertainty weight mechanism enhanced the performance of baseline methods like UA-MT, with a 0.0087 improvement in Dice score when combined with uncertainty weighting.