Generalized Zero-Shot Text Classification for ICD Coding

1 Introduction

In healthcare facilities, clinical records are classified into a set of International Classification of Diseases (ICD) codes that categorize diagnoses. ICD codes are used for a wide range of purposes including billing, reimbursement, and retrieving of diagnostic information. Automatic ICD coding [Stanfill et al., 2010] is in great demand as manual coding can be labor-intensive and error-prone. ICD coding is a multi-label text classification task with a long-tailed class label distribution and noisy document inputs. Majority of ICD codes only have a few or no labeled data due to the rareness of the disease. In the medical dataset MIMIC III [Johnson et al., 2016], among 17,000 unique ICD-9 codes, more than 50% of them never occur in the training data. It is extremely challenging to perform fine-grained multi-label classification on both seen codes (codes with labeled data) and unseen (zero-shot) codes at the same time. Automatic ICD coding for both seen and unseen codes fits into the generalized zero-shot learning (GZSL) paradigm [Chao et al., 2016], where test examples are from both seen and unseen classes and we classify them into the joint labeling space of both types of classes.

∗ Shanghang Zhang accomplishes this work at Petuum Inc.

Modern automatic ICD coding models [Mullenbach et al., 2018; Rios and Kavuluru, 2018] can accurately classify frequent ICD codes while performing poorly on zero-shot codes. To resolve this discrepancy, we propose to generate the synthetic latent features for zero-shot codes and fine-tune the classification model with these generated features. The official ICD guidelines provide each code a short text description and a hierarchical tree structure on all the ICD codes [ICD-9 Guidelines, 2011]. To generate semantically meaningful features, we exploit those domain knowledge about ICD codes. We propose AGM-HT, an Adversarial Generative Model conditioned on code descriptions with Hierarchical Tree structure to generate synthetic features. As illustrated in Figure 1, AGM-HT consists of a generator to synthesize code-specific latent features based on the ICD code descriptions, and a discriminator to decide how realistic the generated features are. To guarantee the semantic consistency between the generated and real features, AGM-HT reconstructs the keywords in the input documents that are relevant to the conditioned codes. To further facilitate the feature synthesis of zero-shot codes, we take advantage of the hierarchical structure of the ICD codes and encourage the zero-shot codes to generate similar features with their nearest sibling code. The generated features are utilized to fine-tune the ICD coding models and achieve more accurate predictions for zero-shot codes.

Our contributions: 1) To the best of our knowledge, we propose the first adversarial generative model for the generalized zero-shot learning on multi-label text classification. AGM-HT generates latent features conditioned on the code descriptions and fine-tunes the zero-shot ICD code assignment classifiers to achieve higher accuracy. 2) AGM-HT exploits the hierarchical structure of ICD codes to generate semantically meaningful features for zero-shot codes without any labeled data. 3) AGM-HT has a novel pseudo cycle generation architecture to guarantee the semantic consistency between the synthetic and real features by reconstructing the relevant keywords in input documents. 4) Extensive experiments demonstrate the effectiveness of our approach. On MIMIC-III dataset, our methods improve the F1 score from nearly 0 to 20.91% for the zero-shot codes and AUC score by 3% (absolute improvement) from previous state of the arts.

2 RelatedWork

Automated ICD coding. Several approaches explore automatic assigning ICD codes on clinical text data [Stanfill et al., 2010]. [Mullenbach et al., 2018] proposed to extract per-code textual features with attention mechanism for ICD assignments. [Shi et al., 2017] explored character based short-term memory (LSTM) with attention and [Xie and Xing, 2018] applied tree LSTM with ICD hierarchy information for ICD coding. Most existing works either focused on predicting the most common ICD code or did not utilize the ICD hierarchy structure. [Rios and Kavuluru, 2018] proposed to utilize ICD hierarchy information for improving the performance on the rare and zero-shot codes. The model hardly assigns rare codes in its final prediction as we show in Section 4, making it impractical to deploy in real applications.

Feature generation for GZSL. The idea of using generative models for GZSL is to generate latent features for unseen classes using generative adversarial networks (GANs) [Goodfellow et al., 2014] and train a classifier on the generated and real features for both seen and unseen classes. [Xian et al., 2018] proposed using conditional GANs to generate visual features given the semantic feature for zero-shot classes. [Felix et al., 2018] added a cycle-consistent loss on generator to ensure the generated features captures the class semantics by using linear regression to map visual features back to class semantic features. [Ni et al., 2019] further improved the semantics preserving using dual GANs formulation instead of a linear model. Previous works focused on vision domain where the features are extracted from well-trained deep models on large-scale image dataset. We introduce the first feature generation framework tailored for zero-shot ICD coding by exploiting existing medical knowledge from limited data.

