Meta Learning for Few-Shot Medical Text Classification

Our primary approach was to first tackle the general NLP meta-learning problem by meta-learning on textual representations on the CLINC150 dataset, an intent classification corpus, and then use the successful modelling strategies across CLINC150 towards the MIMIC III medical notes dataset. Our validated architecture consists of a large pre-trained language model, which generates textual representations, or embeddings, from raw text. With these representations, we automatically construct tasks into a MAML model and Prototypical Network to learn across these tasks.

Our embeddings are retrieved using BERT (Bidirectional Encoder Representations from Transformers) and its variants Devlin et al. (2018), which are pre-trained language models which jointly condition on both left and right context in all layers. These models had achieved state of the art on NLP benchmarks in 2019, and more importantly, are effective in representing text as word vectors (in 768 dimensions) that can be used for downstream tasks. As an example, we can input a CLINC150 meaning_of_life entry ”what do you think life is really about” into BERT and retrieve a vector $e_{1}\in R^{768}$. This vector would be closer to $e_{2}\in R^{768}$ generated from another meaning_of_life entry ”why are humans on earth”, but far from $e_{3}\in R^{768}$ the find_phone entry ”help me find my phone please”.

These embeddings are then constructed into tasks. Each task consists of a support set $D^{tr}$ with $K$ example across $N$ classes and a query set $D^{ts}$ with $Q$ query examples. See figure 1.

The tasks are used to train a meta-learning models: Model Agnostic Meta Learning model and Prototypical Network. Model Agnostic Meta Learning (MAML) Finn et al. (2017) reformulates the training procedure into two steps, (1) an inner loop which is traditional optimization procedure across the task and (2) an outer loop to train across tasks. Combined these methods create a trained network that can easily adapt to multiple tasks.

Prototypical Networks (ProtoNets) also learn a representation across tasks,Snell et al. (2017) but do so by learning a metric space where an encoder function $f_{\theta}$ projects raw examples into an embedding space. Instead of a set of weights learned from MAML, ProtoNets calculate prototypes $c_{n}$ which can be used to classify query examples via the $l_{2}$ distance from the query example to each class prototype:

This approach provides a more memory efficient representation with an easier inductive bias as it is limited to classification problems, making ProtoNets an ideal model for our problem of classifying ICD codes from medical notes.

Finally, we investigate a distributionally robust loss function to our architecture. In Distributionally Robust Optimization (DRO) models are trained to minimize the worst case loss over a set of pre-defined groups. Shiori Sagaw et al. Sagawa et al. (2019) showed that by strongly regularizing the the group DRO with $l_{2}$ regularization can give substantially better higher worst group accuracies. In our experiments we combined DRO and regularization with Meta-Learning to study impact on worst case group predictions (in this case, disease codes) during testing. The sampling was modified to include group codes (disease codes) in $D^{ts}$. These codes were used in the outer loop adaption for MAML to optimize for the worst case loss and track the worst performing groups.

Two datasets were used for experiements. The first CLINC150 dataset was used to establish a baseline to validate the meta-learning model. CLINC150 is used for Intent Classification and Out-of-Scope Prediction and is relatively simple Li and Zhang (2021). The CLINC150 dataset was already processed and split into 6 sets including in/out of domain data for train, test, validation accordingly. For our baseline model, we used in domain data only, the data includes 150 ”in-scope” intent classes, each with 100 records, 20 validations, and 30 test samples. Each record include a piece of text and its label (Figure 3). The table below (Figure 4) shows the distribution of train and test set which is uniform distribution. The diversity of classes indicates that CLINC150 is a good dataset to establish a baseline, however the distribution of intent classes is much more constrained than the distribution of disease codes in MIMIC-III, where there are thousands of diseases and the most challenging to classify are the ones that are rare and often obscured by more common diseases.

After establishing the baseline, we performed experiments using the MIMIC-III dataset, which is a large, freely-available database comprising de-identified health-related data associated with over 40,000 patients who stayed in critical care units of the Beth Israel Deaconess Medical Center between 2001 and 2012. For our scope, we extracted only specific data required for disease classification, which included patient IDs, associated medical notes across time, and associated diagnosis disease codes (or ICD codes). After selection, the dataset ended up containing around 2 millions records of more than 46,000 patients with 6,984 distinct diseases. See 5. The distribution of codes is largely skewed towards heart- or respiratory-related diseases, with a long tail of rare disease codes. Given time and processing constraints, we decided to take sampled 1000 records from each of the top 10 ICDs. We chose to equalize the ICD code samples so that the variance in distribution is due to properties of the medical notes themselves rather than the number of data samples. The selected data distribution is shown in (Figure 6).

