Efficient Document Ranking with Learnable Late Interactions

Authors: ziwei  ji
himanshu  jain
andreas  veit
sashank j reddi
sadeep  jayasumana
ankit singh rawat
aditya krishna menon
felix  yu
sanjiv  kumar
Venue: ICML_cc_2024_Workshop_WANT
Published: 2024-05-14 22:58:02

Abstract: Cross-Encoder (CE) and Dual-Encoder (DE) models are two fundamental approaches for predicting query-document relevance in information retrieval. To predict relevance, CE models use joint query-document embeddings, while DE models maintain factorized query-document embeddings; usually, the former has higher quality while the latter has lower latency.
Recently, late-interaction models have been proposed to realize more favorable latency-quality trade-offs, by using a DE structure followed by a lightweight scorer based on query and document token embeddings. However, these lightweight scorers are often hand-crafted, and there is no understanding of their approximation power; further, such scorers require access to individual document token embeddings, which imposes an increased latency and storage burden over DE models. In this paper, we propose novel \emph{learnable} late-interaction models (LITE) that resolve these issues.
Theoretically, we prove that LITE is a universal approximator of continuous scoring functions, even for relatively small embedding dimension. Empirically, LITE outperforms previous late-interaction models such as ColBERT on both in-domain and zero-shot re-ranking tasks such as MS MARCO and Natural Questions, and out-of-domain tasks such as BEIR. For instance, experiments on MS MARCO passage re-ranking show that LITE not only yields a model with better generalization, but also lowers latency and requires 0.25 times storage compared to ColBERT.

Main paper: 
 introduction: 
 Transformers (Vaswani et al., 2017) have emerged as a successful model for information retrieval problems, where the goal is to retrieve and rank relevant documents for a * Equal contribution 1 Google Research, New York, USA. Correspondence to: Ziwei Ji <ziweiji@google.com>, Himanshu Jain <himj@google.com>.Accepted to the Workshop on Advancing Neural Network Training at International Conference on Machine Learning (WANT@ICML 2024).given query (Nogueira & Cho, 2019) . Two families of Transformer-based models are popular: cross-encoder (CE) and dual-encoder (DE) models. Given a (query, document) pair, CE models operate akin to a BERT-style encoder (Devlin et al., 2019) : the query and document are concatenated, and sent to a Transformer encoder which outputs a relevance score (cf. Figure 1a ). CE models can learn complex querydocument relationships, as they allow for cross-interaction between query and document tokens.By contrast, DE models apply two separate Transformer encoders to the query and document, respectively, producing separate query and document embedding vectors (Reimers & Gurevych, 2019) . The dot product of these two vectors is used as the final relevance score (cf. Figure 1b ). Compared to CE models, DE models are usually less accurate (Hofstätter et al., 2020) , since the only interaction between the query and document occurs in the final dot product. However, DE models have much lower latency, since all the document embedding vectors can be pre-computed offline.Recently, late-interaction models have provided alternatives with a more favorable latency-quality trade-off compared to CE and DE models. Similarly to DE models, lateinteraction models also use a two-Transformer structure, but they store more information and employ additional nonlinear operations to calculate the final score. In particular, let Q ∈ R P ×L1 and D ∈ R P ×L2 denote the query and document token embeddings output by the two Transformers, i.e., there are L 1 query token embedding vectors and L 2 document token embedding vectors of dimension P . DE models simply pool Q and D into two vectors, and take the dot product. By contrast, ColBERT (Khattab & Zaharia, 2020) calculates the (token-wise) similarity matrix Q D and computes the final score via a sum-max reduction i max j (Q D) i,j . While the sum-max score reduction lets ColBERT achieve better accuracy than DE, it is unclear whether this handcrafted reduction can capture arbitrary complex querydocument interactions. Moreover, ColBERT can have higher latency than DE: calculating the similarity matrix Q D requires L 1 • L 2 dot products, while the DE model only requires one dot product. Additionally, to reduce online latency, ColBERT needs to pre-compute and store the Trans- former embedding matrix D for each document (Hofstätter et al., 2020; Santhanam et al., 2022) . This can entail significant storage space if we decide to store a large number of document tokens, since there can be billions of documents in industry-scale information retrieval systems (Zhang & Rui, 2013; Overwijk et al., 2022) . (See Section 2.3 for a detailed discussion.)To reduce latency and storage cost, one may seek to store fewer document tokens, and/or reduce the dimension of each token embedding vector. However, it is unclear how these influence performance. In fact, such reduction can significantly hurt the accuracy of ColBERT, as we show in Section 4.4.In this work, we propose lightweight scoring with token einsum (LITE), which addresses the aforementioned shortcomings of existing late-interaction models. LITE applies a lightweight and learnable non-linear transformation on top of Transformer encoders, which corresponds to processing the (token-wise) similarity matrix S = Q D via shallow multi-layer perceptron (MLP) layers (cf. Figure 1c and Section 3). In particular, we focus on a separable LITE scorer which applies two shared MLPs to the rows and the columns of S (in that order), and then projects the resulting matrix to a single scalar.Theoretically, we rigorously establish the expressive power of LITE: we show that LITE is a universal approximator of continuous scoring functions in 2 distance, even under tight storage constraints (cf. Theorem 3.1). To our knowledge, this is the first formal result about the approximation power of late-interaction methods. Further, we also construct a scoring function that cannot be approximated by a DE model with restricted embedding dimension (cf. Theorem 3.2).Empirically, we show that LITE can systematically improve upon existing late-interaction methods like ColBERT on both in-domain benchmarks such as MS MARCO and Natural Questions (cf. Table 1 ), and out-of-domain benchmarks such as BEIR (cf. Table 2 ). Moreover, LITE can be much more accurate than ColBERT while having lower latency and storage cost (cf. Table 3 ). 
