[论文解读] HydraViT: Adaptive Multi-Branch Transformer for Multi-Label Disease Classification from Chest X-ray Images

Şaban Öztürk, Mehmet Y. Turali+1 more|arXiv (Cornell University)|Oct 9, 2023

COVID-19 diagnosis using AI被引用 5

一句话总结

HydraViT 提出了一种新颖的自适应多分支变换器模型，用于多标签胸部X光（CXR）疾病分类，结合卷积-空间编码器与基于变换器的上下文编码器，以捕捉长距离依赖关系，并采用带有学习加权的多分支输出模块，提升对单个及共现病灶的敏感性。该模型在多标签分类任务中达到最先进性能，平均优于现有方法1.2%。

ABSTRACT

Chest X-ray is an essential diagnostic tool in the identification of chest diseases given its high sensitivity to pathological abnormalities in the lungs. However, image-driven diagnosis is still challenging due to heterogeneity in size and location of pathology, as well as visual similarities and co-occurrence of separate pathology. Since disease-related regions often occupy a relatively small portion of diagnostic images, classification models based on traditional convolutional neural networks (CNNs) are adversely affected given their locality bias. While CNNs were previously augmented with attention maps or spatial masks to guide focus on potentially critical regions, learning localization guidance under heterogeneity in the spatial distribution of pathology is challenging. To improve multi-label classification performance, here we propose a novel method, HydraViT, that synergistically combines a transformer backbone with a multi-branch output module with learned weighting. The transformer backbone enhances sensitivity to long-range context in X-ray images, while using the self-attention mechanism to adaptively focus on task-critical regions. The multi-branch output module dedicates an independent branch to each disease label to attain robust learning across separate disease classes, along with an aggregated branch across labels to maintain sensitivity to co-occurrence relationships among pathology. Experiments demonstrate that, on average, HydraViT outperforms competing attention-guided methods by 1.2%, region-guided methods by 1.4%, and semantic-guided methods by 1.0% in multi-label classification performance.

研究动机与目标

• 解决胸部X光（CXR）图像中病灶大小、位置及共现性差异带来的多标签疾病分类挑战。
• 克服传统CNN及注意力增强模型在捕捉病灶间长距离上下文与共现模式方面的局限性。
• 通过集成多分支输出模块，分别学习单个疾病标签并联合建模标签共现关系，提升分类性能。
• 开发一种计算效率高且高度敏感的深度学习框架，实现对胸部X光扫描中多种胸腔疾病自动化、精准诊断。

提出的方法

• 采用混合架构，先通过卷积空间编码器从CXR图像中提取初始特征图，再通过基于变换器的上下文编码器建模图像块之间的长距离依赖关系。
• 在变换器编码器中利用自注意力机制，自适应聚焦于任务关键且与病灶相关的区域，增强对小病灶或分散病灶的敏感性。
• 实现多分支输出模块，为14种疾病标签中的每一种配置专用头，并设置一个聚合头以建模病灶间的共现关系。
• 在多分支输出上引入可学习的自适应加权机制，以在训练过程中平衡单个标签性能与共现敏感性。
• 使用标签特定损失项与联合损失项相结合的方式训练模型，并引入可学习权重以缓解类别不平衡并提升泛化能力。
• 应用GradCAM生成显著性热力图以增强可解释性，验证注意力机制在代表性CXR图像中正确聚焦于病灶区域。

实验结果

研究问题

• RQ1 与标准CNN相比，混合卷积-变换器主干网络是否能更有效地提升在小病灶或散在病灶的CXR图像中长距离上下文建模能力？
• RQ2 带有自适应加权的多分支输出模块是否能通过同时优化单个病灶与共现病灶的分类性能，提升多标签CXR分类性能？
• RQ3 在平均与逐标签分类性能方面，HydraViT相较于最先进注意力引导、区域引导与语义引导方法表现如何？
• RQ4 所提出的架构在多大程度上通过注意力可视化提升了定位准确率与模型可解释性？

主要发现

• HydraViT 在CXR数据集上的多标签分类性能平均优于竞争性注意力引导方法1.2%。
• 与区域引导方法相比，其在全部14种疾病标签的平均F1分数上提升1.4%，与语义引导方法相比提升1.0%。
• 模型在预测标签与真实标签之间表现出更优的一致性，真实标签的平均预测分数显著高于非存在标签（例如，0.4 vs. 0.1），尤其在前五名预测中表现明显。
• 消融实验证实，基于变换器的上下文编码器与多分支输出模块均对性能有显著贡献，其中后者使训练速度相比顺序单分支训练提升近14倍。
• GradCAM热力图的视觉检查表明，HydraViT能有效突出与病灶相关的区域，表明注意力机制定位稳健。
• 该模型在各类病灶中均保持高性能，即使在低频出现或共现模式复杂的病灶中亦表现优异，表明其具备更强的泛化能力并减少了偏差。