Zero-shot text classification. [Pushp and Srivastava, 2017] explored zero-shot text classification by learning relationship between text and weakly labeled tags on large corpus. The idea is similar to [Rios and Kavuluru, 2018] in learning the relationship between input and code descriptions. [Zhang et al., 2019a] introduced a two-phase framework for zero-shot text classification. An input is first determined as from a seen or an unseen classes before the final classification. This approach does not directly apply to ICD coding as each input is labeled with a set of codes which can include both seen and unseen codes, thus it is not possible to determine if the data is from a seen or an unseen class.

3 Method

The task of automatic ICD coding is to assign ICD codes to patient's clinical notes. We formulate the problem as a multilabel text classification problem. Let L be the set of all ICD codes and L = |L|, given an input text, the goal is to predict yl ∈ {0, 1} for all l ∈ L. Each ICD code l has a short text description. For example, the description for ICD-9 code 403.11 is "Hypertensive chronic kidney disease, benign, with chronic kidney disease stage V or end stage renal disease." There is also a known hierarchical tree structure on all the ICD codes: for a node representing an ICD code, the children of this node represent the subtypes of this ICD code.

We focus on the generalized zero-shot ICD coding problem: accurately predicting code l which is never assigned to any training text (i.e. yl = 0), without sacrificing the performance on seen codes. We assume a pretrained model as a feature extractor that performs ICD coding by extracting label-wise feature fl and predicting yl by σ(g⊤ l · fl), where σ is the sigmoid function and gl is the binary classifier for code l. For the zero-shot codes, gl is never trained on fl with yl = 1 and thus at inference time, the pretrained feature extractor hardly ever assigns zero-shot codes.

Figure 1 shows an overview of the generation framework. We propose to use generative adversarial networks (GAN) [Goodfellow et al., 2014] to generate ˜ fl with yl = 1 by conditioning on code l. The generator G tries to generate the fake feature ˜ f given an ICD code description. The discriminator D tries to distinguish between ˜ f and real latent feature f from the feature extractor model. After the GAN is trained, we use G to synthesize ˜ fl and fine-tune the binary classifier gl with ˜ fl for a given zero-shot code l. Since the binary code classifiers are independently fine-tuned for zeroshot codes, the performance on the seen codes is not affected, achieving the goal of GZSL.

3.1 Feature Extractor

The pretrained feature extractor model is zero-shot attentive graph recurrent neural networks (ZAGRNN) modified from zero-shot attentive graph convolution neural networks (ZAGCNN), which is the only previous work that is tailored towards solving zero-shot ICD coding [Rios and Kavuluru, 2018]. We improve the original implementation by replacing the GCNN with GRNN and adopting the label-distributionaware margin loss [Cao et al., 2019] for training. Figure 2 shows the architecture of the ZAGRNN. At a high-level, given an input x, ZAGRNN extracts label-wise feature fl and performs binary prediction on fl for each ICD code l.

Label-wise feature extraction. Given an input clinical document x containing n words, we represent it with a matrix X = [w1,w2, . . . ,wn] where wi ∈ Rd is the word embedding vector for the i-th word. Each ICD code l has a textual description. To represent l, we construct an embedding vector vl by averaging the embeddings of words in the description.

The word embedding is shared between input and label descriptions for sharing learned knowledge. Adjacent word embeddings are combined using a one-dimension convolutional neural network (CNN) to get the n-gram text features H = conv(X) ∈ RN×dc . Then the label-wise attention feature al ∈ Rd for label l is computed by:

where sl is the attention scores for all rows in H and al is the attended output of H for label l. Intuitively, al extracts the most relevant information in H about the code l by using attention. Each input then has in total L attention feature vectors for each ICD code.

Multi-label classification. For each code l, the binary prediction ˆ yl is generated by:

We use graph gated recurrent neural networks (GRNN) [Li et al., 2015] to encode the classifier gl. Let V(l) denote the set of adjacent codes of l from the ICD tree hierarchy and t be the number of times we propagate the graph, the classifier gl = gt l is computed by:

