SlowGuess/NAACL_2021_OCR huggingface dataset:
Shang-Chi Tsai* Chao-Wei Huang* Yun-Nung Chen
{d08922014,f07922069}@csie.ntu.edu.tw y.v.chen@ieee.org
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 medical code prediction. In the experiments, our proposed framework is able to improve upon best-performing predictors on the benchmark MIMIC datasets.
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 (Choi et al., 2016).
Considering that manual ICD coding has been shown to be labor-intensive (O'malley et al., 2005), several approaches for automatic ICD coding has been proposed and investigated by the research community (Perotte et al., 2013; Kavuluru et al., 2015). With recent introduction of deep neural networks, the performance of automatic
a set of ICD prediction
...250.61, 357.2, 564.0, 564.5, 401.9...
564.0: Constipation
564.5: Functional diarrhea
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 (Choi et al., 2016; Shi et al., 2017; Mullenbach et al., 2018; Baumel et al., 2018; Xie and Xing, 2018; Li and Yu, 2020; Vu et al., 2020; Cao et al., 2020). 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 high-level 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 high-level class. Prior work considered the hierarchical dependencies between ICD codes by using hierarchical SVM (Perotte et al., 2013) or by introducing new loss terms to leverage the ICD structure (Tsai et al., 2019). However, they borrowed the dependency from domain experts and did not consider
The source code of this project is available at https://github.com/MiuLab/ICD-Correlation.
Proceedings of the 2021 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, pages 4043-4052 June 6-11, 2021. ©2021 Association for Computational Linguistics
the label correlation in the data.
Inspired by the success of reranking techniques on automatic speech recognition (Ostendorf et al., 1991) and dependency parsing (Zhu et al., 2015; Sangati et al., 2009), 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 (Saeed et al., 2011; Johnson et al., 2016). 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:
- This paper is the first attempt to improve multi-label classification with a reranking method for automatic ICD coding.
- The experiments show that the proposed approaches are capable of improving best-performing base predictors on the benchmark datasets MIMIC-2 and MIMIC-3, demonstrating its great generalizability.
- The proposed framework has the great potential of benefiting from extra ICD labels, reducing the demand of paired training data towards scalability in the medical NLP field.
This paper focuses on multi-label medical code prediction; hence, We briefly describe the related background about medical code prediction and multi-label classification.
2.1 Medical Code Prediction
ICD code prediction is a challenging task in the medical domain. It has been studied since 1998 (de Lima et al., 1998) and several recent work attempted to approach this task with neural models. Choi et al. (2016) and Baumel et al. (2018) used recurrent neural networks (RNN) to encode the EHR data for predicting diagnostic results. Li and Yu (2020) recently utilized a multi-filter convolutional layer and a residual layer to improve the performance of ICD prediction. On the other hand, several work tried to integrate external medical knowledge into this task. In order to leverage the information of definition of each ICD code, RNN and CNN were adopted to encode the diagnostic descriptions of ICD codes for better prediction via attention mechanism (Shi et al., 2017; Mullenbach et al., 2018). Moreover, the prior work tried to consider the hierarchical structure of ICD codes (Xie and Xing, 2018), which proposed a tree-of-sequences LSTM to simultaneously capture the hierarchical relationship among codes and the semantics of each code. Also, Tsai et al. (2019) introduced various ways of leveraging the hierarchical knowledge of ICD by adding refined loss functions. Recently, Cao et al. (2020) proposed to train ICD code embeddings in hyperbolic space to model the hierarchical structure. Additionally, they used graph neural network to capture the code co-occurrences.
2.2 Multi-Label Classification
Multi-label classification problems are of broad interest to the machine learning community. The goal is to predict a subset of labels associated with a given object. One simple solution to a multi-label classification problem is to transform the problem into  $n$  binary classification problems, where  $n$  denotes the number of labels.
This approach makes an assumption that the predictions of each label are independent. However, in practice, the labels are usually dependent, making these predictors produce undesired predictions.
Figure 2: Illustration of our proposed framework.
There are numerous methods developed to alleviate this issue. Read et al. (2009) proposed classifier chains (CC), which introduce sequential dependency between predictions by adding the decision of one classifier to the input of the next classifier. Cheng et al. (2010) generalized CC to probabilistic classifier chains (PCC), where the proposed approach estimates the joint probability of labels and provides a proper interpretation of CC. Nevertheless, the recurrence between classifiers makes these methods less efficient and not applicable to tasks with large amount of labels.
Another line of research has leveraged the label dependencies that are known beforehand. Deng et al. (2014) used label relation graphs for object classification. Tsai et al. (2019) utilized the hierarchical structure of ICD codes to improve the ICD code prediction. These methods relied on known structures of the labels, which may not be easily accessible and less general.
Some prior work tried to learn label correlation and dependencies directly from the dataset. Zhang et al. (2018) introduced residual blocks to capture label correlation. This method requires paired training data, while our framework can learn from ICD codes only.
The concept of retrieve-and-erank has been widely used in automatic speech recognition (Ostendorf et al., 1991), natural language processing (Collins and Koo, 2005) and machine translation (Shen et al., 2004). Li et al. (2019) proposed to rerank the possible predictions generated by a base predictor with a calibrator. This method is conceptually similar to our framework, where we both follow the retrieve-and-erank procedure. The main difference between is that they leveraged an extra dataset for training the calibrator, while we train a distribution estimator on the same dataset as our base predictor.
The task of ICD code prediction is usually framed as a multi-label classification problem (Kavuluru et al., 2015; Mullenbach et al., 2018). Given a clinical record  $\mathbf{x}$  in EHR, the goal is to predict a set of ICD codes  $\mathbf{y} \subseteq \mathcal{V}$ , where  $\mathcal{V}$  denotes the set of all possible codes. This subset is typically represented as a binary vector  $\mathbf{y} \in \{0,1\}^{|\mathcal{V}|}$ , 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:
1. Label set candidate generation provides multiple ICD set candidates through a base binary predictor, which is detailed in Section 3.1.
2. Label set candidate reranking estimates the probability by leveraging the label correlation for reranking the candidates, which is detailed in Section 3.2.
3.1 Candidate Generation
In the first stage of the framework, we employ a base predictor to perform probabilistic prediction for all labels, and we use the predicted probabilities to generate top- $k$  most probable label sets. More formally, given a clinical note  $\mathbf{x}$ , we perform a base predictor and obtain the prediction for all labels:
$$
P _ {\text {b a s e}} \left(y _ {i} = 1 \mid \mathbf {x}, \theta_ {\text {b a s e}}\right), \quad i = 1, 2, \dots , | \mathcal {Y} |,
$$
where  $\theta_{\mathrm{base}}$  denotes the parameters of the base predictor. The predicted results are used to generate top-  $k$  probable sets, i.e.,  $\hat{\mathbf{y}}\subseteq \mathcal{V}$  with top-  $k$  highest probability prediction:
$$
P _ {\text {b a s e}} (\hat {\mathbf {y}} \mid \mathbf {x}, \theta_ {\text {b a s e}}) = \prod_ {i = 1} ^ {| \mathcal {Y} |} P _ {\text {b a s e}} (y _ {i} = \hat {\mathbf {y}} _ {i} \mid \mathbf {x}, \theta_ {\text {b a s e}}).
$$
Although there are  $2^{|\mathcal{V}|}$  possible subsets, the top- $k$  sets can be efficiently generated with dynamic
Figure 3: Illustration of the proposed reranking process.
programming as described in the prior work (Li et al., 2016).
3.2 Candidate Reranking
One drawback of the base predictor is the assumption about independent labels. To address this issue, in the second stage of the framework, we introduce a label set reranker to rerank the label set candidates generated in the previous stage. The reranker is designed to capture correlation and cooccurrence between labels. Given a label set candidate  $\hat{\mathbf{y}}$ , a reranker should be able to provide a reranking score  $R(\hat{\mathbf{y}})$ , where higher score indicates that the label set is more probable to appear. Similar to prior work (Zhu et al., 2015), We rerank the candidates according to their new scores defined as
$$
\log P _ {\text {b a s e}} (\hat {\mathbf {y}} \mid \mathbf {x}, \theta_ {\text {b a s e}}) + \alpha \cdot R (\hat {\mathbf {y}}),
$$
where  $\alpha$  is a hyperparameter. We use the label set with the highest score after reranking as the final prediction. Note that that reranking is done on label sets, not individual labels.
We employ two rerankers and describe them in the following subsections. Note that we do not restrict the design of rerankers to those we proposed; one can design their own reranker and plug it into the proposed framework. Our reranking framework is illustrated in Figure 3.
3.2.1 MADE Reranker
One intuitive way to assign scores to  $\hat{\mathbf{y}}$  is using the joint probability  $P(\hat{\mathbf{y}})$ . Higher joint probability indicates that the label set is more probable to appear, which aligns with our requirement to
rerankers. However, the joint probability  $P(\hat{\mathbf{y}})$  is often intractable. Therefore, we can only make an estimation with a density estimator.
Here we employ a masked autoencoder (MADE) (Germain et al., 2015) as the density estimator. MADE estimates the joint probability of a binary vector  $P(\hat{\mathbf{y}})$  by decomposing it in an autoregressive fashion with a random ordering
$$
P _ {\mathrm {M A D E}} (\hat {\mathbf {y}}) = \prod_ {i = 1} ^ {| \mathcal {Y} |} P _ {\mathrm {M A D E}} \left(y _ {i} = \hat {\mathbf {y}} _ {i} \mid \hat {\mathbf {y}} _ {o <   i}, \theta_ {\mathrm {M A D E}}\right),
$$
where  $o$  denotes a random permutation of  $\{1,2,\dots ,|\mathcal{Y}|\}$ $o(i)$  denotes the new ordering of  $i$ ,  $\hat{\mathbf{y}}_{o < i} = \{\hat{\mathbf{y}}_j\mid o(j) < o(i)\}$  denotes the set of all elements precede  $\hat{\mathbf{y}}_i$  in the new ordering, and  $\theta_{MADE}$  denotes parameters of the MADE model. MADE introduces a sequential dependency between labels. It enforces this dependency by masking certain connections in the multi-layer perceptron, making the output of  $y_{i}$  only depends on  $\hat{\mathbf{y}}_{o < i}$ .
The MADE model is trained with the labels from the training set of the ICD code prediction task. We use stochastic gradient descent to optimize the parameters  $\theta_{\mathrm{MADE}}$ , and the training objective is to minimize the binary cross-entropy loss  $\mathcal{L}(\hat{\mathbf{y}})$ :
$$
\begin{array}{l} - \frac {1}{| \hat {\mathbf {y}} |} \sum_ {i = 1} ^ {| \hat {\mathbf {y}} |} \left(\hat {\mathbf {y}} _ {i} \log P _ {\mathrm {M A D E}} \left(y _ {i} = 1 \mid \hat {\mathbf {y}} _ {o <   i}, \theta_ {\mathrm {M A D E}}\right) \right. \\ \left. + \left(1 - \hat {\mathbf {y}} _ {i}\right) \log P _ {\mathrm {M A D E}} \left(y _ {i} = 0 \mid \hat {\mathbf {y}} _ {o <   i}, \theta_ {\mathrm {M A D E}}\right)\right). \\ \end{array}
$$
Because we do not know which ordering performs the best, we can sample  $n$  different orderings and use the ensemble of these orderings to
improve estimation:
$$
\begin{array}{l} P _ {\mathrm {M A D E}} (\hat {\mathbf {y}}) = \\ \frac {1}{n} \sum_ {j = i} ^ {n} \prod_ {i = 1} ^ {| \mathcal {Y} |} P _ {\mathrm {M A D E}} (y _ {i} = \hat {\mathbf {y}} _ {i} \mid \hat {\mathbf {y}} _ {o _ {j} <   i}, \theta_ {\mathrm {M A D E}}). \\ \end{array}
$$
The illustration can be found in the blue box of Figure 3.
Given a label set candidate  $\hat{\mathbf{y}}$ , we define the score  $R_{\mathrm{MADE}}(\hat{\mathbf{y}})$  as
$$
R _ {\mathrm {M A D E}} (\hat {\mathbf {y}}) = \frac {\log P _ {\mathrm {M A D E}} (\hat {\mathbf {y}})}{| \hat {\mathbf {y}} | ^ {\beta}},
$$
where  $|\hat{\mathbf{y}}|$  denotes the size of the subset  $\hat{\mathbf{y}}$ , and  $\beta$  is a hyperparameter.  $|\hat{\mathbf{y}}|^{\beta}$  serves as a length penalty similar to the one used in sequence generation (Wu et al., 2016). We find that this length penalty is crucial to the reranker, and without it the score would favor subsets with smaller size.
3.2.2 Masked Self-Attention Reranker (Mask-SA)
As described in the previous subsection, MADE uses a sequential factorization to estimate the joint probability of a label set. This formulation forces the prediction of  $y_{i}$  to only condition on a subset of inputs  $\hat{\mathbf{y}}_{o < i}$ . With this restriction, the MADE model may fail to capture some crucial dependencies.
Inspired by the masked language modeling objective (Devlin et al., 2019), we propose a masked self-attention reranker (Mask-SA). Mask-SA takes as input a set of predicted labels  $\hat{\mathbf{y}}\subseteq \mathcal{V}$ , which is the set representation of the predicted labels. It employs a cloze-style prediction method, where we mask one input at a time and ask the model to predict the masked input. The advantage of this prediction method is that the output is conditioned on all inputs except for itself, which solves the restriction of the MADE model. This procedure is very similar to a denoising autoencoder (Vincent et al., 2008). The illustration can be found in the green box of Figure 3.
The Transformer architecture (Vaswani et al., 2017) has been shown to be efficient and effective in language modeling (Dai et al., 2019). We use it as the architecture of the Mask-SA model, with a slight modification where we remove the positional encodings due to the fact that the predicted ICD codes have no sequential order.
More formally, Mask-SA estimates a distribution over the label vocabulary for the masked input given all other elements in the set  $P_{\mathrm{MSA}}(\hat{\mathbf{y}}_i \mid \hat{\mathbf{y}} - \{\hat{\mathbf{y}}_i\}, \theta_{MSA})$ , where  $\theta_{MSA}$  denotes the parameters of the Mask-SA model.  $\theta_{\mathrm{MSA}}$  can be optimized with stochastic gradient descent to minimize the cross-entropy loss function. Given a label set candidate  $\hat{\mathbf{y}}$ , we compute the score  $R_{\mathrm{MSA}}(\hat{\mathbf{y}})$  as
$$
R _ {\mathrm {M S A}} (\hat {\mathbf {y}}) = \frac {\log \prod_ {i = 1} ^ {| \hat {\mathbf {y}} |} P _ {\mathrm {M S A}} (\hat {\mathbf {y}} _ {i} \mid \hat {\mathbf {y}} - \{\hat {\mathbf {y}} _ {i} \} , \theta_ {\mathrm {M S A}})}{| \hat {\mathbf {y}} | ^ {\beta}},
$$
where  $\beta$  is a hyperparameter. Note that in this formulation, the product of the conditional probabilities is not an exact estimation of the joint probability of  $\hat{\mathbf{y}}$ , but an analogy to the factorization made in the MADE model.
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.
We evaluate our model on two benchmark datasets for ICD code prediction.
- MIMIC-2 Following the prior work (Mullenbach et al., 2018; Li and Yu, 2020), we evaluate our method on the MIMIC-2 dataset. We follow their setting, where 20,533 summaries are used for training, and 2,282 summaries are used for testing. There are 5,031 labels in the dataset.
- MIMIC-3 The Medical Information Mart for Intensive Care III (MIMIC-3) (Johnson et al., 2016) dataset is a benchmark dataset which contains text and structured records from a hospital ICU. We use the same setting as the prior work (Mullenbach et al., 2018), where there are 47,724 discharge summaries for training, with 1,632 summaries and 3,372 summaries for validation and testing, respectively. There are 8,922 labels in the dataset. We also follow the setting in (Shi et al., 2017) where only the top-50 most frequent codes are considered. This setting has 8,067 summaries for training, 1,574 summaries for validation, and 1,730 summaries for testing.
We follow the preprocessing steps described in Mullenbach et al. (2018) with the provided
Model	Top-50 Dev	Top-50 Test	All Dev	All Test
MacroF	MicroF	MacroF	MicroF	MacroF	MicroF	MacroF	MicroF
CAML (2018)	54.03	61.76	53.46	61.41	7.70	54.29	8.84	53.87
+ MADE	56.91†	62.31†	56.64†	62.40†	7.79†	54.70†	9.11†	54.29†
+ Mask-SA	56.57†	62.14†	56.33†	62.26†	8.06†	54.53†	9.27†	54.09†
MultiResCNN (2020)	60.77	66.98	60.84	66.78	7.38	56.05	8.50	55.31
+ MADE	62.10†	67.13†	62.00†	67.13†	7.75†	57.08†	8.81†	56.21†
+ Mask-SA	61.52†	67.15†	62.06†	67.22†	7.97†	57.12†	9.28†	56.49†
LAAT (2020)	65.53	70.38	65.05	70.01	7.48	57.18	8.74	56.56
+ MADE	65.92†	70.34	65.29†	70.13†	7.92†	57.80†	9.16†	57.26†
+ Mask-SA	66.10†	70.30	65.44†	70.15†	8.08†	57.76†	9.41†	57.23†
Table 1: Results on the MIMIC-3  $(\%)$ .  $^\dagger$  indicates the improvement achieved by the proposed rescoring framework. The best scores for each base predictor are marked in bold.
Model	Macro F1	Micro F1
CAML	4.90	44.79
+ MADE	5.31†	46.11†
+ Mask-SA	5.55†	46.08†
MultiResCNN	5.06	45.89
+ MADE	5.50†	47.49†
+ Mask-SA	5.88†	47.55†
LAAT	6.41	47.54
+ MADE	7.23†	49.15†
+ Mask-SA	7.42†	49.05†
Table 2: Results on the MIMIC-2 test set using all codes  $(\%)$ .  $^\dagger$  indicates the improvement achieved by the proposed framework. The best scores are marked in bold.
scripts  $^{2}$ . All discharge summaries are truncated to a maximum length of 2,500 tokens.
4.2 Base Predictors
In order to validate the generalizability of our proposed framework, we employ three different base predictors that are proposed in prior work:
- CAML Convolutional attention for multi-label classification (CAML) is a method proposed in (Mullenbach et al., 2018). CAML aimed at improving ICD code prediction by exploiting the textual description of codes with attention mechanism (Bahdanau et al., 2014).
- MultiResCNN The multi-filter residual convolutional neural network (MultiResCNN) improved the design of CAML with multiple convolution filters and residual
connections (Li and Yu, 2020).
- LAAT Vu et al. (2020) proposed a label attention model which augments the label attention mechanism in CAML with additional transformations. It achieved state-of-the-art performance on MIMIC-2 and MIMIC-3 datasets.
4.3 Training and Evaluation Details
We train our rerankers with the label sets in the training set for 30 epochs. Adam is chosen as the optimizer with a learning rate of  $2e - 5$ . The batch-size is set to 64. The MADE reranker has one hidden layer with 500 neurons, and we find that using  $n = 10$  different orderings provides good estimation without using too much computation power. The Mask-SA reranker employs the transformer architecture with 6 self-attention layers, each with 8 attention heads. The hidden size is set to 256. For each pair of the base predictor and the reranker, we apply a grid search over possible values of  $\alpha$  and  $\beta$  on the validation set to find the best-performing hyperparameters, and we use them to perform evaluation on the test set. During reranking, we generate top-50 label set candidates to rerank. Note that our approach is to rerank the label set candidates instead of modifying the predicted probabilities from the base predictors. Therefore, common metrics considering the predicted probabilities of each label, such as Precision@K and AUC, are not suitable for our evaluation. Instead, we evaluate our methods with two metrics, macro F1 and micro F1.
$^{2}$ https://github.com/jamesmullenbach/caml-mimic
Figure 4: F1 scores  $(\%)$  with different number of candidates.
The results on the MIMIC-3 and MIMIC-2 datasets are shown in Table 1 and Table 2 respectively. All results are obtained by averaging the scores of 5 different training runs. We list the results before reranking in the first row of each base predictor.