B io M ed BERT : A Pre-trained Biomedical Language Model for QA and IR (2020)

Abstract: The SARS-CoV-2 (COVID-19) pandemic spotlighted the importance of moving quickly with biomedical research. However, as the number of biomedical research papers continue to increase, the task of finding relevant articles to answer pressing questions has become significant. In this work, we propose a textual data mining tool that supports literature search to accelerate the work of researchers in the biomedical domain. We achieve this by building a neural-based deep contextual understanding model for Question-Answering (QA) and Information Retrieval (IR) tasks. We also leverage the new BREATHE dataset which is one of the largest available datasets of biomedical research literature, containing abstracts and full-text articles from ten different biomedical literature sources on which we pre-train our BioMedBERT model. Our work achieves state-of-the-art results on the QA fine-tuning task on BioASQ 5b, 6b and 7b datasets. In addition, we observe superior relevant results when BioMedBERT embeddings are used with Elasticsearch for the Information Retrieval task on the intelligently formulated BioASQ dataset. We believe our diverse dataset and our unique model architecture are what led us to achieve the state-of-the-art results for QA and IR tasks.

## introduction the covid-19 pandemic reminded us of the need for a tool that biomedical researchers can use to sift through existing research to extract novel insights, and ultimately help them make novel drug discoveries. the rate of new publications in the biomedical field is on the rise. pubmed reports that more than 1 million biomedical research papers are published each year, amounting to nearly two papers per minute @xcite . for papers mentioning covid-19 alone, as of june 2020 more than 8000 peerreviewed publications have been published on pubmed. with the rate of scientific papers on covid-19 doubling every fourteen days @xcite , it is imperative to have a language understanding tool that can extract relevant information from credible literature, such as the research methodology, data, authors, results, and citations @xcite . in this paper, we address the problem from an information retrieval perspective, extracting the textual and contextual information from the corpus by taking a hierarchical approach. traditional search approaches such as lucene-based elasticsearch @xcite using bm-25 & jaccard-based matrices are efficient in retrieving objective answers where the primary task is to extract specific parts of the passage. however such methods struggle in the contextual retrieval of documents for which we need latent space representation of the query and the corpus of passages. our work leverages the bert language model architecture to pre-train a large-scale biomedical language representation model, named biomedbert. the work is inspired by the research of @xcite from korea university & clova ai research group. in said work, we use the new breathe dataset which combines full text articles and abstracts from ten data sources in the biomedical domain and use it to train our biomedbert model. we use the bert large model as a pre-training backbone to achieve new state-of-the-art results for question answering in the biomedical domain. in addition, we obtained impressive results by combining biomedbert embeddings with elasticsearch to obtain highly relevant results for information retrieval. this is achieved by using a neural passage re-ranking mechanism, which learns the inherent structural dependencies in the query and the research articles. we validated our search algorithm by formulating bioasq as a retrieval dataset. ## related work latent space representation learning and vector space modeling have proven to be extremely successful in the natural language processing domain, where they have shown to efficiently encapsulate the hidden meaning and context of sentences or passages. the journey of distributed word representation learning began with the efficient word2vec @xcite , glove @xcite and fast-text @xcite which outperformed the traditional bag of words model by significant margins in linguistic tasks, as these methods "largely" ignored the word context of the tokens. with the advent of sequential modeling and recurrent neural models, there was a significant enhancement in information content of the latent representation learning, with lstms as the state of the art at the time @xcite . non-recurrent sequence-to-sequence encoder-decoder models, also known as transformers, were introduced by @xcite . transformers stacked attention layers into a multi-headed attention architecture. the attention mechanism makes it possible to learn long-running word dependencies between input and output sequences by computing a context vector of the encoder for each token in the sequence @xcite . the bert architecture, which is primarily the encoder part of the transformer, achieved notable model performances for several linguistic tasks. pre-training and fine-tuning the bert network on a domain specific corpus has also shown to outperform many language models @xcite . our work combines ideas from @xcite and @xcite in using deep bidirectional transformers to learn contextual embeddings of word features from the large breathe corpora. ## breathe dataset biomedical research extensive archive to help everyone, or breathe, is a new dataset collection of biomedical research articles from leading medical archives. it is a combination of both full body texts and abstracts. the development and recent availability of this dataset significantly inspired the work in this paper. to the best of our knowledge, breathe is the largest diverse collection of publicly available, machine readable biomedical corpus for advanced language modeling @xcite . the dataset collection process was done in line with "ethical" principles for scraping, along with public apis that were used when available. the primary advantage of the breathe dataset for our model is its source diversity. breathe contains full text articles and abstracts from nine sources including bmj, arxiv, medrxiv, biorxiv, cord-19, springer nature, ncbi, jama, and bioasq @xcite . we performed our experiments with breathe v1.0 dataset. breathe v1.0 contains more than 6m articles and about 4 billion words. breathe v2.0 is the most recent version @xcite and the results reported in this paper are from training biomedbert on breathe v1.0 dataset. ## methodology biomedbert was built on the foundation of the bert architecture @xcite . in our work, we leverage both the pre-training and fine-tuning aspects of the bert architecture which enhances accuracy while also ensuring the robustness of the model. bert leverages a transformer architecture @xcite and uses the stacked encoder where each encoder consists of a multi-headed attention layer and a feed-forward network. bert builds on the transformer architecture to train large unlabeled data over two self-supervised tasks, and they are the masked language model (mlm) and next sentence prediction (nsp). mlm (also referred to as the cloze task) @xcite works by randomly masking 15% of the sequence and then attempting to predict the masked tokens, where as the nsp task helps the model to understand the relationship between two sentences. it works by predicting if a sentence is the actual "following" sentence or if it is just a random sentence. for more details on the training procedure for bert, the reader is referred to @xcite . in training the bert model, the total loss is computed as the sum of the masked-lm loss and next sentence prediction loss, both of which are cross-entropy, with the former multi-class log-loss and the later binary log-loss. ## neural information retrieval & biomedical search in this phase, the primary objective was to retrieve the most relevant research papers and articles in a ranked order based on the query made by the researcher(q i ). there are two major aspects that enhance the complexity of this problem. first, the computational complexity associated with retrieval from a biomedical corpus of millions of records(n ). second, the ability to accurately retrieve the relevant documents in a proper rank-ordered fashion, and to do so the process needs to account for a deep contextual understanding of the text. we approach the problems in a 2-step hierarchical fashion where in the first step, we primarily reduce the search space complexity by using the elasticsearch and bm-25 algorithm, a well-known metric in retrieval scenarios @xcite . elasticsearch itself operates in two stages. the first stage creates an extremely efficient data structure in textual searching scenarios, an inverted index, of the entire corpus of biomedical research papers. the second stage queries the entire corpus using the inverted index and scores the search results using the similarity (relevancy) ranking function bm25, and returns the top k search results based on that bm 25 score, where k is a tunable parameter @xcite . the output of this stage is an ordered state of retrieved biomedical papers sorted by bm 25 scores and let (a r1 , a r2 , a r3 ...a rk ) be the top k retrieved documents. here, f (t, d) represents the raw frequency of the term t in the document d, idf (t, d) is a function of number of documents, the term t has occurred in, given the universal set of documents d , |d| is the number of words in the document and k, b are parameters. the results of the elasticsearch based retrieval mechanism did not incorporate the contextual aspects of the query and the specific biomedical aspects of the corpus. as mentioned earlier, for effective ir in this domain, contextual consideration is highly relevant. when biomedical researchers search for relevant studies, they search for specific combinations of topics and scenarios. for instance, they might search for "how does coronavirus impact the lungs and to what extent" rather than a generic search such as "what is coronavirus". in order to best answer these types of specific queries, the system needs robust context for both the question and the possible answers. to solve this challenge, in the second step of the hierarchical search we use biomedbert to project the query(q i ), retrieve the top k papers(a r1 , a r2 , a r3 ...a rk ) in a d dimensional latent space and extract the embeddings @xcite . let the query embedding vector of q i be represented by q @xmath0, a @xmath1. we compute the cosine similarity as a vector-matrix product between the normalized query (q @xmath2, wherea @xmath3. let z be the output after computing the cosine similarity between the query and the papers, where @xmath4. finally, we sorted the output cosine scores z , while re-ranking and returning the most relevant biomedical research papers conditioned on the query(q i ). the primary intent of selecting cosine similarity as the preferred metric is that the vector cosine scores are normalized on each of the dimensions and hence are robust to scaling @xcite . ## discussion we observe that the biomedbert model achieves state-of-the-art results in the qa tasks for both of the datasets in the biomedical domain (see current challenges related to research into covid-19 directed us to build a language mining tool to support question-answering and information retrieval for the biomedical domain. with respect to the main purpose of this research, our results are the current state-of-the-art given the model performances on bioasq question-answering datasets. our model outperformed numbers from biobert @xcite on qa tasks irrespective of the fact that biobert was trained on a total of 18b words from the biomedical domain (13.5b from pmc full-text articles and 4.5b from pubmed abstracts), while our biomedbert model was trained on just over 4.1b words. the key reason for our better performance was the diversity in the datasets on which the biomedbert model was trained. diversity helped in enhancing both the performance and robustness for the biomedbert model. another important and novel aspect of our work was how we framed the bioasq dataset to validate our information retrieval methodology. specifically, we debiased reflect reality and achieved, robust and relevant results. biomedbert embeddings coupled with elasticsearch outperformed the retrieval performance of the other models based on the mean reciprocal rank values as shown in ## conclusion in this paper, we present the biomedbert model pre-trained on the breathe v1.0 dataset, one of the largest and most diverse datasets of biomedical research literature. biomedbert achieves stateof-the-art results when fine-tuned on question and answering datasets, and also produces impressive performances on other language tasks such as named entity recognition and relation extraction. biomedbert embeddings coupled with elasticsearch gives state-of-the-art performance on the re-framed bioasq dataset. moreover, the biomedbert model achieves state-of-the-art results for multiple tasks even when only pre-trained on the breathe v1.0 dataset, which contains just over 6 million articles. work is in progress to train an improved biomedbert model on the breathe v2 dataset with over 16 million articles. we believe continued enhancements of the biomedbert model will help biomedical researchers discover meaningful insights from literature faster, and make significant improvements in their field. we would like to thank daniel goncharov and 42 school silicon valley for their contributions to building the breathe dataset. we would also like to thank dave elliott with google cloud platform and google tensorflow research cloud for providing infrastructure and guidance throughout the project.