where g0 l = vl and GRU is gated recurrent units [Chung et al., 2014]. The weights of the binary code classifier is tied with the graph encoded label embedding gl so that the learned knowledge can also benefit zero-shot codes since label embedding is computed from a shared word. The loss function for training is multi-label binary cross-entropy:

As mentioned above, the distribution of ICD codes is extremely long-tailed. To counter the label imbalance issue, we adopt label-distribution-aware margin (LDAM) [Cao et al.,

2019], where we subtract the logit value before sigmoid function by a label-dependent margin ∆l:

where function 1(·) outputs 1 if y1 = 1 and ∆l = C n1/4 l and nl is the number of training data labeled with l and C is a constant. The LDAM loss is thus: LLDAM = LBCE(y, ˆ ym).

3.2 Zero-shot Latent Feature Generation

For a zero-shot code l, the code label yl for any training data example is yl = 0 and the binary classifier gl for code assignment is never trained with data examples with yl = 1 due to the dearth of such data. Previous works have successfully applied GANs for GZSL in the vision domain [Xian et al., 2018; Felix et al., 2018]. We propose to use GANs to improve zero-shot ICD coding by generating pseudo data in the latent feature space for zero-shot codes and fine-tuning the codeassignment binary classifiers using the generated features.

More specifically, we use the Wasserstein GAN [Arjovsky et al., 2017] with gradient penalty (WGAN-GP) [Gulrajani et al., 2017] to generate code-specific latent features conditioned on the textual description of each code. To condition on the code description, we use a label encoder function C : L ↦→ C that maps the code description to a lowdimension vector c. We denote cl = C(l). The generator, G : Z × C ↦→ F, takes in a random Gaussian noise vector z ∈ Z and an encoding vector c ∈ C of a code description to generate a latent feature ˜ fl = G(z, c) for this code. The discriminator or critic, D : F×C ↦→R, takes in a latent feature vector f (either generated by WGAN-GP or extracted from real data examples) and the encoded label vector c to produce a real-valued score D(f, c) representing how realistic f is. The WGAN-GP loss is:

where (·, c) ∼ P ·,c S is the joint distribution of latent features and encoded label vectors from the set of seen code labels S, ˆ f = α· f +(1-α) · ˜ f with α ∼ U(0, 1) and λ is the gradient penalty coefficient. WGAN-GP can be learned by solving the minimax problem: minG maxD LWGAN.

Label encoder. The function C is an ICD-code encoder that maps a code description to an embedding vector. For a code l, we first use a LSTM [Hochreiter and Schmidhuber, 1997] to encode the sequence of M words in the description into a sequence of hidden states [e1, e2, . . . , eM]. We then perform a dimension-wise max-pooling over the hidden state sequence to get a fixed-sized encoding vector el. Finally, we obtain the eventual embedding cl = el||gl of code l by concatenating el with gl which is the embedding of l produced by the graph encoding network. cl contains both the latent semantics of the description (in el) as well as the ICD hierarchy information (in gl).

Keywords reconstruction loss. To ensure the generated feature vector ˜ fl captures the semantic meaning of code l, we encourage ˜ fl to reconstruct the keywords extracted from the clinical notes associated with code l.

For each input text x labeled with code l, we extract the label-specific keyword set Kl = {w1, w2, . . . , wk} as the set of most similar words in x to l, where the similarity is measured by cosine similarity between word embedding in x and label embedding vl. Let Q be a projection matrix, K be the set of all keywords from all inputs and π(·, ·) denote the cosine similarity function, the loss for reconstructing keywords given the generated feature is as following:

Discriminating zero-shot codes using ICD hierarchy. In the current WGAN-GP framework, the discriminator cannot be trained on zero-shot codes due to the lack of real positive features. In order to include zero-shot codes during training, we utilize the ICD hierarchy and use f sib , the latent feature extracted from real data of the nearest sibling lsib of a zeroshot code l, for training the discriminator. The nearest sibling code is the closest code to l that has the same immediate parent. This formulation would encourage the generated feature ˜ f to be close to the real latent features of the siblings of l and thus ˜ f can better preserving the ICD hierarchy. More formally, let csib = C(lsib ), we propose the following modification to LWGAN for training zero-shot codes:

where c ∼ P cU is the distribution of encoded label vectors for the set of zero-shot codes U and (·, c) ∼ P ·,c U is defined similarly as in Equation 6. The loss term is weighted by the cosine similarity π(c, csib ) to prevent generating exact nearest sibling feature for the zero-shot code l. After adding zeroshot codes to training, our full learning objective becomes:

