Meta-learning Pathologies from Radiology Reports using Variance Aware Prototypical Networks (2022)

Abstract: Large pretrained Transformer-based language models like BERT and GPT have changed the landscape of Natural Language Processing (NLP). However, fine tuning such models still requires a large number of training examples for each target task, thus annotating multiple datasets and training these models on various downstream tasks becomes time consuming and expensive. In this work, we propose a simple extension of the Prototypical Networks for few-shot text classification. Our main idea is to replace the class prototypes by Gaussians and introduce a regularization term that encourages the examples to be clustered near the appropriate class centroids. Experimental results show that our method outperforms various strong baselines on 13 public and 4 internal datasets. Furthermore, we use the class distributions as a tool for detecting potential out-of-distribution (OOD) data points during deployment.

## introduction pretrained transformer-based language models (plms) have achieved great success on many nlp tasks @xcite , but still need a large number of in-domain labeled examples for finetuning @xcite . learning to learn @xcite @xcite ) from limited supervision is an important problem with widespread application in areas where obtaining labeled data can be difficult or expensive. to that end, metalearning methods have been proposed as effective solutions for few-shot learning @xcite . current applications of such meta-learning methods have shown improved performance in fewshot learning for vision tasks such as learning to classify new image classes within a similar dataset. namely, on classical few-shot image classification benchmarks, the training tasks are sampled from * equal contribution a "single" larger dataset (for ex: omniglot @xcite and miniimagenet @xcite ), and the label space contains the same task structure for all tasks. there has been a similar trend of such classical methods in nlp as well @xcite . in contrast, in text classification tasks, the set of source tasks available during training and target tasks during evaluation can range from sentiment analysis to grammatical acceptability judgment @xcite . in recent works @xcite , the authors use a range of different source tasks (different not only in terms of input domain, but also their task structure i.e. label semantics, and number of labels) for meta-training and show successful performance on a wide range of downstream tasks. in spite of this success, meta-training on various source tasks is quite challenging as it requires resistance to overfitting to certain source tasks due to its few-shot nature and more task-specific adaptation due to the distinct nature among tasks @xcite . however, in medical nlp, collecting large number of diverse labeled datasets is difficult. in our institution, we collect high quality labeled radiology reports (which are always used as held out test data) and use it to train our internal annotators who then annotate our unlabeled data. this training process is expensive and time consuming. our annotation process is described in section a. thus a natural question is: if we have a large labeled dataset consisting of a lot of classes, can we use it to meta-train a model that can be used on a large number of downstream datasets where we have little to no training examples? this is a challenging problem as the reports can be structured differently based on the report type and there can be a substantial variation in writing style across radiologists from different institutions. our main goal is to build out a set of extensible pipelines that can generalize to new pathologies typically in new sub-specialties while also generalizing across different health systems. in addition, the exact definition of the pathologies and their severity change can change depending on the clinical use case. this makes fully supervised approaches that rely on large labeled datasets expensive. having few-shot capabilities allows us to annotate a handful of cases and rapidly expand the list of pathologies we can detect and classify. in addition, we can use our approach to generate pseudo labels for rare pathologies and enrich our validation and test sets for annotation by an in-house clinical team. lastly our approach can be extended to support patient search and define custom cohorts of patients. our contributions in this work are the following: (1) we develop a novel loss function that extends the vanilla prototypical networks and introduce a regularization term that encourages tight clustering of examples near the class prototypes. (2) we meta-train our models on a large labeled dataset on shoulder mri reports (single domain) and show good performance on 4 diverse downstream classification tasks on radiology reports on knee, cervical spine and chest. in addition to our internal datasets, we show superior performance of our method on 13 public benchmarks over well-known methods like leopard. our model is very simple to train, easy to deploy unlike gradient based methods and just requires a few additional lines of codes to a vanilla prototypical network trainer. (3) we deploy our system and use the dataset statistics to inform out-of-distribution (ood) cases. ## related work there are three common approaches to metalearning: metric-based, model-based, and optimization-based. model agnostic meta-learning (maml) @xcite is an optimizationbased approach to meta-learning which is agnostic to the model architecture and task specification. over the years, several variants of the method have shown that it is an ideal candidate for learning to learn from diverse tasks @xcite @xcite . however, to solve a new task, maml type methods would require training a new classification layer for the task. in contrast, metric-based approaches, such as prototypical networks @xcite , being non-parametric in nature can handle varied number of classes and thus can be easily deployed. given the simple nature of prototypical networks, a lot of work has been done to improve them @xcite @xcite . prototypical networks usually construct a class prototype (mean) using the support vectors to describe the class and, given a query example, assigns the class whose class prototype is closest to the query vector. in @xcite , the authors use a mixture of gaussians to describe the class conditional distribution and in @xcite ; the authors try to model an unknown general class distribution. in @xcite , the authors use spherical gaussians and a kl-divergence type function between the gaussians to compute the function d in equation 2. however, the function used by the above authors is not a true metric, i.e. does not satisfy the triangle inequality. triangle inequality is implicitly important since we use this metric as a form of distance which we optimize, so it makes sense to use a true metric. in this work we replace it by the wasserstein distance which is a true metric and add in a regularization term that encourages the l 2 norm of the covariance matrices to be small, encouraging the class examples to be clustered close to the centroid. one of our main reasons to work with gaussians is due to the closed form formula of the wasserstein distance. few shot learning (fsl) in the medical domain has been mostly focused in computer vision @xcite . there are only a few works that have applied fsl in medical nlp @xcite but most of those works have only focused on different tasks on mimic-iii @xcite which is a single domain dataset (patients from icu and one hospital system). to the best of our knowledge, ours is the first study to successfully apply fsl on a diverse set of medical datasets (diverse in terms of tasks and patient populations). ## datasets all our internal datasets are mri radiology reports detailing various pathologies in different body parts. our models are meta-trained on a dataset of shoulder pathologies which is collected from 74 unique and de-identified institutions in the united states. 60 labels are chosen for training and 20 novel labels are chosen for validation. the number of training labels is similar to some well-known image datasets @xcite @xcite . this diverse dataset has a rich label space detailing multiple structures in shoulder, granular pathologies and their severity levels in each structure. the relationship between the granularity/severity of these pathologies at different structures can be leveraged for other pathologies in different body parts and may lead to successful transfer to various downstream tasks. the labels are split such that all pathologies in a given structure appear at either training or validation but not both. more details about the label space can be found in section b. the figure 1 and the for our downstream tasks, we use a body-part specific custom data processor to collect sentences related to a given structure (acl in knee, different vertebrae in the cervical spine, the entire impression section for lung reports) and concatenate them together to create a paragraph describing all the pathologies in the structure of interest. detailed description of preprocessing for different body parts, is presented in appendix c. the concatenated text from the validation sets of each task is passed to our trained meta-learner to generate the relevant class statistics (mean and variance). we then perform pathology classification on the test set by using our trained meta-learner and the saved class statistics. the downstream tasks are similar to the shoulder task in the sense that the pathology classification is performed on a sequence of sentences that all pertain to the same anatomical structure. thus our approach needs to learn the language that describes the severity of the pathology for a specific anatomical structure. we would like to shed some light on the complexity of the language we encounter. since our dataset is sourced from multiple health systems, and not all reports follow a standard structure, there is a large amount of variation in the language describing the same diagnosis. for example: a severe tear can be referred to as a rupture, or only the size of the nodule is mentioned without specifying that it is low risk (see appendix c for more examples). furthermore, most of our pipelines attempt to classify the different severities for a given pathology and the language describing severity can vary. while it might be possible to construct a rule based system to extract the diagnoses and severities we are interested in, it will be difficult to generalize as we expand to more diagnoses as well as to new health systems. ## prototypical networks prototypical networks or protonets @xcite use an embedding function f θ to encode each input into a m -dimensional feature vector. a prototype is defined for every class @xmath0, as the mean of the set of embedded support data samples (s c ) for the given class, i.e. the distribution over classes for a given test input x is a softmax over the inverse of distances between the test data embedding and prototype vectors. protonets are simple and easy to train and deploy. the mean is used to capture the entire conditional distribution p (@xmath1, thus losing a lot of information about the underlying distribution. a lot of work @xcite @xcite has focused on improving protonets by taking into account the above observation. we extend protonets by incorporating the variance (2nd moment) of the distribution and use distributional distance, i.e. 2-wasserstein metric, directly generalizing the vanilla protonets. ## experiments all our experiments are run on 4 v100 16 gb gpu using pytorch @xcite and hugging-face libraries @xcite . bert-base @xcite , clinical bert @xcite and pubmedbert @xcite are used as our backbone models. adapters @xcite are applied to each of these backbone models. while training adapter based models, the bert weights are frozen and only the adapter weights are updated, thus requiring less resources to train. this idea is similar to @xcite in the sense that we are reusing the features from these deep pretrained models. we compare our methods to leopard @xcite , vanilla protonets and big protonets @xcite . additional results with bert-base and clinical bert backbones can be found in to prevent overfitting on the test set, we choose the best model from each of these experiments based on the meta-validation accuracy and apply it to our downstream classification tasks. we note that these downstream tasks are significantly different from the few shot regime these models are trained in. moreover for these downstream tasks, we train bert models on each task and a multitasking model to provide additional baselines. in all our experiments, pubmedbert consistently outperforms bert-base and clinical bert by an average of 5 points and 3 points respectively. we believe the reason behind the improved performance is the domain specific vocabulary. even though clinical bert is pre-trained on mimic-iii @xcite , it still shares the same vocabulary as bert-base. protonet-bert shows better performance and faster convergence rates during training and validation ( backbone methods accuracy vanilla protonet 89.1 ± 1.1 big protonet 90.8 ± 1.2 leopard 85.1 ± 9.2 pubmedbert protonet w/ isotropic gaussian 90.2 ± 1.4 protonet w/ isotropic gaussian + reg 92.1 ± .8 variance aware protonet (ours) 91.5 ± 1.3 variance aware protonet + reg (ours) 92.9 ± .9 vanilla protonet 88.3 ± 1.4 big protonets 89.4 ± 1.2 pubmedbert w/adapters protonet w/ isotropic gaussian 89.8 ± 1.4 protonet w/ isotropic gaussian + reg 90.9 ± .7 variance aware protonet (ours) 90.5 ± 1.3 variance aware protonet + reg (ours) 91.2 ± .8 the entire validation set to compute the class distribution, we also experiment with choosing a k shots from the validation set to compute the class distribution (figure 12 in section g). our regularized variance aware protonets with bert-base + adapter is also validated on 13 public datasets. for the models and datasets marked with * in ## deployment based on the results described in it is well-known that the bert embeddings are highly anisotropic @xcite . we observe the same phenomenon in our meta-learned models as well (figure 3 ) which we use to our advantage to monitor ood cases. for each class in a dataset, we pick top k-dimensions (a hyperparameter) of maximum variance. we then take the union of these indices that we call the set of dataset indices i.e. the indices that explain the variance among all classes in the dataset. for any given query example, we compute the absolute difference ( ⃗ d j ) between its embedding vector (⃗ q) and class centroids ( ⃗ v j ), i.e. the i-th coordinate ⃗ d j : we then select top k dimensions of the each of these d j . we propose an ood metric called average variance indices (avi_k) by the overlap between the topk difference vector indices and the top-k dataset indices, i.e. av i_k := ## conclusion we extend prototypical networks by using wasserstein distances instead of euclidean distances and introduce a regularization term to encourage the class examples to be clustered close to the class prototype. by training our models on a label rich dataset (shoulder mri reports), we show successful performance on a variety of tasks. since the model weights are reused for all tasks, a single model is deployed enabling us to cut inference costs. moreover, adapters are used allowing us to tune smaller number of parameters (∼ 10 million) resulting in huge training cost savings. our model is also benchmarked on 13 public datasets and outperforms strong baselines like leopard. current work is underway to make our training dataset more diverse so that our models are more generalizable.