Title: Rethinking the Role of Token Retrieval in Multi-Vector Retrieval
Venue: NeurIPS 2023
Accepted: yes
Keywords: information retrieval

document retrieval

natural language processing
Abstract: Multi-vector retrieval models such as ColBERT [Khattab et al., 2020] allow token-level interactions between queries and documents, and hence achieve state of the art on many information retrieval benchmarks. However, their non-linear scoring function cannot be scaled to millions of documents, necessitating a three-stage process for inference: retrieving initial candidates via token retrieval, accessing all token vectors, and scoring the initial candidate documents. The non-linear scoring function is applied over all token vectors of each candidate document, making the inference process complicated and slow. In this paper, we aim to simplify the multi-vector retrieval by rethinking the role of token retrieval. We present XTR, ConteXtualized Token Retriever, which introduces a simple, yet novel, objective function that encourages the model to retrieve the most important document tokens first. The improvement to token retrieval allows XTR to rank candidates only using the retrieved tokens rather than all tokens in the document, and enables a newly designed scoring stage that is two-to-three orders of magnitude cheaper than that of ColBERT. On the popular BEIR benchmark, XTR advances the state-of-the-art by 2.8 nDCG@10 without any distillation. Detailed analysis confirms our decision to revisit the token retrieval stage, as XTR demonstrates much better recall of the token retrieval stage compared to ColBERT.

Rethinking the Role of Token Retrieval in Multi-Vector Retrieval

Abstract

Multi-vector retrieval models such as ColBERT [Khattab and Zaharia, 2020] allow token-level interactions between queries and documents, and hence achieve state of the art on many information retrieval benchmarks. However, their nonlinear scoring function cannot be scaled to millions of documents, necessitating a three-stage process for inference: retrieving initial candidates via token retrieval, accessing all token vectors, and scoring the initial candidate documents. The non-linear scoring function is applied over all token vectors of each candidate document, making the inference process complicated and slow. In this paper, we aim to simplify the multi-vector retrieval by rethinking the role of token retrieval. We present XTR, Conte X tualized T oken R etriever, which introduces a simple, yet novel, objective function that encourages the model to retrieve the most important document tokens first. The improvement to token retrieval allows XTR to rank candidates only using the retrieved tokens rather than all tokens in the document, and enables a newly designed scoring stage that is two-to-three orders of magnitude cheaper than that of ColBERT. On the popular BEIR benchmark, XTR advances the state-of-the-art by 2.8 nDCG@10 without any distillation. Detailed analysis confirms our decision to revisit the token retrieval stage, as XTR demonstrates much better recall of the token retrieval stage compared to ColBERT.

1 Introduction

The performance of a dense retrieval model is largely affected by how it defines expressive representations over queries and documents, and whether it can efficiently retrieve and score a document using these vector representations. For example, dual encoder models [Yih et al., 2011, Lee et al., 2019, Karpukhin et al., 2020, Ni et al., 2021] encode queries and documents into single vectors and compute query-document similarities using dot products. While these models are very efficient for retrieval, their expressivity is limited due to the absence of token-level modeling for scoring. In contrast, multi-vector models such as ColBERT [Khattab and Zaharia, 2020, Santhanam et al., 2022b] are directly designed to capture token-level interactions. By utilizing a (non-linear) scoring function over all query and document token representations, multi-vector models enjoy much better model expressivity and often achieve superior results across various benchmarks [Thakur et al., 2021].

The enhanced model expressivity, however, comes at a great cost of inference complexity. Unlike the case in dual encoders, the non-linear scoring function in multi-vector retrieval models prohibits the use of efficient Maximum Inner Product Search (MIPS) [Ram and Gray, 2012, Shrivastava and Li, 2014, 2015, Shen et al., 2015] for finding the maximum scoring documents. As a result, models such as ColBERT adopt an intricate and resource-intensive inference pipeline, which typically consists

of three stages: 1) token retrieval: using each query token to retrieve document tokens, with their source documents becoming candidates; 2) gathering: collecting all the token embeddings from each candidate document, including those that are not retrieved in the first stage (most document tokens are not retrieved); and 3) scoring: ranking candidates using a non-linear function based on all the token embeddings per document.

This procedure leads to two major issues. First, compared to the token retrieval stage, gathering all document token embeddings and re-scoring the documents can introduce orders of magnitude additional data loading and floating operation cost, making multi-vector models extremely expensive to deploy. Secondly, while the candidate documents are decided in the token retrieval stage, previous training objectives are designed for the scoring stage. This creates a significant training-inference gap causing multi-vector models achieve sub-optimal (and often poor) recall performance. Clearly, the three-stage pipeline has largely limited the potential of multi-vector models, raising an interesting research question can the token retrieval stage alone be sufficient for great performance?

We present XTR, Cont X extualized T oken R etriever: a simplified and efficient method for multivector retrieval, through re-thinking the role of token retrieval. The key insight of XTR is that the token retrieval in multi-vector models should be trained to retrieve the most salient and informative document tokens, so that the score between a query and document can be computed using only the retrieved information, just like how single-vector retrieval models work. By doing so, the gathering step can be completely eliminated, and the cost of scoring is significantly reduced as only a fraction of the tokens need to be considered and the dot products from the token retrieval can be reused. To improve the quality of the token retrieval, XTR proposes a novel, yet simple, training objective, which dramatically improves retrieval accuracy, doubling the chances of a gold token being retrieved in the topk results. Furthermore, despite the improved token retrieval, some relevant tokens may still be missed (i.e., not retrieved). To address this issue, we propose a simple method, called missing similarity imputation, which accounts for the contribution of the missing tokens to the overall score.

XTR streamlines the inference process, bringing it closer to the straightforward procedure of dual encoders, while maintaining and enhancing the expressive scoring function of multi-vector retrieval models. On the BEIR [Thakur et al., 2021] and LoTTE [Santhanam et al., 2022b] benchmarks, XTR attains state-of-the-art performance, requiring neither distillation nor hard negatiave mining. Notably, our model surpasses state-of-the-art dual-encoder GTR [Ni et al., 2021] by 3.6 nDCG@10 on BEIR without any additional training data. On the EntityQuestions benchmark [Sciavolino et al., 2021], XTR outperforms the previous state-of-the-art by 4.1 points on top-20 retrieval accuracy. XTR also does not require any secondary pre-training for retrieval and greatly outperforms mContriever [Izacard et al., 2022] on MIRACL, which contains multilingual retrieval tasks in 18 languages [Zhang et al., 2022b]. Our analysis supports that XTR indeed benefits from retrieving more contextualized tokens in relevant contexts, while making the scoring stage two-to-three orders of magnitude cheaper.

2 Background

2.1 Multi-vector Retrieval

Single-vector retrieval models, also known as dual encoders, encode an input text sequence as a single dense embedding and define the similarity of a query and a document based on the dot product [Lee et al., 2019, Karpukhin et al., 2020]. Multi-vector retrieval models, on the other hand, make use of multiple dense embeddings for each query and document, typically leveraging all contextualized word representations of the input to gain improved model expressivity.

