[论文解读] Dual Multi-scale Mean Teacher Network for Semi-supervised Infection Segmentation in Chest CT Volume for COVID-19

Liansheng Wang, Jiacheng Wang+7 more|arXiv (Cornell University)|Nov 10, 2022

COVID-19 diagnosis using AI被引用 5

一句话总结

本论文提出了一种双多尺度均值教师网络（DM²T-Net），用于胸部CT影像中半监督3D COVID-19肺部感染分割，通过多维注意力CNN（MDA-CNN）捕捉分层多尺度特征，并在多个预测头之间强制学生与教师网络之间的一致性。该方法在COVID-19-P20数据集上取得72.59%的Dice分数，在MosMedData数据集上取得60.19%的Dice分数，展现出在低数据量场景下的优越准确率与泛化能力。

ABSTRACT

Automated detecting lung infections from computed tomography (CT) data plays an important role for combating COVID-19. However, there are still some challenges for developing AI system. 1) Most current COVID-19 infection segmentation methods mainly relied on 2D CT images, which lack 3D sequential constraint. 2) Existing 3D CT segmentation methods focus on single-scale representations, which do not achieve the multiple level receptive field sizes on 3D volume. 3) The emergent breaking out of COVID-19 makes it hard to annotate sufficient CT volumes for training deep model. To address these issues, we first build a multiple dimensional-attention convolutional neural network (MDA-CNN) to aggregate multi-scale information along different dimension of input feature maps and impose supervision on multiple predictions from different CNN layers. Second, we assign this MDA-CNN as a basic network into a novel dual multi-scale mean teacher network (DM${^2}$T-Net) for semi-supervised COVID-19 lung infection segmentation on CT volumes by leveraging unlabeled data and exploring the multi-scale information. Our DM${^2}$T-Net encourages multiple predictions at different CNN layers from the student and teacher networks to be consistent for computing a multi-scale consistency loss on unlabeled data, which is then added to the supervised loss on the labeled data from multiple predictions of MDA-CNN. Third, we collect two COVID-19 segmentation datasets to evaluate our method. The experimental results show that our network consistently outperforms the compared state-of-the-art methods.

研究动机与目标

• 为解决基于2D的方法在捕捉胸部CT中3D空间上下文与层间关系方面的局限性，以实现更精确的感染分割。
• 通过设计多维注意力机制，弥补现有3D分割模型中多尺度感受野的不足，实现分层特征学习。
• 通过开发一种半监督学习框架，有效利用未标注数据，缓解3D CT数据标注有限的挑战。
• 通过集成多尺度一致性损失与多预测监督，提升模型在多样化数据集上的泛化能力。

提出的方法

• 设计了一种多维注意力CNN（MDA-CNN），用于在3D CT影像的不同空间维度上聚合多尺度特征，增强分层表征学习能力。
• 将MDA-CNN集成到双多尺度均值教师网络（DM²T-Net）中，使学生与教师网络均能从不同网络层输出多尺度预测。
• 引入多尺度一致性损失，强制在未标注数据上，学生与教师网络在多个尺度与网络深度之间达成预测一致。
• 训练目标结合了标注数据上的监督损失与未标注数据上的多尺度一致性损失，实现具有改进泛化能力的半监督学习。
• 通过学生网络权重的指数移动平均更新教师网络参数，稳定训练过程并提升预测一致性。
• 收集并使用了两个新的3D胸部CT数据集——COVID-19-P20与MosMedData，对所提方法进行全面评估。

实验结果

研究问题

• RQ1 与单尺度或2D方法相比，具有维度特异性注意力的多尺度3D CNN是否能提升感染分割的准确性？
• RQ2 所提出的多尺度一致性损失在利用未标注数据进行半监督3D分割方面是否高效？
• RQ3 与现有SOTA方法相比，双多尺度均值教师框架在多样化CT成像协议与感染模式下是否具有更好的泛化能力？
• RQ4 来自多个网络层的多预测监督在多大程度上提升了模型的鲁棒性与分割性能？

主要发现

• 所提出的DM²T-Net在COVID-19-P20数据集上取得72.59%的Dice分数，显著优于当前最先进方法。
• 在MosMedData数据集上，该方法取得60.19%的Dice分数，表明即使在更具挑战性与多样性的数据集上也表现出强大性能。
• 多尺度一致性损失在60个周期后收敛稳定，早期训练阶段呈现逐步上升趋势，表明未标注数据被有效利用。
• 可视化对比显示，教师网络的预测始终比学生网络更准确，证实了多尺度监督下知识蒸馏的有效性。
• 泛化性分析表明，当在某一数据集上训练并在另一数据集上测试时，DM²T-Net仍保持优越性能（Dice: 48.81±23.51，HD95: 35.34±30.57），优于2D U-Net++与3D UA-MT。
• 该模型在成像差异与感染模式多样性方面表现出更强鲁棒性，跨数据集评估中HD95值更低、Dice分数更高，优于基线模型。