Rethinking the Role of Token Retrieval in Multi-Vector Retrieval

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. The 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. To 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.

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). First, 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. To the best of my knowledge, this appears, however, unsupported or otherwise inflated on a few dimensions: 1) 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. 2) 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. 3) 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. Second, 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.

Questions

What is the intuition for the normalizer Z during training? How essential is this for the results? How 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?

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

4 excellent

Contribution

4 excellent

Limitations

See Weaknesses & Questions.

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.

Strengths

- Results seem to be strong enough. - Do a decent job in ablations, qualitative analysis, complexity analysis.

Weaknesses

- The paper overcomplicates a simple concept (Eq 4 can be further simplified from what I understand). Max over j would collapse then. - Given 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. - Very bad writing. - Very handwavy at some places - See Questions for detailed list of weakness.

Questions

- 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)?? - Line 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? - Lines 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$. - Lines 134-137. Define top-k_train - 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? - Lines 226-241: what is token at rank k? are these top k retrived tokens?

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

2 fair

Presentation

1 poor

Contribution

2 fair

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.

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. Post rebuttal: I have read the author's response and the rebuttal sufficiently addressed my concerns.

Strengths

Originality: The 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. Significance: Given 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.

Questions

None

Rating

8: Strong Accept: Technically strong paper, with novel ideas, excellent impact on at least one area, or high-to-excellent impact on multiple areas, with excellent evaluation, resources, and reproducibility, and no unaddressed ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

2 fair

Contribution

4 excellent

Limitations

There is sufficient discussion about the limitations after the conclusion section.

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. The 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

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. The 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 - 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. - Figure 4 should report the recall of gold tokens (rather than the precision)

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? - What is the real latency time of such a model when the number of top-K vectors increases - is $k_{train}$ the same thing as $m$ in eq. 2?

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

5: You are absolutely certain about your assessment. You are very familiar with the related work and checked the math/other details carefully.

Soundness

3 good

Presentation

3 good

Contribution

2 fair

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.

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. In 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. Finally, analysis was presented in section 5 to shed lights on why XTR presents a better token retrieval.

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 - The proposed method achieved SOTA results on zero-shot retrieval benchmarks and multilingual retrieval benchmarks - 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. - 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. - The analysis on equation (4) can be more elaborate (see Q1).

Questions

Q1: Could you provide more insights on why the upperbound of m_i in equation (4) works? Have you tried other bounds/estimations? Q2: In table 1, how does the estimated FLOPs/query looks like for single vector retrieval methods such as DPR?

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

- The authors listed that XTR is trained on MSMARCO dataset which may have license issue. - The proposed method only trained on T5 models, and no other type of architecture has been explored.