[论文解读] Semi-supervised Learning for Quantification of Pulmonary Edema in Chest X-Ray Images

Ruizhi Liao, Jonathan Rubin+6 more|arXiv (Cornell University)|Feb 27, 2019

COVID-19 diagnosis using AI参考文献 13被引用 7

一句话总结

该论文提出了一种半监督贝叶斯模型，通过联合训练变分自编码器（VAE）与回归器，仅使用33万张临床数据集中6,000张标注图像，量化胸部X光片中的肺水肿严重程度。通过从标注和未标注数据中学习紧凑的潜在表征，该方法在根均方误差方面显著优于监督学习、熵最小化或多VAE方法，并与真实值具有更高的相关性。

ABSTRACT

We propose and demonstrate machine learning algorithms to assess the severity of pulmonary edema in chest x-ray images of congestive heart failure patients. Accurate assessment of pulmonary edema in heart failure is critical when making treatment and disposition decisions. Our work is grounded in a large-scale clinical dataset of over 300,000 x-ray images with associated radiology reports. While edema severity labels can be extracted unambiguously from a small fraction of the radiology reports, accurate annotation is challenging in most cases. To take advantage of the unlabeled images, we develop a Bayesian model that includes a variational auto-encoder for learning a latent representation from the entire image set trained jointly with a regressor that employs this representation for predicting pulmonary edema severity. Our experimental results suggest that modeling the distribution of images jointly with the limited labels improves the accuracy of pulmonary edema scoring compared to a strictly supervised approach. To the best of our knowledge, this is the first attempt to employ machine learning algorithms to automatically and quantitatively assess the severity of pulmonary edema in chest x-ray images.

研究动机与目标

• 开发一种机器学习方法，以准确量化胸部X光片中的肺水肿严重程度，这是心力衰竭中一个关键但细微的临床标志。
• 通过利用大量未标注的胸部X光片，解决标注医学图像有限的挑战，以提升模型性能。
• 证明通过贝叶斯框架联合学习图像分布与标签预测，可提升回归精度，优于纯监督或自监督方法。
• 实现患者状态的定量表型分析，以支持心力衰竭及其他容量状态相关疾病的临床决策与研究。

提出的方法

• 使用贝叶斯生成模型，单个VAE从全部330,000张胸部X光片中学习共享的潜在表征，包括标注和未标注数据。
• VAE与回归器联合训练，后者将学习到的潜在特征映射至肺水肿严重程度评分（0–3），同时优化重建和预测目标。
• 模型使用重参数化技巧，实现通过随机潜在变量的反向传播，并通过KL散度损失将后验分布正则化为标准正态先验。
• 图像重建损失按图像尺寸的倒数加权，VAE的潜在空间经过归一化以减少尺度偏差。
• 训练过程采用Adam优化器，学习率为0.001，批量大小为4，并基于验证损失采用早停策略。
• 在整个图像集上训练单一VAE，而非按每个标签类别分别训练VAE，作者认为这种方法更适合连续且细微的严重程度分级。

实验结果

研究问题

• RQ1 与完全监督方法相比，利用大规模未标注胸部X光片的半监督学习方法是否能提升肺水肿严重程度回归的准确性？
• RQ2 在标注和未标注数据上联合训练VAE与回归器，是否能生成优于按标签类别分别训练模型的特征表征？
• RQ3 在标注数据有限的情境下，所提方法与熵最小化和多VAE基线相比表现如何？
• RQ4 单一共享的VAE表征是否能有效捕捉与连续肺水肿严重程度评分相关的细微图像特征？
• RQ5 结合全局图像分布在多大程度上提升了肺水肿严重程度预测的泛化能力和鲁棒性？

主要发现

• 所提出的VAE-回归器方法在测试集上实现了0.57的均方根误差（RMSE），优于监督基线（RMSE：0.63）、基于EM的自学习方法（RMSE：0.65）和多VAE方法（RMSE：0.68）。
• 该方法与真实严重程度评分的皮尔逊相关系数（CC）达到0.78，显著高于监督基线（CC：0.72）、EM方法（CC：0.70）和多VAE方法（CC：0.69）。
• 潜在表征与标签预测的联合学习使所有严重程度类别中的预测更加准确且一致，如预测评分分布图所示。
• 该模型通过有效利用未标注数据，展现出更强的泛化能力，减少了在验证集和测试集上的过拟合现象，提升了鲁棒性。
• 结果表明，联合建模图像分布与有限标注数据，比为每个标签类别分别训练VAE更有效，尤其适用于肺水肿严重程度等连续且细微的特征。
• 本研究首次成功将机器学习应用于胸部X光片中肺水肿的自动化、定量评分，为临床决策支持和表型分析提供了有力支持。