Consider a query Q GLYPH<0> ρ q i ξ n i GLYPH<0> 1 and a document D GLYPH<0> ρ d j ξ m j GLYPH<0> 1 where q i and d j denote the d -dimensional query token vector and the document token vector, respectively. Multi-vector retrieval models compute the query-document similarity as follows: f GLYPH<0> Q,D GLYPH<6> GLYPH<0> < n i GLYPH<0> 1 < m j GLYPH<0> 1 A ij P ij where P ij GLYPH<0> q  i d j and A ∀ ρ 0 , 1 ξ n GLYPH<18> m denotes the alignment matrix with A ij being the token-level alignment between the query token vector q i and the document token vector d j . The sum-of-max operator of ColBERT [Khattab and Zaharia, 2020] sets A ij GLYPH<0> ✶ GLYPH<18> j GLYPH<0> argmax j ← GLYPH<0> P ij ← GLYPH<6>GLYPH<24> where the argmax operator is over 1 & j ← & m (i.e., tokens from a single document D ) and ✶ GLYPH<18> GLYPH<152> GLYPH<24> is an indicator function. Then, f ColBERT GLYPH<0> Q,D GLYPH<6> is defined as follows:

4,000× Cheaper

…

…

ColBERT

inference

XTR

training & inference

f

ColBERT

Query

Score

1

Query

f

XTR

Query

Doc

1

Doc

2

Score

1

(b) Gathering

Doc

1

Doc

2

Score

2

Score

2

f

XTR

(a) Token Retrieval

(c) Scoring (training)

(a) Token Retrieval

f

ColBERT

f

XTR

FLOPs

/query (Scoring)

ColBERT

XTR

Recall@100

on BEIR

68.0

64.5

(b) Scoring

new

new

new

Here, we include the normalizer n , which was not included in the original sum-of-max, as it stabilizes training while not affecting the ranking during inference. After computing the query-document similarity, multi-vector retrieval models are typically trained with a cross-entropy loss over in-batch negatives [Santhanam et al., 2022b, Qian et al., 2022]. Specifically, given a positive document D GLYPH<17> for Q and a set of mini-batch documents D 1 GLYPH<2> B GLYPH<0> GLYPH<18> D 1 , . . . , D B ] where D GLYPH<17> ∀ D 1 GLYPH<2> B , they minimize the cross-entropy loss defined as: L CE GLYPH<0> GLYPH<14> log exp f GLYPH<0> Q,D GLYPH<17> GLYPH<6> < B b GLYPH<0> 1 exp f GLYPH<0> Q,D b GLYPH<6> .

2.2 Three-stage inference of Multi-vector Retrieval

Unlike dual encoder models, finding the maximum scoring document-the document that maximizes eq. (1)-cannot be directly handled by MIPS as the scoring function uses a non-linear, sum-of-max operation. Instead, a multi-vector retrieval model typically takes the following steps for the inference. 1) Token Retrieval : for each of the n query token vectors, it first retrieves k ← document token vectors, which is simply used to form initial candidate document set by taking the union of source documents of retrieved tokens. The total number of candidate documents is up to nk ← if each token is coming from a unique document. 2 2) Gathering : since the scoring function eq. (1) requires the computation over all document tokens, multi-vector models need to load all of the token vectors of the candidate documents. To optimize the loading process, a RAM-based index is often employed. 3) Scoring : to provide final ranks of candidate documents, multi-vector models score all the candidate documents with eq. (1). This stage is also called refinement . Note that the training of typical multi-vector models only takes care of the scoring stage with mini-batch documents. Finally, topk documents are returned based on the computed scores. The three-stage inference is illustrated in the top of Figure 1.

3 XTR: Contextualized Token Retriever

Unlike existing multi-vector models that follow the retrieve-gather-score stages, XTR directly scores documents utilizing the tokens retrieved from the token retrieval stage. In this section, we start by showing why the existing cross entropy loss with the sum-of-max scoring function would fail on the first-stage token retrieval. Then, we introduce simple but important modifications for XTR.

Given a positive document D GLYPH<17> and a set of negative documents D GLYPH<14> 1 GLYPH<2> r GLYPH<0> GLYPH<18> D GLYPH<14> 1 , . . . , D GLYPH<14> r GLYPH<24> for a query Q , the first-stage token retrieval needs to retrieve the tokens of D GLYPH<17> , but not the tokens of negative documents. However, the following example shows that the sum-of-max operator used by ColBERT is not specifically designed to retrieve tokens of relevant documents.

Failure case Assume that f ColBERT GLYPH<0> Q,D GLYPH<17> GLYPH<6> GLYPH<0> 0 . 8 where all the individual max token similarity (i.e., q  i d GLYPH<17> j where A ij GLYPH<0> 1 ) is 0.8. On the other hand, assume f ColBERT GLYPH<0> Q,D GLYPH<14> GLYPH<6> GLYPH<0> 0 . 2 for all D GLYPH<14> ∀ D GLYPH<14> 1 GLYPH<2> r where each D GLYPH<14> has a highly peaked token similarity greater than 0.8 but others close to zero (i.e., there exists q  i d GLYPH<14> j % 0 . 8 where A ij GLYPH<0> 1 while other q  i d GLYPH<14> j GLYPH<0> 0 ). Since the sum-of-max operator only cares about the document-level scores, the cross entropy loss would be close to zero during training. 3 However, for each of n query tokens, if there exists at least one negative document token that has a high token similarity greater than 0.8, the token retrieval with topk ← GLYPH<0> 1 would fail to retrieve any tokens of D GLYPH<17> . As a result, multivector retrieval model with the sum-of-max operator will not be able to lower the high scores of some negative tokens. Figure 2 shows that the sum-of-max training causes many document tokens to have unreasonably high scores regardless of their actual relevance to the query tokens.

3.1 In-Batch Token Retrieval

To train multi-vector retrieval models to directly retrieve tokens of relevant documents, we simulate the token retrieval stage during training. This can be simply achieved by employing a different alignment strategy ˆ A . Specifically, we set the alignment ˆ A ij GLYPH<0> ✶ GLYPH<18> j ∀ topk j ← GLYPH<0> P ij ← GLYPH<6>GLYPH<24> where the topk operator is applied over 1 & j ← & mB (i.e., tokens from B mini-batch documents) returning the indices of k largest values. During training, we use a hyperparameter k train for the topk operator. Then, we simply modify eq. (1) as follows:

The intuition is that we consider the token similarities within D only when they are high enough to be retrieved within topk train from a mini-batch. Here, we use a normalizer Z GLYPH<0> ∂ρ i ∂ ↵ j, s.t. ˆ A ij % 0 ξ∂ , which is essentially the number of query tokens that retrieved at least one document token of D . 4 If all ˆ A ij GLYPH<0> 0 , we clip Z to a small number and f XTR GLYPH<0> Q,D GLYPH<6> becomes 0. As a result, our model cannot assign a high token similarity to negative documents as it blocks tokens of positive documents to be retrieved. With the previous failure case where f ColBERT assigned a high score on D GLYPH<17> even though it cannot be retrieved, our similarity function incurs a high loss as f XTR GLYPH<0> Q,D GLYPH<17> GLYPH<6> GLYPH<0> 0 during training (since tokens of D GLYPH<17> were not retrieved). For training, we use the same cross entropy loss defined in §2.1 with our new scoring function. Note that the training data only contains document-level annotations, but XTR encourages important tokens from positive documents to be retrieved.

3.2 Scoring Documents using Retrieved Tokens

During inference, multi-vector retrieval models first have a set of candidate documents ˆ D 1 GLYPH<2> C from the token retrieval stage:

0.5

0.6

0.7

0.8

0.9

1.0

Token Retrieval Score

0.00

0.01

0.02

