Retrieval-Augmented Domain Adaptation of Language Models
Year: 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.

heading: 1 Introduction
text: 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 (Devlin et al., 2019; Brown et al., 2020). However, as such representations are learned from general domain distributions, their
*Equal contribution. †Correspondence to: jiangwenbin@baidu.com
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 (Lee et al., 2020), LegalBERT (Chalkidis et al., 2020), SciBERT (Beltagy et al., 2019), 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 (Chronopoulou et al., 2022). 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 RetrievalAugmented 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
54
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 (Zhang et al., 2019; Peters et al., 2019; Liu et al., 2020), 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:
• Conceptual Contribution We novelly pro-
pose a retrieval-augmented framework for domain-adaptation of PLMs. The framework is designed as a pluggable module, which not only saves the costs to construct and maintain specialized PLMs for each respective domain, but also provides domain specialty at flexible task-level granularity with no concern of the trade-off between domain scale and coverage.
• Technical Contribution We propose a twolevel selection scheme to achieve fine-grained integration of retrieved passages while alleviate noise distraction. Specifically, we introduce a differentiable sampling module for passage-level selection and an attention mechanism to further assign word-level importance.
• Experimental Contribution We achieve consistent improvements across different domains including Science, Biomedicine and Legal. When combined with further pretraining, the improvements of the proposed framework are much more significant, surpassing competitive baseline of vanilla further pretraining.

heading: 2 Related Work
text: Domain Adaptation (Ramponi and Plank, 2020) is an important and widely investigated problem in the area of NLP. In the current pretrain-finetune paradigm, it becomes even more practical as most of NLP systems start with a general-purpose PLM but applied into various downstream scenarios (Guo and Yu, 2022). Many early works have found that continued pretraining of PLM on domain specific corpus could bring benefits (Alsentzer et al., 2019; Lee et al., 2020). 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
𝑄
𝐾&𝑉 𝐾&𝑉
Sh ar ed
Post-Norm
RA-Interface
0.12 0.40 0.31
0.12 0.40 0.31 0.25 0.10 0.22
Passage-level selection
word-level selection
Passage A Passage B
Retriever
𝑥!, 𝑥", … , 𝑥# , … , 𝑥$
Adapted Representation
Figure 2: The overall retrieval-augmented framework of RADA. The RA-Interface serves as a pluggable module to augment task-agnostic input representations with domain-specialized background passages.
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 (Reid et al., 2022). 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 (Chronopoulou et al., 2022; Gururangan et al., 2022; Li et al., 2022) 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, Yu et al. (2021) studied DAPT for the specific task of abstractive summarization. Xu et al. (2021) propose gradual finetuning for domain adaptation. Bai et al. (2021) 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. Guu et al. (2020) and Lee et al. (2019) propose to learn differentiable retriever to better select helpful passages. Lewis et al. (2020) propose a nonparametric retriever and combine it with a generator model. Borgeaud et al. (2022) use attention mechanism to integrate retrieved passages and also explore the datastore in a very large scale. Izacard et al. (2022) 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.

heading: 3 Methodology
text: 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.

heading: 3.1 Datastore Construction
text: We consider textual domain corpus as a necessary recipe to adapt a general-domain language model. From their related sources, we collect domain specialized passages {pi}D, and refer to them as the datastore. In practical implementation, we truncate the length of each passage to 256 tokens for simplicity. We then build an Elasticsearch service†, for any task-specific training or test input sequence x, we can retrieve for its K most relevant background passages {pk}Kk=1. Although Elasticsearch only serves coarse-grained retrieving purpose, we can employ it as an efficient first-stage retriever.
For each specific downstream task, the datastore is constructed at task-level granularity. And the retrieving of most relevant passages further achieves instance-level granularity. This is in contrast with many existing works where datastore is constructed and utilized at domain-level granularity, once a target domain is determined, a fixed-scale datastore is collected and injected into PLM as an integral through further pretraining. As a result, RADA brings better flexibility and save the difficulty to deliberate the trade off between domain specialty and domain coverage.

