CITADEL : Conditional Token Interaction via Dynamic Lexical Routing for Efficient and Effective Multi-Vector Retrieval (2023)

Abstract: Multi-vector retrieval methods combine the merits of sparse (e.g. BM25) and dense (e.g. DPR) retrievers and have achieved state-of-the-art performance on various retrieval tasks. These methods, however, are orders of magnitude slower and need much more space to store their indices compared to their single-vector counterparts. In this paper, we unify different multi-vector retrieval models from a token routing viewpoint and propose conditional token interaction via dynamic lexical routing, namely CITADEL, for efficient and effective multi-vector retrieval.CITADEL learns to route different token vectors to the predicted lexical keys such that a query token vector only interacts with document token vectors routed to the same key. This design significantly reduces the computation cost while maintaining high accuracy. Notably, CITADEL achieves the same or slightly better performance than the previous state of the art, ColBERT-v2, on both in-domain (MS MARCO) and out-of-domain (BEIR) evaluations, while being nearly 40 times faster. Source code and data are available at https://github.com/facebookresearch/dpr-scale/tree/citadel .

## introduction the goal of information retrieval @xcite is to find a set of related documents from a large data collection given a query. traditional bagof-words systems @xcite calculate the ranking scores based on the query terms appearing in each document, which have been widely adopted in many applications such as web search @xcite and open-domain question answering @xcite . recently, dense retrieval @xcite based on figure 1 : gpu latency vs ranking quality (mrr@10) on ms marco passages with an a100 gpu. the circle size represents the relative index storage on disk. all models are trained without hard-negative mining, distillation, or further pre-training. in contrast, multi-vector retrieval has shown strong performance on both in-domain and out-ofdomain evaluations by taking into account tokenlevel interaction. among them, colbert @xcite ) is arguably the most celebrated method that has been the state of the art on multiple datasets so far. however, its wider application is hindered by its large index size and high retrieval latency. this problem results from the redundancy in the token interaction of colbert where many tokens might not contribute to the sentence semantics at all. to improve this, coil @xcite imposes an exact match constraint on colbert figure 2 : a unified token routing view over different multi-vector and single-vector retrieval models. the cylinder in (d) represents the lexical key for token routing. the grey tokens in (c) and (d) represent the tokens that do not contribute to the final similarity. for (d), citadel routes "light", "theory", and "relativity" to the "einstein" key to avoid term mismatch using a learned routing function. for conditional token interaction, where only token embeddings with the same token id could interact with each other. although reducing the latency, the word mismatch problem reoccurs and the model may fail to match queries and passages that use different words to express the same meaning. in this paper, we first give a unified view of existing multi-vector retrieval methods based on token routing (section 2), providing a new lens through which we expose the limitations of current models. under the token routing view, colbert could be seen as all-to-all routing, where each query token exhaustively interacts with all passage tokens (figure 2b ). coil, on the other hand, could be seen as static lexical routing using an exact match constraint, as each query token only interacts with the passage tokens that have the same token id as the query token (figure 2c ). in contrast, we propose a novel conditional token interaction method using dynamic lexical routing called citadel as shown in figure 2d . instead of relying on static heuristics such as exact match, we train our model to dynamically moderate token interaction so that each query token only interacts with the most relevant tokens in the passage. this is achieved by using a lexical router, trained end-to-end with the rest of the model, to route each contextualized token embedding to a set of activated lexical "keys" in the vocabulary. in this way, each query token embedding only interacts with the passage token embeddings that have the same activated key, which is dynamically determined during computation. as we shall see in section 5.1, this learning-based routing does not lose any accuracy compared to all-to-all routing while using fewer token interactions than coil (section 3.4), leading to a highly effective and efficient retriever. experiments on ms marco passages @xcite and trec dl show that citadel achieves the same level of accuracy as colbert-v2. we further test citadel on beir @xcite and citadel still manages to keep up with @xcite which is the current state of the art. as for the latency, citadel can yield an average latency of 3.21 ms/query on ms marco passages using an a100 gpu, which is nearly 40× faster than colbert-v2. by further combining with product quantization, citadel's index only takes 13.3 gb on ms marco passages and reduces the latency to 0.9 ms/query as shown in figure 1 . ## a unified token routing view of multi-vector retrievers we outline a unified view for understanding various neural retrievers using the concept of token routing that dictates token interaction. ## conclusion this paper proposes a novel multi-vector retrieval method that achieves state-of-the-art performance on several benchmark datasets while being 40× faster than colbert-v2 and 17× faster than the most efficient multi-vector retrieval library to date, plaid, on gpus. by jointly optimizing for the token index size and load balancing, our new dynamic lexical routing scheme greatly reduces the redundancy in the all-to-all token interaction of col-bert while bridging the word-mismatch problem in coil. experiments on both in-domain and outof-domain datasets demonstrate the effectiveness and efficiency of our model.