0.03

0.04

0.05

0.06

Density

MS MARCO

T5-ColBERT

XTR (ours)

Da

Db

Qtest

Da

Db

Qtest

m

s

1

f

ColBERT

f

XTR′

m

s

2

s

1

s

2

Here, topk ← GLYPH<0> q GLYPH<152> GLYPH<6> is a union of topk ← document tokens (from the entire corpus) based on the inner product scores with each query vector (i.e., q  d ). Given the n query token vectors, there are C ( & nk ← ) candidate documents. Previous methods load the entire token vectors of each document and compute eq. (1) for every query and candidate document pair, which takes O GLYPH<0> n 2 k ← ¯ md GLYPH<6> computation per query ( ¯ m = average document length). Instead, we propose to score the documents solely using the retrieved token similarity . This significantly reduces the computational cost for the scoring stage since re-using the token retrieval scores removes computing redundant inner products and unnecessary (non-max) inner products. Furthermore, the expensive gathering stage (which requires loading all the document token vectors for computing eq. (1)) can be removed completely. Unlike previous work [Macdonald and Tonellotto, 2021] that leverages token retrieval to sort first-stage candidate documents before the scoring stage, we aim to directly provide the final scores of documents.

Missing similarity imputation During inference, we retrieve k ← document tokens for each of n query tokens. Assume that each document token belongs to a unique document, providing C GLYPH<0> nk ← candidate documents in total. This leaves us with a single token similarity to score each document in the absence of the gathering stage. However, during training-either with eq. (1) or eq. (2)-each positive document has up to n (max) token similarities to average, which mostly converges to n as training proceeds. Hence, during inference, we impute the missing similarity for each query token treating each of candidate documents as if it were positive with n token similarities.

For every candidate document ˆ D , we first define the following scoring function for the inference:

This is similar to eq. (2), but introduces m i ∀ R , which estimates the missing similarity for each q i . ˆ A is defined similar to the one described in eq. (2) except that it uses k ← for the topk operator. Each q i would take the missing similarity m i as the maximum value if ˆ A i GLYPH<152> GLYPH<0> 0 and m i ∋ 0 . Importantly, f XTR ← removes the need of recomputing any q  i d j since when ˆ A ij GLYPH<0> 1 we already know the retrieval score from the token retrieval stage, and when ˆ A ij GLYPH<0> 0 we simply don't need to compute it as ˆ A ij q  i d j GLYPH<0> 0 . Note that when every ˆ A ij GLYPH<0> 1 , the equation becomes the sum-of-max operator. On the other hand, when no document tokens of ˆ D were retrieved for q i (i.e., ˆ A i GLYPH<152> GLYPH<0> 0 ), we fall back to the imputed score m i , which provides an approximated sum-of-max result.

≡ : cross-encoder distillation ∂ : model-based hard negatives

In fact, we can find the upper bound of the missing similarity. For every token retrieval with q i , the missing similarity of the query token for ˆ D will be upper bounded by its last topk ← score. Specifically, for each query token q i , we have the following topk ← token similarity during inference: GLYPH<18> q  i d GLYPH<0> 1 GLYPH<6> , . . . q  i d GLYPH<0> k ← GLYPH<6> GLYPH<24> . Here, each d GLYPH<0> GLYPH<152> GLYPH<6> could come from a different document. Since the missing similarity would have a score less than equal to the score of the last retrieved token, we know that m i & q  i d GLYPH<0> k ← GLYPH<6> . With a larger k ← , the upper bound becomes tighter. In our experiments, we show that simply choosing m i GLYPH<0> q  i d GLYPH<0> k ← GLYPH<6> works well especially when a model is trained with f XTR. 5 While we also tried more complicated imputation methods based on regression, our method was competitive enough despite its simplicity. The imputation process is illustrated in Figure 3.

Table 1 shows the estimated FLOPs of ColBERT and XTR (see Appendix B for more details). Due to the differences in hardware and infrastructure, we mainly compared the theoretical FLOPs. XTR reduces the FLOPs at the scoring stage by 4000 GLYPH<18> making multi-vector retrieval more efficient.

4 Experiments

Experimental Setting Following Ni et al. [2021], we fine-tune XTR on MS MARCO with a fixed set of hard negatives from RocketQA [Qu et al., 2021]. Then, we test XTR on MS MARCO (MS; in-domain) and zero-shot IR datasets. For the zero-shot evaluation, we use 13 datasets from BEIR [Thakur et al., 2021] (see Appendix C for acronyms), 12 datasets from LoTTE [Santhanam et al., 2022b], and 4 datasets on open-domain QA passage retrieval (EQ: EntityQuestions [Sciavolino et al., 2021], NQ, TQA: TriviaQA, SQD: SQuAD). We also train multilingual XTR (mXTR) and evaluate it on MIRACL [Zhang et al., 2022b], which contains retrieval tasks in 18 languages. The performance gap between T5-ColBERT [Qian et al., 2022] and XTR shows the improvement with our methods on a multi-vector retrieval model. For implementation details and baselines, see Appendix C. For the relationship between hyperparameters (e.g., k train and k ← ), see §5.3.

≡ : sparse component ∂ : retrieval pre-training

4.1 In-domain Document Retrieval

MS MARCO The first column of Table 2 (top) shows nDCG@10 on MS MARCO (see Table D.1 for recall@100). XTR outperforms most models and remains competitive with T5-ColBERT. This is encouraging since XTR significantly reduces the cost of the gathering-scoring stage. Note that MS MARCO may fail to reflect the actual improvement of state-of-the-art [Arabzadeh et al., 2022].

4.2 Zero-shot Document Retrieval

BEIR & LoTTE Table 2 (top; except the first columns) shows nDCG@10 on BEIR (see Table D.1 for recall@100). XTRxxl achieves the new state-of-the-art performances significantly outperforming both per-domain models and single model state-of-the-art. Simply scaling XTR removes the needs of designing distillation or hard negative mining pipelines [Santhanam et al., 2022b, Formal et al., 2021]. Results on LoTTE (Table 2 bottom) also show that XTRbase is better than ColBERT and competitive with distillation-based models while XTRxxl advances the state-of-the-art.

Passage retrieval for open-domain QA Table 3 shows results on four open-domain QA datasets. While previous work often includes sparse retrievers (e.g., BM25) [Chen et al., 2021] or contrastive pre-training [Ram et al., 2022, Sachan et al., 2022a,b] to achieve better performances on EntityQuestions, XTR simply fine-tuned on MS MARCO achieves the state-of-the-art performance.

4.3 Multilingual Document Retrieval

MIRACL Since XTR does not need any secondary pre-training, we expect it to be better at multilingual retrieval by better utilizing the multilingual language models. We train a multilingual version of XTR with mT5 [Xue et al., 2021] and test it on multilingual retrieval tasks in 18 languages. Table 4 shows that mXTR greatly outperforms mContriever that uses expensive contrastive pretraining, as well as the hybrid model, BM25 + mDPR.

0

200

400

600

800

1000

Token Rank k

0.00

0.02

0.04

0.06

0.08

0.10

P(Gold|Token@k)

MS MARCO

T5-ColBERT-xxl

XTR-xxl

0

200

400

600

800

1000

Token Rank k

0.0

0.2

0.4

0.6

0.8

1.0

P(Gold|Token@k)

TREC-COVID

T5-ColBERT-xxl

XTR-xxl

0

200

400

600

800