No data pre-processing was performed on the CLINC150 dataset as the data is clean and well organized. Unlike CLINC150, MIMIC-III medical notes contain many typos, errors, and abbreviations. Furthermore, many notes could be associated with a single patient across time, creating redundant information with too many tokens for our text encoder (BERT) to process into embeddings.

As a result, we designed an end-to-end pipeline to extract and clean this MIMIC data. There were other efforts to process MIMIC-III text data Wang et al. (2020) Nuthakki et al. (2019) for classification, we used these pipelines as references and built our own procedure. For data selection, as mentioned above, our pipeline extracts most popular ICD codes and most rare ICD codes from schema tables depending on the experiment settings, example of the extracted data distribution of top 10 ICD is shown in (Figure 6). For each disease, we sampled 1000 medical notes while limiting the number of records from a same patient ID. By doing this, we can increase the variance of the data and minimize the redundant information since notes for a patient in a specific period usually contain same content with small updates.

For each medical note, we applied the standard text processing steps (lower case, remove stop words, remove special characters, lemmatize) and finally solved the input length problem by summarizing the text into a string of maximum 512 tokens using BERT and BART Summarizer Devlin et al. (2018) Lewis et al. (2019). The example below (Figure 7) demonstrates the processing steps for one specific example, the final text output will be used at input data for BERT to get the embedding vector.

For experimentation (Figure 2) we generated embeddings from clinical notes using BERT (bidirectional encoder representations from transformers) and trained two meta-learning models: ProtoNets and MAML. We also extended MAML model by implementing a Distributionally Robust Optimization (DRO) to minimize the worst-case loss during adaptation.

The results of the baseline with CLINC150 using 512 token BERT embedding using the ProtoNet model is shown in Figure 8. For a 10 class 5 shot model, we achieved a meta accuracy of 98% and for a 150 class 3 shot model, we achieved an accuracy of 96.7%

Our purpose was to study the impact of predictions by meta learning models once they are optimized for worst case expected loss due to atypical groups of data. We used four variations of two meta-learning models each to compare test time performance.

1. 1.

   Baseline meta-learning model: MAML and Protonet

2. 2.

   Model with DRO

3. 3.

   Model with Count Based or Group Adjusted DRO

4. 4.

   Model with Count Based or Group Adjusted DRO and $l_{2}$ regularization

For our experiments disease codes were considered groupings with the assumption that there are distribution shifts due to temporal (written over 10 years) and spatial (written by different individuals within each group) shifts in the medical notes. We introduce the notation G which is a set of all the disease codes and g for the codes that are being used in the adaptation loop. When sampling tasks we sample the triplet of (BERT Embeddings, labels, disease codes) as compared to (BERT Embeddings, labels) with a baseline model. The disease codes are only used in the outer adaptation loop and ignored in the inner loop loss calculation. Batch sizes were limited to 16 which results in high variance in the graphs. We also used batch size of 64 with less variance but achieved similar results.
For MAML we apply DRO in the outer loop adaptation only. We perform gradient descent based on the max group loss. We did not implement DRO for the inner loop, hence we do not use $g$ notation for $D^{tr}$. We wanted to focus on how well does the model ”adapt” to distribution shifts. We also track the losses and counts per group globally and use it to determine the worst case and best case disease codes and determine the performance during test time. The count is also used for group adjusted DRO. We first established a baseline with CLINC150 (Figure: 10) and then used the MIMIC-III dataset for final results (Figure: 11).

The following is the DRO objective without group adjustment (not count based) used with MAML ( Finn et al., 2017 Finn et al. (2017) and Sagaw et al., 2020 Sagawa et al. (2019))

The following is the DRO objective with group adjustment (count based) used with MAML. The count term $n_{g}$ acts as a regularizer.( Finn et al., 2017 Finn et al. (2017) and Sagaw et al., 2020 Sagawa et al. (2019))

The following is the DRO objective without group adjustment (count based) used with ProtoNet. The count term $n_{g}$ acts as a regularizer.( Snell et al., 2017 Snell et al. (2017)) and Sagaw et al., 2020 Sagawa et al. (2019))

The following is the DRO objective with group adjustment (count based) used with ProtoNet. The count term $n_{g}$ acts as a regularizer.( Snell et al., 2017 Snell et al. (2017)) and Sagaw et al., 2020 Sagawa et al. (2019))

Results

