Multi-Vector Attention Models for Deep Re-ranking (2021)

Abstract: Large-scale document retrieval systems often utilize two styles of neural network models which live at two different ends of the joint computation vs. accuracy spectrum. The first style is dual encoder (or two-tower) models, where the query and document representations are computed completely independently and combined with a simple dot product operation. The second style is cross-attention models, where the query and document features are concatenated in the input layer and all computation is based on the joint query-document representation. Dual encoder models are typically used for retrieval and deep re-ranking, while cross-attention models are typically used for shallow re-ranking. In this paper, we present a lightweight architecture that explores this joint cost vs. accuracy trade-off based on multi-vector attention (MVA). We thoroughly evaluate our method on the MS-MARCO passage retrieval dataset and show how to efficiently trade off retrieval accuracy with joint computation and offline document storage cost. We show that a highly compressed document representation and inexpensive joint computation can be achieved through a combination of learned pooling tokens and aggressive downprojection. Our code and model checkpoints are open-source and available on GitHub.

## introduction classical information retrieval systems used weighted sparse keyword matching to retrieve relevant documents for incoming search queries. a common approach has been to re-rank these documents with a neural network model which takes the concatenation of the query text and document text as input, and emits a relevance score. we refer to this as a cross-attention network, depicted in figure 1a . in modern nlp, these systems are typically pre-trained using a technique such as bert @xcite and then fine-tuned on humanlabeled relevance data @xcite @xcite . however, even if the re-ranking model is a stateof-the-art neural system, it is still limited by the documents that were produced by the retrieval system. because of this, end-to-end neural approaches to retrieval have become popular to improve the relevance of the documents produced in the retrieval stage. these models typically take the form of a dual encoder network (also called a "two-tower network," "siamese network," and "dssm" @xcite , as depicted in figure 1b ), which emits a "query vector" and "document vector" conditioned on the query text and document text respectively. the relevance score is typically defined as the dot product (or cosine distance) between these vectors. with a dual encoder model, each document vector in the corpus can be precomputed offline, and the query vector only needs to be computed once per incoming query. since scoring each querydocument pair is a simple dot product, a commodity cpu can score thousands of document in a few milliseconds. moreover, retrieving the highest scoring documents for a query can be done in sub-linear time using approximate nearest neighbor algorithms such as hierarchical k-means clustering @xcite . in practice, modern ir systems are not restricted to this simple retrieve-then-rerank framework, and often use multi-stage re-ranking. for example, the system might first retrieve the top 1000 documents using a combination of a dual encoder and keyword retrieval, then re-rank with a very cheap cross-attention model, then finally re-score the top 50 with a more expensive cross-attention model. there are two key costs to consider for these type of deep re-ranking models. the first is the number of bytes necessary to store the precomputed encoder pooler score output (a) cross-attention model. query encoder document encoder pooler pooler score (b) dual encoder model. document representation, since this must be stored in fast-access memory (typically ram) for every document in the corpus. the second is the cost of the "joint computation," which is the part of the model that combines the query and document representation in order to generate a relevance score. in this work, we explore the multi-vector attention (mva) architecture as an extension of dual encoder networks. the mva network produces a query matrix (rather than vector) which attends to a document matrix to produce a query-dependent document representation. the scalar relevance score is then computed in a manner similar to a standard dual encoder model. a mathematical description is given in section 2.1. related work. several retrieval approaches have been proposed that apply lightweight querydocument scoring on last-layer transformer features. these consist of multi-vector dual encoders @xcite @xcite ) that emit multiple query and document vectors which interact via dot products, and multi-layer attention architectures @xcite @xcite . the work presented here can be thought of as an extension to colbert @xcite , where we explore various aspects of the output layer in order to compress the document representation even further. we found the max() operation to be unstable when used in conjunction with the more aggressive pooling and downsampling, so we instead used a differentiable attention operation. 2 multi-vector attention network ## experimental setup datasets. we pre-train all of our models on the colossal clean crawled corpus (c4). for retrieval evaluation, we use the ms-marco @xcite passage re-ranking dataset. we truncate query sequences to length 32 and documents to length 112. we observe that the average lengths of queries and documents in ms-marco are around 6 and 70 respectively. training setup. we use a 12-layer transformer model with 12 attention heads and hidden size 768, equivalent in architecture to bert base . we pretrain dual encoder and cross-attention models on c4 for 100,000 iterations on a v3-128 cloud tpu, with batch size 8,192 and adam with learning rate 3e-4. our dual encoder is pre-trained directly on the mlm and nsp tasks rather than initialized as two identical bert models (which is the conventional practice), and is therefore a stronger baseline than is typically found in the literature. we initialize mva models from the dual encoder checkpoint and pre-train on the same mlm and nsp tasks for 20,000 iterations before fine-tuning to downstream tasks. we fine-tune on ms-marco using a batch size of 256 triples and adam with learning rate 3e-5. we reserve a 10% split of the ms-marco training set as a validation set for ## ablation experiments our results center around two important considerations that impact model deployment: (1) the amount of joint computation needed to score every query-document pair and (2) the cost to store the document representations offline. we first conduct ablation experiments to evaluate how independent architecture variations affect downstream accuracy.