1000

Token Rank k

0.00

0.02

0.04

0.06

0.08

0.10

P(Gold|Token@k)

ArguAna

T5-ColBERT-xxl

XTR-xxl

0

200

400

600

800

1000

Token Rank k

0.0

0.2

0.4

0.6

0.8

1.0

P(Lexical|T oken@k)

MS MARCO

T5-ColBERT-xxl

XTR-xxl

0

200

400

600

800

1000

Token Rank k

0.0

0.2

0.4

0.6

0.8

1.0

P(Lexical|T oken@k)

TREC-COVID

T5-ColBERT-xxl

XTR-xxl

0

200

400

600

800

1000

Token Rank k

0.0

0.2

0.4

0.6

0.8

1.0

P(Lexical|T oken@k)

ArguAna

T5-ColBERT-xxl

XTR-xxl

5 Analysis

5.1 Towards Better Token Retrieval

Gold token retrieval If the tokens of gold documents are not retrieved at all, multi-vector retrieval models would fail to retrieve the gold documents. Hence, a better token retrieval would contain these gold tokens more often in their top results. In Figure 4 (top), we show the probability of a token at the rank k coming from the gold documents of a query. To compute the probability for the rank k , we simply count the number of an event where a token at rank k belongs to the gold document and divide it by the number of tokens at rank k . While this is measuring the precision of the token retrieval, we observed a similar trend for the recall of gold tokens. Compared to T5-ColBERT, XTR retrieves gold tokens with higher probability, even on MS MARCO. This shows that the training objective of XTR encourages it to retrieve tokens from more relevant context.

Lexical token retrieval In Figure 4 (bottom), we show the probability of a token at the rank k being the same as its query token (e.g., ' insulin ' retrieving ' insulin 's). T5-ColBERT has very high probability of retrieving the same token across different ranks and datasets. However, it is unclear to what extent the token retrieval stage should behave as sparse retrieval, as it might suffer from the vocabulary mismatch problem. XTR effectively lowers the reliance on the lexical matching while preserving a good amount of lexical precision so that it would achieve a high retrieval accuracy on the entity-centric dataset (§4.2). In fact, Table 6 in Appendix shows that having lower lexical matching doesn't necessarily mean a lower retrieval quality, but often means better contextualization.

5.2 Efficient Scoring

In Table 5, we show how we can employ the efficient scoring function f XTR ← in XTR with minimal performance losses. We apply f XTR ← on both T5-ColBERT and XTR, and show their performances on MS MARCO. With T5-ColBERT, even if we use the topk ← score for the imputation, the performance

10 2

10 3

10 4

10 5

k

0.70

0.80

0.90

Recall@100

MS MARCO

T5-ColBERT-base (f_ColBERT)

T5-ColBERT-base (f_XTR )

XTR-base

10 1

10 2

10 3

10 4

10 5

k

0.25

0.30

0.35

0.40

MRR@10

MS MARCO

XTR-base (k_train = 32)

XTR-base (k_train = 64)

XTR-base (k_train = 128)

XTR-base (k_train = 256)

XTR-base (k_train = 320)

T5-ColBERT-base (f_XTR )

0

50

100

150

200

250

300

k_train

0.43

0.44

0.45

0.46

nDCG@10

MS MARCO

XTR-base (batch_size = 128)

XTR-base (batch_size = 256)

XTR-base (batch_size = 320)

0

50

100

150

200

250

300

k_train

0.36

0.38

0.40

0.42

0.44

nDCG@10

ArguAna

XTR-base (batch_size = 128)

XTR-base (batch_size = 256)

XTR-base (batch_size = 320)

is much worse than the original sum-of-max scoring. With XTR, the performance greatly improves as it has better token retrieval. Figure 5 shows how Recall@100 improves with larger k ← 's as it provides more exact upper bound for the missing similarity imputation. Table D.2 shows that even if we use smaller k ← , XTR still maintains high performances on BEIR.

5.3 Relationship between Hyperparameters

k train vs. k ← In Figure 6, we show MRR@10 of XTR trained with different k train and evaluated with different k ← on the MS MARCO development set. While all variants of XTR prefer larger k ← , ones trained with smaller k train show higher performances than others under small k ← settings. XTR with larger k train exhibits better performances than ones with smaller k train as k ← becomes larger.

Training batch size vs. k train In Figure 7, we show the relationship between the training batch size and k train during training XTR. In this experiment, we use k ← GLYPH<0> 40 , 000 . While it is evident that XTR mostly favors large training batch sizes, the optimal topk train can be different for different datasets. While most datasets including MS MARCO favored a large enough k train, ArguAna prefers smaller k train. We hypothesize that this is due to the longer query length in ArguAna, which makes multi-vector models fall short compared to dual-encoders (see GTR vs. T5-ColBERT in Table 2).

5.4 Qualitative Analysis

Table 6 shows a prediction sample from MS MARCO. For T5-ColBERT, all of the top retrieved tokens are exact lexical matches. Surprisingly, none of the retrieved passages are about the query, demonstrating T5-ColBERT's failure to retrieve tokens from the correct context. In contrast, XTR retrieves fewer exact lexically matching tokens, but the contexts of the retrieved tokens are much more related to the query. This example explains the lower lexical token retrieval probability of XTR compared to T5-ColBERT in Figure 4 (bottom), but higher gold token retrieval performance in Figure 4 (top). For more qualitative examples, please see Appendix E.

6 Related Work

One of the main limitations of dense retrieval models is that encoding the query and document into a single vector constrains the representational power of the models. Polyencoder [Humeau et al., 2020], MEBERT [Luan et al., 2021], and MVR [Zhang et al., 2022a] propose to use multiple embeddings, instead of one, to represent the query or the document. A more recent approach is token-level multi-vector retrieval, which stores and retrieves with every token embedding. ColBERT [Khattab and Zaharia, 2020] is probably the most renowned model in this family. ALIGNER (i.e. T5ColBERT) [Qian et al., 2022] extends ColBERT by scaling up the backbone langauge model and studying various strategies for aggregating the token-level alignment scores. These token-level retrieval models show strong effectiveness and out-of-domain generalization ability.

Efforts for reducing serving costs of multi-vector models have been mostly focused on the token-level retrieval stage. COIL [Gao et al., 2021] accelerates token-level retrieval by confining retrieval within exact match tokens, sharing the spirit of classic inverted indexing. CITADEL [Li et al., 2022] relaxes COIL with a lexical routing mechanism where a query token vector only retrieves from a subset of

document token vectors routed to the same key. PLAID [Santhanam et al., 2022a] optimizes the speed of ColBERT by pruning weaker candidates in the earlier stages of retrieval and using better vector quantization. ColBERT-v2 [Santhanam et al., 2022b] further adopts residual representations with cluster centroids to improve the efficiency of ColBERT. On the other hand, how to accelerate the scoring stage remains under-explored. To the best of our knowledge, XTR is the first work to simplify the scoring stage and remove the gathering stage in multi-vector retrieval.

7 Conclusion

Multi-vector retrieval leverages query and document token representations for effective information retrieval. In this paper, we propose XTR that simplifies the existing three-stage inference of multivector models by improving the initial token retrieval stage. Specifically, XTR scores documents solely based on the retrieved tokens, which is also optimized during training with in-batch document tokens. As a result, XTR achieves state-of-the-art performances on zero-shot information retrieval benchmarks while greatly reducing the FLOPs of the scoring stage. We further show that our objective function indeed encourages better token retrieval, retrieving more tokens from gold documents, whose contexts are better aligned with the query.

