Rethinking the Role of Token Retrieval in Multi-Vector Retrieval
Authors: Jinhyuk Lee
Zhuyun Dai
Sai Meher Karthik Duddu
Tao Lei
Iftekhar Naim
Ming-Wei Chang
Vincent Y Zhao
Venue: NeurIPS 2023 poster (2023)
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
Jinhyuk Lee"
Zhuyun Dai
Sai Meher Karthik Duddu
Tao Lei
Iftekhar Naim
Ming-Wei Chang
Vincent Y. Zhao
Google DeepMind
Abstract
Multi-vector retrieval models such CoIBERT (Khattab and Zaharia 2020 allow token-level interactions between queries and documents, and hence achieve many information retrieval benchmarks. However, their non-
state of the 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 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 veetor retrieval by rethinking the role of token retrieval. We present XTR, ConteXtualized Token Retriever, which introduces 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 the document, and enables newly designed scoring stage that is two-to-three orders of magnitude cheaper than that COlBERT. On the popular BEIR benchmark, XTR advances the state-of-tee ofrtth-artttt 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.
Introduction
The performance of dense retrieval model is largely affected by how it defines expressive repre- sentations over queries and documents, and whether it can efficiently retrieve and score a document using these vector representations. For example, dual encoder models (Yih al.] 2011 Lee al.) 2019 Karpukhin et al., 22200 Ni et al. [021] 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 KKattab and Zaharia 2020, Santhanam et al. 2022bb are directly designed to capture token-level interactions. By utilizing a (non-lineat 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 great cost of inference complexity. Unlike the case in dual encoders, the non-linear scoring function in multi-vector retrieval models prohibits the 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 result, models such CoIBERT adopt intricate and resource-intensiv inference pipeline, which typically consists
'Correspondence: jinhyuklee@goggle.com
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

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-lineat 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-inferene 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, ContXextualized Token Retriever: simplified and efficient method for multi- vector 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 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-vecton 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, XTT proposes a novel, yet simple, training objective, which dramatically improves retrieval accuracy, doubling the chances of a gold token being retrieved in the top-k 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 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.] [202]) and LoTTE [Santhanam et al. 2022b] benchmarks. XTR attains state-of-the-att performance, requiring neither distillation nor hard negatiave mining. Notably, our model surpasses state-of-the 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 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.
by 4.1 points on top-20 retrieval accuracy. XTR also
Background
2.1 Multi-vector Retrieval
Single-vector retrieval models, also known as dual encoders, encode input text sequence single dense embedding and define the similarity of query and a document based on the dot product (Lee) al. 2019 Karpukhin al. 22200. Multi-vector retrieval models, on the other hand, make use of muluple 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 = {q) and a document D = {d,}"- where q; and d, denote the d- dimensional query token vector and the document token vector, respectively. Multi-vector retrieval models compute the query- query-doeument similarity as follows: S(Q, D) = Σ ΣS-, A, where
P.j = q. d, and A E (0, 1)" nmm denotes the aligmment matrix with A.j being the token-level alignment between the query token vector q, and the document token vector d,. The sum-of-max operator of ColBERT [Khattab and Zaharia 2020 sets A, [(--argmax where the argmax
operator is over 1 s j' s m (i.c., tokens from single document D) and [[x] is an indicator function. Then, SCOBERT ( Q, D) defined as follows:
n
m A.
SCOBERT (Q, n
T max q dy. 1=1 15jsm
(1)
i=1 j=1
2

---

(b) Gatbering
ColBERT inforence
Doc
(c) Scoring (training)
Retrieval
new
Score
Query Q
Doc
new
Score,
XTR training inference
Doc
Scoring
F1.OPsquerry (Scoring)
()) Token Retrieval
Score,
4,000 Cheaper
Query Q
Recalla
CoIBERT
64.5
Scure,
XTR
68LD
Figure 1: Overview of XTR. ColBERT has the stree-stage inference combining (a) the token retrieval, (b) the gathering and (c) the scoring stages ($2.2). XTR leverages the token retrieval for both training and inference. XTR efficiently obtains the score of each candidate document by applying /XtR (or fXTR') the retrieved tokens, completely removing the gathering stage ($..2
Here, we include the normalizer nı, which was not included the original sum-of-max, as it stabilizes training while not affecting the ranking during inference. After computing the query dery-ddo similarity, multi-vector retrieval models are typically trained with a cross-entropy loss over in in-batch negatives [Santhanam et al. 2022b Qian et al. 2022). Specifically, given positive document D* for Q and set of mini-batch documents D1:B [....... DB] where D* Di:B, they minimize the cross-entropy loss defined as: LCu log Σ. exp F(Q.D.)
exp f(Q.D")
2.2 Three-stage inference of Multi-vector Retrieval
Unlike dual encoder models, finding the maximum scoring document- the document that maximizes eq. (I) -tannot be directly handled by MIPS the scoring function uses non-linear, sum-of- of--oo---m- 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 of retrieved tokens. The total number of candidate documents is up nk' if each token coming from a unique document 2) Gathering: since 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. (T). This stage also called refinement. Note that the training of typical multi-vector models only takes of the scoring stage with mini-batch documents. Finally, top-k documents returned based on the computed scores. The three-stage inference is illustrated in the top of Figure]
taking the union of source documents
XTR: Contextualized Token Retriever
Unlike existing multi-vector models that follow the Setrieeegaather-score stages, XTR directly scores documents utilizing the tokens retrieved from the token retrieval stage. 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.
'In fact, each candidate document of a based ColBERT is retrieved by 1.48 tokens per average. meaning that the most of the candidate documents unique.

---

Given a positive document D* and a set of negative documents D Q. the first-stage token retrieval needs to retrieve the tokens of D 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.
D,] for query
0.8 where all the individual max token similarity (i.e., q. dj where Asj 1) is 0.8. On the other hand, assume fCOBERT (T) = 0.2 for all D E Disr where each D has highly peaked token similarity greater than 0.8 but others close to zero (i.e., there exists q, d 0.8 where
Failure case Assume that fCoBERT fc.mERT(Q.
MS MARCO
0.06
TS-CWBERT toun))
0.05
Dender
0.04
A, while other T d 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 However, for each of query tokens, if there exists at least one negative document token that has a high token simi- larity greater than 0.8, the token retrieval with top-k' = would fail to retrieve any tokens of D As result, multi- vector retrieval model with the sum-of-max operator will not be able to lower the high scores of some negative to- kens. 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.
0.5
0.6
5.7
0.8
0.9
3.0
Toien Retrieval Scorc
Figure 2: Density histogram of 4,000 token retrieval scores (cosine similarity). Training with SCOBERT (T5-CoIBERT; $4) causes many document tokens to have extremely high scores regardless of their actual relevance with respect to the input query tokens. XTR mitigates this problem with better training objective.
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 different alignment strategy Å. Specifically, we set the alignment Ai [(ew)p-) where the top-k
operator is applied over 1 s mB (i.e., tokens from B mini-batch documents) returning the indices of k largest values. During training, we use a hyperparameter kimin for the top-k operator. Then, simply modify eq. ( follows:
1 Σ
fXTR EXTR(Q,D) D)
(Q,
max A.j dy.
7
1sysm
The intuition is that consider the token similarities within D only when they are high enough to be retrieved within top-kinin from mini-batch. Here, we normalizer Z [{i] S.t.A 0}|, which is essentially the number of query tokens that retrieved at least one document token of D all À. O, we clip Z to a small number and fXTR (Q. D) becomes 0. As a result, our model cannot assign a high token similarity negative documents as it blocks tokens of positive documents be retrieved. With the previous failure case where fCaBERT assigned high score D" even though it (*) 0 during training
cannot be retrieved, our similarity function incurs a high loss (since tokens of D* were not retrieved). For training, we use the same cross entropy loss defined $2. 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 DEc from the token retrieval stage:
Dio0= = (Dld E E ED top-k (q-))
(3)
Indeed, our derivative analysis in Appendix A shows that the token-level similarity would not change if the document-level scores are already well discriminated.
We tried different normalizers such and found that Z works the best while stabilizing the training.

---

D
D3
D
D3
CoIBERT
XTR
Figure 3: Comparison of SCOBBER in eq. ① and fXIR' in eq. (4) Assume that Da and Do were selected as initial candidate documents from the token retrieval stage. SCOBERT loads all token vectors of Da and Do and exhaustively recomputes pairwise token similarity to obtain the max values (red boxes). On the other hand, fXtr' does not load any token vectors and reuses retrieval scores from the first-stage token retrieval. Assume that, with the top-2 token retrieval results, the first query token retrieved each max score of Da and D,, but the second query token retrieved two tokens only from D, but not Dr. We impute the missing similarity m for D, (denoted as yellow dashed box) by finding its upper bound using the top-2 score (denoted 82) of the second query token (i.e., $2 81).
Scoring
Estimated FLOPs/query
Setting
SCOBERT JXIR!
3 (2md m + n kF +
0.36 10
16, = 128, k' 100.
10
0.09 10
55, 2.5
Table 1: FLOPs comparison of ColBERT and XTR for the scoring stage. XTR only adds minimal complexity for scoring each candidate document. The setting is derived from MS MARCO.
Here, top-k (qx) is union of top-k' document tokens (from the entire corpus) based the inner product scores with each query vector (i.e., q d). Given the query token vectors, there are C (s n.k') 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(n2 (inkkmmd mm) computation per query (m average document length). Instead, we propose to score the documents solely using the reirieved 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. (I) can be removed completely. Unlike previous work [Macdonald and Tonellotto| [022]] 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 document tokens for each of query tokens. Assume that each document token belongs to unique document, providing C' = nu candidate documents in total. This leaves us with single token similarity to score each document in the absence of the gathering stage. However, during training either with eq. (I) 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 were positive with n token similarities.
For every candidate document D, we first define the following scoring function for the inference:
(XTR(Q,,))
max 1=1 Isjsm
[
T d (1 A..)m..
(4)
This is similar to eq. 2 but introduces mn E R. which estimates the missing similarity for each gi- Å is defined similar to the one described in ( except that it uses k' for the top-k operator. Each would take the missing similarity m, the maximum value if A, = 0 and m, 2 0. Importantly, fxtR' removes the need of recomputing any q, d, since when A. we already know the retrieval score from the token retrieval stage, and when A, Å.
we simply don't need to compute as T d, = 0. Note that when every A.j = 1, the equation becomes the sum-of-max operator. On the other hand, when no document tokens î were retrieved for qi (i.e., A = 0), we fall back to the imputed score my, which provides approximated sum-of-max result.
5

---

MS
AR
TO
CF
SF
CV
NQ
HQ
FQ
SD
DB
QU
Avg.
One Retriever per Domain
GenQ
40.8
49.3
18.2
66.9
17.5
64.4
61.9
31.9
35.8
53.4
30.8
14.3
32.8
83.0
43.1
58.8
25.6
76.2
23.5
63.8
70.2
33.7
45.6
61.7
43.0
18.3
34.4
87.5
49.4
One Retriever for All
BM25
22.8
31.5
36.7
75.3
21.3
66.5
65.6
32.5
32.9
60.3
23.6
15.8
31.3
78.9
44,0
ColBERT
40.1
23.3
20.2
77.1
18.4
67.1
67.7
30.5
52.4
59.3
31.7
14.5
39.2
85.4
45.1
GTR
42.0
51.1
21.3
66.0
24.1
60.0
53.9
30.8
49.5
53.5
34,9
14.9
39.2
88.
45.2
IS-ColBERT
45.6
28.8
31.1
72.4
18.1
70.4
68.3
34.0
52.2
61.7
33.4
14.1
41.6
82.3
46.8
XTR
45.0
40.7
20.7
71.0
34.0
53.0
64.7
14.5
40.9
86.1
Splade.3
43.3
47.9
27.2
78.6
23.5
69.3
71.0
33.4
52.1
68.4
33.6
15.8
43,5
83.8
49.9
CoIBERT,
46.3
26.3
78.5
17.6
69.3
73.8
33.8
56.2
66.7
35.6
15.4
44.6
85.2
49.9
GTR.
44.2
54.0
23.3
74.0
26.7
66.2
50.1
34.2
56.8
59.9
46.7
16.1
40.8
89.2
49.1
T5-ColBERTs>
47.3
33.8
31.0
74.2
19.7
73.1
75.8
35.2
60.5
65.2
43.5
17.1
45.0
86.0
50.8
XTR
46.6
44.2
30.9
77.0
24.5
74.3
78.9
35,3
60.9
66.2
43.8
17.1
44.3
88.1
52.7
LOTTE Search
LOTTE Forum
Writing
Rec.
Sci.
Tech.
Life.
Pooled
Writing
Rec.
Sci.
Tech.
Life.
Pooled
BM25
60.3
56.5
32.7
41.8
63.8
48.3
64.0
55.4
37.1
39.4
60.6
47.2
ColBERT
74.7
68.5
53.6
61.9
80.2
67.3
71.0
65.6
41.8
48.5
73.0
58.2
GTR
74.1
65.7
49.8
58.1
82.0
65.0
69.2
62.0
33.7
47.6
72.2
54.9
XTR
77.0
69.4
63.2
82.1
68.7
74.4
60.1
Splade,2
*+
77.1
69.0
55.4
62.4
82.3
68.9
73.0
67.1
43.7
50.8
74.0
60.1
ColBERT
72.3
56.7
66.1
70.8
76.9
GTR,a 83,9
78.0
60.0
69.5
87.4
76.0
79.5
73.5
43.1
62.6
81.9
66.9
XTR,a
83.3
79.3
60.8
73.7
89.1
77.3
83.4
78.4
51.8
64.5
83,9
71.2
cross encoder di distillation
model based hard negatives
Table 2: (top) nDCG@ 10 on MS MARCO (in-domain) and BEIR (zero-shot). The last column shows the average over 13 BEIR datasets. (bottom) Top-5 retrieval accuracy LoTTE datasets (zero-shot).
In fact, can find the upper bound of the missing similarity. For every token retrieval with q. the missing similarity of the query token for D will be upper bounded by its last top-k' score. Specifically, for each query token q,, we have the following top-k' token similarity during inference: [q. d..).... q. dow)]. Here, each d(.) could come from a different document. Since the missing similarity would have score less than equal to the score of the last retrieved token, we know that m, q, d(z). With a larger K., the upper bound becomes tighter. In our experiments, we show that simply choosing m, = q. d(w) works well especially when model is trained with fx1R 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 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 4000x 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 QQu al. 202T]. Then, test XTR MS MARCO (MS; in-domain) and zero-shot IR datasets. For the zero-shot evaluation, we use 13 datasets from BEIR [Thakur al., 2021] (see Appendix Cfor acronyms), 12 datasets from LoTTE (Santhanam et al. 2022b]. and 4 datasets on open-domain QA passage retrieval (EQ: EntityQuestion [Sciavolino 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 Co-CoT [Qian al. 2022) and XTR shows the improvement with our methods on multi-vecton vector retrieval model. For implementation details and baselines, see Appendix C
For the relation relationship between hyperparameters (e.g., kuain and k see $5.3
'e found that directly training with fXTR' instead of fXTR fails to converge, which we leave future work.

---

EQ
NQ
TQA
SQD
Top-20
Top-100
Top-20
Top-100
Top-20
Top-100
Top-20
Top-100
BM25
71.4
80.0
62.9
78.3
76.4
83.2
DPR BM25
73.3
82.6
82.6
88.6
82.6
86.5
75.1
84.4
MARCO
75.3
81.9
78.0
GTR
73.3
80.6
78.5
86.5
76.2
83.4
65.9
77.6
GTR,di
75.3
82.5
83.5
89.8
81.7
86.6
70.4
80.6
DPR
56.7
70.0
79.5
86.1
78.9
84.8
52.0
67.7
COlBERT
79.1
80.3
76.5
XTR
79.0
85.2
79.3
88.1
80.3
85.5
78.2
85.9
XTR
79.4
spense component
retrieval pro-training
Table 3: Zero-shot passage retrieval accuracy on open- open-domain question answering datasets. In-domain performances are underlined and all the other performances are based on the zero-shot evaluation. For EntityQQestion we report macro-averaged performances over different relations.
j6
2h
yo
Avg.
HMZ5
43.1
50.8
35.1
31.9
33.3
25.1
18.3
45.8
44.9
36.9
41.9
33.4
38.3
19.4
48.4
15.0
mDPR
49.9
44.3
39.4
47.8
47.2
43.5
38.3
27.2
43.9
41.9
23.9
35.6
35.8
51.2
RM25
52.6
English
mCnntieser (an)
40.6
36.8
27.4
37.9
WXTR (ei)
66.1
64.7
40.4
40.5
47.9
62.2
37.5
51.4
46.9
568
64.11
49.
43.0
67.7
69.2
47.2
34.5
40.6
52.2
mXTR, (en)
74.1
75.5
$6.0
52.4
56.1
75.1
51.4
61.8
$2.0
68. 7
67.4
61.3
69.7
76.0
76.9
56.9
51.7
60.3
63.5
Trained Engiosh MS MARCO MIRACI, longuages)
mConiriever
61.6
66.4
41.2
00.3
46.3
61.9
42.9
41.9
44.6
55.6
55.4
48.1
65.5
77.6
69.3
45.9
39.6
41.9
52.7
IXTR....
73.0
73.9
76.1
58.2
mXTR.,
77.8
78.4
80.7
Table 4: nDCG@ 10 on 18 multilingual retrieval tasks from MIRACL. Each row shows the perfor- mance of a single multilingual retrieval model. The last two surprise languages (de and yo) are not included in the training dataset of MIRACL. The last column shows the average over 18 languages.
4.1 In-domain Document Retrieval
MS MARCO The first column of Table 2 (top) shows nDCG@ 10 on MS MARCO (see Table D. 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- seaerring-scoringgg 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 (top: except the first columns) shows nDCG@ 10 on BEIR (see Table D. for recall 100). XTR achieves the new state-of-the-art performances significantly outperforning 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 bottom) also show that XTR better than ColBERT and competitive with distillation- based models while XTRxx 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. [202]] or contrastive pre-training [Ram et al.) 2022 Sachan et al. 2022a b) to achieve better performances on EntityQues tions, XTR simply fine-tuned on MS MARCO achieves the state-of-the-an performance.
4.3 Multilingual Document Retrieval
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 m'55 [Xue et al.) 2021) and test it on multilingual retrieval tasks in 18 languages. Table shows that mXTR greatly outperforms mContriever that uses expensive contrastive pre- training. well as the hybrid model, BM25 + mDPR.
MIRACL

---

H45 MAVCC
CRED
ArguAne
PiGele Salttbeeeeee 1g
5.00
XT
OiRERTTesl
0.0
T5- B--C-BER-OO
XTR
3.50
XT Rull
2.06
0.04
5.04
02
0.02
203
400
603
$50
203
400
620
D03
1000
200
600
930
1000
Tiken Rank K
Tiken Rank K
Token Ranc
MARCO
TREC-COVID
ArguAna
0.8
XTB-oo
0.8
XTFloll
TS-Ca RCatio
0.6
9.5
0.6
02 XTBaod
0.2
0.0
0.0
Toben Bankk
Tcler Bank
Tollen Rank
Figure 4: (top) Gold token retrieval performances of T5-CoIBERT and XTR. We plot the probability of each retrieved document token at rank & coming from the gold document. (bottom) Lexical token retrieval performances of T5-ColBERT and XTR. We plot the probability of each retrieved document token rank k being lexically identical its query token.
Model
Imputation
MRR@) 10
R@ 1000
0.30
T5- 55-CoIBERT base
None
0.0
0.0
top-k' score
27.7
91.8
Recalig100
0.80
XTRease
None
22.6
88.7
m; = 0
36.2
97.3
0.70
T5 Biffbasse Co CRIEETT) TF-aa
m; 0.2
36.4
97.3
XTRaaee
top-k' score
37.4
98.0
307
10'
10
16
Table 5: Impact of training objectives and imputation methods comparing T5-ColBERT and XTR. For both models, we apply fXTR' during inference. We report MRR @ 10 and Recall 1000 on the MS MARCO development set.
Figure 5: Recall@ 100 of XTR and T5-CoIBERT with different k.. For T5-ColBERT, we use ci- ther fXTR' or fcanERT.
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 retrieve the gold documents. Hence, better token retrieval would contain these gold tokens more often in their top results. In Figure 4 (top), we show the probability of token at the rank l coming from the gold documents of query. To compute the probability for the rank we simply count the number of an event where token rank & belongs to the gold document and divide by the number of tokens at rank k. While this measuring the precision of the token retrieval, we observed similar trend for the recall of gold tokens. Compared to T5-CoIBERT, XTR retrieves gold tokens with higher probability, even on MS MARCO. This shows that the training objective XTR encourages it to retrieve tokens from more relevant context.
Lexical token retrieval In Figure 4 (bottom), we show the probability of 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 the lexical matching while preserving a good amount of lexical precision so that it would achieve a high retrieval accuracy on the entity-centie dataset ($4.2). In fact, Table 6 in Appendix shows that having lower lexical matching doesn't 't necessarily mean lower retrieval quality, but often means better contextualization.
5.2 Efficient Scoring
In Ta we show how we can employ the efficient scoring function fXTR in XTR with minimal performance losses. We apply fXTR' on both T5-ColBERT and XTR, and show their performances on MS MARCO. With T5-ColBERT, even if we use the top-k' score for the imputation, the performance
8

