Reducing Inference Time of Biomedical NER Tasks using Multi-Task Learning (2022)

Abstract: Recently, fine-tuned transformer-based models (e.g., PubMedBERT, BioBERT) have shown the state-of-the-art performance of a number of BioNLP tasks, such as Named Entity Recognition (NER). However, transformer-based models are complex and have millions of parameters, and, consequently, are relatively slow during inference. In this paper, we address the time complexity limitations of the BioNLP transformer models. In particular, we propose a Multi-Task Learning based framework for jointly learning three different biomedical NER tasks. Our experiments show a reduction in inference time by a factor of three without any reduction in prediction accuracy.

## introduction transformer-based large language models (llms) have made it much easier to perform various nlp tasks with high accuracy. however, due to their large size, they take a lot of time and money to perform inference on large datasets. to give some perspective, one forward pass through pub-medbert @xcite ) takes 8-50ms on an aws g4dn.xlarge instance 1 (which comes with an nvidia t4 gpu). running one pubmedbert model on 1 million biomedical paragraphs would take roughly 9 hours. given the deluge of biological information daily, using fine-tuned pubmed-bert models for each biomedical ner task separately would be too time-consuming and expensive. when it comes to deep learning models, there are generally two variables that are optimized before deployment. these are size (the space occupied by the model's weight on disk and ram) and inference time (the time taken for one prediction). model size tends to matter more when deployed on edge devices and mobile phones since these devices have storage and ram constraints. several techniques, such as knowledge distillation, have 1 https://aws.amazon.com/ec2/instance-types/g4/ been proposed to address this issue, and some of the prominent models which have achieved a significant decrease in model size without much decrease in accuracy are distilbert @xcite , squeezebert @xcite , and mobilebert @xcite . however, size is usually not an issue for models running on servers. for example, a pubmedbert model has a size of only 400 mb. instead, the main concern is inference time, which is what we focus on in this paper. generally, a reduction in the model size naturally leads to a reduction in the inference time. however, in this work, we focus on reducing the inference time without reducing the model size. multi-task learning (mtl) primarily aims to improve the accuracy of multiple prediction tasks that are related to each other by leveraging commonly useful information. many of the previous works have shown the effectiveness of multi-task learning-based models for bioner tasks. the first work to apply mtl for biomedical named entities was attempted by @xcite . they used pre-trained word embeddings with cnn-based neural networks to extract named entities from biomedical texts. @xcite used a combination of bilstm and crf-based model, adapted from @xcite , to extract the entities and further used character and word-based embeddings that were shared by different datasets. a slightly different approach was proposed by @xcite , where they trained a dataset-aware mtl model and showed that their model was able to discriminatively exploit information from all of the related training datasets. the recent developments of large language models, such as bert @xcite , have demonstrated the effectiveness of better contextualized representation of various nlp tasks. @xcite developed biobert using the bert language model and pre-trained it on biomedical abstracts and papers. they achieved state-of-the-art results on several biomedical named entity recognition datasets. khan et al. (2020) and mehmood et al. (2019) incorporated mtl in bert-based models and showed promising results to extract biomedical named entities. although the previous works have shown the importance of multi-task learning when incorporated with either neural network-based models or transformer-based models, none of them have targeted optimizing these large models. while deploying these models for prediction, inference time matters; hence, it is equally important to develop models that reduce the inference time without any significant drop in performance. to this end, we develop a multi-task learning model for three different entities (cell-line, tissue, and strain) and show that we can reduce the inference time by a factor of 3 without any drop in performance when compared with a single-task model for each entity. our main contributions are as follows: the rest of the paper is organized as follows. we provide the details of the dataset in section 2. the experiments, results and their analysis are shown in section 3 and 4 respectively. finally, in section 5, we summarize all the results and provide pointers for future research. ## dataset for the bioner task, there are several publicly available annotated datasets but the most widely 2 the dataset and the source code of our experiments can be found here. used datasets for benchmarking are jnlpba @xcite , ncbi-disease @xcite , bc5cdr, @xcite bc2gm @xcite , and lin-neaus @xcite . these datasets cover mostly cell line, cell type, chemical, disease, gene, protein, and species type entities, and most of them rely on pubmed articles as a source. one of the significant concerns regarding most of the bioner datasets is the data quality, which is not only limited to the biomedical domain. @xcite mentioned annotation quality as one of the major challenges in the field of ner. an updated version of the 2004 jnlpba challenge was released in 2019 to address the flaws in the original corpus @xcite ). another issue is the source and the entity type, which is generally targeted in these datasets. these benchmark datasets lack entities such as tissue and strain that can help create meaningful cohorts across experiments. this information can be used to control the genetic variability in datasets. to address the issues mentioned above, we created a gold-standard corpus manually annotated with cell-line, strain, and tissue on abstracts extracted from the gene expression omnibus (geo) @xcite database. it is a public repository established by national center for biotech- @xcite . the corpus was manually annotated by the domain experts, and the annotation guidelines followed can be found in appendix a. the corpus consists of 9226 english paragraphs, and the number of mentions of strain (3476) is more than cell line (1541) and tissue (1804). despite the less number of total mentions of cell line, the number of unique mentions of cell line (483) is far greater than the number of unique mentions of tissue (338). in the data, wherever the strain entity is tagged, the cell line and tissue are not found, and vice-versa. this is due to the nature of the abstracts (extracted from geo) where we find either the texts contained mention of cell line, tissue, strain or both cell line and tissue in the same text. this makes the corpus unique and more reasonable to perform a multi-task learning model instead of building a multi-class token classification model. ## experiments in this section, we describe our experiments in detail about the model architecture, the training procedures, and the evaluation metrics followed. ## results and analysis the results of the multi-class token classification model built over the paragraphs containing only cell line and tissue are 0.62 and 0.46, respectively. since our data is unique in terms of the entities annotated and their mentions in the paragraphs, deploying a multi-class token classification model to learn the properties of the entities in the text is not a good choice in our case as it gives poor results. it might be possible that the underlying pub-medbert model learns the same features while fine-tuning for different ner tasks; hence, the mtl model is performing well. to rule out this possibility, we fine-tune the models after freezing the encoder layers of the pubmedbert model. the fourth and fifth columns of ## conclusion and future work in this study, we demonstrated how multi-task learning may be used to speed up model inference for complementary tasks that must be performed simultaneously on the same input. in particular, we compared our multi-task model to a single-task model and demonstrated that while the multi-task learning model's performance remained constant, the inference time was reduced by three. moreover, for our experiments, we created a gold-standard corpus, manually tagged with cell-line, tissue and strain. this corpus is the first of its kind where three different entities are manually curated by domain experts. when compared to the other entity types, the models' performance in identifying tissue names was incredibly poor, demonstrating how challenging it is to extract accurate tissue names from the text in the right context. for tissue ner, we must either discover a more suitable auxiliary task or develop some rule-based methods that will enhance the entity's overall performance. to increase the accuracy of tissue, we intend to carry out these actions in the future. investigating the mtl model for inference time on benchmark datasets would be another interesting project.