Limitations

In most of our experiments, XTR was trained on MS MARCO, a large-scale retrieval dataset in English. While our experiments were conducted in a fair setting where most baseline models also utilize MS MARCO, future use cases might need to remove its dependency on MS MARCO due to the license or language-specific issue. We believe that LLM-based retrieval dataset generation [Dai et al., 2022] would be able to mitigate the problem in the future.

References

Negar Arabzadeh, Alexandra Vtyurina, Xinyi Yan, and Charles LA Clarke. Shallow pooling for sparse labels. Information Retrieval Journal , 25(4):365-385, 2022.

Xilun Chen, Kushal Lakhotia, Barlas O˘ guz, Anchit Gupta, Patrick Lewis, Stan Peshterliev, Yashar Mehdad, Sonal Gupta, and Wen-tau Yih. Salient phrase aware dense retrieval: Can a dense retriever imitate a sparse one? arXiv preprint arXiv:2110.06918 , 2021.

Zhuyun Dai, Vincent Y Zhao, Ji Ma, Yi Luan, Jianmo Ni, Jing Lu, Anton Bakalov, Kelvin Guu, Keith B Hall, and Ming-Wei Chang. Promptagator: Few-shot dense retrieval from 8 examples. arXiv preprint arXiv:2209.11755 , 2022.

Thibault Formal, Carlos Lassance, Benjamin Piwowarski, and Stéphane Clinchant. Splade v2: Sparse lexical and expansion model for information retrieval. arXiv preprint arXiv:2109.10086 , 2021.

Luyu Gao, Zhuyun Dai, and Jamie Callan. COIL: revisit exact lexical match in information retrieval with contextualized inverted list. In Kristina Toutanova, Anna Rumshisky, Luke Zettlemoyer, Dilek Hakkani-Tür, Iz Beltagy, Steven Bethard, Ryan Cotterell, Tanmoy Chakraborty, and Yichao Zhou, editors, Proceedings of the 2021 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, NAACL-HLT 2021, Online, June 6-11, 2021 , pages 3030-3042. Association for Computational Linguistics, 2021.

Ruiqi Guo, Philip Sun, Erik Lindgren, Quan Geng, David Simcha, Felix Chern, and Sanjiv Kumar. Accelerating large-scale inference with anisotropic vector quantization. In International Conference on Machine Learning , pages 3887-3896. PMLR, 2020.

Samuel Humeau, Kurt Shuster, Marie-Anne Lachaux, and Jason Weston. Poly-encoders: Architectures and pre-training strategies for fast and accurate multi-sentence scoring. In 8th International Conference on Learning Representations, ICLR 2020, Addis Ababa, Ethiopia, April 26-30, 2020 . OpenReview.net, 2020.

Gautier Izacard, Mathilde Caron, Lucas Hosseini, Sebastian Riedel, Piotr Bojanowski, Armand Joulin, and Edouard Grave. Unsupervised dense information retrieval with contrastive learning. Transactions on Machine Learning Research , 2022.

Vladimir Karpukhin, Barlas O˘ guz, Sewon Min, Patrick Lewis, Ledell Yu Wu, Sergey Edunov, Danqi Chen, and Wen tau Yih. Dense passage retrieval for open-domain question answering. ArXiv , abs/2004.04906, 2020.

Omar Khattab and Matei Zaharia. Colbert: Efficient and effective passage search via contextualized late interaction over bert. In Proceedings of the 43rd International ACM SIGIR conference on research and development in Information Retrieval , pages 39-48, 2020.

Kenton Lee, Ming-Wei Chang, and Kristina Toutanova. Latent retrieval for weakly supervised open domain question answering. In Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics , pages 6086-6096, Florence, Italy, July 2019. Association for Computational Linguistics.

Minghan Li, Sheng-Chieh Lin, Barlas Oguz, Asish Ghoshal, Jimmy Lin, Yashar Mehdad, Wen-tau Yih, and Xilun Chen. Citadel: Conditional token interaction via dynamic lexical routing for efficient and effective multi-vector retrieval. arXiv preprint arXiv:2211.10411 , 2022.

Yi Luan, Jacob Eisenstein, Kristina Toutanova, and Michael Collins. Sparse, Dense, and Attentional Representations for Text Retrieval. Transactions of the Association for Computational Linguistics , 9:329-345, 04 2021.

Craig Macdonald and Nicola Tonellotto. On approximate nearest neighbour selection for multi-stage dense retrieval. In Proceedings of the 30th ACM International Conference on Information & Knowledge Management , pages 3318-3322, 2021.

Jianmo Ni, Chen Qu, Jing Lu, Zhuyun Dai, Gustavo Hern'andez 'Abrego, Ji Ma, Vincent Zhao, Yi Luan, Keith B. Hall, Ming-Wei Chang, and Yinfei Yang. Large dual encoders are generalizable retrievers. In Conference on Empirical Methods in Natural Language Processing , 2021.

Yujie Qian, Jinhyuk Lee, Sai Meher Karthik Duddu, Zhuyun Dai, Siddhartha Brahma, Iftekhar Naim, Tao Lei, and Vincent Y Zhao. Multi-vector retrieval as sparse alignment. arXiv preprint arXiv:2211.01267 , 2022.

Yingqi Qu, Yuchen Ding, Jing Liu, Kai Liu, Ruiyang Ren, Wayne Xin Zhao, Daxiang Dong, Hua Wu, and Haifeng Wang. Rocketqa: An optimized training approach to dense passage retrieval for open-domain question answering. In Proceedings of the 2021 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies , pages 5835-5847, 2021.

Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. The Journal of Machine Learning Research , 21(1):5485-5551, 2020.

Ori Ram, Gal Shachaf, Omer Levy, Jonathan Berant, and Amir Globerson. Learning to retrieve passages without supervision. In Proceedings of the 2022 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies , pages 2687-2700, 2022.

Parikshit Ram and Alexander G Gray. Maximum inner-product search using cone trees. In Proceedings of the 18th ACM SIGKDD international conference on Knowledge discovery and data mining , pages 931-939, 2012.

Devendra Singh Sachan, Mike Lewis, Mandar Joshi, Armen Aghajanyan, Wen-tau Yih, Joelle Pineau, and Luke Zettlemoyer. Improving passage retrieval with zero-shot question generation. arXiv preprint arXiv:2204.07496 , 2022a.

Devendra Singh Sachan, Mike Lewis, Dani Yogatama, Luke Zettlemoyer, Joelle Pineau, and Manzil Zaheer. Questions are all you need to train a dense passage retriever. arXiv preprint arXiv:2206.10658 , 2022b.

Keshav Santhanam, Omar Khattab, Christopher Potts, and Matei Zaharia. Plaid: an efficient engine for late interaction retrieval. In Proceedings of the 31st ACM International Conference on Information &Knowledge Management , pages 1747-1756, 2022a.

Keshav Santhanam, Omar Khattab, Jon Saad-Falcon, Christopher Potts, and Matei Zaharia. Colbertv2: Effective and efficient retrieval via lightweight late interaction. In Proceedings of the 2022 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies , pages 3715-3734, 2022b.

Christopher Sciavolino, Zexuan Zhong, Jinhyuk Lee, and Danqi Chen. Simple entity-centric questions challenge dense retrievers. In Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing , pages 6138-6148, 2021.