where β is the balancing coefficient for keyword reconstruction loss.

Fine-tuning on generated features. After WGAN-GP is trained, we fine-tune the pretrained classifier gl from baseline model with generated features for a given zero-shot code l. We use the generator to synthesize a set of ˜ fl and label them with yl = 1 and collect the set of fl from training data with yl = 0 using baseline model as feature extractor. We finally fine-tune gl on this set of labeled feature vectors to get the final binary classifier for a given zero-shot code l.

4 Experiments

Dataset. We evaluate our approach using the public medical dataset MIMIC-III [Johnson et al., 2016], which contains approximately 58,000 hospital admissions of 47,000 patients who stayed in the ICU of the Beth Israel Deaconess Medical Center between 2001 and 2012. Each admission record has a discharge summary that includes medical history, diagnosis outcomes, surgical procedures, discharge instructions, etc. Each admission record is assigned with a set of most relevant ICD-9 codes by medical coders. The dataset is preprocessed as in [Mullenbach et al., 2018]. Our goal is to accurately predict the ICD codes given the discharge summary.

We split the dataset for training, validation, and testing by patient ID. In total we have 46,157 discharge summaries for training, 3,280 for validation and 3,285 for testing. There are 6916 unique ICD-9 diagnosis codes in MIMIC-III and 6090 of them exist in the training set. We use all the codes for training while using codes that have more than 5 data examples for evaluation. There are 96 out of 1,646 and 85 out of 1,630 unique codes are zero-shot codes in validation and test set, respectively.

ICD-9 code information. We extract the ninth version of the ICD code descriptions and hierarchy from the CDC website [ICD-9 Guidelines, 2011]. In addition to the official description, we extend the descriptions with medical knowledge, including synonyms and clinical information, crawled from online resources [ICD-9 Data, 2006].

Baseline methods. We compare our method with ZAGRNN modified from previous state of the art approaches on zero-shot ICD coding [Rios and Kavuluru, 2018] as described in Section 3.1, meta-embedding for long-tailed problem [Liu et al., 2019] and WGAN-GP with classification loss LCLS [Xian et al., 2018] and with cycle-consistent loss LCYC [Felix et al., 2018] that were applied to GZSL classification in computer vision domain. We also experiment with the recent popular Transformer architecture [Vaswani et al., 2017] as the feature extractor.

Training details. For the ZAGRNN model, we use 100 convolution filters with a filter size of 5. We use 200 dimensional word vectors pretrained on PubMed corpus [Zhang et al., 2019b]. We dropout word embedding layer with rate 0.5. We set C = 2 in LLDAM. We use ADAM [Kingma and Ba, 2015] for optimization with batch size 8 and learning rate 0.001. The final feature size and GRNN hidden layer size are both set to 400. We train the ZAGRNN model for 40 epochs.

For WGAN-GP based methods, the real latent features are extracted from the final layer in the ZAGRNN model. Only features fl for which yl = 1 are collected for training. We use a single-layer fully-connected network with hidden size 800 for both generator and discriminator. We set gradient penalty coefficient λ = 10. For the code-description encoder LSTM, we set the hidden size to 200. We train the discriminator 5 iterations per each generator training iteration. We optimize the WGAN-GP with ADAM [Kingma and Ba, 2015] with mini-batch size 128 and learning rate 0.0001. We train all variants of WGAN-GP for 60 epochs. We set the weight of LCLS to 0.01 and LCYC, LKEY to 0.1. For LKEY, we predict the top 30 most relevant keywords given the generated features.

After the generators are trained, we synthesize 256 features for each zero-shot code l and fine-tune the classifier gl using ADAM and set the learning rate to 0.00001 and the batch size to 128. We fine-tune on all zero-shot codes and select the best performing model on validation set and evaluate the final result on the test set.

Evaluation metrics. We report both the micro and macro precision, recall, F1 and AUC scores on the zero-shot codes for all methods. Micro metrics aggregate the contributions of all codes to compute the average score while macro metrics compute the metric independently for each code and then take the average. All scores are averaged over 10 runs using different random seeds.