background: 
 Given a query q ∈ Q, the goal of information retrieval (Mitra & Craswell, 2018) is to identify the set of relevant documents from some corpus D. Typically, |D| is large (e.g., O(10 9 )), while the number of relevant documents is small (e.g., O(10)). A classical strategy employs a two-phase approach: in the retrieval phase, for moderate K (e.g., O(10 3 )), one retrieves the top-K documents based on a scoring function s ret : Q × D → R. These retrieved documents may potentially include some irrelevant documents. In the re-ranking phase, one applies s rr : Q × D → R to re-score the K documents, and keep the top scoring ones.While s ret and s rr both score query-document relevance, they are often implemented via fundamentally different techniques. Efficiency is more important for s ret since we need to evaluate it over all documents; models such as TF-IDF and BM25 (Robertson & Zaragoza, 2009) and approximate nearest neighbor search (Guo et al., 2016b; Johnson et al., 2019; Guo et al., 2020) are used for this purpose. On the other hand, in the second phase we usually only need to rescore a few (K |D|) documents, and thus we can usually get higher accuracy by using more expensive models for s rr . In this work, we focus on re-ranking.Transformers (Vaswani et al., 2017) have been explored for both retrieval and re-ranking. Given a finite set X, a Transformer is a function T : X L → R P ×L , where L is the sequence length and P is the embedding size of each token in the sequence. A simplified Transformer network is introduced in Section 3.1 and used in our universal approximation results; for more details, we refer the readers to (Vaswani et al., 2017; Devlin et al., 2019) .To estimate query-document relevance via Transformers, one first tokenizes the query and document (e.g., using a SentencePiece tokeniser (Kudo & Richardson, 2018) ) into q = (q 1 , . . . , q L1 ) and d = (d 1 , . . . , d L2 ). There are then two basic strategies. In cross-encoder (CE) models (Nogueira & Cho, 2019) , we apply a single Transformer to the concatenation of q and d, and estimate relevance with learned weights w:EQUATIONwhere pool denotes a pooling strategy by which we reduce a sequence of Transformer token embeddings into a single vector. CE models can often achieve high accuracy since they can take into account interactions between the query and document tokens in every Transformer layer. However, they can also be expensive at inference time: we need to compute (1) for all retrieved documents, each of which involves an expensive Transformer inference (see Section 4.4 for concrete evaluations).By contrast, in dual-encoder (DE) models (Karpukhin et al., 2020) , we apply separate Transformers T 1 , T 2 to the query and document, and then computeEQUATIONIn practice, DE is usually less accurate than CE for reranking (Hofstätter et al., 2020) , since the only interaction between the query and document is the final dot product; however, since all document embeddings pool(T 2 (d)) can be pre-computed offline, DE has much lower latency than CE.Another idea is to apply an MLP to the concatenation of pool(T 1 (q)) and pool(T 2 (d)) (He et al., 2017) . However, Rendle et al. (2020) claim that it may not be better than dot-product DE, partly because it is non-trivial to learn the dot-product operation with an MLP given the concatenated query-document embedding as the input.Recently, there has been interest in late-interaction models. Similarly to DE models, such models also embed queries and documents separately into T 1 (q) and T 2 (d); however, they do not use pooling operations, but instead calculate dot products between all pairs of query and document token embeddings, and perform a non-linear score reduction. Formally, let us define query and document Transformer embeddings Q = (q 1 , . . . , q L1 ) := T 1 (q) ∈ R P ×L1 and D = (d 1 , . . . , d L2 ) := T 2 (d) ∈ R P ×L2 , and let S := Q D denote the similarity matrix. ColBERT (Khattab & Zaharia, 2020) then performs a non-linear sum-max reduction of S:s(q, d) = i∈[L1] max j∈[L2] q i d j .This non-linearity allows ColBERT to achieve better accuracy than DE. See (Luan et al., 2021) for a related model. Another similar approach is CEDR (MacAvaney et al., 2019) , which uses multiple query-document similarity matrices (one for each layer) from pre-trained Transformers.For each query token, instead of only using the most aligned document token, Qian et al. (2022) suggest considering the top-k aligned document tokens.Instead of using similarities between all pairs of query and document token embeddings, COIL (Gao et al., 2021) only considers pairs of query and document tokens that have the same token ID, while CITADEL (Li et al., 2022) further implements a dynamic lexical routing. Li et al. (2023) use sparse token representations that can achieve competitive accuracy compared to ColBERT while being much faster. Mysore et al. (2021) suggest using co-citations as supervision for training.Late-interaction models have precedent in the classical IR literature. For example, DRMM (Guo et al., 2016a) scores (query, document) relevance using a feedforward network on top of count histogram features. On top of the querydocument token similarity matrix based on Word2Vec, MatchPyramid (Pang et al., 2016) applies a convolutional network, while KNRM (Xiong et al., 2017) performs kernelbased pooling. ConvKNRM (Dai et al., 2018) further uses a convolutional network on top of learned token embeddings to produce contextual embeddings. There are also relevant models from the collaborative filtering literature, such as Dziugaite & Roy (2015) .Late-interaction scorers such as ColBERT may be used in both the retrieval and re-ranking phases. In this paper, we focus on the latter, which has been considered in several previous works, e.g., (Hofstätter et al., 2020; Santhanam et al., 2022; Ren et al., 2021) . While ColBERT can yield a more favourable latency versus quality trade-off compared to DE and CE models, there are two important limitations for its use in re-ranking.Limited expressivity of hand-crafted reductions. Although prior late-interaction models include more non-linearity compared with DE, they rely on hand-crafted score reductions, such as sum-max in ColBERT. It is unclear if these operations can capture arbitrary complex interactions among query and document tokens that define the true relevance.Latency and storage overhead. Compared with CE, both DE and late-interaction models reduce latency by relying on pre-computed document (token) embeddings. For DE, this requires storing a single document embedding vector (after proper pooling, cf. ( 2)), and during online inference, we need to take one dot product. Unfortunately, for lateinteraction models, the latency and storage cost can be much higher: suppose we use L 1 query embedding vectors and L 2 document embedding vectors to calculate the similarity matrix, then the storage cost is L 2 times larger than that of DE models1 , and we need to take L 1 L 2 dot products to obtain the similarity matrix. It is unclear how various ways to reduce the latency and storage cost affect the model performance.In the next section, we present LITE, a novel late-interaction scorer that addresses both aforementioned shortcomings:(1) LITE can provably approximate a broad class of ground truth scoring functions (cf. Theorem 3.1); and (2) it is more accurate than prior late-interaction methods on both in-domain and zero-shot tasks, and is amenable to latency and storage reduction with graceful degradation in model performance (cf. Section 4). 
