Augmentation-Adapted Retriever Improves Generalization of Language Models as Generic Plug-In (2023)

Abstract: Retrieval augmentation can aid language models (LMs) in knowledge-intensive tasks by supplying them with external information. Prior works on retrieval augmentation usually jointly fine-tune the retriever and the LM, making them closely coupled. In this paper, we explore the scheme of generic retrieval plug-in: the retriever is to assist target LMs that may not be known beforehand or are unable to be fine-tuned together. To retrieve useful documents for unseen target LMs, we propose augmentation-adapted retriever (AAR), which learns LM’s preferences obtained from a known source LM. Experiments on the MMLU and PopQA datasets demonstrate that our AAR trained with a small source LM is able to significantly improve the zero-shot generalization of larger target LMs ranging from 250M Flan-T5 to 175B InstructGPT. Further analysis indicates that the preferences of different LMs overlap, enabling AAR trained with a single source LM to serve as a generic plug-in for various target LMs. Our code is open-sourced at https://github.com/OpenMatch/Augmentation-Adapted-Retriever .

## introduction large language models (lms) that possess billions of parameters are able to capture a significant amount of human knowledge, leading to consistent improvements on various downstream tasks @xcite @xcite . however, the undeniable drawback of large lms lies in their high computational cost, which negatively impacts their efficiency @xcite . furthermore, the knowledge memorized from pretraining and the implicit reasoning process of lms can be inaccurate and intractable sometimes, hindering their applications on knowledge-intensive tasks @xcite @xcite . instead of leveraging the knowledge and reasoning abilities embedded within the parameters of the lms, retrieval augmentation @xcite @xcite enhances the lm with a retriever that can retrieve knowledge from an external corpus. on the other hand, prior retrieval augmentation methods @xcite necessitate fine-tuning the backbone lm to adjust to the retriever and tackle specific downstream tasks. this kind of fine-tuning can be expensive when more and more unique demands emerge @xcite . more importantly, many toptier lms can only be accessed through black-box apis @xcite . these apis allow users to submit queries and receive responses but typically do not support fine-tuning. in this paper, we introduce augmentation-adapted retriever (aar) to assist black-box lms with downstream tasks as generic plug-in. to retrieve valuable documents for many unseen lms, we propose to leverage a small source lm to provide lm-preferred signals for retriever's training. the retriever after training (i.e., aar) can be directly utilized to assist a large target lm by plugging in the retrieved documents. specifically, we choose a small encoder-decoder lm as the source lm and utilize its fusion-in-decoder attention scores @xcite to annotate lm-preferred documents. the lm-preferred documents are then combined with human-preferred documents to form the positive document set. negative documents are mined by the retriever itself using the ance @xcite technique. after fine-tuning the retriever with lm's preferences, it can directly assist unseen target lms in the zero-shot task generalization. we evaluate aar on a multi-task language understanding dataset mmlu @xcite and an entity-centric question answering dataset popqa @xcite . for the target lms, we choose flan-t5 @xcite series as our backbone for encoder-decoder lms and instructgpt @xcite as our backbone for decoder-only lms. figure 1 shows that assisted with a generic aar, lms of different sizes and architectures can consistently outperform the standalone lms; the performance of smaller lms can sometimes surpass the standalone counterparts of significantly larger sizes (e.g., flan-t5 large w/ aar outperforms standalone flan-t5 xl by 0.6%). aar also demonstrates advantages over other augmentation approaches such as few-shot prompting and adaptive retrieval @xcite . further analysis reveals that the preferences obtained from different-sized source lms are similar, and lms with near capacities tend to yield closer preferred document sets. as a result, our aar model trained from a small source lm can be considered as a generic plug-in to enhance the zeroshot generalization of a significantly larger target lm. we also discover that the documents preferred by lms can provide assistance to the model from alternative perspectives, rather than relying solely on the full information favored by search users. recent researches adopt zero-shot retrieval augmentation that does not fine-tune the lm on in-structgpt @xcite . it can benefit entity-centric question answering @xcite , chain-of-thought reasoning @xcite , and multi-hop question answering @xcite . parallel work @xcite uses lm likelihood to train the retriever for satisfying blackbox lm's preferences, and they adopt gpt-3 curie @xcite to provide the supervision signals. in this work, we devise the retriever that can be used as a generic plug-in to assist a variety of unseen lms. zero-shot learning and reasoning. largescale unsupervised pre-trained lms like gpt-3 @xcite , gpt-4 (openai, 2023), and palm @xcite are able to perform zero-shot learning on many downstream tasks with a task description provided at inference time. instruction-finetuned lms @xcite @xcite , which are pre-trained on multiple supervised tasks using human instructions, also also exhibit robust zeroshot learning capabilities. @xcite propose a new scheme of zero-shot reasoning, which first prompts large lms to generate relevant documents and then perform reading comprehension on the generated contents. recently, there has been a growing trend of utilizing plug-and-play knowledge injection to enhance the zero-shot performance of lms, which is achieved through mapping network @xcite or document encoding @xcite . our work improves the zero-shot generalization of lms by utilizing the retrieved information. we demonstrate that identifying lms' preferences to train the retriever can in turn bring additional evidence texts for lms. ## method in this section, we first introduce the preliminaries of the dense retrieval and the retrieval-augmented lm ( § 3.1), then propose our augmentationadapted retriever ( § 3.2). ## experimental methodologies in this section, we discuss our main experimental setup. more details can be found in appendix a. ## evaluation results in this section, we discuss our main results on mmlu and popqa datasets ( § 5.1) and conduct comprehensive studies about how ( § 5.2, § 5.3, § 5.4) and when ( § 5.5, § 5.6) aar helps. ## conclusion and future work this paper introduces generic retrieval plug-in that utilizes a generic retriever to enhance target lms that may be unknown in advance or are unable to be fine-tuned jointly. our proposed retriever, aar, can directly support black-box lms without requiring any fine-tuning of the lms. this is accomplished by building the aar's training data with preferred documents from a small source lm together with the ground truth. empirical results on mmlu and popqa demonstrate that aar-assisted lms greatly outperform the standalone ones in zero-shot scenarios, and aar generalizes well to lms of different sizes 2 https://github.com/openai/chatgpt-retrieval-plugin and structures. analytical results reveal that lmpreferred and human-preferred documents complement each other; lm-preferred documents from different lms overlap significantly, and lms with similar sizes tend to yield closer document sets. we leave a more detailed explanation of how different lms interact with augmentation documents and a more reasonable selection of lm-preferred documents for future work. we hope our work shed light on a path to a generic way of treating large lms as black boxes and adapting retrievers to augment them.