heading: 3.2 Retrieval-Augmented PLM
text: In a standard PLM-based framework, we first take the task input sequence x = {xi} and use the pretrained encoder to produce contextualized representations:
Hx = encode(x) (1)
whereHx ∈ Rl∗d is the sequence level representation for l words and d dimensions. We accordingly refer to hxi as the contextualized representation at each specific position i. The representations are then fed into a task-tailored head module, such as classification, regression or sequence tagging, etc.
In RADA, accompanied with each task input, we similarly construct the representations for retrieved passages {pk}:
Hpk = encode(pk), 1 ≤ k ≤ K (2)
and we similarly get Hpk and hpkj . The purpose of the proposed RADA framework can be formulated as a universal interface that augments and updates task representations hxi usingH
pk , where †https://github.com/elastic/
elasticsearch
Hpk serves as informative, relevant and domainspecialized knowledge source:
h̃i = RA-Interface(hxi , {Hpk}Kk=1) (3)
h̃i is then unchangeably used in the rest of the task pipeline. The overall framework remains taskagnostic and can be readily integrated into both pretraining and finetuning stage. The interface can also be easily extended to a more broad scenario with different backbone networks other than Transformer.

heading: 3.3 Retrieval Augmentation Interface
text: We elaborate the RA-Interface in this section. Different from many existing methods that resort to integrate structured knowledge (Zhang et al., 2019; Peters et al., 2019), the retrieved textual domain corpus are much less knowledge-intensive and thus might contain irrelevant noise distractions. However, Elasticsearch can only satisfy coarse-grained retrieving purpose, leaving it an important problem to effectively select and integrate actually helpful domain knowledge.
To address such challenge, we propose a twolevel selection scheme which performs discrete sampling at passage-level and attentively aggregate at word-level. Such selection strategy strengthens the proposed interface with fine-grained selection capability to integrate helpful knowledge while filter out its noisy parts. More specifically, we introduce a gumbel sampling mechanism to select the most helpful passage, along with an attention mechanism to assign distinguishable weights for each word insides the selected passage.

heading: 3.3.1 Cross Attention
text: In order to cover both sequence-level and tokenlevel downstream tasks, we would need to augment the representationHx at each specific position. For the i-th token representation hxi , we calculate its relevancy with respect to all other words from the retrieved passages using a canonical cross attention mechanism:
qxi =W Qhxi
kpki =W Khpkj vpki =W V hpkj
(4)
αkij = qxi k pk j√ d
(5)
where αkij denotes the attentive weight of j-th token from passage pk for current position hxi . Intuitively, this attentive weight can naturally measure each retrieved word its importance and relevance. To represent the usefulness of passages and their tokens at integral sequence level, we simply summing over all task input positions:
αkj = ∑
i
αkij (6)

heading: 3.3.2 Differentiable Sampling
text: To alleviate the noise distraction in the retrieved passages, we first sample the most relevant passage, then integrate all its words according to respective attentive weights. Intuitively, the sampling operation should be an argmax operation across all candidate passages:
αk = ∑
j
αkj (7)
k̂ = argmax k∈{1,2,...,K}
αk (8)
Note that here we sum αkj up to α k only temporarily for sampling most relevant passage, and will re-use αkj later in Equation 12 to aggregate representations. The passage level selection operation in Equation 8, although intuitive, is not differentiable, and are not allowed in an end-to-end trainable neural network. To enable such a learning-to-select ability inside the proposed framework, we novelly introduce the Gumbel-Softmax estimator (Jang et al., 2017) to replace Equation 8:
βk = exp((αk + g)/τ)∑ k exp((α k + g)/τ) (9)
where g is the noise sampled from Gumbel distribution: g ∼ Gumbel(0, 1) (10) This is practically implemented as:
g =− log(− log(u)) where u ∼ Uniform(0, 1) (11)
The resulting β distributed over k ∈ {1, 2, ...,K} is an approximate categorical distribution when the temperature τ approaches 0. We follow the Straight-Through procedure (Jang et al., 2017) to discretize β into k̂ in the forward pass but use gumbel approximation in the backward pass.

heading: 3.3.3 Aggregation
text: For the selected passage k̂, we apply a postnormalization on the attentive weights:
α̃k̂j = αk̂j∑ j α k̂ j , ∑ j α̃k̂j = 1 (12)
We then integrate the representations of passage k̂ accordingly:
ĥ = ∑
j
α̃k̂jv pk̂ j (13)
Note that in practical implementation, we use multihead attentions and concatenate all outputs together. We then wrap ĥ within a residual connection and LayerNorm at each position i to produce the final output representation:
h̃i = LayerNorm(Wĥ+ h x i ) (14)
whereW is another feedforward layer that merges representations from multiple attention heads.

