BERT for Long Documents: A Case Study of Automated ICD Coding (2022)

Abstract: Transformer models have achieved great success across many NLP problems. However, previous studies in automated ICD coding concluded that these models fail to outperform some of the earlier solutions such as CNN-based models. In this paper we challenge this conclusion. We present a simple and scalable method to process long text with the existing transformer models such as BERT. We show that this method significantly improves the previous results reported for transformer models in ICD coding, and is able to outperform one of the prominent CNN-based methods.

## introduction the international classification of diseases (icd) codes provide a standard way of keeping track of diagnoses and procedures during a patient visit. these codes are used worldwide for epidemiological studies, billing and reimbursement, and research in health care. the codes are maintained by the world health organization (who) and are revised and updated periodically. as of 2022 the icd codes are in the 11th revision. assigning icd codes to a clinical note, such as a discharge summary, is done by professional medical coders. human coders require extensive training, and the process of coding is often timeconsuming, costly, and error-prone. due to these challenges there is an incentive to automate the coding process. therefore in recent years this problem has gained interest among machine learning researchers in health care @xcite ; @xcite ; @xcite and references therein). on the surface, the problem can be considered as a multi-label document classification problem. however, there are aspects of the problem that make it particularly challenging. the primary challenge is that there are tens of thousands of classes. for instance, billable icd-10-cm codes consist of approximately 73,000 codes. in addition, the distribution of the codes is not uniform. many of the codes are related to rare conditions and are mentioned infrequently in text, which makes it difficult to train a reliable classifier for them. transformer-based language models developed based on self attention @xcite have become the state-of-the-art across many nlp problems by outperforming previous solutions that were mostly based on recurrent neural networks (rnn) and convolutional neural networks (cnn). so one would expect that they perform well in icd coding too. however, examining the literature of icd coding methods reveals that transformer-based solutions fail to outperform cnn-based models. many studies have applied the bert language model @xcite to this task, for example pascual et al. (2021); singh et al. (2020); biseda et al. (2020); amin et al. (2019). more recently, ji et al. ( 2021 ) performed a comprehensive quantitative study to compare bert and some of its variants pre-trained on medical text against cnnbased models such as @xcite and @xcite to answer the question of whether the magic of bert (as observed across many nlp problems) also applies to automated icd coding. they concluded that bert cannot outperform cnn-based models in the full icd code case. unlike rnn or cnn models, which in theory can process sequences of arbitrary length, transformers' computational complexity scales quadratically with sequence length. this means that most of these models can handle limited size sequences. for instance, bert models usually are pre-trained and fine-tuned on sequences with at most 512 tokens. clinical notes normally contain long snippets of text beyond the sequence limit of transformers. we hypothesize that this constraint could explain the poor performance of transformers in this task, and will present empirical evidence for that. we emphasize that we do not claim to achieve state-of-the-art performance in icd coding, or that our design is the most efficient transformer architecture for processing long text. for a review of efficient transformers see @xcite and references therein. our goal is to provide new empirical evidence that shows even the standard transformer models can outperform some of the previous prominent methods and are a viable solution for icd coding. with the success of deep learning in nlp tasks, many researchers focused on using rnn and cnn models for icd coding. cnn models provide a convenient way to learn a contextual representation of text in nlp problems @xcite . for example, @xcite proposed the caml model: a convolutional layer on word2vec embedding vectors to learn a contextual representation for each word. the word representations are combined into a class-specific document representation using the attention mechanism. they also suggested a method to leverage code descriptions via a regularization term. li and yu (2020) proposed multi-filter residual cnn (multirescnn) that uses convolutional layers with different kernel sizes to capture patterns with different lengths. additionally, they used residual blocks on top of the convolutional layer. similar to @xcite they employed a per-class attention mechanism to make the document representation attend to different parts of the input for each code. recurrent neural networks (rnn) are also studied extensively for icd coding. @xcite applied lstm at character and word level to encode both the clinical note and the code description. @xcite employs a two-layer bidirectional gated recurrent unit (gru) model, where the first layer encodes individual sentences, and the second layer encodes the document. with the success of transformer architectures across many nlp tasks, researchers focused their attention to designing such models for icd coding. bert-xml @xcite with access to a large corpus of private data managed to pre-train the model with sequence length of 1024. most of the work in this area, however, considered the standard bert model and its variants pre-trained on medical text to encode the document @xcite @xcite . one observation with these models was that they were unable to outperform cnn-based models. @xcite performed a comprehensive study to answer a few research questions on the suitability of bert models for icd coding. they studied and compared different variants of bert pre-training. they also proposed a hierarchical attention method so that long clinical notes can be processed with a bert model with a limit of 512 tokens. most importantly, they compared different bert variants against traditional cnn-based models, and through extensive experiments showed that bert-based models are not capable of outperforming cnn-based models in icd coding. in the next sections we show that a simple method that enables processing of long text with transformers will attain results that contradict the findings of @xcite . ## method in this section we explain our method for building a model to predict medical codes. as illustrated in figure 1 , our model consists of an encoder that calculates token-level representation of the input text. this can be done in various ways, e.g. mullenbach et al. ( 2018 ) used word2vec and a cnn layer to calculate word-level representations. we choose the bert language model for this purpose. a class-specific representation of the document is then calculated using class-specific attention vectors, similar to @xcite . for ddimensional token representations and k classes, this layer requires d × k parameters. linear binary classifiers are built on top of the document representation to produce the probability that the document belongs to any of the k classes. this layer requires (d + 1) × k parameters (one scalar for the offset). let @xmath0, . . . , x s ] denote the tokenized input sequence with s tokens. let e 1 (x), . . . , e s (x) denote the representation of tokens 1, . . . , s obtained from an encoder. that is, ## evaluation we evaluate the accuracy of the proposed method with several sequence lengths and compare it against the caml method @xcite , which is one of the prominent cnn-based methods for icd coding. ## discussion most of the existing bert models pre-trained on generic or medical text can take input segments of up to 512 tokens. clinical notes, however, are much longer than this limit. to deal with this limitation, much of the existing works in automated icd coding that use bert limit the input to the model by truncating the text or selecting specific spans of text. this results in loss of information and poor performance. in this paper we proposed a simple method to apply bert models to sequences longer than 512 tokens. our method is simple and consists of two key components: (i) apply bert sequentially to (optionally overlapping) segments of 512 tokens, and (ii) concatenate token-level representations from all segments, and combine them using a class-specific attention layer. we demonstrated that processing long text sequences minimizes information loss and is critical for achieving high performance in automated icd coding. we also showed that contrary to previous findings, this method with even a small vanilla bert model outperforms cnn-based methods, and achieves competitive performance. future steps include evaluating medical variants of bert, and exploring other transformer-based architectures that were designed to handle long sequences.