---

NS MAACO
0.40
NS MARCO
Anguline
XCH tbsich stze 1221
XTB bace tharr site 2501
W000070
1000810
9.45
0.42
XTB bace thaice site 3201
6.40
Amhass 203)
278 a7aase man 2001
278 124)
2.44
buse (butch buse bbutch
hone (harth size
0.38
15 Bettbba 201r
0.83
50
250
200
500
1.0
150
250
k Krart
K.twin
Figure 6: MRR@10 of XTR with different kuain and k.. For T5-ColBERT, we also use /XTR' with the top-k' score imputa- tion method for the inference.
Figure 7: Effect of training XTR with different batch sizes and Amiin- For each point of the graph, we train XTReme with the specified training batch size (128. 256, 320) and Kirain (32, 64, 128, 256) and evaluate on each dataset (MS MARCO and ArguAna). nDCG@ 10 of each model is reported.
is much worse than the original cum-of-max scoring. With XTR, the performance greatly improves as it has better token retrieval. Figure shows how Recall @ 100 improves with larger k's as it provides more exact upper bound for the missing similarity imputation. Table shows that even if we use smaller k,, XTR still maintains high performances on BEIR.
5.3 Relationship between Hyperparameters
Kirain vs. k' In Figure 6 we show MRR @ 10 of XTR trained with different kyain and evaluated with different on the MS MARCO development set. While all variants of XTR prefer larger k', ones trained with smaller kinain show higher performances than others under small k settings. XTR with larger Aimii exhibits better performances than ones with smaller kirain as k' becomes larger.
Training batch size vs. krain In Figure 7 we show the relationship between the training batch size and hivain during training XTR. In this experiment, we use k' 40, 000. While it evident that XTR mostly favors large training batch sizes. the optimal top-kurii can be different for different datasets. While most datasets including MS MARCO favored large enough kimans ArguAna prefers smaller kmain 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 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 the query. This example explains the lower lexical token retrieval probability of XTR compared to T5-ColBERT in (bottoo) but higher gold token retrieval performance in Figure (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 single vector constrains the representational power the models. Polyencoder [Humeau et al. 2020). MEBERT (Luan et al.) 202T], and MVR (Zhang et al. [202aa] propose 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 Zaharial 2020] is probably the most renowned model in this family. ALIGNER (i.e. T5-
ColBERT) 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-on-domain generalization ability.
Efforts for reducing serving costs 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 al.) 2022 relaxes COIL with a lexical routing mechanism where query token vector only retrieves from subset of

