Retrieval as Attention: End-to-end Learning of Retrieval and Reading within a Single Transformer

Proceedings of the 2022 Conference on Empirical Methods in Natural Language Processing, pages 2336 - 2349 December 7-11, 2022 ©2022 Association for Computational Linguistics Retrieval as Attention: End-to-end Learning of Retrieval and Reading within a Single Transformer Zhengbao Jiang♡∗, Luyu Gao♡∗, Jun Araki♢, Haibo Ding♢†, Zhiruo Wang♡, Jamie Callan♡, Graham Neubig♡ ♡Language Technologies Institute, Carnegie Mellon University ♢Bosch Research North America {zhengbaj,luyug,zhiruow,callan,gneubig}@cs.cmu.edu {jun.araki,haibo.ding}@us.bosch.com Abstract Systems for knowledge-intensive tasks such as open-domain question answering (QA) usu- ally consist of two stages: efficient retrieval of relevant documents from a large corpus and detailed reading of the selected documents to generate answers. Retrievers and readers are usually modeled separately, which neces- sitates a cumbersome implementation and is hard to train and adapt in an end-to-end fash- ion. In this paper, we revisit this design and eschew the separate architecture and training in favor of a single Transformer that performs Retrieval as Attention (ReAtt), and end-to-end training solely based on supervision from the end QA task. We demonstrate for the first time that a single model trained end-to-end can achieve both competitive retrieval and QA performance, matching or slightly outperform- ing state-of-the-art separately trained retriev- ers and readers. Moreover, end-to-end adap- tation significantly boosts its performance on out-of-domain datasets in both supervised and unsupervised settings, making our model a simple and adaptable solution for knowledge- intensive tasks. Code and models are available at https://github.com/jzbjyb/ReAtt. 1 Introduction Knowledge-intensive tasks such as question an- swering (QA), fact checking, and dialogue gener- ation require models to gather relevant informa- tion from potentially enormous knowledge corpora (e.g., Wikipedia) and generate answers based on gathered evidence. A widely used solution is to first retrieve a small number of relevant documents from the corpus with a bi-encoder architecture which en- codes queries and documents independently for efficiency purposes, then read the retrieved docu- ments in a more careful and expansive way with a cross-encoder architecture which encodes queries ∗The first two authors contributed equally. †Haibo Ding is now at Amazon. and documents jointly (Lee et al., 2019; Guu et al., 2020; Lewis et al., 2020; Izacard et al., 2022). The distinction between retrieval and reading leads to the widely adopted paradigm of treating retrievers and readers separately. Retrievers and readers are usually two separate models with heterogeneous architectures and different training recipes, which is cumbersome to train. Even though two models can be combined in an ad-hoc way for downstream tasks, it hinders effective end-to-end learning and adaptation to new domains. There have been several attempts to connect up reader and retriever training (Lee et al., 2019; Guu et al., 2020; Lewis et al., 2020; Sachan et al., 2021; Lee et al., 2021a; Izacard et al., 2022). How- ever, retrievers in these works are not learned in a fully end-to-end way. They require either initial- ization from existing supervisedly trained dense retrievers (Lewis et al., 2020), or expensive un- supervised retrieval pretraining as warm-up (Lee et al., 2019; Guu et al., 2020; Sachan et al., 2021; Lee et al., 2021a; …

Content truncated to 500 words