[论文解读] K-Diag: Knowledge-enhanced Disease Diagnosis in Radiographic Imaging

Chaoyi Wu, Xiaoman Zhang+3 more|arXiv (Cornell University)|Feb 22, 2023

AI in cancer detection被引用 7

一句话总结

K-Diag 通过对比学习整合医学知识图谱，并利用可学习提示向量来适应视觉表征，提出了一种增强知识的放射科疾病诊断框架。该方法在长尾和零样本疾病识别任务中达到最先进性能，在零样本评估中，106个未见放射学发现中有79个的AUC >0.600。

ABSTRACT

In this paper, we consider the problem of disease diagnosis. Unlike the conventional learning paradigm that treats labels independently, we propose a knowledge-enhanced framework, that enables training visual representation with the guidance of medical domain knowledge. In particular, we make the following contributions: First, to explicitly incorporate experts' knowledge, we propose to learn a neural representation for the medical knowledge graph via contrastive learning, implicitly establishing relations between different medical concepts. Second, while training the visual encoder, we keep the parameters of the knowledge encoder frozen and propose to learn a set of prompt vectors for efficient adaptation. Third, we adopt a Transformer-based disease-query module for cross-model fusion, which naturally enables explainable diagnosis results via cross attention. To validate the effectiveness of our proposed framework, we conduct thorough experiments on three x-ray imaging datasets across different anatomy structures, showing our model is able to exploit the implicit relations between diseases/findings, thus is beneficial to the commonly encountered problem in the medical domain, namely, long-tailed and zero-shot recognition, which conventional methods either struggle or completely fail to realize.

研究动机与目标

• 为解决传统监督学习在医学影像中将疾病标签视为独立正交的问题，忽略其内在语义关系。
• 通过利用相关医学概念间的共享知识，提升长尾疾病识别的泛化能力。
• 通过将领域知识注入视觉表征学习过程，实现对训练中未见疾病的零样本诊断。
• 通过疾病查询与视觉特征之间的交叉注意力机制，构建可解释的诊断框架。

提出的方法

• 通过在医学知识图谱上进行对比学习训练知识编码器，将疾病和发现概念嵌入到连续的、关系感知的表征空间中。
• 视觉编码器在固定知识编码器参数的前提下进行训练，同时通过可学习提示向量模块实现对下游任务的高效适应。
• 基于Transformer的疾病查询模块通过在疾病名称（查询）与视觉特征之间执行交叉注意力，生成可解释的、注意力加权的预测结果。
• 框架使用从知识图谱派生的连续标签嵌入，而非独热向量，使模型能够利用病理之间的语义关系。
• 模型通过对比目标端到端训练，将视觉特征与知识嵌入的疾病描述对齐。
• 系统通过在推理过程中直接将未见疾病名称编码到知识空间，支持少样本和零样本推理。

Figure 1: In the conventional training scheme (left), manual annotations are often converted into discrete one-hot vectors, that are orthogonal in the embedding space, thus ignoring the implicit relations between labels. While in our proposed knowledge-enhanced classification framework (right), the

实验结果

研究问题

• RQ1 将医学领域知识注入视觉表征学习是否能提升放射影像中长尾疾病识别的性能？
• RQ2 增强知识的框架是否能实现训练中未见疾病的零样本诊断？
• RQ3 与独热编码相比，使用连续的、知识嵌入的标签表示在捕捉疾病间语义关系方面表现如何？
• RQ4 知识增强模型在不同解剖结构和标注粒度各异的数据集之间，其泛化能力在多大程度上得以体现？
• RQ5 疾病查询模块中的注意力机制能否为模型预测提供可解释的洞察？

主要发现

• 在三个X光数据集上，与使用独热标签的基线模型相比，K-Diag在ResNet上AUC提升2.53%，在ViT上提升2.15%。
• 在11个公开胸部X光数据集的联合数据集上进行训练时，K-Diag的平均AUC比基线高出2.14%，表明其具有更好的数据利用能力。
• 在PadChest数据集的零样本评估中，K-Diag在106个未见放射学发现中有79个的AUC达到至少0.600，这是传统监督学习无法实现的结果。
• 在零样本设置下，模型在14种发现上的AUC达到至少0.800，在46种发现上的AUC达到至少0.700，表明其对罕见或新发病理具有强大的泛化能力。
• 在不增加目标类别数据量的前提下，通过增加类别多样性（未见疾病）显著提升了所有11个数据集的性能，尤其在GoogleNIH和SIIM-ACR等小数据集上表现更优。
• 消融实验确认，知识编码器和提示学习组件至关重要，当任一组件被移除时，性能显著下降。

Figure 2: Overview of the knowledge-enhanced disease diagnosis workflow. The knowledge encoder (left) is first trained to learn a neural representation of the medical knowledge graph via contrastive learning, and then used to guide the visual representation learning in our knowledge-enhanced classif