---

T5-ColBERT token retrieval for "what is the usual pay for stock associales michael?"
Rank
Token
Context of Token
Relevance
usual
routine passport services: the usual waiting time in logan to get your passport is four (4) cight (8) weeks for routine applications.
No
usual
the usual pay days the 1st and 16th of each month. for annual educational paraprofessionas there is no payroll lag.
No
5
usual
the usual part xiii tax rate is 25% (unless a tax treaty between canada and your home country reduces the rate).
No
50
usual
this where one challenge the judgment debtor's claim, option creditors have to try and make deal with the debtor take less than 25% (the usual amount of wage levy).
No
100
usual
the usual maximum inventory is talisman, 26 elemental runes, and 26 pure essence. the ingredients must he brought to opposing altar.. from the runes being crafted.
No
XTR token retrieval for "what is the usual pay for stock associates at michael?"
Rank
Token
Context of Token
Relevance
usual
store manager. salary: the usual salary a store manager receives can be anywhere around $52,000 to $1 15,000 annually.
No
2
usual
salary: the usual salary a store manager receives can be anywhere around $52,000 to $115,000 annually. 2 bonuses: publix provide bonuses that could reach up to $40,000.
No
5
average
average salaries for michacls stores stock associate: $9. michaels stores hourly pay trends based on salaries posted anonymously by michaels stores employees.
Yes
50
think the avg starting pay is closer 30k for mgr trainees. it is hourly position until you are fully trained (40 hours per week). average macys salaries. the average salary for macys johs is $32,000. av- erage macys salaries can vary greatly due to company, location, industry, experience and benefits.
No
100
average
No
Table 6: Token retrieval example from MS MARCO. Among the top 100 retrieved tokens, 100% of T5-ColBERT tokens lexically identical the query token usual while only 8% of XTR tokens are lexically identical. XTR retrieves the relevant passage by retrieving average for usual.
document token vectors routed the same key. PLAID [Santhanam 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 improve the efficiency of ColBERT. On the other hand, how to accelerate the scoring stage remains under-cxplored. To the best of our knowledge, XTR the first work to simplify the scoring stage and remove the gathering stage in multi-vector retrieval.
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 multi- vector 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 result, XTR achieves state-of-the-arr 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-speciff issue. We believe that L.L.Mbased retrieval dataset generation Dai al. 2022 would be able to mitigate the problem in the future.
10

---

Acknowledgements
We would like to thank the anonymous reviewers for their helpful feedback. We also thank Nicholas Monath, Raphael Hoffmann, Kelvin Guu, Slav Petrov, and others at Google DeepMind for their helpful comments and discussion.
References
Negar Arabzadeh, Alexandra Vtyurina, Xinyi Yan, and Charles LA Clarke. Shallow pooling for sparse labels. Information Retrieval Journal, 25(4))365-338) 2022.
Xilun Chen, Kushal Lakhotia, Barlas Oğuz, Anchit Gupta, Patrick Lewis, Stan Peshterliev, Yashar Mehdad, Sonal Gupta, and Wen-tau Yih. Salient phrase aware dense retrieval: Can dense retriever imitate sparse one? arXiv preprint
2021.
Zhuyun Dai, Vincent 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. 1755, 2022.
Thibault Formal, Carlos Lassance, Benjamin Piwowarski, and Stéphane Clinchant. Splade v2: Sparse lexical and expansion model for information retrieval. arXiv preprint arXiv:2 109.06609 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-Tii 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 Linguisties: Human Language Technologies, NAACL-HLT 2021, Online, June 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 Machine Learning, pages 3887-3896. PMLR, 2020.
Samuel Humeau, Kurt Shuster, Marie-Anne Lachaux, and Jason Weston. Poly-encoders: Architec- tures 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, Sehastian Riedel, Piotr Bojanowski, Armand Joulin, and Edouard Grave. Unsupervised dense information retrieval with contrastive learning. Transactions on Machine Learning Research, 2022.
Vladimir Karpukhin, Barlas Oguz, 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 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 NNXi:2211.104111 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 1021.

