Knowledge Injected Prompt Based Fine-tuning for Multi-label Few-shot ICD Coding (2022)

Abstract: Automatic International Classification of Diseases (ICD) coding aims to assign multiple ICD codes to a medical note with average length of 3,000+ tokens. This task is challenging due to a high-dimensional space of multi-label assignment (tens of thousands of ICD codes) and the long-tail challenge: only a few codes (common diseases) are frequently assigned while most codes (rare diseases) are infrequently assigned. This study addresses the long-tail challenge by adapting a prompt-based fine-tuning technique with label semantics, which has been shown to be effective under few-shot setting. To further enhance the performance in medical domain, we propose a knowledge-enhanced longformer by injecting three domain-specific knowledge: hierarchy, synonym, and abbreviation with additional pretraining using contrastive learning. Experiments on MIMIC-III-full, a benchmark dataset of code assignment, show that our proposed method outperforms previous state-of-the-art method in 14.5% in marco F1 (from 10.3 to 11.8, P<0.001). To further test our model on few-shot setting, we created a new rare diseases coding dataset, MIMIC-III-rare50, on which our model improves marco F1 from 17.1 to 30.4 and micro F1 from 17.2 to 32.6 compared to previous method.

## introduction multi-label learning has many real-word applications in natural language processing (nlp), including but not limited to academic paper labeling @xcite , news framing @xcite , waste crises response @xcite , amazon product labeling @xcite , and medical coding @xcite . in contrast to multi-class classification, an instance in multi-label learning is frequently linked with more than one class label, making the task more challenging due to the combination of potential class labels. in real-world tasks, there are often insufficient training data for rare class labels. taking automatic international classification of diseases (icd) coding as example, given discharge summaries notes as input, the task is to assign multiple icd disease and procedure label codes associated with each note. the assigned codes need to be accurate and complete for the billing purposes. as an example, the mimic-iii dataset @xcite ) contains 8,692 unique icd-9 codes, among which 4,115 (47.3%) codes occur less than 6 times and 203 (2.3%) occur zero times. clinical practice requires a high accuracy, hence, it is not acceptable for a multi-label classifier to fail a disease diagnosis (or code assignment) because it is rare, since such a diagnosis may be of the most clinical importance for the patient. therefore, the classifier is required to perform with high precision even for infrequent codes. this translates to data sparsity due to availability of few training examples. to mitigate the data sparsity problem, additional structured knowledge could be applied. icd codes are organized with an ontological/hierarchical structure where a text description is associated to each code. for instance, icd 250 (diabetes mellitus), shown in figure 1 , is the parent of several child codes including 250.0 (diabetes mellitus without mention of complication), 250.1 (diabetes with ketoacidosis), and 250.2 (diabetes with hyperosmolarity). such child icd codes are more semantically different from each other than their parent code 250. synonyms including acronyms and abbreviations are common in medical notes. for instance, the description of code 250.00 is disease "type ii diabetes mellitus". however, this code can be described in different text forms such as "insulinresistant diabetes", "non-insulin dependent diabetes", "dm2", and "t2dm". therefore, one naive way to assign icd codes is to identify matching between candidate code descriptions and their syn-figure 1 : an illustration of self-alignment pretraining from medical knowledge umls, including the usage of (a) hierarchy, (b) synonym, (c) abbreviation. pink region is the dynamic margin ranges from π/2 to π where we wish to pull negatives apart with a dynamic distance. in this paper, we present a simple but effective knowledge enhanced prompt (kept) framework. we implement and evaluate kept using a lm based on longformer because clinical notes are typically more than 500 tokens. specifically, we first pretrain mimic a longformer lm on mimic-iii dataset. then, we further pretrain umls on structured medical knowledge umls (unified medical language system) using self-alignment learning with contrastive loss to inject medical knowledge into pretrained lm. for the downstream icd-code assignment fine-tuning, we add a sequence of icd code descriptions (label semantics) as prompts in addition to each clinical note as kept lm input. this allows early fusion of code descriptions and the input note. experiments on full disease coding (mimic-iii-full) and common disease coding (mimic-iii-50) show that our keptlongformer outperforms previous sota msmn @xcite . in order to test its few-shot ability, we create a new few-shot rare diseases coding dataset named mimic-iii-rare50, and results show significant improvements compared between msmn and our method. to facilitate future research, we publicly release the code and trained models foot_0 . ## conclusions in this paper, we investigate pretrained clinical language model on auto icd coding task for both common and rare diseases, the latter of which has received limited attention in the past. built on recent advances in contrastive learning, entity representation training and prompt based fine-tuning, our keptlongformer easily achieves a competitive performance over state of the art system in common code assignment, and significantly outperforms baseline model in rare code assignment tasks. finally, our novel hierarchical self-alignment pretrain could be easily applied to other multi-label classification problems such as tumor detection using other ontology such as oncotree.