An Empirical Study of Multi-Task Learning on BERT for Biomedical Text Mining (2020)

Abstract: Multi-task learning (MTL) has achieved remarkable success in natural language processing applications. In this work, we study a multi-task learning model with multiple decoders on varieties of biomedical and clinical natural language processing tasks such as text similarity, relation extraction, named entity recognition, and text inference. Our empirical results demonstrate that the MTL fine-tuned models outperform state-of-the-art transformer models (e.g., BERT and its variants) by 2.0% and 1.3% in biomedical and clinical domain adaptation, respectively. Pairwise MTL further demonstrates more details about which tasks can improve or decrease others. This is particularly helpful in the context that researchers are in the hassle of choosing a suitable model for new problems. The code and models are publicly available at https://github.com/ncbi-nlp/bluebert .

## introduction multi-task learning (mtl) is a field of machine learning where multiple tasks are learned in parallel while using a shared representation @xcite . compared with learning multiple tasks individually, this joint learning effectively increases the sample size for training the model, thus leads to performance improvement by increasing the generalization of the model @xcite . this is particularly helpful in some applications such as medical informatics where (labeled) datasets are hard to collect to fulfill the data-hungry needs of deep learning. mtl has long been studied in machine learning @xcite and has been used successfully across different applications, from natural language processing @xcite @xcite , computer vision @xcite @xcite , to health informatics @xcite @xcite . mtl has also been studied in biomedical and clinical natural language processing (nlp) such as named entity recognition and normalization and the relation extraction. however, most of these studies focus on either one task with multi corpora @xcite or multi-tasks on a single corpus @xcite @xcite . to bridge this gap, we investigate the use of mtl with transformer-based models (bert) on multiple biomedical and clinical nlp tasks. we hypothesize the performance of the models on individual tasks (especially in the same domain) can be improved via joint learning. specifically, we compare three models: the independent single-task model (bert), the model refined via mtl (called mt-bert-refinement), and the model fine-tuned for each task using mt-bert-refinement (called mt-bert-fine-tune). we conduct extensive empirical studies on the biomedical language understanding evaluation (blue) benchmark @xcite , which offers a diverse range of text genres (biomedical and clinical text) and nlp tasks (such as text similarity, relation extraction, and named entity recognition). when learned and fine-tuned on biomedical and clinical domains separately, we find that mtl achieved over 2% performance on average, created new state-of-the-art results on four blue benchmark tasks. we also demonstrate the use of multi-task learning to obtain a single model that still produces state-of-the-art performance on all tasks. this positive answer will be very helpful in the context that researchers are in the hassle of choosing a suitable model for new problems where training resources are limited. our contribution in this work is three-fold: (1) we conduct extensive empirical studies on 8 tasks from a diverse range of text genres. (2) we demonstrate that the mtl fine-tuned model (mt-bert-fine-tune) achieved state-of-the-art performance on average and there is still a benefit to utilizing the mtl refinement model (mt-bert-refinement). pairwise mtl, where two tasks were trained jointly, further demonstrates which tasks can improve or decrease other tasks. (3) we make codes and pre-trained mt models publicly available. the rest of the paper is organized as follows. we first present related work in section 2. then, we describe the multi-task learning in section 3, followed by our experimental setup, results, and discussion in section 4. we conclude with future work in the last section. ## related work multi-tasking learning (mtl) aims to improve the learning of a model for task t by using the knowledge contained in the tasks where all or a subset of tasks are related @xcite . it has long been studied and has applications on neural networks in the natural language processing domain @xcite . @xcite proposed to jointly learn six tasks such as part-of-speech tagging and language modeling in a time-decay neural network. @xcite summarized recent studies on applying mtl in sequence tagging tasks. @xcite and martínez alonso and plank (2017) focused on conditions under which mtl leads to gain in nlp, and suggest that certain data features such as learning curve and entropy distribution are probably better predictors of mtl gains. in the biomedical and clinical domains, mtl has been studied mostly in two directions. one is to apply mtl on a single task with multiple corpora. for example, many studies focused on named entity recognition (ner) tasks @xcite . @xcite al. (2018), khan et al. (2020), and mehmood et al. (2019) the other direction is to apply mtl on different tasks, but the annotations are from a single corpus. @xcite proposed a joint model extract biomedical entities as well as their relations simultaneously and carried out experiments on either the adverse drug event corpus @xcite or the bacteria biotope corpus @xcite . @xcite also jointly extract entities and relations but focused on the biocreative/ohnlp 2018 challenge regarding family history extraction @xcite . @xcite integrated the bert language model into joint learning through dynamic range attention mechanism and fine-tuned ner and relation extraction tasks jointly on one in-house dataset of coronary arteriography reports. different from these works, we studied to jointly learn 8 different corpora from 4 different types of tasks. while mtl has brought significant improvements in medicine tasks, no (or mixed) results have been reported when pre-training mtl models in different tasks on different corpora. to this end, we deem that our model can provide more insights about conditions under which mtl leads to gains in bionlp and clinical nlp, and sheds light on the specific task relations that can lead to gains from mtl models over single-task setups. ## multi-task model the architecture of the mt-bert model is shown in figure 1 . the shared layers are based on bert @xcite . the input x can be either a sentence or a pair of sentences packed together by a special token @xcite . if x is longer than the allowed maximum length (e.g., 128 tokens in the bert's base configuration), we truncate x to the maximum length. when x is packed by a sequence pair, we truncate the longer sequence one token at a time. similar to @xcite , two additional tokens are added at the start ([cls]) and end ([sep]) of x, respectively. similar to @xcite , in the sequence tagging tasks, we split one sentence into several sub-sentences if it is longer than 30 words. in the shared layers, the bert model first converts the input sequence to a sequence of embedding vectors. then, it applies attention mechanisms to gather contextual information. this se- text similarity relation extraction named entity recognition 𝑠𝑖𝑚(𝑋 1 , 𝑋 2 ) 𝑃(𝑐|𝑋) 𝑃(𝑦 𝑖 |𝑋) [cls] tok 1 tok 2 [sep] ## experiments we evaluate the proposed mt-bert on 8 tasks in blue benchmarks. we compare three types of models: (1) existing start-of-the-art bert models fine-tuned directly on each task, respectively; ## https://github.com/ncbi-nlp/bluebert to evaluate the models on different domains, we multi-task learned various mt-bert on blue biomedical tasks and clinical tasks, respectively. bluebert clinical is the base bluebert model pretrained on pubmed abstracts and mimic-iii clinical notes, and fine-tuned for each blue task on task-specific data. mt-model are the proposed models described in section 3. we used the pre-trained bluebert clinical to initialize its shared layers, refined the model via mtl on the blue tasks (mt-bluebert-refinement clinical ). we keep fine-tuning the model for each blue task using task-specific data, then got mt-bluebert-fine-tune clinical . the mednli task, the mt model gained improvement by 2.4%. on the remaining tasks, the mt model also performed well by reaching the stateof-the-art performance with less than 1% differences. when compared the models with and without fine-tuning on single datasets, in biomedical tasks, we used bluebert biomedical as the baseline because it achieved the best performance on the blue benchmark. ## conclusions and future work in this work, we conduct an empirical study on mtl for biomedical and clinical tasks, which so far has been mostly studied with one or two tasks. our results provide insights regarding domain adaptation and show benefits of the mtl refinement and fine-tuning. we recommend a combination of the mtl refinement and task-specific fine-tuning approach based on the evaluation results. when learned and fine-tuned on a different domain, mt-bert achieved improvements by 2.0% and 1.3% in biomedical and clinical domains, respectively. specifically, it has brought significant improvements in 4 tasks. there are two limitations to this work. first, our results on mtl training across all blue benchmark show that mtl is not always effective. we are interested in exploring further the characterization of task relationships. for example, it is not clear whether there are data characteristics that help to determine its success @xcite . in addition, our results suggest that the model could benefit more from some specific examples of some of the tasks in while our work only scratches the surface of mtl in the medical domain, we hope it will shed light on the development of generalizable nlp models and task relations that can lead to gains from mtl models over single-task setups.