---

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-Wo Wei Chang, and Yinfei Yang. Large dual encoders are generalizable retrievers. In Conference on Empirical Methods 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.01260 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 1835-584,400 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 Learming Research, 21 (1)5:55855-551, 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 Proceed- ings 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 Zeutlemoyrr. Improving passage retrieval with zero-shot question generation. arXiv preprint arXiv:2204.0744 077496, 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 31 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- questions challenge dense retrievers. In Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing, pages 6138-6144 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 2015.
Computer Vision, pages
Anshumali Shrivastava and Ping Li. Asymmetric Ish (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.
12

---

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-23660 2022.
Linting Xue, Noah Constant, Adam Roberts, Mihir Kale, Rami AI-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. 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 Anmual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers). ACL 2022, Dublin, Ireland, May 22-27, 2022, pages 5990 -0990-6000. Association for Computational Linguistics, 2022a.
Xinyu Zhang, Nandan Thakur, Odunayo Ogundepo, Ehsan Kamalloo, David Alfonso-Hermelee Xiaoguang Li, Qun Liu, Mehdi Rezagholizadeh, and Jimmy Lin. Making miracl: Multilingual information retrieval across continuum of languages. arXiv preprint arXiv:2210 09984, 2022b.
13

---

A Derivatives w.r.t. Similarity Scores
Sum-of-max Here, we use cross-entropy loss LCE with the sum-of-max operator SCOBERT and analyze the derivatives with respect to the token similarity scores.
exp ((Q.D")
B
CCB log
-COBERT (Q, D*) + log
exp (Q, (5)
B Σ exp f(Q, D,)
SCOBERT (Q,
A P
P E
n
1=1
Here, we denote j as the index of the row-wise maximum value, dependent on each (i.e., Aij 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 for positive document D* E D1:B:
S(Q,D*) OCEE 0P
B
exp f()
OP
E b=1 exp f(Q,D)) O
1/
i
P
exp f(Q, D,)
ap+
77
1=1
1sjsm
2 b=1
B
exp f(Q, D
ap* b-
ij
Σ
B
xpp/(Q/ D₂) as(Q, D.)
exp
2b=1 exp S(Q, Db) b=1
B
aP"
exp f(Q,
1
- [1- 10, D..n)].
7
h- exp f(Q D.)
Similarly, the gradient w.r.t. a maximum token similarity P for a negative document D D1:B computed as follows:
Ocrr 0P-
f(Q.D")
1
B exp f(Q, D.)
E exp f(Q, D₂)
OP
1j
ij
exp DT)
- IQ, D1B).
77
B Σ exp f(Q, D.)
Hence, the positive token-level score P will gradually increase until P(D* IQ, Dı:n) + and the negative token-level score P will decrease until P(D D1:B):))))) This shows that the token- level scores are trained based on the document-level scores, which might stagnate the token-level hith-later
scores. For instance, even P ij is very high later causing d to be retrieved instead of ones from
positive documents-it will not be penalized as long as P(D IQ, D1:B) is low enough.
In-batch token retrieval Compared to the sum-of-max operator, our in-batch sum-of-max fxIR considers the max values only when they are retrieved over other negative tokens in the mini-batch.
SxTRQ,,D:1 = 7 E A A, = Z
E P
[-1
Here, we denote the index of the row-wise maximum value that is also within the mini-batch top-kuain given 4i (i.e., satisfies both A, and A, = 1). If there is such j. we simply use Py O. We also use a normalizer Z, which is the number of non-zero Pij. In this analysis, we assume Z > 0 since if every P& is zero, the gradient is undefined.
14

---

The gradient w.r.t. the maximum token similarity P (non-zero) for positive document D* E Du:B computed as follows:
OCCE
1(Q,D")
B
exp f(Q, D₂)
aP 0
E exp f(Q, D₂)
OP
= z- [1 exp f(Q, D")
Σ exp f(Q, D.)
-- [1-1 P(D")Q,
This is very similar result compared the sum-of-max operator except that 1) the gradient is defined only when P is non-zero (i.e. retrieved) and 2) it is dependent on Z", which means that the gradient will be large whenever there is small number of retrieved tokens from the positive document. If only a handful of tokens are retrieved for D*, our objective function increases P...
For negative similarity score P.j. we have the following:
UcC
S(Q.D")
B exp f(Q, D.)
OP; 1j
OP.
I exp f(Q, Da)
O 15
exp S(Q,
-
7.
Σ 1=1 exp S(Q, D&)
7- P(D D1.B)
Again, it is similar to the sum-of-max result, but it depends Z_ In this case, even when P(D QQ, D1:B) is low, if there is small number of retrieved tokens from D (iee,, small Z ) P. will be decreased significantly. Note that when ZT is large, zt naturally becomes smaller as they compete for in-batch token retrieval, which causes positive tokens have higher scores.
B Inference Complexity
We compare the complexity of CoIBERT 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 jet al. 2022a b. Guo et al.) 2020].
For the scoring stage, both COlBERT and XTR have O(nk') candidate documents. Here, we assume the worst case nk' where each document token comes from unique document. For each candidate document, ColBERT loads set of document vectors of md floating points (m = average document length) and computes eq. (T) with the query vectors of nd floating points. Computing eq. (L) per candidate document requires 2nmn F1.OPs s for token-level inner products, nma for finding the row- wise max, and n for the final average. In total, ColBERT requires n²2 K (2m2dd m + 1) FLOPs for the scoring stage. Note that this does not include the latency of loading the Onkkmmd) floating points onto the memory, which amounts up 450MB per query when 16, x 1000,
55, d 128.
On the other hand, XTR first imputes the missing similarity, which simply done by caching the t'tth token retrieval score for each query token. Then, each of nk'' candidate documents requires nf FLOPs for finding row-wise max and n for the average where is the average number of retrieved tokens per each candidate document. In total, we have n² k + 1) FLOPs. Table shows the estimated FLOPs of the two models. XTR reduces the FLOPs at the scoring stage by 4000x making multi-vector retrieval more efficient and practical.
15

---

C Implementation Details
XTR uses Kirain for retrieving in-batch document tokens. Since we retrieve over mini-batches, the size of mini-batch affects the performance for different kiiin which is shown in $5.3 In our experiments, we tried Kiraio {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' = 40, 000, which is possible due to the efficient scoring stage of XTR We analyze the effect of using different k's as well as its relationship to Aigain in 5.3 We initialize XTR from the base and xxl versions of the T5 encoder [Raffel al. 20200 and provide XTReaee and XT For multilingual XTR, we initialize XTR from mT5 [uue et al. 2021]. We fine-tune XTR for 50,000 iterations with the learning rate 1c-3. Up 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-COVIIE NF: NFCorpus. NQ: Natural Questions. HQ: HotpotQA. FQ: FiQA-2018. SD: SCIDOCS. DB: DBPedia. QU: Quora).
Baselines There are two main paradigms training retriever models for the out-of- eval- uation. 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 indepen- dent models for N different domains, For this emc-retriever-pe--domace 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-typicat
trained on large-scale IR dataset such as MS MARCO-an and directly applies it on the out-of-domain corpora and queries. For this one-retriever- for-all approaches, we present results of state-of sfate-of-the-art retrievers including Splade,2 [Formal et al.) 2021), Co (Santhanam et al. 2022b], and GTR,sl (Ni et al. 2021). We also show the results of T5-ColBERT [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-fort baselines, as well as XTR, are trained on English MS MARCO, unless otherwise stated.
'In fact. XTR with &' 40, 000 has still two-to-three orders of magnitude cheaper scoring stage than ColBERT with k'' 000 and T5-CoIBERT with k' 4, 000.
16

---

D Additional Results
In Table D.I
we show
Recall @ 100 BEIR.
MS
AR
TO
SF
CV
NF
NQ
HQ
PQ
SD
D3B
QU
Avg.
One Retriever per Domain
GenQ
97.8
45.1
45.6
PTRarirwr
98.9
47.5
94,1
53.1
91.8
55.9
30.6
89.8
74.6
76.5
41.6
46.3
99.6
69.2
One Retriever for All
BM25
65.8
94.2
53.8
93.1
43.6
90.8
49.8
25.0
76.0
74.0
53.9
35.6
39.8
97.3
63.6
64.5
89.8
97.4
44.3
92.3
52.2
87.2
41.1
27.5
89.3
67.6
67.0
34,0
41.8
99.6
64.7
TS ColBERT
91.8
76.0
49.9
90.4
46.2
91.3
55.4
27.6
90.5
78.3
63.0
34.2
50.5
97.9
65.5
XTR 91.0
92.1
50.8
90.5
57.3
28.0
80.7
63.5
34.8
52.0
68.0
GTR 91.6
98.3
46.6
94.7
55.6
90.0
40.7
30.0
94.6
75.2
78.0
36.6
49.4
99.7
68.4
TS-ColBERT
93.3
81.4
50.1
91.7
49.8
94.6
60.3
29.0
95.5
81.6
72.5
38.5
54.6
99.1
69.1
XTR.
Table D.I: Recall 100 MS-MARCO and BEIR. The last column shows the average over 13 BEIR benchmarks. Compared to GTR, T5-ColBERT only marginally improves the recall. On the other hand, XTR greatly improves the recall showing the importance of having better token retrieval.
In Table D.2
we show nDCG@ @ 10 and Recall @ 100 BEIR with different k'.
MS
AR
TO
FR
CF
SF
CV
NF
NQ
HQ
FQ
SD
DB
QU
Avg.
nDCG@IN 10
40,000
45.0
40.7
31.3
73.7
20.7
71.0
73.6
34.0
53.0
64.7
34.7
14.5
40.9
86.1
49.1
1,000
43.2
44.6
29.0
72.1
20.4
71.7
67.5
34,2
49.8
61.3
33.0
15.9
37.0
86.3
47.9
Recall @
40,000
91.0
92.1
50.8
92.5
51.6
$7.3
28.0
91.6
80.7
63.5
34.8
52.0
98.9
68.0
1,000
48.0
92.5
53.3
48.1
28.6
88.8
62.5
37.0
47.0
67.1
Table D.2: nDCG@ and Recall 100 of XTR on MS-MARCO and BEIR with different k'. The last column shows the average over 13 BEIR benchmarks.
17

---

E
Qualitative Analysis
In Table L5.
show token retrieval results from T5-ColBERT and XTR.
T5-ColBERT token retrieval for "lauren london age?"
Rank
Token
Context of Token
Relevance
la
laura bush laura lane welch bush (born november 4, 1946) is the wife of the 43rd president of the united states, george bush. is laura branigan dead? laura branigan died august 26, 2004 the age
No
2
la
No
5
la
laika death in space. laika died within hours from overheating. her body temperature got way too hot for her to survive. the heat in her spacecraft had risen to 40 degrees celsius (104 degrees fahrenheit).
No
50
singer laura branigan dies at singer laura branigan dies at laura branigan, grammy-noninated pop singer best known for 1982 platinum hit gloria, has died.
No
100
la
lauren bacall lauren bacall born betty joan perske; september 16, 1924
No
august)
XTR token retrieval for "lauren london age?"
Rank
Token
Context of Token
Relevance
lauren london birthday, age, family & biography 33 years, I month, 23 days old age lauren london will turn 34 on 05 december, 2018.
Yes
2
la
lauren london current age 33 years old. lauren london height 5 feet 7 inches (1.5 m/ 157 cm) and her weight 119 lbs (54 kg).
Yes
until now, lauren taylor's age is 28 year old have gemini constellation. count down 363 days will come next birthday of lauren taylor!
No
50 la
if dwayne johnson. 43, and his longtime girlfriend, lauren hashian, 31. have baby, would they have pebble? the furious 7 star and his bac reportedly expecting their first child together.
No
100
laura bush biography after his defeat, bush returned oil business and laura became a housewife, but soon returned to polities to help her father-in-law, george h.w. bush's 's presidential campaign in 1980.
No
Table E.1: Token retrieval example from MS MARCO for the token "la in the query "lauren london age". Among the top 100 retrieved tokens, 100% of T5-ColBERT tokens are lexically identical as the query token and 1 100% of XTR tokens are also lexically identical. However, top retrieved results from XTR contain the correct entity (Lauren London) while those from T5-ColBERT are about wrong entities (Laura Bush, Laura Branigan, etc.).
18

---

15-ColBERT token retrieval for "temple university student population?"
Rank
Token
Context of Token
Relevance
temple
about temple university tuition, cost, financial aid, scholarships, and admission rates
No
2 temple
overview the application fee at temple university is SSS. it selective, with acceptance of 61.7 percent and early acceptance rate of 78 percent.
No
5 temple
the application fee at temple university S55. is selective, with an acceptance rate of 61.7 percent and an early acceptance rate of 78 percent.
No
50
temple
temple university staff accountants hour, which 14% higher than the national average all staff accoun- tants at $45,000 annually and 16% lower than the national salary average for all working americans browse expedia's 's selection and check the best hotels close temple university the world-class spas and restaurants, snatch up one the cheap hotel deals near temple university
$52,000 annually, or $25 per
No.
100
temple
No
XTR token retrieval for "temple university student population?"
Rank
Token
Context of Token
Relevance
temple
by gender, the school has 18,009 male and 19,476 female students. by race/ethnicity, 20,664 white, 4,466 black, and 3,819 asian students are attending at temple university. below tables and charts represent enrollment statistics including school degree, gender, race/cthnicity, and tranfer-in students at the school. temple university, 37,485 students enrolled.
Yes
temple
Yes
5
temple
temple university the big picture: how many students were on campus in fall 2015? the 28,886 new freshman applicants. 56% were admitted and 3 of the admitted students enrolled at temple university in fall 2015.
Yes
50
temple
temple university
founded in 1884 by russell conwell, yale-
No
educated boston lawyer, orator, and ordained baptist minister kaiser said temple's endowment fund is low because the university is late to the idea of fundraising.
100
temple
No
Table E.2: Token retrieval example from MS MARCO for the token "temple" " in the query "temple university student population?". Among the top 100 retrieved tokens, 100% of T5-ColBERT tokens lexically identical the query token temple and 100% of XTR tokens are also lexically identical. However, top retrieved results from XTR are of the correct context (student population) while those from T5-ColBERT are off-topic (e.g., tuition. salary, etc.).
19

---

T5-ColBERT token retrieval for "aire expressed some skin tumors"
Rank
Token
Context of Token
Relevance
aire
acids: structures, properties, and functions (university science books, sausalito, ca, 2000). humans expressing defective form of the transcrip- tion factor aire (autoimmune regulator) develop multiorgan autoimmune disease.
No
2 aire
the primary biochemical defect in apeced is unknown. we have isolated novel gene, aire, encoding for putative nuclear protein featuring two phd-type fingefinger motifs, suggesting its involvement in transcriptional regulation.
No
5 aire
control of central and peripheral tolerance by aire. the negative selection of self-reactive thymocytes depends the expression of tissue-specific antigens by medullary thymic epithelial cells.
No
50
aire
we found that a human patient mice with defects aire develop similar lung pathology, demonstrating that the sire-deficien model of autoimmunity is suitable translational system in which to unravel fundamental mechanisms of ild pathogenesis.
No
100
air
cool air initiates just downstream of the major sense transcript poly(a) site terminates either early or extends into the promoter region.
No
XTR token retrieval for "aire expressed in skin lumors"
Rank
Token
Context of Token
Relevance
aire
keratin deaetdd- regulation of aire and gene expression in skin tumor keratinocytes expression of the intermediate filament protein keratin 17 (k17) robustly upregulated inflammatory skin diseases and in many tumors...
Yes
2 aire
the thymic transcription factor autoimmune regulator (aire) prevents autoimmunity in part by promoting expression of tissue-specific self- antigens, which include many cancer antigens. for example, aire- deficient patients are predisposed to vitiligo, autoimmune disease of melanocytes that is often triggered by efficacious immunotherapies against melanoma.
Yes
5 aire
aire regulates negative selection of organ-specific t cells autoimmunc polyendocrinopathy syndrome type is recessive mendelian disorder resulting from mutations in novel gene, aire, and characterized by spectrum of organ-specific autoimmune diseases.
No
50
aire
here we demonstrate novel role for cd4+3- inducer cell population, previously linked to development of organized secondary lymphoid structures maintenance of cell memory the functional regulation of aire-medien mediated promiscuous gene expression in the thymus.
No
100
air
this localization is dependent on the presence of sperm in the spermath- eca. after fertilization, air-2 remains associated with chromosomes during each meiotic division.
No
Table E.3: Token retrieval example from MS MARCO for the token "aire" in the query "aire is expressed in some skin tumors". Among the top 100 retrieved tokens, 77% of T5 tokens lexically identical as the query token aire and 77% of XTR tokens are also lexically identical. Top retrieved results from XTR relevant to the query (about cancer, tumor, skin, and melanocyte), while those from TS-ColBERT are off-topic.
20

---

15-ColBERT for "women with higher birth weight are more likely develop breast cancer later life"
Rank
Token
Context of Token
Relevance
later
context exposure cardiovascular risk factors during childhood and adolescence may be associated with the development of atherosele- rosis later in life.
No
2
later
despite the high incidence of febrile scizures, their contribution to the development of epilepsy later in life has remained controversial. prospectively collected data from intervention studies in adults with severe malaria were analysed focusing on laboratory features on presentation and their association with later requirement for rrt. they did have limited amount of proteolytic activity and were able to kill aureus. with time. the nuclear envelope ruptured, and dna filled the eytoplasm presumably for later lytie net production finally, we address the need for a careful consideration of potential benefits of bisphosphonate therapy and the risk for osteonecrosis of the jaw, recently recognized late-toxicity of their use.
No
5
later
No
50
later
No
100
late
No
XTR for "women with a higher birth weight are more likely to develop breast cancer later in life."
Rank
Token
Context of Token
Relevance
later
life course breast cancer risk factors and adult breast density (united kingdom) objective to determine whether risk factors childhood and carly adulthood affect later mammographic breast density. exposure to cardiovascular risk factors during childhood and ado- lescence may be associated with the development of atherosclerosis later life.
Yes
2
later
No
5
subsequent
emerging evidence suggests an association between female prenatal experience and her subsequent risk of developing breast cancer. our nested casc-control study of ch progression included 138 cases, who diagnosed with eh and then with carcinoma (1970-2003) at least year (median, years) later, and 241 controls... obesity and being overweight during adulthood have been consis- tently linked to increased risk for development of dementia later in life, especially alzheimer's 's discase.
Yes
50
later
No
100
during
No
Table E.4: Token retrieval example from Scifact for the token "later" in the query "women with higher birth weight more likely develop breast cancer later in life". Among the top 100 retrieved tokens, 72% T5-CoIBERT tokens are lexically identical as the query token later while only 33% of XTR tokens are lexically identical. Top retrieved results from XTR can retrieves synonyms (sebsequent) from relevant context, while those from T5-ColBERT are off-topic.
21

---

T5-ColBERT for "venules have a thinner or absent smooth layer compared to arterioles."
Rank
Token
Context of Token
Relevance
1
thinner
platelet cd401 is associated with smaller plaques and thinner caps, while p-selectin is associated with smaller core size. conclusions: blood cell activation significantly associated with atherosclerotic changes of the carotid wall.
No
2
thin
the periosteum is a thin, cellular and fibrous tissue that tightly ad- heres the outer surface of all but the articulated surface of bone and appears to play a pivotal role driving fracture pain.
No
5
thin
immunohistological scoring showed significantly (p<0.0001) higher median 5hmc levels in ben and den than in thin ssm, thick ssm, and emd.
No
50
weak
subarachnoid haemorrhage (1-43 [1-25-1-63)), and stable angina (1-41 [1-3 36-))) and weakest ahdominal aortic aneurysm (1-08 [1 11-00-1-17)..
No
100
slight
the ucp-2 gene expression was widely detected in the whole body with substantial levels in the wat and with slight levels the skeletal muscle and bat.
No
XTR for "venules have thinner absent smooth layer compared to arterioles."
Rank
Token
Context of Token
Relevance
1
thinner
platelet cd401 is associated with smaller plaques and thinner caps, while p-selectin is associated with smaller core size. conclusions: blood cell activation significantly associated with atherosclerotic changes the carotid wall.
No
2
thin
the periosteum a thin, cellular and fibrous tissue that tightly ad- heres to the outer surface of all hut the articulated surface of bone and appears to play a pivotal role in driving fracture pain.
No
5
thick
in dense fibrotic zones, thickening of the arterial and venous wall with severe luminal narrowing was present each patient. we assessed vasomotor function of the adipose microvasculature using vidcomicroscopy of small arterioles isolated from different fat compartments.
No
50
small
No
100
particle
context circulating concentration lipoprotein(a) (Ipfa)) large glycoprotein attached a low-density lipoprotein-like particle, may be associated with risk of coronary heart disease (chd) and stroke.
No
Table E.5: Token retrieval example from Scifact for the token "thinner" in the query "vanules have thinner or absent smooth later compared to arterioles". Among the top 100 retrieved tokens, only 1% of T5-COBBERT tokens are lexically identical as the query token thinner and only 1 1% of XTR tokens are also lexically identical.
22


Reviews:
[{"review_id": "Q0lXVVp7tA", "reviewer": "Reviewer_VrQp", "summary": "This 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", "questions": "Q1: 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?", "limitations": "- 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.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- 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", "weaknesses": "- 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).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your detailed review on our work.\n\n**A1. The XTR method is mainly built on T5 models, without exploration on other architectures.**\n\nWe mainly used T5 models (encoder-only) since it is easier to scale the architecture (from base to xxl) and has been shown to work well for the initialization of dual encoders such as GTR and Sentence-T5. In particular, GTR and Sentence-T5 have already been shown to work equal or better than BERT-based encoders such as DPR or Sentence-BERT. In the future, we are planning to use stronger and larger pre-trained language models such as PaLM or LLAMA, but using decoder-only LMs as retrievers is out-of-scope in this work.\n\n**A2. The analysis on equation (4) can be more elaborate (also related to Q1).**\n\nPlease see our response A1 to R4 on a more formal description of the alignment matrices. Specifically, we defined the alignment matrix in a more mathematical way giving a clear distinction between Eq 1 and Eq 4. \n\n**A3. Could you provide more insights on why the upperbound of m_i in equation (4) works? Have you tried other bounds/estimations?**\n\nThe success of our upper bound-based approximation tells us that estimating missing value in a query-dependent manner is more important than exactly imputing missing value. As reported in Table 5, constant imputation does not work well. On the other hand, more sophisticated imputation methods such as power-law based imputation, which was omitted from the paper for brevity, did not provide significant improvement despite its complexity. As a result, we use the upper bound-based estimation, which is simple to utilize while robust enough to provide good estimation. We will report the power-law based imputation method in the final version.\n\n**A4. In table 1, how does the estimated FLOPs/query looks like for single vector retrieval methods such as DPR?**\n\nFor a single vector retrieval, FLOPs/query is much lower since it does not have multiple representations of each query and document. Note that single-vector retrieval does not have the scoring stage since the retrieval directly provides final scores and rankings. Assuming sublinear MIPS, dual encoders like DPR/GTR would have $2d \\log L$ FLOPs/query for retrieval where $L$ is the number of documents (e.g. $d=768, L=5\\times10^6$ for DPR). On the other hand, naive multi-vector retrieval models would have $2nd \\log M$ where n is the number of query tokens and $M$ is the number of document tokens in a corpus (e.g. $d=128, n=16, M=1\\times 10^9$ for XTR). While our focus in this work is to better optimize the scoring stage in terms of FLOPs/query and the memory usage, previous works such as PLAID or ColBERTer are worth mentioning, which reduce the complexity of the first-stage token retrieval as described in our general response G2."}, {"review_id": "6XwGXZzBVh", "reviewer": "Reviewer_BRdC", "summary": "The 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 ", "questions": "- 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?", "limitations": "The limitations section is a bit too generic (not really related to the proposed approach). See weaknesses for possible issues to report at this level.", "rating": 6, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "The 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. ", "weaknesses": "- 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)", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your review.\n\n**A1. No experiment measuring the observed latency are reported.**\n\nSince it is difficult to reimplement every baseline within our hardware and infrastructure, apples-to-apples latency comparisons are a bit tricky. Admitting the differences in implementations, we report the latency of XTR in our general response G1. To summarize, our implementation of XTR runs reasonably fast while significantly improving the speed of T5-ColBERT, both of which are tested in the same environment. \n\n**A2. No comparison with PLAID or ColBERTer**\n\nPLAID and ColBERTer optimize the first-stage token retrieval, which is orthogonal to the contribution of XTR (better scoring stage without the gathering stage) and they can be applied to XTR as well. For instance, sparse optimization techniques of AligneR (Qian et al., 2022), which can be more directly compared to ColBERTer as it optimizes the first-stage token retrieval, can be applied to XTR without minimum performance losses. Please also refer to our general response G2 for the detailed comparisons.\n\n**A3. Figure 4 should report the recall of gold tokens (rather than the precision)**\n\nWhile Figure 4 reported the precision of token at the k-th rank, we agree that the recall of gold tokens would be interesting to see since it would demonstrate how well the token retrieval serves as the first stage of the multi-vector retrieval. We found that indeed, the token retrieval of XTR has higher token recall (i.e., how many tokens of all gold documents are retrieved?), achieving 0.2517 recall@1000 while T5-ColBERT achieves 0.1047 recall@1000 (TREC-COVID). Note that the overall recall would be very low since the model doesn’t have to retrieve the entire set of gold tokens, which include many stopwords as well. Similar results were observed on MS-MARCO and other BEIR benchmarks, which translate to the superior document-level recall@100 of XTR (Table E.1 in Appendix). We also added figures of the token retrieval recall in the pdf attached in the general response, which we will include in our final draft.\n\n\n**A4. 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\nExploring the sparsity loss was extensively studied in the prior work. In particular, AligneR (Qian et al., 2022) tested the Optimal Transport and L1 regularization techniques to sparsify the pairwise as well as unary alignments. However, while these sparsification techniques make the model decide which tokens of a document are salient, we found that existing sparsification methods do not encourage those salient tokens to be retrieved across other documents. This is primarily due to the fact that sparsification methods work independent of the query representations.\n\n**A5. What is the real latency time of such a model when the number of top-K vectors increases**\n\nFor the real latency of each stage (retrieval and scoring), please refer to the general response G1 where we show some fair comparisons with increasing top-k vectors.\n\n**A6. is k_train the same thing as m in eq. 2?**\n\nk_train is a hyperparameter that decides which top tokens in a mini-batch should be taken into considerations for the scoring of mini-batch documents. On the other hand, m is the length of a document $D$ (i.e., number of tokens in $D$)."}, {"review_id": "OWiMsyazNn", "reviewer": "Reviewer_5SZT", "summary": "This 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.", "questions": "None", "limitations": "There is sufficient discussion about the limitations after the conclusion section.", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 4, "strengths": "Originality:\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.", "weaknesses": "Clarity: 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your review.\n\n**A1. Clarity: I feel this paper needs little fine-tuning in clarity.**\n\nWe will definitely add a few more sentences to make the problem statement and the experiment sections clearer. For instance, the third paragraph of the introduction has the problem statement where we can add a few sentences for better understanding. Additionally, for a better description of our method, we provided more formal definitions of the alignment matrices in the response A1 to R4."}, {"review_id": "EXPR8EOG2e", "reviewer": "Reviewer_HcB5", "summary": "Authors 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.", "questions": "-\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?", "limitations": "The authors discuss the limitations. Given that the work focusses on the training mechanism, it would have been nice to see trainings with other datasets.", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 1, "contribution": 2, "strengths": "-\tResults seem to be strong enough.\n-\tDo a decent job in ablations, qualitative analysis, complexity analysis.", "weaknesses": "-\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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your detailed review. Please read our clarification for any misunderstanding you might have had while reading the paper.\n\n**A1. The paper overcomplicates a simple concept (Eq 4 can be further simplified from what I understand). Max over j would collapse then.**\n\nFirst of all, the definition of the alignment matrix $A$ is different in Eq 1 and Eq 4 as described in our paper. Eq 1 shows the baseline sum-of-max operator and how the alignment matrix $A$ can be used to describe its algorithm. Here $A_{ij} = 1_{j=\\text{argmax}(P_{ij^\\prime})}$ where argmax is over $1\\leq j^\\prime \\leq m$ (i.e., single document) and $P_{ij}=q_i^\\top d_j$, meaning that $A_{ij}$ is 1 when $P_{ij}$ is the maximum value among $P_{ij^\\prime}$ otherwise 0. Hence, we can replace max over $j$ with $A_{ij}$ in Eq 1. On the other hand, Eq 4 shows the approximated sum-of-max for XTR. In XTR, $A_{ij} = 1_{j \\in \\text{top-k}(P_{ij^\\prime})}$ where the top-k is over $1 \\leq j^\\prime \\leq mB$, meaning that $A_{ij}$ will be 1 when $P_{ij}$ is within the top-k retrieved values among all the $P_{ij^\\prime}$ within a mini-batch of $B$ documents. The top-k operator here returns the top-k indices (which will be $k_\\text{train}$ during training and $k^\\prime$ for inference) among the $j^\\prime$s. The main difference with Eq 1 is that the $j^\\prime$ spans over all tokens from mini-batch documents in Eq 4, not just from a single document as in Eq 1. As a result, unless all $A_{ij}=1$ (which is unrealistic), Eq 4 will be different from Eq 1 and cannot be collapsed.\n\nWe hope that our mathematical formulation can provide proper clarification for your concerns. We will add this clarification in the paper and use a different notation for the alignment matrix of XTR (e.g., $A^\\text{XTR}$) to prevent any confusion.\n\n**A2. it would be interesting to see how this method performs with training on NQ, TriviaQA, etc.**\n\nIt is a common practice to use MS-MARCO for training retrieval models in many recent works. Indeed, Ni et al., 2021 showed that using MS-MARCO transfers better than using NQ when evaluated on the BEIR benchmarks. Most of the recent papers on neural retrieval were all trained on MS-MARCO, such as ColBERT, ColBERT v2, SPLADE v2, ColBERTer, PLAID, etc. \n\n**A3. For 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\nIn our paper (176-177), we discuss that Eq 4 would degenerate to Eq 1 when every $A_{ij}$ of Eq 4 is $1$, but this is an unrealistic special case. Based on our description of $A_{ij}$, Eq 4 does not collapse to Eq 1. Please see our response A1 on the description of alignment matrices, which we believe would resolve some of the reviewer’s concerns.\n\n**A4. Line 120. How is f_colbert(Q,D−)=0.2 when qiTdj−>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\nit is possible to have $f_\\text{colbert}(Q,D−)=0.2$ when $q_i^\\top d_j \\rightarrow 0.8$ because the document-level score (e.g. $f_\\text{colbert}$) is an average score of token-level pairwise scores (i.e., $q_i^\\top d_j$). For instance, imagine there are 6 tokens with $q_i^\\top d_j=0.1$ and one token with $q_i^\\top d_j=0.8$. $f_\\text{colbert}(Q,D-)$ would become (0.8 + 0.1*6)/7=0.2.\nThe failure case example is presented for the sake of intuitive understanding of the motivation.  For a more formal understanding of the failure case, we also provided a theoretical analysis of the sum-of-max operator in Appendix A. We will bring some of these results into the main text to make it less confusing.\n\n**A5. Lines 134-137. Very handwavy again. Authors should use proper math format to denote in-batch tokens. It is very difficult to understand when Aij=1.**\n\nWe believe that our description in A1 would resolve this concern.\n\n**A6. Lines 134-137. Define top-k_train**\n\n$k_\\text{train}$ is a hyperparameter that decides how many in-batch tokens should be retrieved. We will clarify this in the paper.\n\n**A7. For 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\nIf the token vectors are already loaded on RAM, computing all $A_{ij}$ can be done pretty fast  But we cannot simply assume token vectors are already in RAM. In fact, loading/storing the many token vectors in RAM has been one of the most expensive parts of multi-vector retrieval, and how to make it cheaper is the main problem studied by XTR as well as many prior works (ColBERT v2, PLAID, etc).\nThe main advantage of XTR is to completely remove the need of storing any token vectors on RAM and score each document solely based on the retrieved tokens and their scores. While this can significantly reduce the RAM requirement for ColBERT scoring, our general response (G1) on the actual latency also shows how much time we can save with RAM vs Disk-based similarity computation as well as with the XTR-style scoring.\n\n**A8. Lines 226-241: what is token at rank k? are these top k retrived tokens?**\n\nYes, the token at rank k means the k-th retrieved token. We will clarify this in the paper."}, {"review_id": "WBbPq8tMv1", "reviewer": "Reviewer_XbZR", "summary": "The 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.", "questions": "What 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?", "limitations": "See Weaknesses & Questions.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 4, "strengths": "The 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).", "weaknesses": "Despite 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your detailed review.\n\n**A1. This comparison doesn't count the actual FLOPs used. It also doesn't measure the latency of the scoring stage or the full pipeline.**\n\nAs the reviewer summarized, our goal of the paper is to simplify the three-stage inference of ColBERT while making the scoring stage very efficient. Also, our claim on the FLOPs improvement is over the scoring stage, not over the entire pipeline. The apples-to-apples comparison of all baselines is difficult due to the differences in libraries (e.g. FAISS vs ScaNN, pytorch vs jax) and implementations depending on hardware constraints (e.g., multiprocessing, RAM usage). Nevertheless, for the sake of understanding the actual latency, we summarize the latency of our current implementation of XTR in the general response G1. We also show how much memory is needed to compute the full sum-of-max for T5-ColBERT. Please note that optimizing the latency of the entire pipeline is not our immediate goal as it requires lots of engineering efforts, which will depend on the given resources and algorithms. We will make this clearer in the paper. For comparisons against PLAID and ColBERTer, please see our general response G2.\n\n**A2. 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.**\n\nIn the abstract, we will properly tone down the state-of-the-art message and focus more on the contributions of XTR. We use the encoder part of T5-XXL so it is about 5B.\n\n**A3. What is the intuition for the normalizer Z during training? How essential is this for the results?**\n\nThe normalizer Z provides stable losses during training. Since each document has a different number of mini-batch tokens retrieved, removing the normalization would give very high scores on positive documents, making the training process unstable. Not using Z (or setting Z to a constant) does not perform well compared to our proposed version of Z.\n\n**A4. How robust are the analyses with the T5-ColBERT, compared to the original ColBERT and ColBERTv2 models?**\n\nWhile most of our analyses were based on T5-ColBERT, we do acknowledge that there is a difference between T5-ColBERT and ColBERT (or v2) including query expansion and cross-encoder distillation. Based on our theoretical analyses in Appendix A, which uses the sum-of-max operator, we believe that our findings would hold for existing ColBERT variants even if they use the query expansion. More empirical comparisons against ColBERT could be needed to better support our hypothesis in the future.\n"}]

Author rebuttal:
We thank all the reviewers for their thoughtful comments and feedback.

Most of our reviewers agree that XTR effectively mitigates the problem of three-stage inference of multi-vector models and provides significant improvements over strong baseline models. Some of the main concerns include latency reports (with comparison against the existing baselines) (R2, R5) and the presentation quality of the paper (R1, R3, R4).

**G1. Actual Latency of XTR**

Our initial manuscript did not compare the actual latency due to the differences in libraries and infrastructure used for implementing each baseline and XTR. Specifically, XTR uses ScaNN (Guo et al., 2020), which has different optimization techniques and a different distributed system compared to Faiss (Johnson et al., 2019) used by various baselines. XTR also does not use any centroid-based approximation used by ColBERTv2 or PLAID. Nevertheless, as many reviewers (R2, R5) asked for, the actual latency might provide a clearer sense of the contribution of XTR. Under the same environment, we report the latency of the token retrieval as well as the scoring stage, which uses naive CPU-based Numpy without any multi-processing. The token retrieval stage uses CPU-based distributed ScaNN (same for T5-ColBERT and XTR). For T5-ColBERT, document token vectors are loaded for the exact sum-of-max, either from Disk or RAM. MS-MARCO dev set queries were used to average per-query latencies.

| Top-k’ tokens | Token retrieval | T5-ColBERT-Disk scoring | T5-ColBERT-RAM scoring | XTR-scoring | XTR total latency |
| --- | :---: | :---: | :---: | :---: | :---: |
| 1,000 | 0.25 sec | 11.56 sec | 3.13 sec | 0.38 sec | 0.63 sec |
| 4,000 | 0.31 sec | 47.05 sec | 12.67 sec | 1.64 sec | 1.95 sec |
| 40,000 | 0.47 sec | 8 min 2.09 sec | 2 min 7 sec | 17.17 sec | 17.64 sec |  

Note that these results are mainly for the fair comparisons against the full sum-of-max baseline (i.e., T5-ColBERT) so we did not use multi-processing. Our implementation of the XTR scoring with multiprocessing actually takes about < 0.3 sec even when $k^\prime=40,000$, hence giving < 0.8 sec in total. The scoring stage of XTR also does not require loading $\mathcal{O}(nk^\prime\bar{m}d )$ floating points from Disk or RAM as T5-ColBERT does, which can range from 450MB ($k^\prime=1,000$) to 18GB ($k^\prime=40,000$) per query.

**G2. Comparison against PLAID and ColBERTer**

Prior work such as PLAID and ColBERTer focus on different aspects of multi-vector retrieval. PLAID and ColBERTer focus on improving the efficiency of the first stage token retrieval, while XTR aims to greatly simplify the gathering and scoring stages as well as improve the training objective.  More specifically, PLAID 1) uses cluster-centroid based approximated search to provide a small number of candidates and 2) looks up the candidate vectors to apply the full sum-of-max over them. ColBERTer proposes to use whole-word representations with more aggressive dimensionality reduction, reducing the cost of first-stage retrieval. XTR, on the other hand, 1) uses a naive MIPS library that provides a large number of candidates and 2) uses the light-weight scoring function, which eliminates index lookup. XTR additionally modifies the training objective offering better token retrieval. We believe that these paradigms are very orthogonal and hope to explore combining these efforts in the future. For instance, it would be interesting to see if the training objective of XTR could improve the candidate generation part of PLAID so that it can further have a smaller number of candidates.

**G3. Presentation Quality**

While many reviewers (R1, R2, R5) believe that the presentation of our paper is either good or excellent, we acknowledge that there is room for improvement as R3 and R4 suggested. In particular, we provided detailed clarification on each concern of R4 regarding the equations and examples (please see R4-A1), which will be added in our manuscript. For a better introduction, we will add clarifying sentences as R3 suggested (please see R3-A1).

We also attached figures related to the question raised by R2-A3.

Meta-review: The paper is a clear accept. It identifies a practical problem and provides a principled solution, based on a neat generalization of the multi-vector retrieval objective, with strong empirical results. The reviews are positive and the author response adequately addressed the reviewers' questions.
Decision: Accept (poster)