heading: 3.4 Retrieval-Augmented Training
text: The proposed RA-Interface serves as a pluggable module that interacts with the representations in-place. It is also end-to-end trainable and does not modify the final training objective. In practical usages, we can either directly apply such a retrievalaugmented framework in downstream finetuning, or first adapt it with a domain-specific further pretraining stage. We investigate both settings in the experiments.

heading: 4 Experiments

heading: 4.1 Setup
text: Dataset We investigate the proposed framework on a variety of domain-specialized downstream tasks. Specifically, biomedical domain tasks including QIC (Intent Classification), QQR (QueryQuery Relevance), CMeEE (Named Entity Recognition), CMeIE (Information Extraction), scientific domain tasks including SCIERC (Relation Classification) (Luan et al., 2018), ACL-ARC (Citation Intent Classification) (Jurgens et al., 2018) and legal domain task CAIL2019-SCM (Similar Case Matching) (Xiao et al., 2019). And these 4 biomedical tasks are selected from the well established benchmark CBLUE (Zhang et al., 2022). Detailed statistics are listed in Table 4.
Datastore We collect domain specialized corpus at scale as datastore. For scientific corpus, we use the S2ORC corpus constructed from semantic scholar (Lo et al., 2020), for biomedical and legal corpus, we use in-house data regarding medical reviews or legal documents crawled online. Detailed statistics are listed in Table 3.
Training For retrieval-augmented pretraining, we split 10% from the entire domain corpus as pretraining pretext task data while the rest 90% as pretraining datastore. The pretraining steps are set to 10k for both DAPT and RADA. For retrievalaugmented finetuning, all corpus are considered as datastore, but we only retrieval and keep the most relevant passages for each task instances. In ElasticSearch, we set the number of retrieved passages
for each input data as K = 10. For each task, we search batch size through {16, 32}, learning rate through {1e − 5, 2e − 5, 5e − 5}†, and set epoch to 10†. We also run with 3 different random seeds, and accordingly report the average and standard deviation. We use BERT (Devlin et al., 2019) as PLM for English tasks, and RoBERTa (Cui et al., 2020; Liu et al., 2019) for Chinese ones.

heading: 4.2 Main Results
text: Table 1 gives the main results for RADA. We can observe two conclusions: 1) RADA outperforms standard finetuning on 4 out of 7 selected tasks; 2) When combined with further pretraining, the proposed framework is further improved, significantly surpassing standard finetuning and DAPT. Specifically the absolute benefits over FT baseline are +1.59.
We further investigate a low-resource setting by subsampling 30% of the training set. This is a very practical scenario as we often need to deal with domain tasks at a relatively low annotation cost. Results in Table 2 also demonstrate the effectiveness of the proposed framework in such settings. On 6 out of 7 tasks, RADA achieves the best performance.

heading: 4.3 Ablations
text: We investigate the impact of various components in this section.
†For CMeIE, we search for {1e-5, 2e-5}. †For CMeIE and CMeEE, we set to 35 and 15 respectively.

heading: 4.3.1 Scale of Datastore Corpus
text: One essential component of the proposed framework is the datastore. We first look into the effects of its scale. At pretraining stage, we fix the pretraining steps to 5,000, and accordingly set the scale of datastore to 128, 256, 512 and 1024. The results are illustrated in Figure 3. We observe clear trends of increasing performance w.r.t. increased datastore scale. As a consequence, the proposed RADA framework can always benefit from a larger datastore.

heading: 4.3.2 Training Data as Datastore
text: We further investigate the feasibility of directly using downstream task data as datastore corpus. As previous study has demonstrated, training data itself can also provide useful background information to an extent (Wang et al., 2022). Similarly, Gururangan et al. (2020) have also used task training data to perform further pretraining, which they referred to as TAPT. In Table 6, we provide the results when using training data in the proposed retrievalaugmented framework, and also reproduce TAPT for more comprehensive comparison. The results show that both TAPT and RADA with Trainset can provide considerable benefits, but are outperformed by RADA (full). As for these two methods, they perform comparably on the investigated three datasets. We also find that training data is more useful in sentence-level tasks (as included in the table), but less helpful in other tasks such as classification, sequence tagging, etc.

