[论文解读] Vision Transformer using Low-level Chest X-ray Feature Corpus for COVID-19 Diagnosis and Severity Quantification

Sang Joon Park, Gwanghyun Kim+7 more|arXiv (Cornell University)|Apr 15, 2021

COVID-19 diagnosis using AI参考文献 34被引用 8

一句话总结

本文提出了一种新颖的视觉Transformer（ViT）框架，利用大规模数据集预训练的低层次胸部X光（CXR）特征嵌入，以提升COVID-19诊断与严重程度量化性能。通过使用在超过22万张CXR图像上预训练的主干网络提取常见异常（如实变、密度增高），该模型在多种外部数据集上实现了最先进性能，泛化能力显著优于现有的CNN与ViT基线模型。

ABSTRACT

Developing a robust algorithm to diagnose and quantify the severity of COVID-19 using Chest X-ray (CXR) requires a large number of well-curated COVID-19 datasets, which is difficult to collect under the global COVID-19 pandemic. On the other hand, CXR data with other findings are abundant. This situation is ideally suited for the Vision Transformer (ViT) architecture, where a lot of unlabeled data can be used through structural modeling by the self-attention mechanism. However, the use of existing ViT is not optimal, since feature embedding through direct patch flattening or ResNet backbone in the standard ViT is not intended for CXR. To address this problem, here we propose a novel Vision Transformer that utilizes low-level CXR feature corpus obtained from a backbone network that extracts common CXR findings. Specifically, the backbone network is first trained with large public datasets to detect common abnormal findings such as consolidation, opacity, edema, etc. Then, the embedded features from the backbone network are used as corpora for a Transformer model for the diagnosis and the severity quantification of COVID-19. We evaluate our model on various external test datasets from totally different institutions to evaluate the generalization capability. The experimental results confirm that our model can achieve the state-of-the-art performance in both diagnosis and severity quantification tasks with superior generalization capability, which are sine qua non of widespread deployment.

研究动机与目标

• 为应对疫情期间COVID-19 CXR标注数据有限的挑战，利用大量未标注CXR数据中常见的异常，提升模型性能。
• 通过将特征学习与疾病分类解耦，提升COVID-19诊断与严重程度量化任务中的泛化能力。
• 开发一种稳健且可迁移的模型，可在多种临床环境中广泛部署，最大限度减少领域偏移。
• 证明预训练模型提取的低层次CXR特征可作为视觉Transformer在医学影像任务中的最优语料。

提出的方法

• 在大规模公开CXR数据集（如CheXpert）上预训练主干网络，以检测常见的低层次异常（如实变、密度增高、肺水肿）。
• 从预训练主干网络中提取特征嵌入，作为视觉Transformer的结构化语料。
• 使用ViT架构关注这些学习到的特征嵌入，通过建模CXR图像中长距离依赖关系实现高层级诊断。
• 应用ROI最大池化将每个patch的严重程度图转换为全局严重程度评分数组，提升与临床标注的一致性。
• 在有限的标注COVID-19数据上微调ViT，同时保持主干网络冻结，以防止过拟合。
• 采用均方误差（MSE）进行严重程度量化，二元交叉熵损失进行诊断任务，结合数据增强与机构间数据多样性，以增强模型鲁棒性。

实验结果

研究问题

• RQ1 预训练模型提取的低层次CXR特征嵌入是否可作为视觉Transformer在COVID-19诊断中的有效语料？
• RQ2 与标准ViT或CNN基线相比，使用预训练特征语料是否能提升在外部未见数据集上的泛化性能？
• RQ3 ROI池化方式选择（最大池化 vs. 平均池化）对严重程度量化性能有何影响？
• RQ4 仅使用CXR图像，该模型能否在严重程度预测中达到与临床专家相当的性能？
• RQ5 在训练过程中引入多样化来源的机构数据是否显著提升严重程度量化任务中的泛化能力？

主要发现

• 所提模型在三个外部数据集上的平均严重程度预测MSE为1.655，优于直接使用patch嵌入的基线ViT及采用ResNet特征的混合ViT。
• 在CNUH、YNU和KNUH的外部测试集中，模型分别取得1.682、1.677和1.607的MSE，展现出优越的泛化性能。
• 与平均池化相比，采用ROI最大池化进行严重程度图聚合使MSE降低0.191，证明其在捕捉局部异常方面的优越性。
• 若在训练中排除机构数据，性能显著下降（CNUH上MSE = 2.503），证实多样化数据对稳定严重程度预测的必要性。
• 该模型在诊断与严重程度量化任务中均达到最先进性能，严重程度评分性能与临床专家相当。
• 模型在多个机构间表现出的鲁棒性，证实其在低资源或资源匮乏医疗环境中的实际部署潜力。