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

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.

## introduction automated diagnosis prediction @xcite is the challenging task of classifying different diseases based on a patient's ehr for applications such as treatment recommendation (e.g., selecting ems protocols @xcite @xcite ) or medical billing (e.g., assigning icd-9 codes) @xcite . diagnosis prediction based on the free-text medical notes is known as multi-label text classification (mltc) @xcite , which is the task of assigning the most relevant labels to a text instance. mltc is more complex than the traditional multi-class text classification because the number of possible label combinations grows exponentially with the number of classes @xcite . another challenge in diagnosis prediction is the imbalanced distribution of diagnoses as some medical conditions happen more frequently than others, causing a long-tail data distribution. for example, the total number of chest pain-related reports in a real-world ems dataset is ten times more than overdose/poisoning reports @xcite . training on such imbalanced datasets, also called "powerlaw datasets" @xcite , introduces bias in model predictions towards head label classes while ignoring the few-shot or tail classes. most existing diagnosis prediction solutions @xcite are task-agnostic and rely on integrating biomedical domain knowledge with transformer models in the pre-training stage. for example, core-bert @xcite uses clinical outcome pre-training to learn relations among symptoms, risk factors and clinical outcomes by incorporating wikipedia and pubmed knowledge bases. recent pre-trained large language models (llms) @xcite @xcite have demonstrated superior performance by leveraging large external clinical corpora and huge number of parameters. however, these models only incorporate uncurated knowledge in pre-training, neglect task-specific knowledge and label relations, and are costly to fine-tune and deploy on resource-constrained devices @xcite . to solve the class-imbalance problem in mltc, the convolutional attention network and its variants @xcite @xcite were proposed to extract meaningful document representations that cover different ranges of clinical text. other works @xcite integrated hierarchical information by graph convolutional neural networks to select label-models encoder attention mechanism knowledge integration knowledge source datasets @xcite bert self-attention pre-training wikipedia, pubmed mimic-iii @xcite megatronbert self-attention pre-training wikipedia, pubmed mimic-iii @xcite gpt2 self-attention pre-training pubmed medmcqa @xcite cnn label-wise attention mimic-iii @xcite cnn label-wise attention icd-9 hierarchy graph icd-9 description mimic-iii @xcite multi-filter residual cnn label-wise attention mimic-iii @xcite multi relevant features. some follow-up studies @xcite @xcite have also proposed to incorporate label co-occurrence graphs, along with hierarchical structures to capture label concurrent and mutual exclusive relations for icd-9 code classification. however, most of these works neglected the potential benefits of incorporating expert knowledge from medical guidelines. external domain knowledge can provide additional information for training with few-shot labels to compensate for data scarcity and model size or be applied as constraints in training based on label relations. this paper presents dkec (figure 1 ), a knowledge and data-driven approach to class-imbalanced mltc by (i) automated extraction of label-specific semantic relations from online sources and (ii) integrating them as heterogeneous knowledge graphs with different encoders using a label-wise attention mechanism. our contributions are as follows: pre-trained transformers for diagnosis prediction one avenue explored in prior work is focused on large-scale pre-training from clinical admissions, discharge summaries, and other biomedical texts, such as biobert @xcite , @xcite and gatortron @xcite ( ## heterogeneous knowledge graph construction our goal is to construct a heterogeneous knowledge graph g, which for every disease diagnosis code d k in a set of diagnosis codes d : {d k } l @xmath0, represents the corresponding sets of medical concepts such as signs and symptoms s : and hierarchy : {h k } |h| @xmath1. as shown in figure 1 , the heterogeneous graph of medical concepts is constructed as @xmath2, e), with n as the set of nodes and e as the set of edges. there are four different types of nodes in the graph, diagnosis codes d, signs and symptoms s, treatment t , and hierarchy h, and three types of bidirec-tional edges, has/indicates ←→ e ds between d and s, suggests/administers ←→ e dt between d and t , and children/parent ←→ e dh between d and h. for example, the "injury -crush syndrome" diagnosis code d i is connected to the signs and symptom "muscle mass" s j using an edge of type "has/indicates" ←→ e ds . "injury -head" (d j ) and d i are the children of ( ←→ e dh ) node "adult trauma emergencies". next, we present a systematic method for automated extraction of disease-relevant medical entities from external knowledge sources and mapping them into normalized concepts for unique representation in the graph. for every diagnosis code in a training dataset, we extract the relevant symptoms and treatments for the disease textual descriptions in online sources and generate the @xmath3, relation, t k /s k >. the hierarchy information is given by the label coding in each dataset. ## knowledge-enhanced classification figure 1 shows the overall dkec pipeline, consisting of (i) a text branch that extracts features from the input ehr notes using a long document encoder (ii) a graph branch that takes the heterogeneous knowledge graph g as input and incorporates it for label co-occurrence extraction and multi-label classification via a heterogeneous label-wise attention (hla) mechanism. ## conclusion this paper proposes a domain knowledge-enhanced multi-label text classification method for diagnosis prediction. we present an approach for automatic knowledge graph construction from online sources based on medical entity extraction using chain-ofthought prompting with gpt-4. we also introduce a heterogeneous label-wise attention mechanism that incorporates relations among diverse medical entities in the knowledge graph to capture labelrelated text features for classification. we evaluated our methods on two real-world datasets. experiments show the accuracy of knowledge graph construction based on three knowledge bases and improved performance over several sota methods. we also demonstrated the applicability of our approach to different language models sizes and its scalability to large number of labels.