Results on seen codes. Table 1 shows the results of ZAGRNN models on all the seen codes. With ZAGRNN, almost all metrics slightly increased from ZAGCNN except for micro precision. With LLDAM loss, our final feature extractor can improve more significantly from ZAGCNN especially for macro metrics and achieve better precision recall tradeoff. The p-value from Student's t-test for comparing micro F1 scores between LDAM and baseline is 10-8 , indicating the improvement of using LDAM is significant. Nonetheless, these modifications are not enough to get reasonable performance zero-shot codes due to the lack of positive example for zero-shot codes during training. Results of Transformer as feature extractor is shown in the second row of the table. Due to the O(ℓ2 ) memory of Transformer for inputs with ℓ words, we can only train a small Transformer model on long clinical notes, instead of a large model [Devlin et al., 2018], yielding worse performance than CNN feature extractor.

Note that fine-tuning the zero-shot code classifiers with meta-embedding or WGAN-GP will not affect the classification for seen codes since the code assignment classifiers are independently fine-tuned.

Results on zero-shot codes. Table 2 summarizes the results for zero-shot codes. For the baseline ZAGRNN and metaembedding models, the AUC on zero-shot codes is much better than random guessing. LLDAM improves the AUC scores and meta-embedding can achieve slightly better F1 scores. However, since these methods never train the binary classifiers for zero-shot codes on positive examples, both micro and macro recall and F1 scores are close to zero. In other words, these models almost never assign zero-shot codes at inference time. For WGAN-GP based methods, all the metrics improve from ZAGRNN and meta-embedding except for micro precision. This is due to the fact that the binary zeroshot classifiers are fine-tuned on positive generated features which drastically increases the chance of the models assigning zero-shot codes.

Ablation studies. We next examine the detailed performance of each component of AGM-HT as shown in Table 2. Adding LCLS hurts the micro metrics, which might be counter-intuitive at first. However, since the LCLS is computed based on the pretrained classifiers, which are not wellgeneralized on infrequent codes, adding the loss might provide bad gradient signal for the generator. Adding LCYC, LKEY and LWGAN-Z improves LWGAN and achieves comparable performances in terms of both micro and macro metrics. At a closer look, LWGAN-Z improves the recall most, which matches the intuition that learning with the sibling codes enables the model to generate more diverse latent features. The performance drops when combing LWGAN-Z with LCLS and LCYC. We suspect this might be due to a conflict during optimization as the generator tries to synthesize ˜ f close to the sibling code lsib and simultaneously maps ˜ f back to the exact code semantic space of l. Using LKEY resolves the conflict as it captures more generic semantics from the words instead of from the exact code descriptions. Our final model that uses the combination of LWGAN-Z and LKEY achieves the best performance on both micro and macro F1 and AUC score.

We also conduct Student's t-test on performance scores from the 10 runs for demonstrating the significance of our final model (LWGAN-Z + LKEY) compared to the adopted LCLS and LCYC. The p-values are 0.01% and 1.39%, indicating the improvement of our final model is significant.

T-SNE visualization of generated features. We plot the TSNE projection of the generated features for zero-shot codes with LWGAN and LWGAN-Z respectively in Figure 3. Dots with lighter color represent the projections of generated features and those with darker color correspond to the real features from the nearest sibling codes. Features generated for zeroshot codes using LWGAN-Z are closer to the real features from the nearest sibling codes. This shows LWGAN-Z learns to generate latent features that better preserve the ICD hierarchy.

Keywords reconstruction from generated features. We next qualitatively evaluate the generated features by examining their reconstructed keywords. We first train a keyword predictor using LKEY on the real latent features and their keywords extracted from training data. Then we feed the generated features from zero-shot codes into the keyword predictor to get the reconstructed keywords.

Table 3 shows some examples of the top predicted keywords for zero-shot codes. Although the keyword predictor is never trained on zero-shot code features, the generated features are able to find relevant words that are semantically close to the code descriptions. In addition, features generated with LWGAN-Z can find more relevant keywords than LWGAN. For instance, for zero-shot code V10.62, the top predicted keywords from LWGAN-Z include leukemia, myelogenous, CML (Chronic myelogenous leukemia) which are all related to myeloid leukemia, a type of cancer of the blood and bone marrow.

5 Conclusion

We proposed the first feature generation framework, AGMHT, for generalized zero-shot multi-label classification. We incorporated the ICD tree hierarchy structure and a cycle reconstruction architecture to latent feature generation models that significantly improved zero-shot ICD coding without compromising the performance on seen codes. Extensive experiments demonstrate the superior performance of AGMHT on public datasets for both seen and unseen codes. We also qualitatively demonstrated the generated features from our framework can better preserve the class semantics and the ICD hierarchy compared to existing methods.