Improving Referring Ability for Biomedical Language Models (2024)

Abstract: Existing auto-regressive large language models (LLMs) are primarily trained using documents from general domains. In the biomedical domain, continual pre-training is a prevalent method for domain adaptation to inject professional knowledge into powerful LLMs that have been pre-trained in general domains. Previous studies typically conduct standard pre-training by randomly packing multiple documents into a long pre-training sequence. Recently, some existing works suggest that enhancing the relatedness of documents within the same pre-training sequence may be advantageous. However, these studies primarily focus on general domains, which cannot be readily applied in the biomedical domain where the distinction of fine-grained topics is harder. Is it possible to further improve the pre-training for biomedical language models (LMs) using exactly the same corpus? In this paper, we explore an improved approach to continual pre-training, which is a prevalent method for domain adaptation, by utilizing information from the citation network in this challenging scenario. Empirical studies demonstrate that our proposed LinkLM data improves both the intra-sample and inter-sample referring abilities of auto-regressive LMs in the biomedical domain, encouraging more profound consideration of task-specific pre-training sequence design for continual pre-training.

## introduction pre-trained language models (plms) benefit from large-scale, readily accessible, unsupervised texts. particularly in the biomedical domain, numerous studies conducted pre-training on academic papers and abstracts to enhance representations and professional knowledge @xcite @xcite . most of them are encoder-based language models @xcite . with the development of auto-regressive language models (lms), numerous studies have demonstrated their superior generalization ability and performance compared to encoder-based plms when the models are sufficiently large @xcite @xcite . they can not only understand instructions or background information provided in the context, which can be considered as the intra-sample referring ability (as shown in figure 1 ), but also adapt to new tasks by referring several provided demonstrations, which can be regarded as the inter-sample referring ability. moreover, with the advent of remarkable open-sourced large language models (llms), such as the llama family @xcite , researchers turn to explore the possibility of conducting continual pre-training to develop llms tailored for specific-domains @xcite @xcite . several pre-training methods have been proposed for encoder-based models, including masked language modeling, next sentence prediction @xcite , document relation prediction @xcite , translation language modeling (conneau and lample, 2019). these methods have effectively helped in learning specific knowledge and significantly promoted the development of encoder-based lms. however, to the best of our knowledge, most auto-regressive lms adhere to a conventional method for preparing input sequences for pre-training or continual pre-training, which involves first shuffling the corpora, followed by the random packing (concatenation) of documents until the concatenated sequence reaches the prescribed maximum input length @xcite @xcite . recently, some studies demonstrate that the standard pre-training method for auto-regressive lms can be further improved by designing appropriate pre-training sequences @xcite @xcite , such as incorporating relevant texts into the preceding context. linkbert @xcite constructs three types of segment pairs based on a citation network to classify whether they are continuous, linked, or random, motivating models to capture the citing relationship between two text segments. considering its success, we consider whether this methodology can be extended to auto-regressive lms, helping them learn to capture relationships between multiple text segments and improving their referring ability. therefore, in this paper, we explore the linking information from the citation network to construct sequences for training an auto-regressive lm, which we call it as linklm. specifically, we design the pre-training sequences by organizing the documents based on their citing relationships. when optimizing the language modeling objective, auto-regressive lms can learn to refer to possible information from the previous context. as illustrated in figure 2 , when predicting the tokens in the abstract d 1 1 (<pmid 37893869>), models can access information from its citing papers, learning from the findings about other detection tools (e.g., enfen battery in d 1 2 ) and different aspects (e.g., neurobiology in d 2 2 ). furthermore, by referring d 1 2 , d 1 3 , and d 1 4 , we can understand attention deficit hyperactivity disor-der (adhd) with a series of related works along the science history. therefore, training with lin-klm data encourages lms to refer to necessary information from the previous context, and therefore enhances models' referring ability, which can be used in tasks such as open-book question answering @xcite and the in-context learning (icl) setting @xcite . though the success of constructing appropriate pre-training sequences has been revealed by some previous works @xcite @xcite , they primarily focus on general domains where the distinction of topics is less challenging than that in the biomedical domain. additionally, they only trained their models from scratch. however, after pre-training with largescale, randomly concatenated documents, lms may tend to avoid breaking document boundaries (i.e., [eos] token) to refer to adjacent concatenated documents. whether the conclusion still holds under the continual pre-training scenario is not clear. since continual pre-training is a prevalent practice for developing biomedical llms, we focus on this setting in our experiments. in summary, our contributions are threefold: ## preliminary experiment all references of a given paper can serve as background information, but their importance towards the given paper is different. therefore, it is necessary to rank them based on their significance. a natural solution is using retrievers. as one of our preliminary experiments, we realize that retrievers are not as reliable as we expect in identifying the most appropriate reference for a given abstract. before using the retriever to select references that provide sufficient background information for the following anchor abstract, we should first understand how well a retriever can find out the reference that provides the most information for predicting a given abstract. we know the information that a reference provides can be measured by i(ref ; anchor) = p (anchor) -p (anchor|ref ) to the best of our knowledge, meditron @xcite is currently the best open-sourced biomedical lm because it is continually pre-trained with biomedical texts on the top of the powerful llm, llama-2 @xcite , so that it can provide a relatively accurate measurement for conditional perplexity. therefore, we use meditron-7b to compute the ranking of references as the ground truth. subsequently, we use some popular models including the contriever foot_1 to rank the references of a given abstract. we selected 1,000 anchor abstracts for this analysis. results are summarized in ## methodology in the scenario of pre-training biomedical lms, we usually collect abstracts or full papers as the pre-training corpus. the key of our methodology is to construct a long input sequence containing relevant information in the context. scientific researchers typically cite pertinent papers to support their conclusions and these citing papers are often previous stages of their research. based on this, we construct the pre-training input sequence with the help of the citation network, which is easy to obtain in the biomedical domain. algorithm 1 shows the procedure of our methodology. to develop biomedical lms, we use one of the most commonly used data sources, pubmed ab- algorithm 1 linklm sequence construction require: @xmath0, l): citation network require: r(d): return the citing references require: retriever 1: p ← [], q ← [] 2: while |d| > 0 do 3: randomly select d i from d 4: q.append(d i ) 5: d.remove(d i ) 6: while r(d i ) ∩ d ̸ = ∅ do 7: q ← [] 17: end while 18: shuffle p 19: return p : list of abstracts stract 3 . after pre-processing the raw data, we extract both textual and citing information, forming a citation network g. we begin with a randomly selected abstract as the anchor (e.g., d 1 1 in figure 2 ). unlike previous works @xcite , we select multiple relevant references at the same time to increase the hit rate of selected references. this approach addresses the limitations of retrievers, which do not always retrieve the most relevant reference from the given candidates, as discussed in section 3. to increase the diversity of our linklm data, we randomly sample the number of selected references, k, following a poisson distribution. according to at the beginning of the data construction, we easily obtain multi-hop long sequences. however, since we delete nodes once they are visited to prevent duplication of pre-training samples, the original citation graph becomes sparse gradually. many sequences will be composed by a single document at the end of the process. therefore, after constructing all sequences, we perform sequence-wise shuffling so that the sequences comprising a single document will be distributed uniformly alongside other longer sequences. in this way, each batch contains linked long sequences, making full use of the constructed linklm data. ## conclusions in this paper, we propose a pre-training sequence construction method for improving the referring ability of biomedical language models. previous studies mostly focus on general domains and they train the lms from scratch with designed pretraining sequences. in contrast, we explore this topic in a more challenging scenario, where the distinction of fine-grained topics is more difficult in the biomedical domain. moreover, we explore it under the continual pre-training setting, since it is a prevalent method for developing domain-specific lms now, filling the gap in this series of work. in this paper, we construct pre-training sequences by concatenating relevant references into the previous context using linking information from a citation network. empirical studies show that compared to the standard pre-training (i.e., randomly packing documents), our method significantly improves the intra-sample referring ability and the inter-sample referring ability on biomedical mcqa tasks, which answers our research question: by carefully designing pre-training sequences, we can still improve the pre-training for biomedical language models by reordering the pre-training documents (using exactly the same corpus). especially, pre-training using linklm data can further improve the performance when using retrieved demonstrations, revealing the future potential of our proposed methodology.