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

Proceedings of the 61st Annual Meeting of the Association for Computational Linguistics Volume 1: Long Papers, pages 11891–11907 July 9-14, 2023 ©2023 Association for Computational Linguistics CITADEL: Conditional Token Interaction via Dynamic Lexical Routing for Efficient and Effective Multi-Vector Retrieval Minghan Li1∗, Sheng-Chieh Lin1, Barlas Oguz2, Asish Ghoshal2, Jimmy Lin1, Yashar Mehdad2, Wen-tau Yih2, and Xilun Chen2† University of Waterloo1, Meta AI2 {m692li,s269lin,jimmylin}@uwaterloo.ca {barlaso,aghoshal,mehdad,scottyih,xilun}@meta.com Abstract Multi-vector retrieval methods combine the merits of sparse (e.g. BM25) and dense (e.g. DPR) retrievers and achieve state-of-the-art performance on various retrieval tasks. These methods, however, are orders of magnitude slower and need more space to store their indexes 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 each token vector 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. 1 Introduction The goal of information retrieval (Manning et al., 2008) is to find a set of related documents from a large data collection given a query. Traditional bag- of-words systems (Robertson and Zaragoza, 2009; Lin et al., 2021a) calculate the ranking scores based on the query terms appearing in each document, which have been widely adopted in many applica- tions such as web search (Nguyen et al., 2016; Noy et al., 2019) and open-domain question answer- ing (Chen et al., 2017; Lee et al., 2019). Recently, dense retrieval (Karpukhin et al., 2020) based on ∗This work is done during Minghan’s internship at Meta. †Xilun and Minghan contributed equally to this work. 100 101 102 GPU Latency (ms/query) 0.30 0.32 0.34 0.36 MRR@10 DPR-768 DPR-128 COIL-tok COIL-full ColBERT ColBERT-PLAID CITADEL CITADEL + PQ 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. pre-trained language models (Devlin et al., 2019; Liu et al., 2019) has been shown to be very effec- tive. It circumvents the term mismatch problem in bag-of-words systems by encoding the queries and documents into low-dimensional embeddings and using their dot product as the similarity score (Figure 2a). However, dense retrieval is less robust on entity-heavy questions (Sciavolino et al., 2021) and out-of-domain datasets (Thakur et al., 2021), therefore calling for better solutions (Formal et al., 2021b; Gao and Callan, 2022). In contrast, multi-vector retrieval has shown strong performance on both in-domain and out-of- domain evaluations by taking into account token- level interaction. Among them, ColBERT (Khattab and …

Content truncated to 500 words