[Paper Review] 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 Technologies4 citations

TL;DR

DKEC proposes a domain knowledge-enhanced multi-label text classification framework for medical diagnosis prediction that integrates heterogeneous medical entity relationships via a label-wise attention mechanism and a group-wise training strategy to improve performance on rare (tail) classes. It achieves state-of-the-art results, especially for few-shot labels, and enables small language models to match the performance of large models.

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.

Motivation & Objective

• To address the long-tail label distribution problem in medical multi-label text classification, where rare diagnoses have insufficient training samples.
• To incorporate structured domain knowledge from medical guidelines into neural models to improve representation learning for low-resource labels.
• To develop a method that enhances small language models to achieve performance comparable to large language models in diagnosis prediction tasks.
• To create a generalizable framework that integrates heterogeneous medical entity relations (e.g., signs, medications, protocols) beyond traditional ontologies.
• To evaluate the method across diverse language models and real-world EHR and ePCR datasets.

Proposed method

• A heterogeneous graph is constructed from medical guidelines, integrating diverse medical entities (e.g., symptoms, drugs, protocols) and their heterogeneous relations (e.g., 'treats', 'indicates').
• A label-wise attention mechanism uses the heterogeneous graph to compute dynamic attention weights, focusing model attention on relevant medical entities for each label.
• Labels are grouped based on semantic similarity using embeddings from the heterogeneous graph, forming clusters of related diagnoses.
• Group-wise training computes binary cross-entropy loss over entire label groups during optimization, increasing effective sample size for rare classes.
• Post-processing applies logic rules (e.g., age-based distinctions) to refine predictions within each group, improving precision for overlapping or ambiguous labels.
• The framework is integrated with various language models, including small models like TinyClinicalBERT, to assess generalization and efficiency.

Experimental results

Research questions

• RQ1 Can incorporating heterogeneous medical entity relations from official guidelines improve representation learning in multi-label medical text classification?
• RQ2 Does group-wise training based on label semantic similarity effectively mitigate data imbalance and improve performance on rare (tail) classes?
• RQ3 Can small language models achieve performance comparable to large language models when enhanced with domain knowledge and group-wise training?
• RQ4 How generalizable is the DKEC framework across different language models and real-world medical datasets?
• RQ5 To what extent does the integration of domain knowledge from guidelines reduce bias toward frequent (head) labels in imbalanced medical datasets?

Key findings

• DKEC improves macro F1 score by 8.2% to 32.8% on the RAA dataset compared to state-of-the-art methods, with the largest gains observed on rare (tail) classes.
• On the MIMIC-III dataset, DKEC-Comb achieves a macro F1 of 0.499 with TinyClinicalBERT (15M parameters), outperforming the baseline BioMedLM (2.7B parameters) with 0.459.
• In MIMIC-III, TinyClinicalBERT with DKEC-Comb achieves 0.516 micro F1, outperforming the baseline GatorTron (325M parameters) at 0.523, indicating strong generalization.
• For the same model (TinyClinicalBERT), DKEC-Comb achieves 0.234 macro F1, slightly better than the baseline BioMedLM (0.248), demonstrating effectiveness on small models.
• The group-wise training and heterogeneous label-wise attention mechanisms significantly improve few-shot learning, particularly for low-resource diagnoses.
• The method enables small models to achieve performance comparable to large language models, making it suitable for deployment on resource-constrained devices.