Accurate and Well-Calibrated ICD Code Assignment Through Attention Over Diverse Label Embeddings (2024)

Abstract: Although the International Classification of Diseases (ICD) has been adopted worldwide, manually assigning ICD codes to clinical text is time-consuming, error-prone, and expensive, motivating the development of automated approaches. This paper describes a novel approach for automated ICD coding, combining several ideas from previous related work. We specifically employ a strong Transformer-based model as a text encoder and, to handle lengthy clinical narratives, we explored either (a) adapting the base encoder model into a Longformer, or (b) dividing the text into chunks and processing each chunk independently. The representations produced by the encoder are combined with a label embedding mechanism that explores diverse ICD code synonyms. Experiments with different splits of the MIMIC-III dataset show that the proposed approach outperforms the current state-of-the-art models in ICD coding, with the label embeddings significantly contributing to the good performance. Our approach also leads to properly calibrated classification results, which can effectively inform downstream tasks such as quantification.

## introduction the international classification of diseases (icd foot_0 ) coding system, proposed by the world health organization, stands as a universal standard for precise documentation in the medical domain @xcite . still, the manual assignment of icd codes to clinical text is a time-consuming, labor intensive, and error-prone task, which has led to the exploration of automated methods, e.g. using deep learning algorithms for text classification. despite many previous efforts, automatic icd coding is still challenging. clinical notes consist of long text narratives that use a specialized medical vocabulary and are associated with a high dimensional, sparse, and imbalanced label space. in addition to accurately classifying individual clinical notes, estimating the prevalence of icd codes within a dataset is also important for many practical applications. this corresponds to a text quantification problem @xcite , for which access to properly calibrated text classification models can be helpful. this paper describes a novel approach to icd coding, combining several ideas from previous work. in particular, we explored two different text encoding strategies using a strong transformerbased model @xcite , explicitly dealing with the lengthy nature of documents like hospital discharge summaries. the resulting representations are combined with a label embedding mechanism inspired by the proposal from @xcite , which explores diverse icd code synonyms. additionally, taking inspiration from the mlp-based quantification approach from @xcite , we present a training setup in which multi-label classification and text quantification are jointly addressed. this additional step aims to improve model calibration while also informing downstream tasks such as text quantification. following previous studies, the proposed model was evaluated on the publicly available mimic-iii dataset @xcite , specifically analyzing results on two subsets of hospital discharge summaries, namely mimic-iii-50 @xcite and mimic-iii-clean @xcite . our approach surpasses standard baselines and previous state-of-the-art models for icd coding across all evaluated metrics, while simultaneously providing interesting results regarding model calibration and icd code quantification. the source code supporting our experiments is available in a github repository foot_1 . the remaining parts of this paper are organized as follows: section 2 reviews the related literature, while section 3 introduces our novel approach for icd coding and quantification. section 4 presents the experimental results, establishing a direct comparison with previous studies. finally, section 5 summarizes our contributions and discusses future research directions. the paper ends with a discussion on limitations and ethical considerations. ## related work several previous studies have addressed the problem of automatic icd coding. for instance, @xcite introduced the convolutional attention for multi-label classification (caml) approach, which is still commonly considered as a baseline. caml employs a label-wise attention mechanism that enables the model to learn distinct document representations for each label, selecting relevant parts of the document for each icd code. the authors conducted experiments on mimic datasets @xcite , and the data splits developed for this work were made publicly available. this study is considered an essential milestone for reproducibility. aiming to address caml's limitations in capturing variable-sized text patterns, @xcite improved the convolutional attention model by introducing a densely connected cnn with multi-scale feature attention (msatt-kg), which produces variable n-gram features and adaptively selects informative features based on neighborhood context. this method also incorporates a graph cnn to capture hierarchical relationships among medical codes. in turn, @xcite proposed mul-tirescnn, i.e. a tailored cnn architecture combining multi-filter convolutions and residual convolutions, capturing patterns of different lengths and achieving superior performance over caml. @xcite introduced laat, i.e., a model that combines an rnn-based encoder with a new label attention mechanism for icd coding. laat aimed to handle the variability in text segment lengths and the interdependence among different segments related to icd codes. additionally, the authors introduced a hierarchical joint learning mechanism to address the class imbalance issue. @xcite put forth the multiple synonyms matching network (msmn) as an alternative approach to icd coding. rather than relying on the icd code hierarchy, the authors leveraged synonyms to enhance code representation learning and improve coding performance. nowadays, transformer-based language models (lms) are becoming the fundamental technology for medical ai systems to process clinical narratives. for instance, @xcite developed a large clinical lm named gatortron, training a transformer encoder model on text narratives from de-identified clinical notes from university of florida health, pubmed articles, and wikipedia. the authors also examined how increasing the number of parameters enhances the performance on various natural language processing (nlp) tasks, including named entity recognition, medical relation extraction, semantic textual similarity, natural language inference, and medical question answering. their results showed that gatortron outperformed previous transformer models across various nlp tasks within the biomedical and clinical domains. @xcite compared transformer models for long document classification, focusing on mitigating the computational overheads associated with encoding large texts. in turn, @xcite investigated limitations associated with using pre-trained transformer-based lms, identifying challenges regarding large label spaces, long input lengths, and domain disparities. the authors proposed plm-icd, i.e., a framework that effectively handles these challenges and achieves superior results on the mimic-iii dataset, surpassing previously existing methods. in a recent study, @xcite argued that the proper assessment of model performance on icd coding had often struggled with weak configurations, poorly designed train-test splits, and inadequate evaluation procedures. the authors pinpointed significant issues with the mimic-iii splits released by @xcite . they proposed a new dataset split using stratified sampling to ensure a complete representation of all classes, referred to as mimic-iii-clean. regarding text quantification, various algorithms have been proposed in recent years @xcite . still, few previous studies have specifically considered multi-label settings @xcite . @xcite explored the use of a multi-layer perceptron (mlp) model for icd code quantification, taking inspiration from under-complete denoising auto-encoders. the mlp was trained to refine estimates provided by the probabilistic classify and count (pcc) method, considering label correlations. experiments with different mimic-iii dataset splits showed that the proposed method outperforms baseline approaches such as classify and count (cc) and pcc. ## proposed approach this work presents a novel approach for icd coding, aiming at strong classification performance together with well-calibrated outputs, which can inform downstream tasks such as text quantification. ## experimental evaluation this section presents the experimental evaluation of the proposed method, establishing a comparison with previously reported results. ## conclusions and future work this work introduced a novel deep learning method for icd coding, achieving state-of-the-art results in tests with two well established mimic-iii dataset splits. the proposed method processes long clinical documents and uses a label embedding mechanism that explores diverse icd code synonyms. besides achieving highly accurate classification results, the proposed approach produces well-calibrated estimates that can effectively inform downstream tasks such as text quantification. despite the strong results, it should be noted that our model does not exploit the hierarchical structure inherent to the icd coding system. thus, a promising avenue for further improvement involves using this structural knowledge, e.g. by implementing dual classification heads. another path worth exploring relates to using alternative methods to improve calibration (e.g., using other classification loss functions besides the bce), since improving calibration is beneficial for classification and essential for accurate results in quantification.