Rethinking the Role of Token Retrieval in Multi-Vector Retrieval (2304.01982v3)

Abstract: Multi-vector retrieval models such as ColBERT [Khattab and Zaharia, 2020] allow token-level interactions between queries and documents, and hence achieve state of the art on many information retrieval benchmarks. However, their 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.

• The paper presents XTR, which simplifies inference by retrieving salient tokens and reducing computation in multi-vector retrieval models.
• It achieves a 2.8 nDCG@10 improvement on BEIR benchmarks, demonstrating enhanced recall and efficiency over traditional methods.
• By bypassing redundant token gathering, XTR offers scalable solutions for real-time search, question answering, and multilingual retrieval.

Rethinking the Role of Token Retrieval in Multi-Vector Retrieval

The paper presents a novel approach to simplifying the inference process in multi-vector retrieval models, particularly focusing on token retrieval's critical role. Multi-vector retrieval models like ColBERT leverage token-level interactions between queries and documents to achieve high performance in information retrieval tasks. However, the complexity of their non-linear scoring function poses scalability challenges when dealing with large document collections, requiring a cumbersome three-stage process involving token retrieval, token gathering, and document scoring.

The authors introduce the ConteXtextualized Token Retriever (XTR), which simplifies the existing multi-vector retrieval process by enhancing the token retrieval stage. XTR is designed to retrieve the most salient document tokens first, enabling effective scoring of candidate documents using only a subset of the token vectors. This innovation leads to a scoring process that is significantly more efficient than its predecessors, reducing computation by orders of magnitude.

On standard retrieval benchmarks like BEIR, XTR advances the state-of-the-art by 2.8 nDCG@10 without distillation, highlighting its improved recall performance over ColBERT. The paper reveals that by rethinking the token retrieval stage, XTR can effectively rely on the retrieved tokens alone for scoring, thus circumventing the need for the gathering stage.

Key Contributions

1. Simplified Inference Process: XTR proposes a new paradigm by focusing the retrieval on contextually important tokens, eliminating the need to score based on all document tokens. This drastically reduces computational overhead and accelerates inference.
2. Improved Scoring Efficiency: XTR's new scoring method avoids redundant computation by reusing token retrieval scores, contrasting with exhaustive pairwise token similarity calculations required by previous models.
3. Superior Benchmark Performance: XTR achieves state-of-the-art results on multiple Information Retrieval benchmarks like BEIR and LoTTE, and it demonstrates better recall capabilities. The authors attribute this to the enhanced focus on retrieving contextually relevant tokens.
4. Cross-Language Efficacy: The multilingual version of XTR surpasses other advanced multilingual retrieval models on MIRACL, showcasing its flexibility across different languages.

Implications and Future Directions

The introduction of XTR has the potential to reshape the deployment of multi-vector models in real-world applications where efficiency is crucial. Its ability to streamline the retrieval process while maintaining or improving performance presents opportunities for more widespread use across diverse domains, including real-time search, question answering, and multilingual retrieval.

Prospective extensions of this work might explore the integration of XTR with emerging LLMs. Further fine-tuning with domain-specific data or zero-shot capabilities could enhance its adaptability and accuracy. Additionally, applying this strategy in semi-supervised or unsupervised contexts may yield further insights into its generalizability across different tasks or datasets.

In summary, XTR provides a substantial efficiency gain for multi-vector retrieval models by strategically optimizing token retrieval and scoring processes. This work not only contributes a practical solution to existing scalability issues but also opens avenues for further exploration in efficient, contextually-aware retrieval systems.