heading: 4.3.3 Pretraining Steps
text: We have equipped RADA with a further pretraining stage with domain specialized data, and proved it to be very effective in Table 1. In Figure 4, we further ablate the effects of pretraining steps. We save and evaluate the intermediate checkpoint at respectively 2.5k, 5k and 7.5k pretraining steps. The results exhibit clear trends that with more pretraining steps, RADA continually brings better adaptation performance.

heading: 4.3.4 Differentiable Sampling
text: One key design of RADA is the two-level selection scheme. Specifically, the differentiable sampling module based on Gumbel Sampling trick enables the model to dynamically learn which passage to integrate. In Table 5 we investigate this choice of design. We provide an alternative baseline, where instead of learning to sample, we simply sum the cross attention weights over all positions as passage-level score, and select the maximum one to integrate into the retrieval-augmented interface. We refer to this as Top 1 selection strategy. The results show that both methods can bring improve-
ments, but the proposed differentiable sampling is much more effective.

heading: 4.3.5 Efficiency Analysis
text: One potential concern for retrieval-augmented methods is their efficiency. We therefore investigate this factor in Table 7. We consider the inference efficiency at deployment time. For each task input, we retrieve 10 background passages using ElasticSearch, then compute and incorporate them into the RA-Interface. Note that at deployment time, it is practical to encode and cache all passages from the datastore in advance, so we only account for the time consumption starting from the interface. The results show that, the retrievalaugmentation framework only brings acceptable time increase. And the overall inference speed is maintained at around 0.02 seconds per instance on a production-level device.

heading: 4.3.6 Qualitative Analysis
text: We provide qualitative analysis in Table 5. The example is sampled from the SCI task, the target is to extract the relation between subject rewriting systems and object grammar formalisms, i.e.,
HYPONYM-OF. We can clearly see that elasticsearch can only provide coarse-grained, shallowsemantic retrieving capability based on keywords, such as grammar, formalisms, HPSG, etc. However, the proposed differentiable sampling module can more effectively identify the most helpful passage by reasoning over deep representations produced by shared encoders. From passage 2 with the highest sampling score, we can indeed reason, understand and accordingly infer the actual relationship between the target entity pair.

heading: 5 Conclusion
text: 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.
Limitation
We summarize two limitations which also serve as promising directions to be explored in our future work. RADA framework only considers textual domain corpus as the datastore, although this has greatly improve the coverage of domain knowledge as texts are always relatively easy to collect. However, it is widely investigated that structured knowledge such as knowledge base can also serve similar purpose. And such resources are generally in higher quality and are easier to match. Therefore, it would be benefiting to further integrate such resources at certain scenario where KB is available.
The other limitation regards to the scale of the RADA implementation. As large language models have becoming increasingly powerful, they have demonstrated quite impressive capability in memorizing and recalling a wide range of background knowledge existed in the massive corpora they have been pretrained on. This trends of development naturally raises question for the proposed framework: will it still be beneficial when scale up to LLMs? and on what kinds of scenario does it brings best improvements? These are very important questions to answer, and we can certainly expect them to be explored in future works.
Ethical Statement
We evaluate the proposed method on established and publicly available datasets. There is also no human evaluation involved. This paper is not concerned with the above ethical risks. When the proposed framework is deployed into domain specific production, the domain adapted language models might express ethical-related outputs, but just as any other language models do (Weidinger et al., 2021), and should be treated with according techniques to eliminate ethical risks such as bias, stereotypes.

heading: Future Work
text: In the future, we hope to further extend the framework to more scaled large language models and also more challenging few-shot prompting setting. Limitation
We summarize two limitations which also serve as promising directions to be explored in our future work. RADA framework only considers textual domain corpus as the datastore, although this has greatly improve the coverage of domain knowledge as texts are always relatively easy to collect. However, it is widely investigated that structured knowledge such as knowledge base can also serve similar purpose. And such resources are generally in higher quality and are easier to match. Therefore, it would be benefiting to further integrate such resources at certain scenario where KB is available. The other limitation regards to the scale of the RADA implementation. As large language models have becoming increasingly powerful, they have demonstrated quite impressive capability in memorizing and recalling a wide range of background knowledge existed in the massive corpora they have been pretrained on. This trends of development naturally raises question for the proposed framework: will it still be beneficial when scale up to LLMs? and on what kinds of scenario does it brings best improvements? These are very important questions to answer, and we can certainly expect them to be explored in future works