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

Abstract: Systems for knowledge-intensive tasks such as open-domain question answering (QA) usually consist of two stages: efficient retrieval of relevant documents from a large corpus and detailed reading of the selected documents. This is usually done through two separate models, a retriever that encodes the query and finds nearest neighbors, and a reader based on Transformers. These two components are usually modeled separately, which necessitates a cumbersome implementation and is awkward to optimize in an end-to-end fashion. 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 (RAA), and end-to-end training solely based on supervision from the end QA task. We demonstrate for the first time that an end-to-end trained single Transformer can achieve both competitive retrieval and QA performance on in-domain datasets, matching or even slightly outperforming state-of-the-art dense retrievers and readers. Moreover, end-to-end adaptation of our model significantly boosts its performance on out-of-domain datasets in both supervised and unsupervised settings, making our model a simple and adaptable end-to-end solution for knowledge-intensive tasks.

## introduction knowledge-intensive tasks such as question answering (qa), fact checking, and dialogue generation require models to gather relevant information 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 encodes queries and documents independently for efficiency purposes, then read the retrieved documents in a more careful and expansive way with a cross-encoder architecture which encodes queries and documents jointly @xcite @xcite . 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 @xcite @xcite @xcite . however, retrievers in these works are not learned in a fully end-to-end way. they require either initialization from existing supervisedly trained dense retrievers @xcite , or expensive unsupervised retrieval pretraining as warm-up @xcite @xcite @xcite . the reliance on retrieval-specific warm-up and the ad-hoc combination of retrievers and readers makes them less of a unified solution and potentially hinders their domain adaptation ability. with the ultimate goal of facilitating downstream tasks, retriever and reader should instead be fused more organically and learned in a fully end-to-end way. in this paper, we focus on one of the most important knowledge-intensive tasks, open-domain qa. we ask the following question: is it possible to perform both retrieval and reading within a single transformer model, and train the model in a fully end-to-end fashion to achieve competitive performance from both perspectives? such a single-model end-to-end solution eliminates the need for retrieval-specific annotation and warm-up and simplifies retrieval-augmented training, making adaptation to new domains easier. based on the analogy between self-attention which relates dif-ferent tokens in a single sequence @xcite and the goal of retrieval which is to relate queries with relevant documents, we hypothesize that self-attention could be a natural fit for retrieval, and it allows an organic fusion of retriever and reader within a single transformer. specifically, we start from a encode-decoder t5 @xcite and use it as both retriever and reader. we use the first b encoder layers as bi-encoder to encode queries and documents independently, and the attention score at layer b + 1 (denoted as retrieval attention) to compute relevance scores, as shown in fig. 1 . we found that directly using self-attention for retrieval underperforms strong retrievers, which we conjecture is because self-attention pretrained on local context is not sufficient to identify relevant information in the large representation space of the whole corpus. to solve this, we propose to compute retrieval attention between a query and a large number of documents and adjust the retrieval attention across documents. for each query, we compute retrieval attention over both close documents that potentially contain positive and hard negative documents, and documents of other queries in the same batch as random negatives. the retrieval attention is adjusted by minimizing its discrepancy from the cross-attention between the decoder and encoder (denoted as target attention), which is indicative of the usefulness of each document in generating answers @xcite . the resulting retrieval as attention model (reatt) is a single t5 trained based on only qa annotations and simultaneously learns to promote useful documents through cross-document adjustment. we train reatt on natural questions dataset (nq) @xcite in a fully end-toend manner. it achieves both competitive retrieval and qa performance, matching or slightly outperforming state-of-the-art retriever colbert-nq @xcite trained with explicit retrieval annotations and strong qa model fid @xcite , demonstrating for the first time end-to-end training can produce competitive retriever and reader within a single model. to further test reatt's generalization and end-to-end adaptation ability, we conduct zero-shot, supervised, and unsupervised adaptation experiments on 7 datasets from the beir benchmark @xcite . in all settings, end-to-end adaptation improves the retrieval performance usually by a large margin, during training, the retrieval attention between a query q 1 and documents d 11,12,13 is adjusted by minimizing its discrepancy from the target attention. for simplicity, we use a single arrow to represent attention of a single head between multiple tokens. ## retrieval as attention (reatt) with the goal of developing a single transformer that can perform both retrieval and reading, and the analogy between retrieval and self-attention, we first introduce architecture changes to allow retrieval as attention ( § 2.2), then examine how well attention as-is can be directly used to perform retrieval ( § 2.3). ## learning retrieval as attention we first approximate attention over the corpus at training time by sub-sampling a manageable number of documents for each query containing both potentially relevant and random documents ( § 3.1). next, we introduce our end-to-end training objective that optimizes a standard qa loss while also adding supervision to promote attention over documents that are useful for the end task ( § 3.2). ## implementation details of reatt reatt is based on t5-large with @xmath0.001 to select the best retrieval head. we retrieve @xmath1, and use a batch size of |q| = 64 queries to obtain in-batch random documents. we use α = 8 to combine crossdocument adjustment loss with qa loss. we use adamw with a learning rate of 5e-5, 10% steps of warmup, and linear decay. we first warmup crossattention's ability to distinguish documents by only using the qa loss for 3k steps, then train with the combined losses (eq. 3) for 4 iterations, where the first iteration uses close documents returned by bm25, and the following 3 iterations use close documents returned by the previous reatt model (denoted as reatt bm25 ). each iteration has 8k update steps and takes ∼ 1.5 days on a single node with 8 × a100 gpus with 80gb memory. since dpr @xcite achieves stronger performance than bm25, training with close doc3 colbert is trained on ms marco, colbert-nq is on nq. ## out-of-domain generalization and adaptation in this section, we examine both zero-shot retrieval performance on out-of-domain datasets and reatt's end-to-end adaptability in supervised (qa, ir) and unsupervised settings. ## related work retrieval-augmented question answering utilizes evidence retrieved from an external knowledge source to facilitate question answering. there have been several attempts to learn retrievers and readers jointly. orqa, realm, rag, emdr 2 , yono, and atlas @xcite @xcite @xcite first warm-up retrievers using unsupervised pretraining methods such as inverse cloze task (ict), salient span masking (ssm), and large-scale contrastive learning, or initialize from supervised retrievers, then fine-tune both retriever and reader on downstream tasks. they either use fixed index @xcite or asynchronously update the index during training @xcite @xcite . recently, retrieval-augmented models are scaled up to very large corpora such as the web @xcite , making them capable of handling information out of the scope of wikipedia. atlas @xcite scales up retrieval-augmented models with t5-11b as the reader and contriever @xcite as the retriever and achieves strong few-shot performance on multiple benchmarks. detailed comparisons of these models can be found in tab. 7. more related works on dense retrieval, unsupervised retrieval learning, and retrieval augmentation for language modeling can be found in appendix c. ## conclusion we propose retrieval as attention (reatt), a single transformer model that can be learned in an endto-end fashion only using end task loss. we demonstrated on nq dataset that reatt can achieve both competitive retrieval and qa performance. we further show that reatt is easy to adapt to other domains in both supervised and unsupervised settings, achieving both boosted retrieval and end task performance. future directions include better endto-end training objectives, efficient training and inference, and transferring our solution to large-scale pretraining.