lite scorers: 
 We now introduce LITE scorers. Let S := Q D ∈ R L1×L2 denote the similarity matrix which consists of the dot products of all query-document Transformer token embedding pairs. LITE models apply MLPs to reduce S to a scalar score. A natural option is to flatten S and then apply an MLP; we call this flattened LITE. On the other hand, in this paper we focus on another MLP model which we call separable LITE, motivated by separable convolution (Chollet, 2017) and MLP-Mixer (Tolstikhin et al., 2021) : we first apply row-wise updates to S, then column-wise updates, and then a linear projection to get a scalar score. Formally, we first calculate S , S ∈ R L1×L2 as follows: for allEQUATIONEQUATIONwhere LN, σ respectively denote layer-norm and ReLU. The final score is given by w vec(S ).Given the above definitions, it is natural to consider the expressivity of LITE. In particular, there are two fundamental questions: (1) Can we always approximate (continuous) scoring functions using LITE, even though LITE only has the similarity matrix as inputs and the original Transformer embeddings are lost? (2) Are LITE models more expressive than simpler models such as DE?We answer these questions in the following: we show that LITE models are universal approximators of continuous scoring functions (cf. Theorem 3.1), while there exists a scoring function which cannot be approximated by a simple dot-product DE (cf. Theorem 3.2).We consider the Transformer architecture described by (Yun et al., 2020) : it includes multiple encoding layers, each of them can be parameterized as A(X) + FF(A(X)), where X ∈ R P ×L denotes the input, FF denotes a feedforward network, and A(X) denotes an attention block:X + H i=1 W i o W i v XSoftmax((W i k X) (W i q X)).HereW i q , W i k , W i v ∈ R C×Pare query, key and value and projection matrices, W i o ∈ R P ×C are output projection matrices, and H, C denotes the number of heads and dimension of each head. The Softmax function is applied to each input column.A Transformer network defined in the above way is permutation-equivariant (Yun et al., 2020 , Claim 1): if we permute the input token sequence, then the output token sequence is permuted in the same way. If we want the network to distinguish between different orders of tokens, we can add a positional encoding matrix E ∈ R P ×L to the input X, and apply a Transformer network to X + E.As discussed in previous sections, in the late-interaction setting, we may need to store the whole Transformer output with shape P × L, which can be expensive. One solution is to apply a pooling function to reduce the number of tokens; we empirically study this method in Section 4.4, and in Theorem 3.1, we apply pooling functions to map the Transformer output in R P ×L to R P ×2 , i.e., a sequence of two token embeddings. We show that two query tokens and two document tokens are enough for universal approximation.Next, we define the scorers. Let F σ,n denote the set of 2-layer ReLU networks with n-dimensional inputs and a scalar output:F σ,n := z → a σ(Wz + b) ,where σ denotes the ReLU activation, z ∈ R n , W ∈ R m×n , a, b ∈ R m , and we allow m to be arbitrarily large. We first consider a class of flattened LITE scorers, including all two-layer ReLU networks on top of S that output a scalar score:F f := {S → f (vec(S))|f ∈ F σ,L1•L2 } .For separable LITE, we consider a simplified version of ( 3) and ( 4), but without loss of generality, as described blow: we first use a 2-layer ReLU network f 1 : R L2 → R to reduce every row of S to a single scalar, and thus transform S into a column vector; and then we apply another 2-layer ReLU network f 2 to reduce this column vector into a scalar. Formally,F s := {S → f 2 (f 1 (S))|f 1 ∈ F σ,L2 , f 2 ∈ F σ,L1 } ,where we let f 1 (S) ∈ R L1 denote the result of applying f 1 to every row of S. Note that F s is a subset of the function class defined by ( 3) and (4) (ignoring layer normalization).Here is our universal approximation result.Theorem 3.1 (Universal approximation with LITE). Let s : R (P ×L1)×(P ×L2) → R denote a continuous scoring function with a compact support Ω and L 1 , L 2 ≥ 2. For any F ∈ {F f , F s } and any > 0, there exist a scorer f ∈ F, and T 1 : R P ×L1 → R P ×2 and T 2 : R P ×L2 → R P ×2 , both of which consist of positional encodings, a Transformer and a pooling function, such thatΩ (f (T 1 (X) T 2 (Y)) -s(X, Y)) 2 d(X, Y) ≤ .The proof is given in Appendix B, and is based on the "contextual mapping" techniques from (Yun et al., 2020) . This result is non-trivial, since the input to LITE scorers is the similarity matrix based on only two query tokens and two document tokens; this means LITE models are universal approximators even under strong constraints on the total embedding size. In contrast, as we show in Theorem 3.2, if the total embedding size is less than P •L, then a dot-product DE can have a large approximation error.In addition to Theorem 3.1, even without positional encodings, in Theorem B.1 we show that LITE scorers are still universal approximators of arbitrary continuous scoring functions if we do not apply pooling. By contrast, without positional encodings, ColBERT can only represent permutationequivariant ground-truth scoring functions, because the summation and maximum operations do not consider the order of input tokens. It is an open question if ColBERT is a universal approximator with positional encodings.If we ask whether a dot-product DE can approximate arbitrary continuous functions, then we give a negative result.Theorem 3.2 (Limitation of DE with restricted embedding dimension). Suppose each query and document both have L ≥ 2 tokens. There exists a continuous ground-truth scoring function s supported on Ω := [0, 1] P ×L × [0, 1] P ×L , such that if O ≤ P L -1, then for any mappings h 1 , h 2 : R P ×L → R O that map queries and documents to Odimensional vectors respectively,Ω (h 1 (X) h 2 (Y) -s(X, Y)) 2 d(X, Y) ≥ 1 20 .Previously Menon et al. (2022) showed that if there is no constraint on the embedding dimension, then dot-product DE is a universal approximator of continuous functions. By contrast, here we show if the DE embedding dimension is less than P L, there could be a constant approximation error. 
