[论文解读] Leveraging Labeling Representations in Uncertainty-based Semi-supervised Segmentation

Sukesh Adiga V, José Dolz+1 more|arXiv (Cornell University)|Mar 10, 2022

Advanced Neural Network Applications被引用 8

一句话总结

该论文提出了一种新颖的不确定性估计方法，用于半监督3D左心房分割，利用预训练的去噪自编码器（DAE）从稀疏标注中学习标注表征。通过将模型预测映射到DAE生成的合理分割掩码，该方法仅需一次前向推理即可计算像素级不确定性，从而在降低计算成本的同时提升分割精度，相较于最先进方法，在10%标注数据下实现了86.58%的Dice分数和11.82mm的Hausdorff距离。

ABSTRACT

Semi-supervised segmentation tackles the scarcity of annotations by leveraging unlabeled data with a small amount of labeled data. A prominent way to utilize the unlabeled data is by consistency training which commonly uses a teacher-student network, where a teacher guides a student segmentation. The predictions of unlabeled data are not reliable, therefore, uncertainty-aware methods have been proposed to gradually learn from meaningful and reliable predictions. Uncertainty estimation, however, relies on multiple inferences from model predictions that need to be computed for each training step, which is computationally expensive. This work proposes a novel method to estimate the pixel-level uncertainty by leveraging the labeling representation of segmentation masks. On the one hand, a labeling representation is learnt to represent the available segmentation masks. The learnt labeling representation is used to map the prediction of the segmentation into a set of plausible masks. Such a reconstructed segmentation mask aids in estimating the pixel-level uncertainty guiding the segmentation network. The proposed method estimates the uncertainty with a single inference from the labeling representation, thereby reducing the total computation. We evaluate our method on the 3D segmentation of left atrium in MRI, and we show that our uncertainty estimates from our labeling representation improve the segmentation accuracy over state-of-the-art methods.

研究动机与目标

• 为解决半监督医学图像分割中不确定性估计的高计算成本问题，特别是如蒙特卡洛Dropout或集成方法等需要多次推理的方法。
• 通过利用学习到的标注表征作为形状先验，提升分割可靠性，且无需对齐图像或额外训练任务。
• 减少对复杂训练方案（如协同训练或对抗训练）的依赖，这些方案常面临收敛性和稳定性问题。
• 开发一种更鲁棒且高效的不确定性估计机制，通过预训练DAE的单次推理指导学生模型。
• 证明基于标注表征的不确定性估计可带来优于基于熵或集成的不确定性估计的分割性能。

提出的方法

• 使用预训练的去噪自编码器（DAE）从可用的分割掩码中学习标注表征，捕捉合理的解剖学形状。
• DAE将教师模型的分割预测映射为一组重建后的合理掩码，有效充当形状先验。
• 通过原始预测与DAE重建结果之间的L2差异计算像素级不确定性，实现单次推理的不确定性估计。
• 将不确定性图作为正则化项集成到均值教师框架中，引导学生模型关注可靠预测区域。
• 采用指数移动平均（EMA）更新教师模型权重，并结合监督损失与一致性损失进行训练。
• 该方法在2018年心房分割挑战赛提供的3D左心房MRI数据上进行评估，与最先进基线方法进行比较。

实验结果

研究问题

• RQ1 从稀疏标注中学习到的标注表征是否能以更低的计算成本改善半监督分割中的不确定性估计？
• RQ2 基于DAE重建掩码的不确定性是否能带来优于基于熵或集成的不确定性估计的分割性能？
• RQ3 在低标注率设置下，该方法在Dice分数和Hausdorff距离方面与最先进半监督分割方法相比表现如何？
• RQ4 所提出的不确定性估计在不同超参数设置（如不确定性权重γ和一致性权重β）下是否具有鲁棒性？
• RQ5 基于标注表征的不确定性估计是否可在无需对齐图像或额外训练任务（如形状预测或对抗损失）的情况下有效应用？

主要发现

• 在10%标注数据设置下，该方法实现了86.58%的Dice分数和11.82mm的Hausdorff距离，优于最佳基线方法（SASSnet）2.88%的Dice分数和5.08mm的HD。
• 在20%标注数据设置下，该方法实现了88.60%的Dice分数和7.61mm的Hausdorff距离，优于最佳基线方法（URPC）0.02%的Dice分数和5.49mm的HD。
• 消融实验表明，结合阈值或熵策略使用标注表征进行不确定性估计，性能优于UAMT，完整方法取得最佳结果。
• 该方法每轮训练仅需一次DAE推理，相比UAMT每步使用8次推理，显著降低了计算成本。
• 对超参数γ和β的消融分析表明，该方法在不同设置下均表现鲁棒且持续优于基线方法，其中γ=1时性能最优。
• 可视化对比显示，所提方法生成的分割结果更平滑、更准确，孔洞更少，边界贴合度更高，优于SASSnet、DTC、UAMT和URPC。