Extremely efficient online query encoding for dense retrieval (2024)

Abstract: Existing dense retrieval systems utilize the same model architecture for encoding both the passages and the queries, even though queries are much shorter and simpler than passages. This leads to high latency of the query encoding, which is performed online and therefore might impact user experience. We show that combining a standard large passage encoder with a small efficient query encoder can provide significant latency drops with only a small decrease in quality. We offer a pretraining and training solution for multiple small query encoder architectures. Using a small transformer architecture we are able to decrease latency by up to \sim12\times , while MRR@10 on the MS MARCO dev set only decreases from 38.2 to 36.2. If this solution does not reach the desired latency requirements, we propose an efficient RNN as the query encoder, which processes the query prefix incrementally and only infers the last word after the query is issued. This shortens latency by \sim38\times with only a minor drop in quality, reaching 35.5 MRR@10 score.

## introduction information retrieval was revolutionized by semantic matching models @xcite @xcite . such models encode the corpus of passages foot_1 and the query in a shared embedding space, where retrieval is performed using an (approximated) nearest neighbors search @xcite . these models increase the quality of search results dramatically @xcite , but suffer from a large computational overhead @xcite . while training a large model and encoding the corpus is costly, this can usually be done offline once (or every couple of days/weeks) and cost is bounded by the size of the corpus. on the other hand, encoding queries is a major part of the retrieval system that is performed frequently and online, making latency an important consideration. 3 hence, cutting the latency of this component directly leads to a cut in the online-latency of the whole system. 4 today, practically all semantic retrieval models use the same architecture to embed both the corpus (passages) and the queries. knowledge distillation @xcite has been used to improve efficiency by creating smaller models. however, mainly transformer-based architectures @xcite of medium size were considered @xcite , putting a bound on the achievable latency of the query encoding. balancing between the latency and cost requirements is challenging; while sophisticated gpu implementations can run bert inference in just a few milliseconds, this hardware is very costly. this is especially problematic in an over-provisioning setting, where utilization is kept low to handle burst of traffic. further, as query encoding is run online, it is often necessary to use a batch size of 1, which also limits the gpu utilization. therefore, it is often necessary to use a cpu for query encoding, which in-turn increases the latency overhead. this challenge calls for a query-embedding solution that can balance cost and latency, while still providing quality embeddings for retrieval. the simple, yet crucial observation we make in this paper is that queries are usually very short; often just 3-5 words, and rarely exceeding 15 words. this is in contrast to passages (or documents), which consist of dozens of words or more in some settings. therefore, we argue that while large, complex models and a vast amount of training data are crucial for quality passage embedding foot_4 , for query embedding it is sufficient to use smaller, simpler models. with this observation, we propose a method to trade-off latency and quality of online query encoding for dense retrieval, reaching low latency while preserving reasonable quality. specifically, we propose training two different variations of small models for query encoding. the first straightforward option is a small, efficient transformer. this leads to impressive results, barely hurting the retrieval quality, but the decrease in online-latency is limited to 12×, reaching 2.1 milliseconds. to extend our solution and also deal with cases where a more significant decrease in online latency is needed, we propose using an rnn-based model. as mentioned above, queries tend to be short, making rnns a viable option. apart from being efficient, the rnn architecture offers another benefit for online latency. since an rnn processes tokens sequentially, the system can feed the model with the prefix of the query as it is typed by the user. when the user issues the query, the model only needs to process its last word. this method is denoted as incremental inference in this paper, and is able to further reduce online-latency. our smallest proposed model reaches a 38× drop in latency compared to the baseline, with an online latency of only 0.7 milliseconds running on a cpu, while also achieving competitive quality results. finally, for cases where online-latency is of utmost importance, we suggest a method to practically nullify the contribution of query-encoding to the online-latency at the cost of ∼ 4× rise in compute. ## design we want to train a dual encoder system composed of a small and efficient query encoder and a standard larger transformer passage encoder. we are not interested in the training procedure of the large encoder, which was already studied thoroughly. therefore, we assume one is available. we denote by t enc the large transformer encoder, and s q as the small query encoder (either a smaller transformer, or an rnn). even though the passage and query encoders cannot share all their weights due to their different sizes and archi-tectures, we opt to keep the token embeddings of both models tied. this ensures that a token has the same "meaning" in both models @xcite . this decision is further discussed in appendix a. to train the efficient query encoder we operate in stages, as detailed in this section. ## incremental inference with rnns when using an rnn to encode a token-sequence, the encoding of the prefix of tokens is independent of the rest of the tokens: this property enables incremental encoding of user queries before they are fully composed. upon query completion, encoding only the remaining part accelerates encoding and minimizes latency. however, while the model can encode tokens incrementally, the tokenization process is not independent of the prefix. for example, while hell is a prefix of hello, their token representations are not. luckily, word boundaries (e.g., a space) are not crossed by the tokenizer, so prefix encoding can immediately be applied when coming across such a boundary. when the user issues the query, only the last word has to be encoded. a single word can span multiple tokens (e.g., 'cephalosporin' consists of 5 tokens), necessitating multiple inference steps. still, in the ms marco dev set queries, the last word's token count percentiles (p50, p90, p95, and p99) are 1, 2, 3, and 4 (respectively). in complete queries the same percentiles correspond to 9, 12, 14, and 18 tokens. this suggests that by only processing the last word of a query on the critical path of inference, we can significantly reduce the latency. note, the current guinness record for fast typing is 212 words per minute foot_5 , or 283ms per average word. as our computation speed per word is significantly smaller, computations of the query prefix are done before the last word is issued. while this approach could increase overall computation time foot_6 , in most cases the latency of the critical path is more important than overall latency. section 4.1 shows that incremental encoding can vastly reduce this measure. ## experiments our models are based on the tevatron framework @xcite and therefore cocondenser is the main baseline we compare to. for completeness we also include the results of bm25, dpr @xcite and ance @xcite . we follow many previous works and train and test our methods on the ms marco dataset @xcite using m rr@10 as the main metric and r@50/1000 as complementary metrics, and on the nq dataset @xcite using r@5/20/100 as metrics. for t enc we use the pretrained version of cocondenser, trained on the mlm task in a retrieval-friendly way. s q is implemented both using an rnn model and a st model as described in section 2. we denote by rnn(ℓ, f ) an rnn model with ℓ layers and f feedforward layers. st(ℓ) is an st model with ℓ layers. ## conclusions in this paper, we point out that queries are significantly shorter and simpler than passages, suggesting that using similar architectures for both passage and query encoders might be wasteful. indeed, we show that small transformer-based query encoders improve latency with only a minor hurt to quality. we also introduce incremental inference with rnn-based encoders, and show they produce an even lower latency, better suited for cases where latency is highly constrained. again, we show this improvement in latency comes with only a small drop in the quality of the generated embeddings.