experiments: 
 We now evaluate the proposed LITE scorer on a few standard information retrieval benchmarks, where we confirm that LITE significantly improves accuracy over existing DE and late-interaction methods on both in-domain and outof-domain tasks. Moreover, we show that LITE remains competitive as we reduce the latency and storage cost, and in particular, LITE can achieve higher accuracy than ColBERT with less latency and 0.25× storage cost.Datasets. We evaluate scorers on both in-domain reranking on the MS MARCO (Nguyen et al., 2016) and Natural Questions (NQ; (Kwiatkowski et al., 2019) ) datasets, and zero-shot re-ranking on the BEIR (Thakur et al., 2021) dataset.Training. For training on MS MARCO, we use the official training set of triplets (q, d + , d -), where document d + is relevant to query q while d -is irrelevant. State-of-the-art methods on MS MARCO also use hard-negative mining (Qu et al., 2021; Santhanam et al., 2022) ; however, in this paper our focus is on comparing different late-interaction scorers, and thus we simply use the original triplet training data.We use labels from a CE teacher model during training, as it has been observed that distillation can significantly improve performance (Santhanam et al., 2022; Menon et al., 2022) .For MS MARCO, we use the scores from the T2 teacher released by Hofstätter et al. (2020) . For the NQ dataset, we use a teacher model trained with 19 hard-negatives mined with BM25, following (Menon et al., 2022) . For loss functions, we try the KL loss and the margin MSE loss (see Appendix A.2 for definitions of loss functions and more details of training).Evaluation. For MS MARCO, we use the standard Dev set and the TREC DL 19 and 20 test sets (Craswell et al., 2020; 2021) . For NQ, we utilize the version of this dataset used in (Karpukhin et al., 2020) , which consists of questions, positive passages containing the correct answer, and a collection of Wikipedia passages. Re-ranking metrics are reported on the Dev query set with 200 passages containing positives, 100 BM25 hard-negatives and up to 100 random negatives, following (Menon et al., 2022) . We report MRR@10 ( Radev et al., 2002) and nDCG@10 (Järvelin & Kekäläinen, 2002) scores.For BEIR, following (Thakur et al., 2021) , we take the scorers trained on MS MARCO and evaluate zero-shot transfer performance. Specifically, we report evaluation results on the 14 public datasets. Thakur et al. ( 2021) evaluate the CE model by first retrieving 100 documents using BM25, and then calculating the nDCG@10 score for CE re-ranking; we use the same procedure.Models. For the Transformer encoder, we start from a pretrained BERT model (Turc et al., 2019) which has 6 layers and 768 token dimension. For DE and late-interaction models, we let the query encoder and document encoder share weights. We use a query sequence length of 30 and a document sequence length of 200 with the Transformer. If we use all 200 document tokens to calculate the similarity matrix S, then S ∈ R 30×200 . In some experiments the document sequence length is reduced in the end to save latency and storage cost; we will specify the details later. More hyperparameter details are given in Appendix A.1.In Table 1 , we report MRR@10 and nDCG@10 scores for different scorers on all datasets. When calculating the similarity matrix for ColBERT and LITE, we use the original sequence length (200) and token embedding dimension (768) of the Transformer encoder. We try both the KL loss and margin MSE loss and report the better results; more details can be found in Appendix A.3.On MS MARCO, the T2 teacher (Hofstätter et al., 2020) has Dev MRR@10 of 0.399. A DE student can only achieve MRR@10 of 0.355. Both ColBERT and separable LITE can significantly reduce this gap, but separable LITE is much better than ColBERT (0.393 vs. 0.383). We also train a 6-layer, 768-dimensional CE student using distillation from the T2 teacher; it has MRR@10 of 0.395, which is only slightly better than separable LITE. Moreover, on TREC DL 19 and 20 datasets, separable LITE also achieves better MRR@10 and nDCG@10 scores than ColBERT.These observations generalize to the NQ dataset as well: we find that late-interaction models are much better than DE, and separable LITE is much better than ColBERT.We also try a few ablations, including using top-k aligned document tokens instead of top-1 in ColBERT, and freezing the backbone and only fine-tuning the scorers. Separable LITE achieves better accuracy than ColBERT in all cases. See Appendix A.4 for details.4.3. Zero-shot re-ranking on BEIR 1 ) on 14 public datasets from BEIR (Thakur et al., 2021) . We also include results for the 6-layer CE model mentioned above, which is trained in the same way as other late-interaction models. We can see that separable LITE achieves better zero-shot transfer than ColBERT on 11 out of 14 datasets. CE still gives better zeroshot transfer than separable LITE, but as we show below, CE has much higher latency (cf. Table 3 ).As discussed previously, late-interaction methods may have higher latency and storage cost than DE. Suppose the Transformer encoders use L 1 query tokens and L 2 document tokens of dimension P , then DE only needs to take one dot product, while calculating the similarity matrix for lateinteraction methods requires L 1 L 2 dot products. Moreover, to save online latency, we need to pre-compute and store one P -dimensional document embedding vector for DE, while for late-interaction methods we might need to store a P × L 2 embedding matrix. This increase in storage cost is significant in industry-scale information retrieval systems, since there can be billions of documents (Zhang & Rui, 2013; Overwijk et al., 2022) .One solution is to reduce P and L 2 to some smaller P and L 2 (by projection, pooling, etc.), and then store a P × L 2 embedding matrix for each document. Correspondingly, for each query we use L 1 embedding vectors of dimension P , and to calculate the similarity matrix, we need L 1 L 2 dot products between P -dimensional vectors. This can reduce both latency and storage; below we analyze how performance drops with such reduction, and show that separable LITE remains competitive compared to ColBERT.Reducing the number of output document tokens.Here, we keep the token dimension at 768 and reduce the number of output document tokens. The Transformer encoder outputs an embedding matrix D ∈ R 768×200 of 200 token embeddings, and we try to reduce the number of tokens either by directly taking average of adjacent columns (average pooling), or by applying a trainable linear projection to every row of D. We try both methods and find that separable LITE prefers learnable projection while Col-BERT prefers average pooling. The results are shown in Reducing token dimension. Next, we fix the number of document tokens at 200, and reduce the dimension of each output token via learnable linear projections. The results are given in Figure 3 . With different token dimension, separable LITE is always more accurate than ColBERT. Achieving lower latency/storage than ColBERT using LITE. If the size of pre-computed document embedding matrix is fixed, then LITE has higher latency than ColBERT since its MLP scorer is slower than sum-max. However, since LITE is robust to embedding size reduction, it can remain more accurate than ColBERT while being more time and space efficient by using fewer document tokens. The result is shown in Table 3 , together with latency of other scorers studied before.In Table 3 , we evaluate the latency of scoring relevance between 1 query and 100 documents. For CE, we use the 6-layer distilled student and evaluate the total time to calculate the joint embeddings between the query and every document. For DE, ColBERT and separable LITE, we use models from Table 1 ; we pre-compute the document embeddings, and evaluate the query embedding generation and scoring time. For the "small separable LITE" model, we only store 50 tokens for each document, and we also use a small MLP (we let W 1 in (3) have shape (768, 50)). In Table 3, small separable LITE only uses 0.25× storage space compared with ColBERT which stores 200 document token embeddings, and it also achieves lower latency while still being much more accurate than ColBERT (0.391 vs. 0.383) . We can also see that the CE latency is 100× of the LITE latency, since CE cannot use offline pre-computation.  
