[论文解读] Self-training with dual uncertainty for semi-supervised medical image segmentation

Zhanhong Qiu, Haitao Gan+3 more|arXiv (Cornell University)|Apr 10, 2023

Radiomics and Machine Learning in Medical ImagingMedicine被引用 3

一句话总结

该论文提出了一种带有双重不确定性估计的自训练框架——样本级和像素级——以提升半监督医学图像分割性能。通过利用多个检查点中模型预测结果的方差来估计样本级不确定性，以及通过使用不同上采样方法的两个解码器输出差异来估计像素级不确定性，该方法选择性地对难样本进行再训练并优化伪标签，在相同设置下于ACDC和Prostate数据集上取得了最先进性能。

ABSTRACT

In the field of semi-supervised medical image segmentation, the shortage of labeled data is the fundamental problem. How to effectively learn image features from unlabeled images to improve segmentation accuracy is the main research direction in this field. Traditional self-training methods can partially solve the problem of insufficient labeled data by generating pseudo labels for iterative training. However, noise generated due to the model's uncertainty during training directly affects the segmentation results. Therefore, we added sample-level and pixel-level uncertainty to stabilize the training process based on the self-training framework. Specifically, we saved several moments of the model during pre-training, and used the difference between their predictions on unlabeled samples as the sample-level uncertainty estimate for that sample. Then, we gradually add unlabeled samples from easy to hard during training. At the same time, we added a decoder with different upsampling methods to the segmentation network and used the difference between the outputs of the two decoders as pixel-level uncertainty. In short, we selectively retrained unlabeled samples and assigned pixel-level uncertainty to pseudo labels to optimize the self-training process. We compared the segmentation results of our model with five semi-supervised approaches on the public 2017 ACDC dataset and 2018 Prostate dataset. Our proposed method achieves better segmentation performance on both datasets under the same settings, demonstrating its effectiveness, robustness, and potential transferability to other medical image segmentation tasks. Keywords: Medical image segmentation, semi-supervised learning, self-training, uncertainty estimation

研究动机与目标

• 为解决分割任务中医学图像标注数据有限的关键挑战。
• 通过减少模型不确定性带来的噪声，稳定半监督学习中的自训练过程。
• 通过基于不确定性估计选择性地再训练难样本，提升分割精度。
• 通过使用双解码器进行像素级不确定性估计，提升伪标签质量。
• 在多种医学影像数据集中展示方法的泛化性与鲁棒性。

提出的方法

• 样本级不确定性通过在预训练阶段对同一未标注样本在多个检查点上的模型预测结果方差进行估计。
• 未标注样本按样本级不确定性由低到高逐步引入训练，即从易到难。
• 像素级不确定性通过使用不同上采样策略（如转置卷积与反卷积）的两个解码器输出差异计算。
• 利用像素级不确定性对伪标签进行优化，对不确定性较低的区域赋予更高置信度。
• 模型通过迭代自训练持续改进，仅对最不确定的样本进行再训练，以减少误差传播。
• 该框架将不确定性感知的伪标签生成集成到标准自训练流程中，无需对网络架构进行大规模修改。

实验结果

研究问题

• RQ1 样本级不确定性估计能否提升半监督医学图像分割中伪标签的可靠性？
• RQ2 基于双解码器的像素级不确定性能否提升伪标签质量并减少噪声？
• RQ3 基于双重不确定性的难样本选择性再训练能否提升分割性能？
• RQ4 所提出方法在基准医学影像数据集上与现有自训练方法相比表现如何？
• RQ5 该方法在不同医学图像分割任务中的鲁棒性与可迁移性如何？

主要发现

• 在与基线方法相同的训练协议下，该方法在2017年ACDC数据集上达到了最先进性能。
• 在2018年Prostate数据集上，该方法在分割精度上优于五种现有的半监督方法。
• 双重不确定性集成显著减少了自训练过程中由噪声伪标签引起的误差传播。
• 通过仅对最不确定样本进行选择性再训练，模型表现出更强的鲁棒性，避免了低质量伪标签导致的性能下降。
• 该方法展现出强大的泛化潜力，在两个不同的医学影像任务中均实现了稳定且一致的性能提升。
• 使用多个模型检查点和双解码器在不增加额外训练或推理成本的前提下，提供了可靠的不确定性估计。