Modeling Diagnostic Label Correlation for Automatic ICD Coding (2021)

Abstract: Given the clinical notes written in electronic health records (EHRs), it is challenging to predict the diagnostic codes which is formulated as a multi-label classification task. The large set of labels, the hierarchical dependency, and the imbalanced data make this prediction task extremely hard. Most existing work built a binary prediction for each label independently, ignoring the dependencies between labels. To address this problem, we propose a two-stage framework to improve automatic ICD coding by capturing the label correlation. Specifically, we train a label set distribution estimator to rescore the probability of each label set candidate generated by a base predictor. This paper is the first attempt at learning the label set distribution as a reranking module for ICD coding. In the experiments, our proposed framework is able to improve upon best-performing predictors for medical code prediction on the benchmark MIMIC datasets.

## introduction clinical notes from electronic health records (ehrs) are free-from text generated by clinicians during patient visits. the associated diagnostic codes from the international classification of diseases (icd) represent diagnostic and procedural information of the visit. the icd codes provide an standardized and systematic way to encode information and has several potential use cases @xcite . considering that manual icd coding has been shown to be labor-intensive @xcite , several approaches for automatic icd coding has been proposed and investigated by the research community @xcite . with recent introduction of deep neural networks, the performance of automatic …250.61, 357.2, 564.0, 564.5, 401.9… a set of icd prediction constipation functional diarrhea 564.0 : 564.5 : figure 1 : an example of conflicting predictions. while these two codes share the same root in the hierarchical icd structure and are semantically similar, they are unlikely to appear together. icd coding has been improved significantly @xcite @xcite @xcite @xcite . prior work on neural models mostly treated the task of automatic icd coding as a multi-label classification problem. these models mostly employ a shared text encoder, and build one binary classifier for each label on top of the encoder. this architecture along side with binary cross-entropy loss make the prediction of each label independent of each other, which might lead to incomplete or conflicting predictions. an example of such error is shown in figure 1 . this issue is especially problematic in icd code prediction, since the icd codes share a hierarchical structure. that is, the low-level codes are more specific, and the highlevel ones are more general. in some cases, the low-level codes under the same high-level category are more likely to be jointly diagnosed. rare codes also have more opportunity to be considered from the frequent codes in the same highlevel class. prior work considered the hierarchical dependencies between icd codes by using hierarchical svm @xcite or by introducing new loss terms to leverage the icd structure @xcite . however, they borrowed the dependency from domain experts and did not consider the label correlation in the data. inspired by the success of reranking techniques on automatic speech recognition @xcite and dependency parsing @xcite , we propose a two-stage reranking framework for icd code prediction, which captures the label correlation without any expert knowledge. in the first stage, we use a base predictor to generate possible label set candidates. in the second stage, a label set reranker is employed to rerank the candidates. we design two rerankers to help to capture the correlation between labels. the experimental results show that our proposed framework consistently improves the results of different base predictors on the benchmark mimic datasets @xcite . the results also show that the proposed framework is model agnostic, i.e., we can use any base predictor in the first stage. data privacy is a major difficulty for medical nlp research. the personal health information (phi) which explains a patient's ailments, treatments and outcomes is highly sensitive, making it hard to distribute due to privacy concerns. in addition, ehrs across multiple hospitals or languages may contain different writing style, typos and abbreviations. it is labor-demanding to train separate models for each hospital with their in-house data only. the advantage of our proposed two-stage framework is that we can train base predictors with in-house data, while enjoying the universality of icd codes to train a reranker on icd labels from different sources. this reranker is able to generally work with various base predictor trained on health records from any specific hospital. the contributions of this paper are 3-fold: ## related work this paper focuses on multi-label medical code prediction; hence, we briefly describe the related background about medical code prediction and multi-label classification. ## proposed framework the task of icd code prediction is usually framed as a multi-label classification problem @xcite . given a clinical record x in ehr, the goal is to predict a set of icd codes y ⊆ y, where y denotes the set of all possible codes. this subset is typically represented as a binary vector @xmath0, 1} |y| , where each bit y i indicates the presence or absence of the corresponding label. the proposed framework is illustrated in figure 2 and consists of two stages: ## experiments in order to evaluate the effectiveness of our proposed framework, we conduct experiments on two benchmark datasets. we employ three different base predictors to validate the generalizability of the proposed framework. ## qualitative analysis after comparing the results produced from two rerankers, we find that both methods have similar tendency of prediction. in other words, the original candidate sets would be improved by adding or deleting similar icd codes after rescoring from both made and mask-sa. to further analyze prediction change, ## conclusions this paper proposes a novel framework to improve multi-label classification for automatic icd coding, which includes candidate generation and candidate reranking modules. in the first stage, a base predictor is performed to generate top-k probable label set candidates. in the second stage, we propose a reranker to capture the correlation between icd labels without any external knowledge. two types of the reranker, made and mask-sa, are employed to rerank the candidate sets. our experiments show that both rerankers can consistently improve the performance of all predictors in mimic-2 and mimic-3 datasets, demonstrating the generalizability of our framework and the great potential of the flexible usage.