Towards BERT -based Automatic ICD Coding: Limitations and Opportunities (2021)

Abstract: Automatic ICD coding is the task of assigning codes from the International Classification of Diseases (ICD) to medical notes. These codes describe the state of the patient and have multiple applications, e.g., computer-assisted diagnosis or epidemiological studies. ICD coding is a challenging task due to the complexity and length of medical notes. Unlike the general trend in language processing, no transformer model has been reported to reach high performance on this task. Here, we investigate in detail ICD coding using PubMedBERT, a state-of-the-art transformer model for biomedical language understanding. We find that the difficulty of fine-tuning the model on long pieces of text is the main limitation for BERT-based models on ICD coding. We run extensive experiments and show that despite the gap with current state-of-the-art, pretrained transformers can reach competitive performance using relatively small portions of text. We point at better methods to aggregate information from long texts as the main need for improving BERT-based ICD coding.

## introduction during patient stays in medical institutions, clinicians generate text notes that record the state of the patient as well as the diagnoses and the treatments administered. typically, a code from the international classification of diseases (icd) is assigned to these clinical notes, in order to provide standardized information about the patient condition. icd codes are used for different purposes, such as billing, computer-assisted diagnosis or epidemiological studies @xcite @xcite . assigning icd codes to medical notes is usually done manually by clinicians. this is an error-prone and time-consuming procedure and therefore, automatic solutions have been studied for over two decades @xcite . however, automatic icd code assignment proves challenging for multiple reasons. first, there exists a very large number of icd codes ( 17.000) and each clinical report may have associated more than one code. to deal with this large multi-label classification problem, it is common to reduce the number of codes to those that appear most frequently @xcite . second, medical text usually lacks structure, includes irrelevant passages, as well as abbreviations, misspellings, numbers and a very specific vocabulary. on top of that, medical notes are long, which makes it difficult for automatic coding models to draw relations between different sections of the reports. current state-of-the-art methods for automatic icd coding from medical notes are based on deep learning @xcite @xcite . these methods use different configurations of convolutional (cnn) and recurrent (rnn) neural networks as well as attention modules @xcite . this stands in contrast to most areas of natural language processing (nlp), where models based on the transformer architecture @xcite dominate the state-of-the-art @xcite . one of the main strengths of transformer models is their ability to deal with long range dependencies. this is a desirable property in icd coding, where an understanding of different parts of the document may be necessary to assign a code. the lack of transformer models for icd coding is surprising, especially since there already exist bert-based models @xcite (a type of bidirectional transformer) that are trained on medical text data @xcite @xcite . these models have achieved state-ofthe-art performance on other tasks such as named entity recognition or question answering on medical documents @xcite . on the other hand, the complexity of transformers scales quadratically with the length of their in-put, which restricts the maximum number of words that they can process at once. this limitation may be critical in icd coding, since clinical notes usually exceed this maximum input length. in this work, we investigate in detail bert-based icd coding, and explore different strategies to overcome the constraint on the input length by using an encoder-decoder architecture. we use the mimic-iii dataset @xcite , a big and widely used dataset for the icd coding task, in order that our results are directly comparable to other existing methods @xcite @xcite . by exposing the limitations and benefits of bert-based models on this task our work sets a solid basis for further research on automatic icd coding systems. ## related work automatic icd coding has been an active area of research for over two decades. already @xcite and @xcite proposed different strategies to extract features from medical documents in order to build classifiers for automatically assigning icd codes to medical notes. more recently, perotte et al. (2014) proposed a multi-level support vector machine (svm) model to predict icd codes from the mimic-ii dataset @xcite , the precursor of the mimic-iii dataset @xcite that we consider in this work. similarly, @xcite presented a method to extract features from structured and unstructured text and evaluated it on the mimic-iii dataset. in the last years, the state-of-the-art of automatic icd coding has been dominated by deep learning models. @xcite proposed an lstm model that operates at the character-level combined with an attention mechanism @xcite . @xcite proposed an embedding model based on glove embeddings @xcite that maps text and labels to the same space, where predictions are made using the cosine similarity. @xcite proposed a model that combined convolutions with a per-label attention mechanism. this model was further improved by @xcite and @xcite . @xcite , proposed a label-attention model that reached the current best performance for icd coding on the mimic-iii dataset. all of these works represent only a portion of the research carried out in this field @xcite @xcite @xcite . since the appearance of the transformer model @xcite , transformer-based architectures @xcite @xcite have become state-ofthe-art in almost every area of natural language processing @xcite thanks to their ability to handle long range dependencies. bert @xcite , a bidirectional transformer, is of particular importance since it is the basis of many other language understanding models. nonetheless, given the specific characteristics of medical text, e.g., specialized vocabulary, models pretrained on generic language, like bert, do not reach high performance on biomedical language understanding tasks. therefore, specialized models, such as biobert @xcite or clinicalbert @xcite , pretrained on medical text have been proposed. in particular, the recent pubmedbert model @xcite is the state-of-the-art in the blurb benchmark @xcite , a benchmark for biomedical language understanding which includes the following tasks: named entity recognition, question answering, document classification, relation extraction, sentence similarity and evidence-based medical information extraction. despite its prominence in medical language understanding, automatic icd coding escapes the set of tasks where bert-based models excel. to the best of our knowledge, no bertbased model has been proposed yet that reaches competitive performance on icd coding on the mimic-iii dataset. in this work, we investigate in detail bert-based icd coding and identify existing limitations and opportunities. ## background in this section we present the bert model used in our experiments as well as the evaluation metrics. ## dataset in this work, we use the widely-used mimic-iii dataset @xcite . this dataset contains medical information in various forms, however, as in previous studies @xcite @xcite , we consider exclusively the discharge summaries for icd coding. discharge summaries are medical notes created by doctors at the end of a stay in a medical facility and contain all the information about the stay. in the mimic-iii dataset, the length of the discharge summaries after tokenization ranges from 78 to 18, 429 tokens with a mean length of 2, 740 tokens and a median of 2, 500. each of these discharge summaries has associated to it one or more icd codes from the icd-9 taxonomy, with an average of 13.15 icd codes per summary. therefore, icd coding is a multi-label classification task. the mimic-iii dataset consists of 52, 722 discharge summaries with a total of 8, 921 unique icd codes. however, most of the codes are very infrequent, and therefore, existing work @xcite @xcite narrows down the task to finding only the 50 most frequent icd codes. we follow this strategy and use the reduced dataset, sometimes referred to as mimic-iii-50. this dataset consists of a training set of 8, 067 samples, a validation set of 1, 574 samples and a test set of 1, 730 samples. this data split is aligned with previous work, and thus, our results are directly comparable to those in the existing literature. ## model discharge summaries are longer than the maximum length accepted by pubmedbert such that it fits in the memory of a modern gpu and thus, we need to split the summaries into pieces of text. in order to process more than one piece of text per summary we adopt an encoder-decoder structure, where the encoder and the decoder are trained separately. this way, the encoder is the bert model that maps the different pieces of text to vector representations. these vector representations are then combined and decoded into icd codes by the decoder, which can be any kind of model. ## results we pose six research questions regarding the different strategies to encode and decode discharge summaries using a bert-based encoder. in our experiments, we fix the random seed so that all the results are comparable. how much does fine-tuning the encoder help decoding? here, we consider only the pubmedbert models fine-tuned on front, back and mixed data, since they were the only ones to learn during fine-tuning, as shown in section 5.1. to investigate the impact of this fine-tuning step on decoding performance, we use a simple linear layer which receives as input the concatenation of the encodings of the front, back and mixed chunks. each of these pieces of text is encoded by the pubmedbert model trained on that piece of text, i.e., we use three different encoders. we study the difference in performance for three different training points of the encoders: not fine-tuned, fine-tuned for three epochs and finetuned for six epochs. the results are detailed in these results show that fine-tuning the encoder significantly improves the decoding performance and that the best performance is obtained after six epochs. in fact, the difference between fine-tuning for six epochs and not fine-tuning is as large as 27.24 points for the macro auc score and 17.43 points for the micro auc score. we observed the same pattern in all of our experiments, and therefore, in the following we will only present results with the encoder fine-tuned for six epochs, unless stated otherwise. which of the three pieces of text, front, back or mixed, contains the most relevant information for icd coding? we experiment with a linear and a flat mlp decoder and apply these models to the encodings of each of the three chunks of text separately, i.e., front, back and mixed. we report the results in figure 2 . we see that front, i.e., the first 512 tokens of the discharge summary yields the best performance, both when the decoder is a linear layer and an mlp. although slightly inferior, the mixed chunk produces competitive scores while when using an mlp the auc scores are more than 3 points lower for back than for front. furthermore, using as decoder an mlp improves the performance significantly over using a linear layer; with the front non-linear model performing comparably to the combination of the three chunks with a linear decoder, as reported in the previous section, this naturally raises the question of whether the combination of the chunks yields an improvement. to study this, we use the same non-linear mlp architecture as in figure 2 (bottom) on 1) the concatenation of the encodings of front and back and 2) the concatenation of the three encodings, front, back and mixed. we report the results in these results show that combining front and back improves performance in comparison to using only front. as it may be expected, adding the mixed paragraph, which contains redundant information, produces only a small improvement. overall, the combination of the three chunks produces an improvement of 2.07 points for macro auc and 1.67 points for micro auc over using only front. given the larger input, these models have more parameters than the ones using only one of the chunks, which could partly explain the improvement, especially when adding redundant information, i.e., the mixed chunk. this result leads us to investigate the influence of the decoder architecture. how does the architecture of the decoder impact performance? here, we consider flat mlp, parallel mlp and transformer decoders on the combination of front, back and mixed. for each of these architectures, we evaluate three different sizes: base, large and x-large, where the difference between these sizes is only the number of layers and the size of the internal representations. this way, our experiments aim at discerning whether the structure of the decoder, the number of parameters, or both, influence the performance of the icd coding model. none of the models considered obtains a performance significantly higher than the others, with the largest difference across macro and micro auc scores being of only 0.28 and 0.57 points, respectively. this result is surprising since, given the complexity of the task, it could be expected that larger and more sophisticated decoders would perform better. notwithstanding, the saturation in performance suggests that all the information available in the input of the decoder is successfully extracted by every model, regardless of its complexity. this in turn indicates that the performance of the whole encoder-decoder model is limited by the reduced amount of text that is given as input (only the beginning and the end of the discharge summaries). therefore, we next consider providing larger portions of text from the discharge summaries as input. is dividing the discharge summaries by paragraphs a good splitting strategy? by splitting the discharge summaries into paragraphs we take into account information from a larger body of text than by using the front and the back. the main disadvantage of this approach is that the encoder fails to converge during fine-tuning. here, we test the hypothesis of whether the decoder can compensate the lack of fine-tuning of the encoder and, by leveraging the larger amount of information available, reach competitive performance. we encode the 200 most frequent paragraphs using the pubmedbert model fine-tuned on paragraph data, although due to lack of convergence during fine-tuning, we observed very similar results when using the not fine-tuned version. since not all the discharge summaries contain the same paragraphs, there is a misalignment between samples. for this reason, here we consider only the transformer decoder; the self-attention modules of the transformer should be able to cope with the misalignment better than the other architectures. we consider the transformer decoders (base, large and x-large) from the previous section. now, the transformer decoder receives 200 encoded representations, one per paragraph. given this large number of input representations or tokens, we aggregate the output of the transformers by taking the mean over the representations produced for all the paragraphs 3 . in figure 3 , we compare these paragraph decoders to the parallel mlp model on the front, back and mixed chunks from the previous section. we see that dividing the discharge summaries into paragraphs greatly under-performs in comparison to using the beginning and end of the summaries encoded by fine-tuned pubmedbert models. this result partly rejects the hypothesis that the decoder can benefit from a larger unstructured input. next, we continue investigating this hypothesis by feeding the decoder with the complete discharge summaries following the all strategy. ## discussion automatic icd coding from discharge summaries using transformer models has proven to be challenging. discharge summaries are very long documents and thus, they need to be divided into chunks in order to be entirely processed by bert-like models.this way, a decoder is necessary to combine the representations of each chunk, which are independently generated by the bert encoder. we have shown that for these representations to be meaningful the encoder needs to be fine-tuned on the icd coding task. it is straight-forward to fine-tune a bert encoder such as pubmedbert using specific parts of the summary, e.g., the beginning or the end. however, in our experiments, fine-tuning pubmedbert on excerpts extracted from different parts of the text, i.e., paragraph and all, prevented convergence due to the lack of alignment between samples, i.e., due to each training sample containing information from a different section of a discharge summary. furthermore, our results show that the decoder, regardless of its architecture, cannot compensate for lack of convergence during the fine-tuning of the encoder. on the other hand, our best bert-icd model reaches competitive performance using only 1, 024 tokens (front and back), which represents a significantly smaller portion of text than state-of-the-art models, based on cnns and rnns. unlike bert, cnn and rnn models can extract information from texts of any length without needing to split them, which allows for end-to-end training over long pieces of text. @xcite found that the performance of their convolutional attention model benefits from longer input texts until a length of between 2, 500 and 6, 500 words, and @xcite use up to 4, 000 words as input. our model combines encodings from the beginning and the end of the discharge summary, and reaches better performance in that case than when it processes either of those portions of text alone. this supports the statement that including more text improves icd coding. all of these results suggest that the difficulty of fine-tuning a bert encoder on long pieces of text is the main bottleneck for performance and the reason for the existing gap with state-of-the-art models for icd coding. one of the main advantages of transformer models over cnns and rnns is that they can handle long-range dependencies. hence, if longer text could be fed at once into a bert encoder it would be possible to find relationships and patterns over longer spans of text. it is therefore likely that advances either in terms of hardware, i.e., larger gpu memories allowing for longer pieces of text to be processed at once; or in compressing bert-like models, e.g., distillation, will progressively close the gap with the state-of-the-art, following the same trend of other areas of nlp. on top of that, we consider that the two most promising directions for future research on bert-based icd coding are: 1) devising strategies to fine-tune the encoder over longer spans of text, e.g., building an ensemble of models where each of them is trained on one section of the text; 2) improving the methods to aggregate encodings from different parts of the document. finally, to deploy automatic icd coding systems in the real world, it is important that their decisions can be explained. explaining transformer models is currently a field of active research, and although there exist important concerns about the interpretability of attention distributions in transformers @xcite , methods based on gradient attribution @xcite or on attention flow (abnar and zuidema, 2020) can provide insights on their decision-making. a bert-based icd coding system could directly benefit from this field of research and eventually provide explanations together with its icd code predictions. ## conclusion contrary to what is common in most nlp tasks, the transformer architecture is not the state-of-the-art in assigning icd codes to discharge summaries. in this work, we have presented a thorough study of the performance of bert-based models on this task and we have identified the length of the discharge summaries as the main obstacle holding back their performance. our work sets a solid foundation for further research on icd coding and suggests that overcoming the exposed limitations of bert-based models is likely to lead to a new state-of-the-art. furthermore, we believe that the interpretability of icd coding models is an interesting avenue for future work, which can benefit from a large body of existing research.