conclusion: 
 In this work, we propose LITE models that can provably approximate any continuous scoring functions. We also show that LITE outperforms prior late-interaction models in both in-domain and zero-shot reranking. Moreover, LITE is still more accurate than prior methods if we try to reduce the storage cost of pre-computed document embeddings, and in particular LITE can achieve higher accuracy with less latency and storage cost. 
 
 Appendix: A.1. Hyper-parametersThe main hyperparameters for LITE are the MLP widths. For Separable LITE (cf. ( 3) and ( 4)), if the input dot-product matrix has shape L 1 × L 2 , then W 1 has shape (m 2 , L 2 ), W 2 has shape (L 2 , m 2 ), W 3 has shape (m 1 , L 1 ), and W 4 has shape (L 1 , m 1 ). In this work, we let m 1 = 360 and m 2 = 2400 in most experiments for simplicity, but we also note that much smaller widths can already give a high accuracy while also reducing the latency (cf. Table 3 ).Here we first define the loss functions used in our experiments.For simplicity, let us first consider the triplet setting, where we are given a query q, a positive document d + , and a negative document d -. Suppose the teacher score is given by t = (t + , t -), and the student score s = (s + , s -). The margin MSE loss is defined as s) ).((t + -t -) -(s + -s -)) 2 , i.e.,In our NQ experiments, we use one positive document and multiple negative documents. In this case the KL loss is defined similarly, while for the margin MSE loss we consider the margins between the positive document and every negative document. Formally, suppose there are N documents, the first one is positive while the remaining ones are negative, and let t i and s i denote the teacher and student scores for the i-th document, then we considerN i=2 ((t 1 -t i ) -(s 1 -s i )) 2 .It is also an interesting open direction to try other training frameworks, such as sRank (Zhu et al., 2023) .On the optimization algorithm, we use AdamW (Loshchilov & Hutter, 2019) with batch size 128, peak learning rate 2.8 × 10 -5 , weight decay 0.01, and 1.5 million steps. We use a linear learning rate warm up of 30000 steps, then a linear learning rate decay.Here we present results on different scorers and loss functions.First, Table 4 includes results on MS MARCO. For context, the T2 teacher (Hofstätter et al., 2020) achieves a Dev MRR@10 of 0.399. Even a CE student (with 6 layers and token dimension 768) cannot match this teacher performance: the best MRR@10 we get is 0.395.We also note that separable LITE get good results for both the KL loss and margin MSE loss, while other scorers seem to prefer only one loss. It is interesting to understand the effects of loss functions. q i d j .In other words, for each query token q i , ColBERT finds the most-aligned document embedding vector and includes their dot-product in the score. Qian et al. (2022) suggest using top-k aligned document tokens for each query token; here we try k = 2, 4, 8 on MS MARCO, but do not notice significant improvement compared with k = 1. k 1 2 4 8 MRR@10 0.383 0.378 0.380 0.382 Table 6 . Dev MRR@10 on MS MARCO with different values of k. We find that k = 1 (i.e., the original ColBERT) is better than other options we try (k = 2, 4, 8).Freezing query and document encoders. Recall that we use pretrained BERT models for query and document encoding, and moreover in all experiments above we also fine-tune the pretrained Transformers on MS MARCO and NQ. Here we explore performance of different scorers when the query and document Transformer encoders are frozen (i.e., pre-trained but not fine-tuned on MS MARCO).When the query and document encoders are frozen, ColBERT does not require any additional fine-tuning since the sum-max function does not include any weights. In this case, ColBERT can achieve Dev MRR@10 score 0.112 on MS MARCO. 3) and ( 4)), then it can achieve Dev MRR@10 score 0.188 on MS MARCO, which is much better than ColBERT.For KNRM, following (Xiong et al., 2017) , we use K = 11 kernels, where µ 1 = 0.9, µ 2 = 0.7, . . ., µ 10 = -0.9 with σ 1 = • • • = σ 10 = 0.1, and µ 11 = 1.0 with σ 11 = 10 -3 . We hold µ k and σ k fixed and only train w.We report MRR@10 and nDCG@10 scores on in-domain tasks in Table 7 . KNRM achieves similar scores to ColBERT overall, while separable LITE is more accurate than KNRM on all benchmarks. Moreover, separable LITE is much better than KNRM on zero-shot transfer: it is better than KNRM on 12 out of 14 datasets, as shown in Table 8 . Here we prove Theorem 3.1. We first restate it here and also include a universal approximation result without positional encodings.Theorem B.1 (Universal approximation with LITE). Let s : R (P ×L1)×(P ×L2) → R denote a continuous scoring function with a compact support Ω and L 1 , L 2 ≥ 2. For any F ∈ {F f , F s } and any > 0, there exists a query Transformer T 1 : R P ×L1 → R P ×L1 , a document Transformer T 2 : R P ×L2 → R P ×L2 , and a scorer f ∈ F, such thatΩ f T 1 (X) T 2 (Y) -s(X, Y) 2 d(X, Y) ≤ .Under the same conditions, there also exist positional encoding matrices E ∈ R P ×L1 and F ∈ R P ×L2 , a query Transformer T 1 : R P ×L1 → R P ×L1 and a pooling function pool 1 : R P ×L1 → R P ×2 , a document Transformer T 2 : R P ×L2 → R P ×L2 and a pooling function pool 2 : R P ×L2 → R P ×2 , and a scorer f ∈ F, such thatΩ f pool 1 (T 1 (X + E)) pool 2 (T 2 (Y + F)) -s(X, Y) 2 d(X, Y) ≤ .Our proof is based on the analysis of (Yun et al., 2020) : they showed that Transformer networks are universal approximators of continuous and compactly-supported sequence-to-sequence functions. In our case, we need to show universal approximation with the dot-product matrix; to this end, we actually need a few technical lemmas from (Yun et al., 2020) , as detailed below.Without loss of generality, we assume the support of the ground-truth scoring function is contained in [0, 1) P ×L1 ×[0, 1) P ×L2 . The first step is to replace the ground-truth scoring function s with a piece-wise constant function: let δ > 0 be small enough, and letEQUATIONThe document (token) index can be stored on disk, or in RAM. Storing in RAM significantly reduces latency, as we do not need to pay the cost of transferring embeddings from disk. Even if one were to store the index on disk, it is still of interest to reduce the total embedding size to reduce the storage and transfer cost/latency (which would scale linearly with embedding size). 


Decision: This manuscript targets the application of predicting query-document relevance and introduces a learnable yet efficient strategy to automatically trade off the latency-quality. The paper is written in good shape: it has sufficient technical novelty and extensive numerical results to fit the interest of the workshop. Authors are encouraged to conduct more ablation studies to justify the effectiveness of the proposed LITE, e.g., in terms of domain generalization.