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

Proceedings of the 61st Annual Meeting of the Association for Computational Linguistics Volume 1: Long Papers, pages 2421–2436 July 9-14, 2023 ©2023 Association for Computational Linguistics Augmentation-Adapted Retriever Improves Generalization of Language Models as Generic Plug-In Zichun Yu1 Chenyan Xiong2 Shi Yu1 Zhiyuan Liu13 1Dept. of Comp. Sci. & Tech., Institute for AI, Tsinghua University, Beijing, China 2Microsoft Research, Redmond, USA 3Beijing National Research Center for Information Science and Technology, Beijing, China {yuzc19, yus21}@mails.tsinghua.edu.cn; chenyan.xiong@microsoft.com liuzy@tsinghua.edu.cn 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 signif- icantly 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. 1 Introduction Large language models (LMs) that possess bil- lions of parameters are able to capture a signif- icant amount of human knowledge, leading to consistent improvements on various downstream tasks (Brown et al., 2020; Kaplan et al., 2020; Roberts et al., 2020). However, the undeniable drawback of large LMs lies in their high compu- tational cost, which negatively impacts their effi- ciency (Strubell et al., 2019; Bender et al., 2021). Furthermore, the knowledge memorized from pre- training and the implicit reasoning process of LMs can be inaccurate and intractable sometimes, hin- dering their applications on knowledge-intensive tasks (Guu et al., 2020; Lewis et al., 2020; Mallen et al., 2022; Wei et al., 2022). Flan-T5Base (250M) Flan-T5Large (780M) Flan-T5XL (3B) InstructGPT (175B) # Parameters 35 40 45 50 55 60 65 MMLU Accuracy Standalone LM LM w/ Few-Shot Prompting LM w/ Adaptive Retrieval LM w/ AAR (Ours) Figure 1: Performance of LM w/ AAR (Ours). Instead of leveraging the knowledge and rea- soning abilities embedded within the parameters of the LMs, retrieval augmentation (Guu et al., 2020; Lewis et al., 2020; Borgeaud et al., 2022) enhances the LM with a retriever that can retrieve knowledge from an external corpus. On the other hand, prior retrieval augmentation methods (Izac- ard and Grave, 2021a; Izacard et al., 2022) necessi- tate 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 (Maronikolakis and Schütze, 2021). More importantly, many top- tier LMs can only be accessed through black-box APIs (Ouyang …

Content truncated to 500 words