[论文解读] 3D Semi-Supervised Learning with Uncertainty-Aware Multi-View Co-Training

Yingda Xia, Fengze Liu+7 more|arXiv (Cornell University)|Nov 29, 2018

AI in cancer detection被引用 24

一句话总结

本文提出了一种不确定性感知的多视角协同训练（UMCT）方法，一种3D弱监督学习框架。该方法通过空间变换生成3D医学影像的多样化视角，并采用从2D预训练模型初始化的非对称3D卷积核，以促进特征多样性。该方法利用贝叶斯深度学习结合Dropout来估计认知不确定性，实现基于不确定性的标签融合，从而提升伪标签的可靠性，在仅使用少量标注数据的情况下，于NIH胰腺和LiTS肝脏肿瘤分割任务中取得了最先进性能。

ABSTRACT

While making a tremendous impact in various fields, deep neural networks usually require large amounts of labeled data for training which are expensive to collect in many applications, especially in the medical domain. Unlabeled data, on the other hand, is much more abundant. Semi-supervised learning techniques, such as co-training, could provide a powerful tool to leverage unlabeled data. In this paper, we propose a novel framework, uncertainty-aware multi-view co-training (UMCT), to address semi-supervised learning on 3D data, such as volumetric data from medical imaging. In our work, co-training is achieved by exploiting multi-viewpoint consistency of 3D data. We generate different views by rotating or permuting the 3D data and utilize asymmetrical 3D kernels to encourage diversified features in different sub-networks. In addition, we propose an uncertainty-weighted label fusion mechanism to estimate the reliability of each view's prediction with Bayesian deep learning. As one view requires the supervision from other views in co-training, our self-adaptive approach computes a confidence score for the prediction of each unlabeled sample in order to assign a reliable pseudo label. Thus, our approach can take advantage of unlabeled data during training. We show the effectiveness of our proposed semi-supervised method on several public datasets from medical image segmentation tasks (NIH pancreas & LiTS liver tumor dataset). Meanwhile, a fully-supervised method based on our approach achieved state-of-the-art performances on both the LiTS liver tumor segmentation and the Medical Segmentation Decathlon (MSD) challenge, demonstrating the robustness and value of our framework, even when fully supervised training is feasible.

研究动机与目标

• 通过利用大量未标注数据，解决3D医学体素数据标注成本过高的问题。
• 通过空间变换引入视角多样性，以及采用非对称3D卷积核设计，提升弱监督3D分割性能。
• 通过使用贝叶斯深度学习估计预测不确定性，提升协同训练中伪标签的质量。
• 证明所提出的框架即使在完全监督设置下，通过多视角训练也能提升性能。
• 探究2D预训练权重对3D网络初始化的影响，以提升训练稳定性和性能。

提出的方法

• 通过在3D数据上应用空间变换（如旋转、排列等）生成多视角输入，为协同训练提供不同视角。
• 在子网络中使用非对称3D卷积核（如3×3×1），并从预训练的2D模型进行初始化，以引入模型层面的视角差异。
• 通过Dropout实现贝叶斯深度学习，以估计认知不确定性，生成每个视角预测的置信度分数。
• 设计不确定性加权标签融合（ULF）模块，根据预测的估计不确定性对不同视角的预测进行加权，生成未标注数据的伪标签。
• 在端到端训练多视角网络过程中，使用其他视角的伪标签和真实标签作为监督信号。
• 该框架支持弱监督和完全监督训练，即使所有数据均被标注，也应用协同训练目标。

实验结果

研究问题

• RQ1 通过数据层面和模型层面的视角多样性，多视角协同训练能否提升3D弱监督分割性能？
• RQ2 通过贝叶斯深度学习估计不确定性，能否提升协同训练中3D医学图像分割的伪标签可靠性？
• RQ3 2D预训练权重对3D网络初始化在体素分割任务中的影响如何？
• RQ4 不确定性加权标签融合是否优于标准平均或无权重融合方法？
• RQ5 所提出的协同训练框架是否能在仅使用少量标注数据的完全监督设置下仍提升性能？

主要发现

• 在NIH胰腺数据集上，所提出的UMCT框架在仅使用10%标注数据的情况下，Dice分数达到78.70%，优于其他弱监督方法。
• 在LiTS肝脏肿瘤数据集上，UMCT在弱监督学习设置下实现了最先进性能，展现出在有限标注下的强大泛化能力。
• UMCT的完全监督变体在NIH数据集上达到78.70%的Dice分数，优于标准的3D ResNet-50和V-Net，凸显了多视角协同训练的优势。
• 消融实验证实，采用2D初始化的非对称3D卷积核相比对称架构，性能提升超过1%。
• 不确定性加权标签融合在三个或更多视角下显著提升性能，表明置信度估计能有效降低不可靠伪标签带来的噪声。
• 关于初始化的消融研究显示，2D预训练权重显著提升了训练稳定性和最终性能，凸显其在3D分割任务中的关键作用。