Multi-stage Retrieve and Re-rank Model for Automatic Medical Coding Recommendation (2024)

Abstract: The International Classification of Diseases (ICD) serves as a definitive medical classification system encompassing a wide range of diseases and conditions. The primary objective of ICD indexing is to allocate a subset of ICD codes to a medical record, which facilitates standardized documentation and management of various health conditions. Most existing approaches have suffered from selecting the proper label subsets from an extremely large ICD collection with a heavy long-tailed label distribution. In this paper, we leverage a multi-stage “retrieve and re-rank” framework as a novel solution to ICD indexing, via a hybrid discrete retrieval method, and re-rank retrieved candidates with contrastive learning that allows the model to make more accurate predictions from a simplified label space. The retrieval model is a hybrid of auxiliary knowledge of the electronic health records (EHR) and a discrete retrieval method (BM25), which efficiently collects high-quality candidates. In the last stage, we propose a label co-occurrence guided contrastive re-ranking model, which re-ranks the candidate labels by pulling together the clinical notes with positive ICD codes. Experimental results show the proposed method achieves state-of-the-art performance on a number of measures on the MIMIC-III benchmark.

## introduction electronic health records 1 (ehrs) contain a comprehensive repository of essential administrative and clinical data pertinent to a person's care within a specific healthcare provider setting. in order to conduct meaningful statistical analysis, these ehr data are annotated with structured codes in a classification system known as medical codes. the international classification of diseases 2 (icd) is 1 https://www.cms.gov/medicare/e-health/ ehealthrecords 2 https://www.who.int/standards/classifications/ classification-of-diseases one of the most widely-used coding systems, and it provides a taxonomy of classes, each uniquely identified by a code assigned to an episode of patient care. the task of medical coding associates icd codes with ehr documents. the status quo of assigning medical codes is a manual process, which is labour-intensive, time-consuming, and error-prone @xcite . to reduce coding errors and cost, the demand for automated medical coding has become imperative. previous deep learning approaches regarded medical coding as an extreme multi-label text classification problem @xcite @xcite @xcite , where an encoder is typically employed to learn the representations of the clinical notes and a label-specific binary classifier is subsequently constructed on top of the encoder for label predictions. however, some remaining difficulties have still posed immense challenges. first, clinical documents are lengthy (containing on average 1596 words in the mimic-iii dataset) and noisy (including terse abbreviations, symbols, and misspellings). second, the label set is extremely large and complex; for instance, in the 10 th icd edition, there are over 130,000 codes foot_0 . third, the distribution of icd codes is extremely long-tailed; while some icd codes occur frequently, many others seldom appear, if at all, because of the rarity of the diseases. for instance, among the 942 unique 3-digit icd codes in the mimic-iii dataset @xcite , the ten most common codes account for 26% of all code occurrences and the 437 least common codes account for only 1% of occurrences @xcite . to address the aforementioned challenges, we propose a novel multistage retrieve and re-rank framework, where the goal is to first generate a curated icd list and then figure 1 : an example of a medical record from the mimic-iii dataset which includes the discharge summary, assigned icd codes and auxiliary knowledge. we colour each code and its corresponding mentions in the discharge summary and auxiliary knowledge. we use the auxiliary knowledge of the notes to retrieve the candidate subset of the label space. to summarize, the major contributions of this paper are: the automatic icd indexing task is well established in the healthcare domain. extensive research using deep learning has been dedicated to icd indexing, including recurrent-based neural networks (rnns), convolution-based neural networks (cnns), and their variations @xcite @xcite . these architectures are able to extract and categorize semantic features, reducing the need for medical domain expertise during the traditional feature selection stage seen in conventional algorithms @xcite . the icd indexing task is formulated as a multi-label classification problem in these approaches. mullenbach et al. (2018) introduced a combination of cnn with an attention mechanism to effectively capture pertinent information within clinical texts for each icd code. building on this foundation, xie et al. (2019) enhanced the cnn attention model by integrating a multi-scale feature attention technique. many cnn variants were subsequently introduced to address the challenges posed by lengthy and noisy clinical texts, including multirescnn (li and yu, 2020), dcan (ji et al., 2020), and ef-fectivecan (liu et al., 2021). rnn-based models, renowned for their capacity to capture contextual information across input texts, have also been widely used for icd indexing. shi et al. (2017) proposed a character-aware long short-term memory (lstm) recurrent network to learn the underlying representations of clinical texts. xie and xing (2018) introduced a tree-of-sequences lstm architecture alongside adversarial learning to capture hierarchical relationships among icd codes. additionally, baumel et al. (2018) presented a hierarchical attention-bidirectional gated recurrent unit (ha-gru) model, facilitating document labeling by identifying sentences relevant to each icd code. laat (vu et al., 2020) used a bidirectional long-short term memory (bilstm) encoder and a customized label-wise attention mechanism to cultivate label-specific vectors across distinct clinical text fragments. to address the hierarchical relationships intrinsic to icd codes, graph convolutional neural networks (gcnns) @xcite have emerged as a powerful tool. @xcite and @xcite used gcnns to capture both the hierarchical interplay among icd codes and the semantic information specific to each code. hypercore @xcite took a comprehensive approach by considering both code hierarchy and code co-occurrence, employing gcnns to learn code representations within the co-graph. incorporating external knowledge beyond icd code information has also gained traction. @xcite introduced a knowledge source integration (ksi) model that integrates external knowledge from wikipedia. this integration calculated matching scores between clinical notes and disease-related wikipedia documents, in order to enrich the available information for icd predictions. additionally, @xcite proposed a multiple synonym matching network (msmn) to use synonyms of icd codes, enhancing the quality of code representation learning. expanding on this, @xcite integrated a pre-trained language model with three domain-specific knowledge sources: code hierarchy, synonyms, and abbreviations. this fusion of knowledge sources contributes significantly to the performance of icd classification. ## results and discussion in order to asses the efficacy of our proposed framework, we compare with the existing state-of-the-art (sota) models, as outlined in as the occurrence frequencies of the icd codes are imbalanced, our focus lies in assessing the efficacy of our model, specifically on the infrequently appearing icd codes. we categorize the icd codes into four groups based on their occurrences in the training set: [0, 10), [10, 50), [50, 500), and [500,). figure 3 illustrates the distribution of icd codes and their occurrence percentages across the four categorized groups in the training set, which show that the distribution of icd frequency is highly biased, conforming to a long-tail distribution. figures 3b and 3c present the performance of our model on mimic-iii-full in com-parison to the caml baseline @xcite across the four icd groups on macro-auc and micro-f1, respectively. our model demonstrates significant improvements for both frequent and infrequent labels on both metrics. to confirm the specific contributions of these modules in terms of enhancing both the effectiveness and robustness of the model, we conduct ablation studies with three different settings: (a) we examine the effectiveness of using auxiliary knowledge in the retrieval stage by removing the retrieval stage and rank the icd codes on the whole label set; (b) we examine the influence of different embedding methods by replacing the clinical-longformer with clinical-bert; and (c) we test the effectiveness of label embedding by replacing the encoding of the label with the average of word embeddings in the label descriptors. the experimental results are shown in ## conclusion in this paper, we regard the medical coding task as a recommendation problem and present a novel multi-stage retrieve and re-rank framework. the primary objective of the proposed framework is twofold: to construct a curated list of icd codes and, subsequently, to further refine the candidate list for a given medical record. specifically, we first conduct a two-step retrieval process, incorporating auxiliary knowledge and the bm25 algorithm. this approach retrieves a concise subset of the candidate list, mitigating the challenges of a very large and imbalanced label distribution. we then use a reranking model to refine the previously obtained candidate list, employing code co-occurrence guided contrastive learning. experimental results demonstrate that our proposed framework outperforms the previous sota, which suggests that it provides more precise and contextually grounded icd recommendations for the given medical records. in the future, our proposed framework may be extended with more external knowledge such as the unified medical language system (umls) and code synonymy.