DKEC: Domain Knowledge Enhanced Multi-Label Classification for Diagnosis Prediction

TL;DR

This work tackles long-tail multi-label diagnosis prediction from free-text EHR notes by introducing DKEC, which automatically builds heterogeneous knowledge graphs from online medical sources and integrates them via a heterogeneous label-wise attention mechanism. The method combines a long-document encoder with a graph transformer to fuse label-specific medical knowledge (Symptoms, Treatments, and Hierarchy) into per-label document representations, enabling improved few-shot performance. Empirical results on EMS and MIMIC-III datasets show that DKEC outperforms state-of-the-art baselines, especially for tail labels, and that the approach helps smaller language models approach the performance of larger LLMs, highlighting practical benefits for deployment on resource-constrained devices. The work demonstrates the potential of external domain knowledge, normalized via UMLS, to enhance MLTC in medical settings, while noting limitations in knowledge coverage and the need for further human evaluation and causal reasoning.

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.

Figures (6)

• Figure 1: (a) DKEC Pipeline includes three main modules: a text branch to derive text embeddings, a graph branch to derive updated diagnosis embeddings, and (b) an HLA module to derive label-attentive document embeddings.
• Figure 2: Knowledge Graph Construction
• Figure 3: DKEC with different pre-trained transformers.
• Figure 4: Performance on subsets of MIMIC-III dataset with varying label sizes. Subsets with 1.0k and 3.7k labels have full knowledge. 6.7k has partial knowledge.
• Figure 5: Spurious relation examples. Label-related keywords are in red, and spurious words are in blue.