Explaining the Effectiveness of Multi-Task Learning for Efficient Knowledge Extraction from Spine MRI Reports (2022)

Abstract: Pretrained Transformer based models finetuned on domain specific corpora have changed the landscape of NLP. However, training or fine-tuning these models for individual tasks can be time consuming and resource intensive. Thus, a lot of current research is focused on using transformers for multi-task learning (Raffel et al., 2020) and how to group the tasks to help a multi-task model to learn effective representations that can be shared across tasks (Standley et al., 2020; Fifty et al., 2021) . In this work, we show that a single multi-tasking model can match the performance of task specific model when the task specific models show similar representations across all of their hidden layers and their gradients are aligned, i.e. their gradients follow the same direction. We hypothesize that the above observations explain the effectiveness of multi-task learning. We validate our observations on our internal radiologist-annotated datasets on the cervical and lumbar spine. Our method is simple and intuitive, and can be used in a wide range of NLP problems.

## introduction since the seminal work by @xcite , transformers have become the main architecture for almost all natural language processing (nlp) tasks. self-supervised pretraining of massive language models like bert @xcite and @xcite has allowed practitioners to use these large language models with little or no finetuning to various downstream tasks. multi-task learning (mtl) in nlp has been a very promising approach and has shown to lead to performance gains even over task specific fine-tuned models @xcite @xcite . however, applying these large pre-trained transformer models to downstream medical nlp tasks is quite difficult. medical nlp has its unique challenges ranging from domain specific corpora, noisy annotation labels and scarcity of high quality labeled data. despite these challenges, a number of researchers and practitioners have successfully finetuned these large language models for various medical nlp tasks. however, there is not much literature that uses multi-task learning in medical nlp to classify and extract diagnoses from clinical text @xcite . moreover, there is almost no work in predicting spine pathologies from the radiologists' notes @xcite . in this article, we are interested in extracting information from radiologists' notes on the cervical and the lumbar spine. in a given note, the radiologist discusses the specific, and often multiple pathologies, present in the medical images and grade their severity. extracting relevant pathologies from these reports can facilitate the creation of structured databases that can be used for a number of downstream use-cases, such as cohort creation, quality assessment and outcome tracking. single-task learning for information extraction in medical nlp has enjoyed much success in deep learning @xcite . however, an ultimate nlp system for a complete understanding of the medical report must be able to perform many diverse information extraction and classification tasks simultaneously and efficiently. such a system can be enabled by mtl, where one model shares weights across multiple tasks and makes multiple inferences in one forward pass. such networks can not only be trained with limited resources, but are more scalable and deployable when compared to several single-task models. moreover, the shared features within these mtl networks can induce more robust regularization and boost performance. thus there is a lot of interest in the academic and industry research communities to understand when multi-task learning improves performance over single-tasking models @xcite , and how to group a diverse set of tasks to encourage the model to learn a representation that can be shared across tasks @xcite @xcite . some of the aforementioned works, most notably in @xcite , define a notion of task similarity via the wasserstein distance and show that a small wasserstein distance between tasks aids in mtl. this work is an extension of our earlier work (sehanobish et al., 2022) where we used parameter efficient mtl models to extract information from cervical spine. in that work, we defined tasks as a conditional distribution over the classes, and we attributed our success of mtl to smaller wasserstein distance between tasks. however, computing wasserstein distance is expensive and suffers from the curse of dimensionality @xcite , which requires the number of samples to be significantly larger than the dimension of the representation (768 for many transformer models) in order for the distance to be accurately estimated. this prevents us from being able to estimate wasserstein distance for some of our minority classes, which have about 200 examples. even for majority classes where we have about 5k samples, our work suffers from large error rates. thus, to alleviate the above drawbacks, in this work, we sought to use methods that are applicable to small data regimes that lie in high dimensional space. inspired by the work of @xcite and @xcite , we hypothesize if the single-task models show similar representations across their hidden layers and the task specific gradients are aligned (see definition 1 in section 4.2), the multi-task model can match or outperform the task-specific, single-task models. we validate this hypothesis on two multi-task settings on our internal datasets: (a) four of the most common pathologies in the cervical spine -cen-tral canal and foraminal stenosis, disc herniation and cord compression, and (b) three pathologies in the lumbar spine -central canal stenosis, disc herniation and nerve root impingement. in this work, we (a) extend our novel pipeline to extract and predict the severity of various pathologies in the lumbar and cervical spine at each motion segment, (b) compute central kernel alignment (cka) and show similarity between the transformer layers trained for individual tasks on a given dataset, (c) compute dot products between the gradients of the task specific loss functions with respect to various parameters and show that most of the gradients flow along a similar trajectory and (d) show how to leverage that information into a simple mtl framework allowing us to achieve significant model compression during deployment and also speed up our inference without sacrificing the accuracy of our predictions. ## datasets we use an internal dataset consisting of radiologists' mri reports on the cervical and the lumbar spine. our dataset is heterogeneous and is diversely sampled from a large number of different radiology practices and medical institutions; the cervical mri data consists of 1578 reports from 97 different radiology practices detailing various pathologies of the cervical spine and our lumbar mri data contains 2004 reports from 170 different practices. we annotate the cervical reports with the 4 following pathologies: spinal stenosis, disc herniation, cord compression, and neural foraminal stenosis, and the lumbar reports with the 3 pathologies: disc herniation, spinal stenosis, and nerve impingement. each of these pathologies is accompanied by an indication of severity. in the cervical reports, the three categories for the central canal stenosis are based on gradation; none/mild are not clinically significant, moderate and severe definitions involve cord compression or flattening. the moderate versus severe gradation refers to the varying degrees of cord involvement. for disc herniation and central canal stenosis, the categories are based on a continuous spectrum and it is a standard practice in radiology for any continuous spectrum to be bucketed in mild, moderate and severe discrete categories. cord compression severity is binary: compression/signal change versus none. this is because both cord compression and signal change can cause symptoms, and are therefore clinically relevant. foraminal stenosis is treated as a binary task as well: severe versus non-severe, as severe foraminal stenosis may indicate nerve impingement, which is clinically significant. similar considerations are taken into account when annotating the lumbar reports. the splits and the details of each category can be found in every report is tagged by annotators with labels for relevant pathologies and severities, along with span information indicating which part(s) of the report mentions each pathology. for example, in a report for the lumbar spine, the sentence "l1-l2: there is no disc herniation. no spinal canal or foraminal narrowing" would be given normal or 0 class for each of the 3 pathologies (central canal stenosis, disc herniation and nerve root impingement). similarly in a cervical spine report, the sentence " c2-3: normal; no disc herniation or bulge. no central canal stenosis or neuroforaminal narrowing" would be given a normal or 0 class for all the 4 pathologies. an example of a full radiology report can be found in appendix a. ## workflow in this section, we will briefly describe our pipeline. the reports are first de-identified according to hipaa regulations. next, a spacy @xcite parser is used to break the report into sentences. a bert based ner model which we call the report segmenter is then used to identify the motion segment(s) referenced in each sentence, and all the sentences containing a particular motion segment are concatenated together. this report segmenter has been shown to achieve an f1 score of .9 on our internal datasets, and the same model is common across both the lumbar and the cervical datasets. more details about the ner model and the hyperparameters used to train it can be found in appendix b and c. all pathologies are predicted using the concatenated text for a particular motion segment. finally, the severities for each pathology are modeled as multi-label classification problem, and a pre-trained transformer is finetuned using the text for each motion segment. for more details about our pipeline and data processing, please see appendix b. figure 1 breaks down how a report looks as it is processed through our spine pipeline. ## similarity of representations between in this section we will describe our methodology to understand the similarity between the representations of various single task models. for all the experiments in this section, we use the pubmed-bert @xcite as the backbone. ## results on multi-task models in this section, we give empirical evidence on the success of mtl for our datasets. the results shown in this section are from our test set. for our classification task, the pubmedbert model is used as the backbone. this bert model is finetuned on the the cervical tasks resulting in 4 task-specific bert sequence classifier models which provides our baseline results. for the lumbar dataset, the pubmedbert model is finetuned on the 3 classification tasks resulting in 3 task-specific bert sequence classifier models. now, instead of finetuning the task specific models for extracting various pathology information from the cervical spine dataset, 4 classifier heads (i.e. 4 linear layers) are added to a single pubmed-bert model to create an output layer of shape [3, 3, 2, 2], where the first 3 outputs correspond to the logits for the stenosis severity prediction, the next 3 for the disc severity, the next 2 for the cord severity and the final 2 logits for the foraminal severity. for the lumbar dataset, 3 classifier heads are added to the pubmedbert model to create an output layer of shape [3, 3, 2], where the first 3 outputs correspond to the logits for the stenosis severity prediction, the next 3 for the disc severity, and the final 2 logits for the nerve severity. for the experiments, with both the datasets, a dropout of .5 is added to the bert vectors before passing them to the classifier layers. each of these classifier heads is trained with a cross entropy loss with the predicted logits and the ground truth targets. all the losses are added up which allows the gradients to backpropagate through the whole model and train these classifier heads jointly. the results for our experiments are shown in backbone model stenosis disc cord foraminal bert base baseline (single tasker) .62 ± .03 .64 ± .03 .70 ± .03 .79 ± .03 multi-tasking .62 ± .02 .65 ± .03 .72 ± .02 .78 ± .01 clinical bert baseline (single tasker) .64 ± .05 .66 ± .02 .71 ± .02 .82 ± .01 multi-tasking .63 ± .02 .67 ± .01 .75 ± .01 .79 ± .03 msr pubmedbert baseline (single tasker) .66 ± .03 .68 ± .04 .73 ± .05 .84 ± .01 multi-tasking .67 ± .01 .69 ± .01 .72 ± .04 .83 ± .03 the hyperparameters and other training and implementation details can be found in appendix c. ## deployment we deploy our spine pipeline system on an aws p3.2x machine with a single nvidia v100 gpu. reports are passed through the pipeline daily and first go through the report segmenter which tags sentences belonging to our set of motion segments. post-processing is done per report to aggregate sentences belonging to each motion segment group and to filter out any reports that do not contain motion segments. each grouping of motion segments is individually classified through our mtl model to predict a severity class per pathology. both the report segmenter and the multi-tasking model are processed in batch mode with latencies of 31ms/report and 56ms/report, respectively. compared to single pathology models, we observe a 3x improvement in latency per study when using the mtl pathology model. the spine pipeline is routinely evaluated in an offline setting for studies that do not produce any motion segment groupings or fail to capture any sentences for a given motion segment, per report. our current deployment only supports the lumbar reports and we are in the process of extending our deployment to also support the cervical pathologies. ## conclusion and future work in this work, a simple multi-tasking model is presented that is competitive with task specific models. instead of training and deploying task specific models, only one model is trained and deployed. this allows us to save significant costs during train-ing and faster inference during deployment while achieving significant model compression, without any loss in the quality of performance. our work opens the possibility of using multi-tasking models to extract information over various different body parts, allowing users to leverage large transformer models using limited compute resources. our novel pipeline is one of the very few works that attempts to extract pathologies and their severities from a heterogeneous source of radiologists' notes on lumbar and cervical spine mris at the level of motion segments. these findings suggest that our approach may not only be more widely generalizable and applicable, but also more clinically actionable. we believe our analysis with cka and gradient alignment sheds more light on the success of mtl. this insight has led to our process change from single-task bert based models to a more cost-effective mtl system. our analysis is widely applicable for other datasets and tasks. it is tempting to ask if one can use one multi-task model for both the lumbar and the cervical datasets. this is a work in progress and we have found strong similarity between single task models in the two datasets (most notably between the lumbar disc and the cervical disc models and the lumbar stenosis and the cervical stenosis models). however, unlike in the above analysis, we see low cka scores between various other task specific models which may make mtl difficult (see appendix d). we are in the process of using our analysis, along with insights borrowed from @xcite to either group tasks from the two datasets or align different task-specific gradients to create an efficient learner. the biggest drawback of our work is the limited amount of data on which our observations are verified. we are actively addressing this issue as we annotate more reports concerning various pathologies in different body parts.