[论文解读] Deep Expectation-Maximization for Semi-Supervised Lung Cancer Screening

Sumeet Menon, David Chapman+4 more|arXiv (Cornell University)|Oct 2, 2020

COVID-19 diagnosis using AI参考文献 16被引用 9

一句话总结

本文提出一种半监督深度学习框架，结合期望最大化（EM）算法，利用标注和未标注的CT扫描数据训练3D卷积神经网络（CNN），用于肺癌筛查。通过将未标注样本的标签视为隐变量，该方法在仅使用较小数据集中的标注数据时，分类准确率超越完全监督基线，于NLST数据集上达到最高81.1%的准确率，较监督基线提升20%。

ABSTRACT

We present a semi-supervised algorithm for lung cancer screening in which a 3D Convolutional Neural Network (CNN) is trained using the Expectation-Maximization (EM) meta-algorithm. Semi-supervised learning allows a smaller labelled data-set to be combined with an unlabeled data-set in order to provide a larger and more diverse training sample. EM allows the algorithm to simultaneously calculate a maximum likelihood estimate of the CNN training coefficients along with the labels for the unlabeled training set which are defined as a latent variable space. We evaluate the model performance of the Semi-Supervised EM algorithm for CNNs through cross-domain training of the Kaggle Data Science Bowl 2017 (Kaggle17) data-set with the National Lung Screening Trial (NLST) data-set. Our results show that the Semi-Supervised EM algorithm greatly improves the classification accuracy of the cross-domain lung cancer screening, although results are lower than a fully supervised approach with the advantage of additional labelled data from the unsupervised sample. As such, we demonstrate that Semi-Supervised EM is a valuable technique to improve the accuracy of lung cancer screening models using 3D CNNs.

研究动机与目标

• 解决肺癌筛查中医学影像标注数据有限的挑战。
• 提升3D CNN在跨域肺癌分类中的泛化能力和准确率。
• 探索利用EM算法联合优化CNN权重并推断未标注数据中隐标签的可行性。
• 评估半监督学习结合EM是否能有效利用大量未标注CT扫描数据以增强模型性能。
• 对比半监督EM方法与完全监督基线及全量标注数据训练的上限性能。

提出的方法

• 该方法采用期望最大化（EM）元算法，在存在观测（标注）标签和未观测（隐变量）标签的未标注数据中训练3D CNN。
• 在E步中，模型使用当前CNN权重预测未标注CT扫描的类别概率。
• 在M步中，CNN通过结合标注数据和E步生成的伪标签预测结果进行微调，以最大化似然函数。
• 该框架整合了两个不同的肺癌筛查数据集：Kaggle数据科学杯2017（Kaggle17）和国家肺癌筛查试验（NLST），其数据量和成像协议各不相同。
• 模型采用两层3D CNN和AlexNet架构进行训练，数据预处理为20层切片的50×50×50体素块。
• 通过在某一数据集的标注数据上进行训练，并引入另一数据集的未标注数据，实现跨域评估，以模拟现实世界中数据稀缺的情况。

实验结果

研究问题

• RQ1 当应用于结合未标注CT扫描的半监督学习时，EM算法是否能有效提升3D CNN在肺癌筛查中的性能？
• RQ2 基于EM的半监督模型性能与仅使用标注数据训练的完全监督模型相比如何？
• RQ3 将更大规模、跨域的数据集（NLST）中的未标注数据引入较小规模、同域数据集（Kaggle17）时，分类准确率提升的幅度有多大？
• RQ4 当未标注数据量相对于标注数据量增加时，半监督EM带来的性能提升是否随之扩大？
• RQ5 当标注数据来自较小数据集而未标注数据来自较大数据集时，EM方法是否表现出更高的有效性？

主要发现

• 当在Kaggle17标注数据上训练并引入NLST未标注数据时，两层3D CNN的准确率从75.95%提升至77.5%，绝对提升1.55%。
• 在相同跨域设置下，AlexNet架构的准确率从79.36%提升至81.1%，提升1.74%。
• 当在NLST标注数据上训练并在Kaggle17上测试时，两层3D CNN的准确率从76.75%提升至78.1%，提升1.35%；而AlexNet从72.9%提升至74.4%。
• 半监督EM带来的性能提升始终约为使用全部标注数据训练的完全监督上限性能提升的一半。
• 当标注数据来自较小的Kaggle17数据集、未标注数据来自更大的NLST数据集时，该方法表现出更大的相对提升，表明在标注数据稀缺的场景下具有更高的数据效率。
• 结果证实，基于EM的半监督学习在3D CNN医学影像任务中是有效的，尤其在标注数据有限但存在大量未标注扫描时更具优势。