Fumin Shen, Wei Liu, Shaoting Zhang, Yang Yang, and Heng Tao Shen. Learning binary codes for maximum inner product search. In Proceedings of the IEEE International Conference on Computer Vision , pages 4148-4156, 2015.

Anshumali Shrivastava and Ping Li. Asymmetric lsh (alsh) for sublinear time maximum inner product search (mips). In Advances in Neural Information Processing Systems , pages 2321-2329, 2014.

Anshumali Shrivastava and Ping Li. Improved asymmetric locality sensitive hashing (alsh) for maximum inner product search (mips). In Conference on Uncertainty in Artificial Intelligence , 2015.

Nandan Thakur, Nils Reimers, Andreas Rücklé, Abhishek Srivastava, and Iryna Gurevych. Beir: A heterogeneous benchmark for zero-shot evaluation of information retrieval models. In Thirty-fifth Conference on Neural Information Processing Systems Datasets and Benchmarks Track (Round 2) , 2021.

Kexin Wang, Nandan Thakur, Nils Reimers, and Iryna Gurevych. Gpl: Generative pseudo labeling for unsupervised domain adaptation of dense retrieval. In Proceedings of the 2022 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies , pages 2345-2360, 2022.

Linting Xue, Noah Constant, Adam Roberts, Mihir Kale, Rami Al-Rfou, Aditya Siddhant, Aditya Barua, and Colin Raffel. mt5: A massively multilingual pre-trained text-to-text transformer. In Proceedings of the 2021 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies , pages 483-498, 2021.

Wen-tau Yih, Kristina Toutanova, John C. Platt, and Christopher Meek. Learning discriminative projections for text similarity measures. In Conference on Computational Natural Language Learning , 2011.

Shunyu Zhang, Yaobo Liang, Ming Gong, Daxin Jiang, and Nan Duan. Multi-view document representation learning for open-domain dense retrieval. In Smaranda Muresan, Preslav Nakov, and Aline Villavicencio, editors, Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), ACL 2022, Dublin, Ireland, May 22-27, 2022 , pages 5990-6000. Association for Computational Linguistics, 2022a.

Xinyu Zhang, Nandan Thakur, Odunayo Ogundepo, Ehsan Kamalloo, David Alfonso-Hermelo, Xiaoguang Li, Qun Liu, Mehdi Rezagholizadeh, and Jimmy Lin. Making a miracl: Multilingual information retrieval across a continuum of languages. arXiv preprint arXiv:2210.09984 , 2022b.

A Derivatives w.r.t. Similarity Scores

Sum-of-max Here, we use a cross-entropy loss L CE with the sum-of-max operator f ColBERT and analyze the derivatives with respect to the token similarity scores.

Here, we denote ˆ j as the index of the row-wise maximum value, dependent on each i (i.e., A ij GLYPH<0> 1 ). Given the cross-entropy loss with the sum-of-max operator, we compute the gradient with respect to one of the maximum token similarities P GLYPH<17> i ˆ j for a positive document D GLYPH<17> ∀ D 1 GLYPH<2> B :

Similarly, the gradient w.r.t. a maximum token similarity P GLYPH<14> i ˆ j for a negative document D GLYPH<14> ∀ D 1 GLYPH<2> B is computed as follows:

Hence, the positive token-level score P GLYPH<17> i ˆ j will gradually increase until P GLYPH<0> D GLYPH<17> ∂ Q,D 1 GLYPH<2> B GLYPH<6> GLYPH<0> 1 and the negative token-level score P GLYPH<14> i ˆ j will decrease until P GLYPH<0> D GLYPH<14> ∂ Q,D 1 GLYPH<2> B GLYPH<6> GLYPH<0> 0 . This shows that the tokenlevel scores are trained based on the document-level scores, which might stagnate the token-level scores. For instance, even if P GLYPH<14> i ˆ j is very high-later causing d GLYPH<14> ˆ j to be retrieved instead of ones from positive documents-it will not be penalized as long as P GLYPH<0> D GLYPH<14> ∂ Q,D 1 GLYPH<2> B GLYPH<6> is low enough.

In-batch token retrieval Compared to the sum-of-max operator, our in-batch sum-of-max f XTR considers the max values only when they are retrieved over other negative tokens in the mini-batch.

Here, we denote ¯ j as the index of the row-wise maximum value that is also within the mini-batch topk train given q i (i.e., satisfies both A ij GLYPH<0> 1 and ˆ A ij GLYPH<0> 1 ). If there is no such ¯ j , we simply use P i ¯ j GLYPH<0> 0 . We also use a normalizer Z , which is the number of non-zero P i ¯ j . In this analysis, we assume Z % 0 since if every P i ¯ j is zero, the gradient is undefined.

The gradient w.r.t. the maximum token similarity P GLYPH<17> i ¯ j (non-zero) for a positive document D GLYPH<17> ∀ D 1 GLYPH<2> B is computed as follows:

This is a very similar result compared to the sum-of-max operator except that 1) the gradient is defined only when P GLYPH<17> i ¯ j is non-zero (i.e. retrieved) and 2) it is dependent on Z GLYPH<17> , which means that the gradient will be large whenever there is a small number of retrieved tokens from the positive document. If only a handful of tokens are retrieved for D GLYPH<17> , our objective function increases P GLYPH<17> i ¯ j .

For negative similarity score P GLYPH<14> i ¯ j , we have the following:

Again, it is similar to the sum-of-max result, but it depends on Z GLYPH<14> . In this case, even when P GLYPH<0> D GLYPH<14> ∂ Q,D 1 GLYPH<2> B GLYPH<6> is low, if there is a small number of retrieved tokens from D GLYPH<14> (i.e., small Z GLYPH<14> ), P GLYPH<14> i ˆ j will be decreased significantly. Note that when Z GLYPH<14> is large, Z GLYPH<17> naturally becomes smaller as they compete for in-batch token retrieval, which causes positive tokens to have higher scores.

B Inference Complexity

Wecompare the complexity of ColBERT and XTR during the scoring stage in terms of FLOPs. We do not measure the complexity for the online query encoding and maximum inner product search (MIPS), which have been extensively studied for both dual encoders and multi-vector retrieval [Santhanam et al., 2022a,b, Guo et al., 2020].

For the scoring stage, both ColBERT and XTR have O GLYPH<0> nk ← GLYPH<6> candidate documents. Here, we assume the worst case nk ← where each document token comes from a unique document. For each candidate document, ColBERT loads a set of document vectors of ¯ md floating points ( ¯ m GLYPH<0> average document length) and computes eq. (1) with the query vectors of nd floating points. Computing eq. (1) per candidate document requires 2 n ¯ md FLOPs for token-level inner products, n ¯ m for finding the rowwise max, and n for the final average. In total, ColBERT requires n 2 k ← GLYPH<0> 2 ¯ md GLYPH<17> ¯ m GLYPH<17> 1 GLYPH<6> FLOPs for the scoring stage. Note that this does not include the latency of loading the O GLYPH<0> nk ← ¯ md GLYPH<6> floating points onto the memory, which amounts up to 450MB per query when n GLYPH<0> 16 , k ← GLYPH<0> 1000 , ¯ m GLYPH<0> 55 , d GLYPH<0> 128 .

