[论文解读] COVID-19 Pneumonia Severity Prediction using Hybrid Convolution-Attention Neural Architectures

Nam V. Nguyen, Chang Jm|arXiv (Cornell University)|Jul 6, 2021

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

一句话总结

本文提出一种结合卷积与注意力机制的混合神经网络，并采用以数据为中心的预训练框架，以提升在有限医学影像数据下的新冠肺炎肺炎严重程度预测性能。通过利用Transformer的自注意力机制以及密集关联记忆（Hopfield网络），该模型在地理范围预测中实现了 R² = 0.85 ± 0.05 和 ρ = 0.92 ± 0.02，显著优于迁移学习基线模型。

ABSTRACT

This study proposed a novel framework for COVID-19 severity prediction, which is a combination of data-centric and model-centric approaches. First, we propose a data-centric pre-training for extremely scare data scenarios of the investigating dataset. Second, we propose two hybrid convolution-attention neural architectures that leverage the self-attention from the Transformer and the Dense Associative Memory (Modern Hopfield networks). Our proposed approach achieves significant improvement from the conventional baseline approach. The best model from our proposed approach achieves $R^2 = 0.85 \pm 0.05$ and Pearson correlation coefficient $ρ= 0.92 \pm 0.02$ in geographic extend and $R^2 = 0.72 \pm 0.09, ρ= 0.85\pm 0.06$ in opacity prediction.

研究动机与目标

• 解决因数据集极小且稀缺而导致的深度学习模型在新冠肺炎严重程度预测中性能低下问题。
• 开发一种以数据为中心的预训练框架，以增强有限医学影像数据的表征学习能力。
• 设计融合自注意力机制的混合卷积-注意力架构，以捕捉医学影像中的全局上下文信息。
• 通过新型神经架构提升新冠肺炎肺炎在地理范围和密度严重程度指标上的预测性能。
• 证明以数据为中心的预训练方法与混合注意力机制优于传统的ImageNet迁移学习方法。

提出的方法

• 提出一种以数据为中心的预训练框架，通过生成合成或增强数据来改善下游严重程度预测任务前的特征学习。
• 将Transformer架构中的自注意力模块以及密集关联记忆（现代Hopfield网络）集成到卷积神经网络主干网络中。
• 采用混合架构（HCT与HCH），结合卷积特征提取与通过注意力机制实现的全局上下文建模。
• 分两个阶段训练模型：首先在预处理数据上使用以数据为中心的框架进行预训练；其次在严重程度预测任务上进行微调。
• 采用交叉验证，损失函数基于公式6计算，以评估模型的泛化能力与训练稳定性。
• 使用DenseNet121、ResNet50和EfficientNet-B4等主干模型，权重初始化来自以数据为中心的预训练阶段。

实验结果

研究问题

• RQ1 以数据为中心的预训练框架是否能在极低数据量场景下显著提升深度学习模型在新冠肺炎严重程度预测中的性能？
• RQ2 融合自注意力与Hopfield网络的混合卷积-注意力架构与标准CNN及迁移学习基线相比，在预测肺炎严重程度方面表现如何？
• RQ3 通过注意力机制实现的全局上下文建模在多大程度上提升了新冠肺炎影像中地理范围与密度的预测准确性？
• RQ4 所提出的框架是否在不同主干网络架构上均表现出一致的性能提升，尤其是在小规模数据集上？
• RQ5 以数据为中心的预训练对模型收敛与泛化能力的影响如何？通过交叉验证损失与相关性度量进行评估。

主要发现

• 所提出的以数据为中心的预训练框架显著优于ImageNet迁移学习，在所有主干网络上均显著提升了R²与皮尔逊相关系数。
• 表现最佳的模型HCT-DenseNet121在地理范围预测中实现了 R² = 0.85 ± 0.05 和 ρ = 0.92 ± 0.02，表现出极高的预测准确性。
• 密度预测性能也显著提升，R² = 0.72 ± 0.09 与 ρ = 0.85 ± 0.06，表明实际值与预测值高度一致。
• 与独立的CNN相比，混合架构平均将均方误差（MSE）降低约0.2个点，平均将平均绝对误差（MAE）降低1.5个点。
• 混合模型的交叉验证损失始终低于独立CNN，表明其泛化能力更强且训练更稳定。
• 在DenseNet121上性能提升最为显著，地理范围预测的R²从0.74提升至0.82，经混合注意力机制整合后。