[论文解读] DKEC: Domain Knowledge Enhanced Multi-Label Classification for Diagnosis Prediction

Xueren Ge, Abhishek, Satpathy+3 more|arXiv (Cornell University)|Oct 10, 2023

Text and Document Classification Technologies被引用 4

一句话总结

DKEC 提出了一种用于医疗诊断预测的领域知识增强型多标签文本分类框架，通过标签级注意力机制和分组训练策略，整合异构医疗实体关系，以提升罕见（尾部）类别的性能。该方法在少样本标签上表现尤为出色，实现了最先进水平的结果，并使小型语言模型能够达到与大型模型相当的性能。

ABSTRACT

Multi-label text classification (MLTC) tasks in the medical domain often face the long-tail label distribution problem. Prior works have explored hierarchical label structures to find relevant information for few-shot classes, but mostly neglected to incorporate external knowledge from medical guidelines. This paper presents DKEC, Domain Knowledge Enhanced Classification for diagnosis prediction with two innovations: (1) automated construction of heterogeneous knowledge graphs from external sources to capture semantic relations among diverse medical entities, (2) incorporating the heterogeneous knowledge graphs in few-shot classification using a label-wise attention mechanism. We construct DKEC using three online medical knowledge sources and evaluate it on a real-world Emergency Medical Services (EMS) dataset and a public electronic health record (EHR) dataset. Results show that DKEC outperforms the state-of-the-art label-wise attention networks and transformer models of different sizes, particularly for the few-shot classes. More importantly, it helps the smaller language models achieve comparable performance to large language models.

研究动机与目标

• 为解决医疗多标签文本分类中的长尾标签分布问题，即罕见诊断缺乏足够的训练样本。
• 将来自医疗指南的结构化领域知识整合到神经网络模型中，以改善低资源标签的表征学习。
• 开发一种方法，提升小型语言模型在诊断预测任务中的性能，使其可与大型语言模型相媲美。
• 构建一个可泛化的框架，整合异构医疗实体关系（如症状、药物、诊疗方案）——超越传统本体。
• 在多种语言模型及真实世界电子健康记录（EHR）和电子病历摘要（ePCR）数据集上评估该方法。

提出的方法

• 从医疗指南构建异构图，整合多样化的医疗实体（如症状、药物、诊疗方案）及其异构关系（如‘治疗’、‘提示’）。
• 标签级注意力机制利用异构图计算动态注意力权重，使模型注意力聚焦于每个标签相关的医疗实体。
• 基于异构图中的嵌入，按语义相似性对标签进行分组，形成相关诊断的聚类。
• 分组训练在优化过程中对整个标签组计算二元交叉熵损失，从而增加罕见类别的有效样本量。
• 后处理应用逻辑规则（如基于年龄的区分）以优化组内预测，提升重叠或模糊标签的精确度。
• 该框架可集成于多种语言模型中，包括小型模型（如 TinyClinicalBERT），以评估其泛化能力与效率。

实验结果

研究问题

• RQ1 将官方指南中的异构医疗实体关系整合到模型中，是否能改善多标签医疗文本分类中的表征学习？
• RQ2 基于标签语义相似性的分组训练是否能有效缓解数据不平衡问题，并提升罕见（尾部）类别的性能？
• RQ3 在引入领域知识和分组训练后，小型语言模型是否能实现与大型语言模型相当的性能？
• RQ4 DKEC 框架在不同语言模型和真实世界医疗数据集上的泛化能力如何？
• RQ5 在数据不平衡的医疗数据集中，整合指南中的领域知识在多大程度上减少了对频繁（头部）标签的偏见？

主要发现

• 与最先进方法相比，DKEC 在 RAA 数据集上的宏平均 F1 分数提升 8.2%，达到 32.8%，其中在罕见（尾部）类别上增益最大。
• 在 MIMIC-III 数据集上，使用 TinyClinicalBERT（1500 万参数）的 DKEC-Comb 实现了 0.499 的宏平均 F1，优于基线模型 BioMedLM（27 亿参数）的 0.459。
• 在 MIMIC-III 中，TinyClinicalBERT 搭配 DKEC-Comb 实现了 0.516 的微平均 F1，优于基线模型 GatorTron（3.25 亿参数）的 0.523，表明其具有强大的泛化能力。
• 对于同一模型（TinyClinicalBERT），DKEC-Comb 实现了 0.234 的宏平均 F1，略高于基线模型 BioMedLM 的 0.248，证明其在小型模型上的有效性。
• 分组训练与异构标签级注意力机制显著提升了少样本学习性能，尤其在低资源诊断任务中表现突出。
• 该方法使小型模型实现与大型语言模型相当的性能，因此适用于资源受限设备的部署。