On the other hand, XTR first imputes the missing similarity, which is simply done by caching the k ← -th token retrieval score for each query token. Then, each of nk ← candidate documents requires n ¯ r FLOPs for finding row-wise max and n for the average where ¯ r is the average number of retrieved tokens per each candidate document. In total, we have n 2 k ← GLYPH<0> ¯ r GLYPH<17> 1 GLYPH<6> FLOPs. Table 1 shows the estimated FLOPs of the two models. XTR reduces the FLOPs at the scoring stage by 4000 GLYPH<18> making multi-vector retrieval more efficient and practical.

C Implementation Details

XTR uses k train for retrieving in-batch document tokens. Since we retrieve over mini-batches, the size of mini-batch affects the performance for different k train, which is shown in §5.3. In our experiments, we tried k train GLYPH<0> ρ 32 , 64 , 128 , 256 , 320 ξ for each batch size and choose the best model based on their performance on the MS MARCO development set. For inference, XTR uses k ← for the token retrieval. We use k ← GLYPH<0> 40 , 000 , which is possible due to the efficient scoring stage of XTR. 6 We analyze the effect of using different k ← 's as well as its relationship to k train in §5.3. We initialize XTR from the base and xxl versions of the T5 encoder [Raffel et al., 2020] and provide XTRbase and XTRxxl. For multilingual XTR, we initialize XTR from mT5 [Xue et al., 2021]. We fine-tune XTR for 50,000 iterations with the learning rate to 1e-3. Up to 256 chips of TPU v3 accelerator were used depending on the size of the model. We use ScaNN [Guo et al., 2020] for the MIPS during the token retrieval stage. For BEIR, we use 13 datasets (AR: ArguAna. TO: Touché-2020. FE: Fever. CF: Climate-Fever. SF: Scifact. CV: TREC-COVID. NF: NFCorpus. NQ: Natural Questions. HQ: HotpotQA. FQ: FiQA-2018. SD: SCIDOCS. DB: DBPedia. QU: Quora).

Baselines There are two main paradigms on training retriever models for the out-of-domain evaluation. The first group trains a single retriever for each dataset (or domain) by generating queries for each out-of-domain corpus. Typically, this approach generates N datasets to train N independent models for N different domains. For this one-retriever-per-domain approaches, we include GenQ [Thakur et al., 2021], GPL [Wang et al., 2022], and Promptagator [Dai et al., 2022]. The second group builds a single retriever-typically trained on a large-scale IR dataset such as MS MARCO-and directly applies it on the out-of-domain corpora and queries. For this one-retrieverfor-all approaches, we present results of state-of-the-art retrievers including Spladev2 [Formal et al., 2021], ColBERTv2 [Santhanam et al., 2022b], and GTRxxl [Ni et al., 2021]. We also show the results of T5-ColBERTxxl [Qian et al., 2022], which is a T5-initialized ColBERT model and shares the same backbone LM and training dataset with XTR. Note that T5-ColBERT uses the heavy scoring stage based on the original sum-of-max. All of our one-retriever-for-all baselines, as well as XTR, are trained on English MS MARCO, unless otherwise stated.

D Additional Results

Table D.2: nDCG@10 and Recall@100 of XTRbase on MS-MARCO and BEIR with different k ← . The last column shows the average over 13 BEIR benchmarks.

E Qualitative Analysis

In Table 6-E.5, we show token retrieval results from T5-ColBERT and XTR.