1. 1.

   MAML with DRO using CLINC150
   During test time, with CLINC150 we saw improved performance on worst case groups but slightly decreased accuracy for best case groups. This is expected as DRO also acts as a regularizer that forces the model to pay attention to worst case groups. See Figure 10 for a summary of results.

   

2. 2.

   MAML with DRO using MIMICIII
   Please see Table 1 and Figure 11 for a summary of results. Baseline MAML showed signs of overfitting with the MIMICIII dataset. It did not perform well on the worst case groups. When DRO was enabled, there was no overfitting even when the $l_{2}$ regularizer was not used. With DRO there is drop in the best case accuracy but a perceptible increase or comparable results in the worst case scenarios. MAML with Count Based DRO (without $l_{2}$) performs well in comparison to other DROs as it show the most improvement on worst case and middle case groups. MAML with Count Based DRO (with $l_{2}$) performs better amongst all the DROs in best case scenarios.

   Type	Avg	Worst Case	Best Case	Middle Case
   Baseline MAML	0.714$\pm$0.029	0.594$\pm$0.020	0.963$\pm$0.007	0.631$\pm$0.021
   MAML DRO	0.725$\pm$0.029	0.587$\pm$0.020	0.892$\pm$0.011	0.659$\pm$0.024
   MAML Group Adjusted DRO (No $l_{2}$)	0.696$\pm$0.031	0.646$\pm$0.018	0.908$\pm$0.010	0.698$\pm$0.030
   MAML Group Adjusted DRO (With $l_{2}$)	0.694$\pm$0.029	0.626$\pm$0.021	0.949$\pm$0.007	0.658$\pm$0.023

   

3. 3.

   ProtoNet with DRO using MIMICIII
   Please see Figure 12 for a summary of results. As with MAML, baseline ProtoNet overfits the MIMICIII data. DRO does regularize and reduce overfitting but does not eliminate it. Count Based or Group Adjusted DRO with $l_{2}$ regularization performs the best in all cases. For our experiment we passed the BERT embeddings through fully connected layers since we did not train BERT.

We compare the performance of ProtoNets on Clinical BERT embedding on random, semi-rare, and popular disease classifications (Figure 9)

• •

  Random: Charts sampled from the full distribution of codes with greater than 10 notes per code

• •

  Semi-rare: Charts sampled with codes with the bottom 50 codes with greater than 10 notes per code

• •

  Popular: Charts among the top 50 codes (See Figure 5).

Results
Surprisingly, randomly distributed codes performed better than the popular and semi-rare codes. Likely the popular codes performed worse due to the skewed distribution of medical codes within the top codes which are focused on heart and lung diseases. This introduces a lower variation in the embeddings, making it difficult for the ProtoNet prototypes to be distinguished via minimum $l_{2}$ distance. Still, the meta-learning approach on rare codes has shown positive results. In contrast, fine-tuning a BERT model for similar performance required ¿500 samples, indicating that meta-learning on this dataset is significantly more memory efficient and useful for rare disease settings.

Our purpose was to understand the performance of Meta-Learning in tackling low-resource classification such as disease code classification using medical notes. We validated our approach on a NLP benchmark and devised two experiments on a medical dataset by coupling DRO with meta-learning algorithms and using ProtoNet to classify rare and popular disease codes.

In our experiment, we did not exploit temporal data such as time-series data of disease change or all notes of a patient over each time period. This temporal data would be very useful to track the progress of each patient and get insight about the conditions development. In the future, processing each patient records as a whole observation would capture the time factor and would allow us to cover a broader spectrum of ICD codes. Another useful improvement could be the relationship between diseases as many diseases are correlated or even belong to a branch of a broader category. Exploiting the temporal data and hierarchy structure of ICD and integrating this information into our current baseline would improve model performance as beside comparing popular and rare diseases, we can also use top-down approach to compare each branch of disease.

We conclude that DRO combined with MAML does improve prediction and does accounts for distribution shift. The choice between DROs with or without $l_{2}$ may end up being domain specific. DRO combined with ProtoNet gives mixed results. We equalized the distribution of medical notes by choosing 1000 notes per disease code for sampling so that the distribution shift is limited to the content of the medical notes. Future experiments would involve removing that constraint (equalized distribution) and evaluate few shot learning on rare disease codes.

In our architecture, BERT was used an embedding generator and was not trained using gradient descent. Future work could involve integrating the BERT loss function with meta-learning algorithms. Furthermore, we can investigate this approach on partial snippets of medical text to understand the limits of NLP meta-learning on these datasets.

The results for ProtoNet predictions for random, popular and semi-rare codes are also encouraging. We believe that more experimentation and fine-tuning our encoder will give improved results.