[论文解读] Co-training with High-Confidence Pseudo Labels for Semi-supervised Medical Image Segmentation

Zhiqiang Shen, Peng Cao+4 more|arXiv (Cornell University)|Jan 11, 2023

Advanced Neural Network Applications被引用 7

一句话总结

该论文提出了一种不确定性引导的协作均值教师网络（UCMT），这是一种半监督医学图像分割框架，通过不确定性引导的区域混合（UMIX）生成高置信度伪标签，并利用协作均值教师（CMT）架构保持模型之间的分歧。UCMT在四个公开的2D和3D医学图像数据集上，使用有限的标注数据，实现了最先进性能。

ABSTRACT

Consistency regularization and pseudo labeling-based semi-supervised methods perform co-training using the pseudo labels from multi-view inputs. However, such co-training models tend to converge early to a consensus, degenerating to the self-training ones, and produce low-confidence pseudo labels from the perturbed inputs during training. To address these issues, we propose an Uncertainty-guided Collaborative Mean-Teacher (UCMT) for semi-supervised semantic segmentation with the high-confidence pseudo labels. Concretely, UCMT consists of two main components: 1) collaborative mean-teacher (CMT) for encouraging model disagreement and performing co-training between the sub-networks, and 2) uncertainty-guided region mix (UMIX) for manipulating the input images according to the uncertainty maps of CMT and facilitating CMT to produce high-confidence pseudo labels. Combining the strengths of UMIX with CMT, UCMT can retain model disagreement and enhance the quality of pseudo labels for the co-training segmentation. Extensive experiments on four public medical image datasets including 2D and 3D modalities demonstrate the superiority of UCMT over the state-of-the-art. Code is available at: https://github.com/Senyh/UCMT.

研究动机与目标

• 为解决协同训练模型在半监督医学图像分割中过早收敛至共识并生成低置信度伪标签的局限性。
• 通过在训练过程中减少不确定性，同时保持子网络之间的模型分歧，提升伪标签质量。
• 开发一个统一框架，同时增强模型分歧和伪标签置信度，以在标注数据有限的情况下实现更好的泛化能力。
• 探究协同训练型半监督学习中模型分歧与伪标签不确定性之间的相互作用。

提出的方法

• UCMT引入了协作均值教师（CMT）组件，通过一个动量平均的教师模型，对两个学生网络实施协同训练，以正则化训练过程并防止共识崩溃。
• 不确定性引导的区域混合（UMIX）模块通过选择性地混合CMT预测中的高不确定性区域来操作输入图像，促进生成高置信度伪标签。
• UMIX利用CMT生成的不确定性图识别并增强模糊区域，确保仅保留低不确定性区域，而高不确定性区域则被主动优化。
• 该框架采用联合训练方案结合CMT与UMIX，其中CMT利用UMIX增强输入生成的伪标签，以提升预测的鲁棒性与多样性。
• 该方法采用学生网络之间的交叉伪监督和均值教师正则化，以稳定训练并提升泛化性能。
• 对UMIX中的超参数（如top-k区域数量和补丁大小）进行了鲁棒性分析，结果表明在不同配置下性能均保持稳定。

实验结果

研究问题

• RQ1 在基于协同训练的半监督医学图像分割中，如何联合优化模型分歧与伪标签置信度？
• RQ2 在伪标签质量和模型性能方面，不确定性感知的数据增强有何影响？
• RQ3 具有动量正则化教师模型的协作均值教师框架能否防止过早收敛至自训练行为？
• RQ4 在低数据场景下，不确定性引导的区域混合与随机增强方法（如CutMix）相比表现如何？
• RQ5 在协同训练设置中，模型分歧、伪标签不确定性与分割性能之间的关系是什么？

主要发现

• UCMT在四个公开的医学图像分割数据集（包括ISIC、LA、ACDC和Polyp）上，于多种2D和3D模态下均实现了最先进性能。
• 在ISIC数据集上，仅使用5%标注数据时，UCMT的Dice分数达到88.22，优于基线模型和先前的SOTA方法。
• 消融实验证实，CMT与UMIX的结合能获得最佳性能，完整UCMT模型显著优于缺少任一组件的变体。
• 在低数据场景（如2%标注数据）下，UMIX始终优于CutMix，证明其在生成高置信度伪标签方面的有效性。
• 超参数敏感性分析表明，UMIX对超参数选择具有鲁棒性，不同补丁大小和top-k区域选择下的性能下降可忽略不计。
• 定性结果表明，UCMT生成的分割边界更平滑、更准确，且对细微结构的恢复能力优于均值教师和交叉伪监督基线方法。