Retrieval-Augmented Domain Adaptation of Language Models (2023)

Abstract: Language models pretrained on general domain corpora usually exhibit considerable degradation when generalizing to downstream tasks of specialized domains. Existing approaches try to construct PLMs for each specific domains either from scratch or through further pretraining, which not only costs substantial resources, but also fails to cover all target domains at various granularity. In this work, we propose RADA, a novel Retrieval-Augmented framework for Domain Adaptation. We first construct a textual corpora that covers the downstream task at flexible domain granularity and resource availability. We employ it as a pluggable datastore to retrieve informative background knowledge, and integrate them into the standard language model framework to augment representations. We then propose a two-level selection scheme to integrate the most relevant information while alleviating irrelevant noises. Specifically, we introduce a differentiable sampling module as well as an attention mechanism to achieve both passage-level and word-level selection. Such a retrieval-augmented framework enables domain adaptation of language models with flexible domain coverage and fine-grained domain knowledge integration. We conduct comprehensive experiments across biomedical, science and legal domains to demonstrate the effectiveness of the overall framework, and its advantage over existing solutions.

## introduction language models pretrained on large scale of unsupervised corpora are capable of producing powerful representations as well as providing satisfactory performance when finetuned for general domain downstream tasks @xcite . however, as such representations are learned from general domain distributions, their generalization performance are deteriorate on specialized domains where distribution are much more different, such as biomedicine, science, legal, etc. many works thus trying to construct domainspecific plms either from scratch or initialized from the original plm, such as biobert @xcite , legalbert @xcite , scibert @xcite , etc. through the pretraining procedure, domain-specific knowledge are memorized and internalized into the parameter of the plms, thus providing a better initialization point for learning downstream tasks. however, such attempts not only consume much more costs, as it requires considerable computation to pretrain a language model, but also can not cover adaptive needs at task-or even instance-level granularity, as domain-specific plms are preliminarily trained with a fixed scale domain corpora but applied for all domain tasks. in fact, domain can be defined at various granularity w.r.t. various applications @xcite . for example, in the biomedical domain, texts might distribute across academic publication, medical record or medical-situated dialog, and exhibit completely different style and background knowledge. in this paper, we novelly propose a retrieval-augmented framework for domain adaptation (rada) to address the above challenges. instead of internalizing a fixed corpora into plm parameters via further pretraining, we construct and retrieve a datastore to dynamically enhance the learning of domain-specific downstream tasks. the flexibility of such a retrieval-augmented framework is in two-fold: 1) on one hand, we no longer need to maintain multiple plms respectively for each application domains, all domains simply share the same checkpoint for initialization but with domainspecialized datastore. 2) on the other hand, the datastore only consists of most relevant background passages tailored for each specific task. such tasklevel granularity provides best trade-off between general domain plm domain a plm domain b plm domain c plm domain a corpus domain b corpus domain c corpus general domain plm domain b corpus domain a corpus domain c corpus domain a task domain b task domain c task domain a task domain b task domain c task scale and coverage, we either do not need to collect a very large corpora in order to effectively cover all sub-domains, nor concerned with its coverage for this specific task. the illustration of the motivated framework compared to existing further pretraining solution is given in figure 1 . the core component of such a retrievalaugmented framework remains how to effectively integrate the retrieved passages. one line of related works is knowledge-enhanced plms @xcite @xcite , they mostly resort to attention mechanism to directly incorporate the matched knowledge triples. differently, we consider textual domain corpus instead of structured knowledge in this paper as the latter only has very limited coverage, especially for domain specialized scenarios. as a consequence, we aim to integrate retrieved passages that are less knowledgeintensive and might be mixed with noisy information that are irrelevant to solve the task. we accordingly propose a two-level selection scheme: passage-level selection and word-level selection. among the multiple retrieved passages given by a coarse-grained but efficient retriever, we first propose a differentiable sampling module that enables the model to dynamically select the most helpful background passage. we then adopt an attention mechanism to weigh all words inside the selected passage to achieve more fine-grained integration. the retrieval-augmentation interface works as a pluggable module that needs no further modification of the main encoding pipeline. we conduct experiments on 7 popular tasks covering three different downstream domains and a wide range of task variety. the results and ablations demonstrate the effectiveness and advantage of the proposed framework. the contribution of this paper is three-fold: gururangan et al. (2020) take a step further to systematically investigate the effect of further pretraining on multiple domains, which they refer to as domain-adaptive plm plm task input retrieval augmentation interface domain corpus cross-attention cross-attention aggregate ⨁ gumbel softmax gradient approximation 𝑄 𝐾&𝑉 𝐾&𝑉 pretraining (dapt). besides, they also propose task-adaptive pretraining (tapt) to perform further pretraining also on task data, which further brings significant improvements. there are mainly two concerns for these further pretraining based methods. one is the cost to construct and maintain considerable number of plms for each specific domain, which can be quite a lot spreading over various application scenarios. the other is the complication to formally define the concept of domain because domains exhibit hierarchical taxonomy in practical scenarios @xcite . the scope of a domain should neither be too broad, which results in loss of domain specialty, nor be too narrow, which results in loss of generalization ability. many recent works @xcite @xcite thus propose to incorporate multiple domains into one shared model, and introduce mixture-ofexperts module to allow different sub-domains also share with each other. by contrast, rada takes a very different retrieval-based road to resolve the above concerns. several other works also try to advance domain adaptation from various perspectives. for example, @xcite studied dapt for the specific task of abstractive summarization. @xcite propose gradual finetuning for domain adaptation. @xcite investigates the trade-off between dapt and data annotation under limited budget for domain adaptation. another very related line of work is retrievalaugmented language models. knn-lm (khan-delwal et al., 2020) retrieves for similar language modeling probability distributions and interpolate them with the current step of token generation. @xcite and @xcite propose to learn differentiable retriever to better select helpful passages. @xcite propose a nonparametric retriever and combine it with a generator model. @xcite use attention mechanism to integrate retrieved passages and also explore the datastore in a very large scale. @xcite investigate the usefulness of retrieval-augmentation in the context of few-shot prompting with llms. although being technically related, most of these works focus on general-purpose language understanding, especially knowledge-intensive tasks. while in this paper, we focus on the challenge of lm domain adaptation, and also novelly propose a two-level selection scheme to further augment the framework. ## methodology rada is a general framework for task-agnostic domain adaptation, and can be readily integrated into the standard pipeline of plm deployment. in the following section, we first brief the construction of domain specialized datastore in section 3.1, and the framework formulation in section 3.2. we then introduce the core component, a pluggable interface for retrieval-augmented domain adaptation in section 3.3, and finally brief the training usage in section 3.4. the overall framework is illustrated in figure 2 . ## conclusion in this paper we propose a retrieval-augmented framework to novelly address the challenge of language model domain adaptation. we use domain specialized corpus as datastore and retrieve from it for informative and helpful domain knowledge. the key module of the framework is a retrievalaugmentation interface, where we design a twolevel selection scheme to integrate the most relevant passage and its words while alleviating the noise. the overall framework enables flexible domain coverage and fine-grained domain knowledge integration. on a variety of downstream domains and tasks, we conduct comprehensive experiments and comparisons to demonstrate the effectiveness of the motivated framework and its components. in the future, we hope to further extend the framework to more scaled large language models and also more challenging few-shot prompting setting.