XTR for 'women with a higher birth weight are more likely to develop breast cancer later in life. '

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe authors improve multi-vector retrieval to move beyond the standard retrieve-gather-score stages of ColBERT. In particular, they modify the objective function during training as well as the scoring mechanism so it doesn't require gathering all token vectors of each candidate document before the final scores are computed.\n\nThe authors begin by showing that, in ColBERT, the standard cross-entropy loss applied over the aggregated scores fails to reduce token-level scores if the average score is low. This can reduce the precision of the initial \"retrieve\" stage, putting a higher burden on the gather/score stages.\n\nTo tackle that, the authors simulate the token retrieval stage during training by masking/skipping document token scores for tokens that aren't among the nearest top_{k_{train}} to a given query token within the training batch. At search time, the authors use exclusively the retrieved (i.e., close) tokens from candidate documents. Scores of \"missing\" tokens are imputed/estimated with a lower bound from the kNN retrieval stage.\n\nSTRENGTHS:\nThe paper is very strong overall. It tackles a well-defined problem with a well-motivated and novel solution. I can see this being widely applicable to multi-vector / ColBERT-like retrievers. The analysis and rich results are very strong as well, considering the paper doesn't use distillation from cross encoders. (The paper does use hard negatives, though --- from RocketQA).\n\nWEAKNESSES:\nDespite the very strong contributions, a number of claims of the paper are needlessly inflated (or unsupported).\n\nFirst, the abstract (and paper) say that XTR \"enables a newly designed scoring stage that is two-to-three orders of magnitude cheaper than that of ColBERT\". The basis of this claim is a theoretical FLOPs analysis of one of the three scoring stages of ColBERT against XTR.\n\nTo the best of my knowledge, this appears, however, unsupported or otherwise inflated on a few dimensions:\n\n1) This comparison doesn't count the actual FLOPs used. It also doesn't measure the latency of the scoring stage or the full pipeline. Even though these measurements may be less idealized than theoretical analysis, not reporting them makes it much harder to assess the proposed methods.\n\n2) While XTR appears to make the scoring stage essentially free (at least based on the FLOPs analysis), it is unclear how much this affects the total computational cost of retrieval. Is the latency now low? Is it lower than existing multi-vector retrievers? The paper offers no empirical insight on this, but it's unlikely (or, rather, impossible?) that XTR is 2-3 orders of magnitude faster than them overall. For instance, ColBERT retriever in PLAID (2022) seems to report latency of 58 ms (or lower) per query on MS MARCO. Two orders of magnitude faster than that would be 0.6 milliseconds per query, which is much faster than even basic BM25 retrieval permits.\n\n3) The comparison is conducted against the original ColBERT (2020), but this would appear to ignore years of optimizations for multi-vector scoring as in ColBERTv2 (2021), PLAID (2022), and other papers like ColBERTer (2022), etc. For instance, the PLAID ColBERTv2 retriever appears to show that candidate scoring (and the corresponding index lookup) is an extremely cheap step. This may be because PLAID moves a lot of the work to earlier stages, making them more expensive instead(?). However, there is very minimal engagement in the paper with the concerns of these developments in terms of faster candidate generation and/or much smaller index. While this isn't the focus of the paper, it's important to (at least) report the index size, how much of it needs to reside in memory, and how fast the overall retrieval process can be for XTR.\n\nSecond, the abstract reports that \"on the popular BEIR benchmark, XTR advances the state-of-the-art by 2.8 nDCG@10 without any distillation.\" It's true this is achieved without any distillation, but such a result is only achieved with the XXL (11B or 5.5B?) T5 encoder of XTR, which is orders of magnitude larger than related methods. This isn't an issue with the evaluation itself, but it shouldn't be glossed over.\n\nQUESTIONS:\nWhat is the intuition for the normalizer Z during training? How essential is this for the results?\n\nHow robust are the analyses with the T5-ColBERT, compared to the original ColBERT and ColBERTv2 models? ColBERT typically has a query expansion stage that appears to be skipped in XTR. How does this affect the analysis and results?\n\nLIMITATIONS:\nSee Weaknesses & Questions.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-05"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nAuthors propose a better document retrieval method. They build on top of ColBert and instead of reranking using all tokens of documents retrieved by stage 1(a query token document token retrieval), they just use retrieved document tokens from stage 1 and perform retrieval using them.\n\nSTRENGTHS:\n-\tResults seem to be strong enough.\n-\tDo a decent job in ablations, qualitative analysis, complexity analysis.\n\nWEAKNESSES:\n-\tThe paper overcomplicates a simple concept (Eq 4 can be further simplified from what I understand). Max over j would collapse then.\n-\tGiven the proposed methodology is a training mechanism, it would be interesting to see how this method performs with training on NQ, TriviaQA, etc. Currently, the authors perform trainning only on MS MACRO and perform testing on NQ, triviaQA, etc. DPR does this.\n-\tVery bad writing.\n-\tVery handwavy at some places\n-\tSee Questions for detailed list of weakness.\n\nQUESTIONS:\n-\tFor the choice of m_i, Eq. 4 would have collapsed to the form of Eq.1 with extra Ai. Why show it in this form and complicate? Isn't Max over j the same as Aij(qi.Tdj)??\n-\tLine 120. How is $f_{colbert}(Q, D^-)=0.2$ when $q_i^Td^-_j>0.8$. Failure case very handwavy. Please explain with proper numbers and reasoning. Is Fig2 different tokens for any +ve/-ve documents or just the retrieved tokens?\n-\tLines 134-137. Very handwavy again. Authors should use proper math format to denote in-batch tokens. It is very difficult to understand when $A_{ij}=1$.\n-\tLines 134-137. Define top-k_train\n-\tFor the document tokens and query tokens already in RAM, why not compute the pairwise similarities computed if Aij=0? This would still retain some speed?\n-\tLines 226-241: what is token at rank k? are these top k retrived tokens?\n\nLIMITATIONS:\nThe authors discuss the limitations. Given that the work focusses on the training mechanism, it would have been nice to see trainings with other datasets.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper deals with the problem statement of document retrieval. First, the paper contrasts and explains the differences between single vector and multi-vector retrieval models. As multi-vectors retrieval models perform better due to their accessibility to more tokens, it involves significant inference costs. The authors propose a new model Contextualized Token Retriever XTR which aims to overcome the disadvantages during inference and bring closer to single vector models. Experimental results on BEIR and LoTTE benchmarks show that XTR achieves State-of-the-art results.\n\n\nPost rebuttal:\nI have read the author's response and the rebuttal sufficiently addressed my concerns.\n\nSTRENGTHS:\nOriginality:\nThe paper brought a lot of theoretical analysis showing why and where multi-vector retrieval models fail. This analysis provided a strong justification to their propose approach. In addition, the authors were able to show the effectiveness of XTR on multiple benchmarks. \n\nSignificance:\nGiven the gap that exists between single and multi-vector retrieval models in terms of performance vs computational efficiency, this paper bridges the gap by not only making the retrieval efficient but also improving the performance. Hence, I feel this paper is quite significant.\n\nWEAKNESSES:\nClarity: I feel this paper needs little fine-tuning in clarity. In the introduction, I couldn't understand the problem statement i.e document retrieval they are dealing with but there was a lot of motivation. I believe adding few sentences about the problem statement makes introduction better. Similarly in the experiments section, there could be separate sub-sections for the datasets, model comparisons and experimental settings. It will makes things more coherent.\n\nQUESTIONS:\nNone\n\nLIMITATIONS:\nThere is sufficient discussion about the limitations after the conclusion section.", "TITLE": "Official Review", "RECOMMENDATION": 8, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-07"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe XTR model extends the ColBERT neural IR model by removing one of the efficiency issues: candidate documents (selected using a dense vector index, e.g. FAISS) have to be re-scored by loading as much vectors as there are tokens in the document.\n\nThe proposed approach simply reduces the number of vectors representing a document. The candidate score is then (supposedly) equal to the final score (since all the vectors are in the FAISS index). The authors show that in practice (on a variety of datasets, include LoTTE, MIRACL and BEIR) performs roughly similarly to ColBERT\n\nSTRENGTHS:\nThe paper presents a simple method that improves over ColBERT by reducing the number of vectors per document. The simple strategy used in the paper is a good alternative to harder to implement pipelines like e.g. PLAID or ColBERTer.  The model performs a little bit worse than ColBERT but improves (in theory) the latency, although the latter is not measured in the paper.\n\nThe experimental work is well conducted on a variety of dataset, showing the robustness of the proposed approach.\n\nWEAKNESSES:\n- No experiment measuring the observed latency are reported. While the estimated FLOPs/query is important, seing actual difference on the same hardware would strengthen the message\n\n- No comparison with PLAID which implements a lot of strategies to improve the efficiency of ColBERT (which is another way to deal with the problem of the number of vectors in documents). More importantly, another very related work, ColBERTer (CIKM 2022) which also reduces the number of vectors per document using a sparsity-inducing loss, is missing.\n\n- Figure 4 should report the recall of gold tokens (rather than the precision)\n\nQUESTIONS:\n- Adding a sparsity loss would have been an alternative to the learning scheme: why was this not used since it might give more stability to the process?\n\n- What is the real latency time of such a model when the number of top-K vectors increases\n\n- is $k_{train}$ the same thing as $m$ in eq. 2?\n\nLIMITATIONS:\nThe limitations section is a bit too generic (not really related to the proposed approach). See weaknesses for possible issues to report at this level.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 5, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-07"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis work proposes XTR, ConteXtualized Token Retriever, which is a method for multi-vector retrieval with a simple and effective objective function. Comparing to prior work on multi-vector retrieval method ColBERT where all the tokens from query and candidate document needs to be computed in order to calculate the final query/document score, XTR only uses retrieved tokens in the document to calculate the score, which greatly reduces inference time. It also provided analysis on the situations when ColBERT training objective may fail and proposed new objective that can mitigate the issue. \nIn experimental session, extensive experiments haven been done to show the effectiveness of the proposed method. XTR achieves competitive performances with current SOTA multi-vector retrieval models, while being much efficient at inference time. It also achieved SOTA in zero-shot retrieval benchmark BEIR and strong results in multi-lingual retrieval benchmarks.\nFinally, analysis was presented in section 5 to shed lights on why XTR presents a better token retrieval.\n\nSTRENGTHS:\n- It proposes a method for multi-vector retrieval which achieved strong performances on multiple benchmarks and can greatly reduce inference time comparing to prior work on multi-vector retrieval\n- The proposed method achieved SOTA results on zero-shot retrieval benchmarks and multilingual retrieval benchmarks\n- It provides insightful analysis on where previous multi-vector retrieval method (ColBERT) failed and proposed new training schema that solve the problem. It also provides examples that validate the proposed methods in qualitative analysis.\n- It is well written overall with thoughtful analysis and extensive experimentation results\n\nWEAKNESSES:\n- The XTR method is mainly built on T5 models, without exploration on other architectures. \n- The analysis on equation (4) can be more elaborate (see Q1).\n\nQUESTIONS:\nQ1: Could you provide more insights on why the upperbound of m_i in equation (4) works? Have you tried other bounds/estimations?\n\nQ2: In table 1, how does the estimated FLOPs/query looks like for single vector retrieval methods such as DPR?\n\nLIMITATIONS:\n- The authors listed that XTR is trained on MSMARCO dataset which may have license issue. \n- The proposed method only trained on T5 models, and no other type of architecture has been explored.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-10"}