Quick Dense Retrievers Consume KALE: Post Training Kullback–Leibler Alignment of Embeddings for Asymmetrical dual encoders
Year: 2023

Abstract:
In this paper, we consider the problem of improving the inference latency of language model-based dense retrieval systems by introducing structural compression and model size asymmetry between the context and query encoders. First, we investigate the impact of pre and post-training compression on the MSMARCO, Natural Questions, TriviaQA, SQUAD, and SCIFACT, finding that asymmetry in the dual-encoders in dense retrieval can lead to improved inference efficiency. Knowing this, we introduce Kullback–Leibler Alignment of Embeddings (KALE), an efficient and accurate method for increasing the inference efficiency of dense retrieval methods by pruning and aligning the query encoder after training. Specifically, KALE extends traditional Knowledge Distillation after bi-encoder training, allowing for effective query encoder compression without full retraining or index generation. Using KALE and asymmetric training, we can generate models which exceed the performance of DistilBERT despite having 3x faster inference.

text: Proceedings of The Fourth Workshop on Simple and Efficient Natural Language Processing (SustaiNLP), pages 59–77 July 13, 2023 ©2023 Association for Computational Linguistics

heading: 1 Introduction
text: A bi-encoder-based retrieval, often called dense retrieval, is a retrieval function that leverages the vector representation of queries and documents as a proxy for relevance. Using two encoders, one for the query and one for the document, the input data is mapped into a common latent space where closeness becomes a proxy for relevance. Dense retrievers have become increasingly popular due to their ability to capture the semantic relationships between query and document terms. However, bi-encoder-based models can also be computationally expensive, particularly when dealing
∗ Corresponding author: dcampos3@illinois.edu
with large datasets. As a result, there has been a growing interest in methods for compressing these models to reduce their computational complexity without sacrificing performance. While the use of smaller models (Wang et al., 2020) has provided a path to improving model performance, compression cannot be adjusted to suit varying latency needs. In other words, a model must match latency requirements before it can be experimented with. Additionally, since bi-encoders require a complete index generation to evaluate performance iteratively compressing models and
59
retraining them can be very expensive. Seeing the bottleneck caused by trying to train compressed models for retrieval we explore approaches to compress models after training. By doing so it becomes cheaper to evaluate the impact of compression of retrieval and generate variants of many sizes. In this paper, we explore the role of asymmetry in the size of query and document encoders that leverage language models. Through experiments on several benchmarks, we demonstrate that our approach can significantly reduce the number of parameters in the bi-encoder model without sacrificing performance. As shown in figure 1, the combination of asymmetric bi-encoders and post-training KALE allows for 3x more QPS than an uncompressed bi-encoder with less than 1% loss in accuracy and nearly 5x with less than 2%. Building on the favorable implications of asymmetry for efficient inference, we introduce a compression mechanism called Kullback-Leibler Allingment of Embeddings (KALE). KALE uses an alignment of representations to compress models without requiring any form of retraining or index regeneration. To ground our approaches, we evaluate the effectiveness of KALE and asymmetry on several benchmark datasets and compare the results to existing efficient inference approaches. The following research questions drive our work:
• Is the performance of dense retrieval methods more driven by the query or document encoder size?
• Is it possible to compress query encoders without retraining and index regeneration?
• How can dense retrieval asymmetry and posttraining alignment be leveraged to improve query encoder latency?
It is in answering these questions that we deliver the following contributions:
• We present the first robust studies on the role of document-query encoder symmetry, demonstrating that the size of the document encoder dominates performance.
• We introduce and demonstrate the effectiveness of KALE, a post-training compression
and alignment approach demonstrating its effectiveness and
• We empirically demonstrate on various benchmarks how Asymmetric Compression can lead to 4.5 better QPS with 1% loss in recall accuracy at 100.

heading: 2 Related Work
text: Transformer Based Language Models such as BERT (Devlin et al., 2019) provide contextual language representations built on the Transformer architecture (Vaswani et al., 2017) which can be specialized and adapted for specific tasks and domains (Lee et al., 2020). Using contextual word representations, it becomes relatively easy to excel at a broad range of natural language processing tasks such as Question Answering, Text Classification, and sentiment analysis. Bi-Encoders, commonly called dual-encoders or dense retrievers, decompose ranking by leveraging the inner product of query and document representations to produce a relevance score for query document pairs. While not as accurate at cross-encoders (Reimers and Gurevych, 2019), they are more efficient for inference and easier to deploy. Bi-encoder document representations are query invariant, allowing them to be pre-computed and loaded into an Approximate Nearest Neighbor (ANN) such as FAISS (Johnson et al., 2019). At runtime, a query is an encoder into a latent space, and the k documents are retrieved using a nearest neighbor algorithm such as HNSW (Malkov and Yashunin, 2016). Since the entire document index has already been created the retrieval latency is limited to a single call of the query encoder. Bi-encoders commonly leverage LLM such as BERT (Devlin et al., 2019) to retrieve short passages of text leading to the task descriptor of Dense Passage Retrievers (DPR) (Karpukhin et al., 2020). Driven by their efficiency in deployment and relevance performance, DPR-based models have rapidly become the building blocks for systems doing product search (Magnani et al., 2022), open domain question answering (Karpukhin et al., 2020) and customer support (Mesquita et al., 2022). Efficient Inference study methods and models which decrease the model execution cost while
2
minimizing the losses to model performance. Knowledge Distillation (Hinton et al., 2015) is a training method where a model, called the student, learns to emulate a teacher model, which is commonly larger or better performing than the student. Unstructured pruning removes individual weights or groups of weights in a model by applying a mask or setting the weight values to 0. When paired with a sparsity-aware inference engine, it is possible to gain 3-5x speedups in inference throughput with little to no loss in accuracy (Kurtić et al., 2022). Structured pruning removes fundamental structural components in a language model, such as individual attention heads (Voita et al., 2019) or entire model layers (Sanh et al., 2019). Removing entire model layers is one of the most pervasive approaches, as latency gains are easy to realize, and pruning is straightforward. While their training regimes may differ, models like DistilBERT (Sanh et al., 2019) and TinyBERT (Jiao et al., 2020), and MiniLM (Wang et al., 2020) leverage structural pruning as ways of generation 2-10x speedups. Methods like quantization (Pouransari and Tuzel, 2020) (Zafrir et al., 2019), early exiting (Xin et al., 2020) or token pruning (Kim et al., 2021) have been effective in other NLP tasks. Still, our work primarily focuses on structured pruning and its relationship with asymmetry. We leave studying the impacts of asymmetry on these compression methods to future work. Asymmetrical deep learning broadly refers to any non-uniformity in shape or attribute of models. Traditional modeling approaches favor uniformity as it is preferable for optimization algorithms (Mihaylova and Martins, 2019), and using models for inference should match training as closely as possible (Ranzato et al., 2015) as improvements in training loss during optimization result in improvements in model performance during inference. However, this does not account for cost or latency asymmetries during usage. Kasai et al. demonstrated how the sequence-to-sequence encoder depth dominates language model performance for machine translation (Kasai et al., 2020). Tay et al. 2021 extend this work by finding a DeepNarrow which shows that for broad language modeling, it is possible to have 50% fewer parameters and a 40% faster inference with no loss in accuracy. Embedding Distillation Concurrent to our work on bi-encoder compression, Kim et al. 2023 study how distillation in embeddings leads to general compression of bi-encoders and cross-encoders (Kim et al., 2023). Our work differs from theirs as we focus on the role of asymmetry between query and document encoders and how to leverage it for improved inference efficiency.

heading: 3 Method
text: The use of representation models for retrieval begins with a document space d and a query space q where each of which is generated by some model m. Models do not need to share the same initialization, shape, or size, but their representation vectors must share size without some projection. These two models learn a notion of relevance by training to minimize the distance of positive query-document pairs as shown in equation 1 where x is a query vector and y is a document vector, and · denotes the dot product of the vectors. L = 1− x · y|x||y| (1) The query and document encoder models are commonly initialized with a pre-trained language model such as BERT. Then, using pairs of labels for positive relevance scores for queries and documents, the models are trained to minimize the distance between queries and their relevant documents (Karpukhin et al., 2020) While it is common practice to initialize the query encoder and document encoder with identical language models, this ignores the cost asymmetry of the usage patterns. The document encoder is usually only used once during a large-scale batch generation of the index. Index generation happens in a latency-insensitive environment and can easily leverage many GPUs and large batch sizes to improve efficiency. The query encoder runs every time a user issues a query, which can be irregular and sporadically. The query encoder responds to each user query independently. Thus, query encoders often use a batch size of 1 and commonly leverage small inference-optimized hardware like the T4 GPU or small CPUs.
3
Since the document encoder does not run very often, any improvement in latency produces a single fixed gain utterly dependent on the corpus size and index refresh cycle. The query encoder’s userfacing nature means latency improvements occur whenever a user queries.

heading: 3.1 Role of model symmetry with Bi-encoders
text: Since the query encoder runs many times online and the document encoder runs once, offline, we question: Is there some form of asymmetry between the query encoder and the document encoder that can be exploited? Do the two encoders need to be compressed symmetrically? To answer this question, we explore the impact on the performance of pruning the query and document encoders on the NQ passage retrieval dataset (Kwiatkowski et al., 2019). Using a BERT-base uncased model with 12 transformer encoder layers, we generate structurally pruned models with 9,6,3,2 and 1 layer. We also further pre-train the three and six-layer models using knowledge distillation, represented as 6KD and 3KD, from a 12- layer model on the Wikipedia-book corpus similar to distilBERT (Sanh et al., 2019). Then, using each of these seven models, we train dense retrieval models on the NQ passage retrieval dataset with variations of query and document models resulting in 72 variants. With each of these
models, we generate a full index and evaluate retrieval performance on the development portion of the dataset. We do not tune any parameters to avoid overfitting and to explore asymmetry without overoptimizing. Each model’s retrieval accuracy is evaluated with retrieval sets of depth 20, 100, and 200. We compare the impact of varying the encoders to the uncompressed baseline and a distilBERT model (denoted by 6db). Looking at the impact of symmetric compression
as shown in table 1, we see that the impact of compression is more pronounced with a small recall set as retrieval accuracy impact at 20 is 3x that of at 200. As shown in table 1 we observe major accuracy gains by fine-tuning the pruned model with a 4% gap between 6 and 6KD and a 8% gap between 3 and 3KD with a 4% gap for recall at 20 on the NQ dataset. Looking at the impact of asymmetry of the depth
4
of encoders as shown in table 2 and figure 2 we find there is the size of the query and document encoders cause similar impacts on retrieval accuracy. A retriever with 3 layers in the query encoder and 12 in the document encoder loses 11.9% of its retrieval accuracy and 12.55% when the sizes of the document encoder and query encoders are flipped. These asymmetric retrievers perform better than the symmetric 3-layer models, which lose 16.2% which highlights the ability to improve retrieval performance by having non-uniform compression. It is worth noting that having a larger document encoder is preferable to a larger query encoder which supports the notion that the document encoder is more important than the query encoder (Li and Lin, 2021).// Similar results can be seen with the introduction of fine-tuned three and 6-layer models as shown in table 6. Unsurprisingly, KD-optimized language models outperform non-distilled models, and any asymmetrical variant that leverages a distilled model outperforms the un-distilled variant. Without further optimization, a model with a distilled 3-layer query encoder and a 12-layer document encoder will outperform a model with symmetrical 6-layer models despite being 2x faster.

heading: 3.2 Inference Benchmarks
text: To evaluate the impact of structural pruning, we benchmark inference speeds of query encoding while varying the number of transformer layers. We perform benchmarking using an Intel Xeon Gold 6238R Processor and a T4 Nvidia GPU. For each model, we evaluate the performance on encoding 6500 queries with a batch size of one and a max context length of 32. For CPU inference, we evaluate the performance of models using the ONNX library 1, and for GPU inference, we evaluate native Pytorch inference. We repeat each run five times to ensure consistency and report the mean. Summary statistics can be found in table 3 and full results, including percentile, standard deviation, and confidence intervals, can be found in the appendix .5.
Table 2: Impact of Structural pruning before fine-tuning on Retrieval Accuracy on NQ passage retrieval dataset
layersq layersd Top 20 Impact Top 100 Impact Top 200 Impact 12 12 79.86% 0.00% 85.84% 0.00% 88.42% 0.00% 9 12 74.27% -7.00% 84.40% -1.67% 86.95% -1.66% 6 12 73.63% -7.80% 84.27% -1.83% 86.79% -1.85% 3 12 69.83% -12.55% 82.58% -3.80% 85.35% -3.48% 2 12 69.67% -12.76% 82.19% -4.25% 84.68% -4.23% 1 12 59.00% -26.12% 75.37% -12.19% 81.00% -8.39% 12 9 74.21% -7.07% 84.40% -1.67% 87.06% -1.53% 9 9 73.41% -8.08% 83.68% -2.51% 86.51% -2.16% 6 9 71.63% -10.30% 83.05% -3.25% 85.98% -2.76% 3 9 67.89% -14.98% 80.94% -5.71% 84.79% -4.10% 2 9 67.15% -15.92% 80.53% -6.19% 83.66% -5.39% 1 9 56.04% -29.83% 73.35% -14.55% 78.12% -11.65% 12 6 72.22% -9.57% 83.41% -2.83% 85.84% -2.91% 9 6 71.61% -10.33% 83.30% -2.96% 85.93% -2.82% 6 6 71.69% -10.23% 83.30% -2.96% 86.04% -2.69% 3 6 66.93% -16.20% 80.28% -6.48% 83.96% -5.04% 2 6 66.12% -17.20% 80.33% -6.42% 83.49% -5.58% 1 6 59.53% -25.46% 75.37% -12.19% 79.83% -9.71% 12 3 70.36% -11.90% 81.72% -4.80% 84.60% -4.32% 9 3 68.67% -14.01% 80.47% -6.25% 84.46% -4.48% 6 3 67.92% -14.95% 80.06% -6.74% 83.85% -5.17% 3 3 66.93% -16.20% 80.61% -6.09% 84.49% -4.45% 2 3 63.30% -20.74% 78.37% -8.71% 83.02% -6.11% 1 3 59.53% -25.46% 75.68% -11.84% 80.08% -9.43% 12 2 69.56% -12.90% 81.33% -5.25% 84.49% -4.45% 9 2 67.92% -14.95% 80.75% -5.93% 84.32% -4.64% 6 2 67.53% -15.43% 80.33% -6.42% 83.82% -5.20% 3 2 66.90% -16.23% 80.36% -6.38% 84.24% -4.73% 2 2 66.87% -16.27% 80.42% -6.32% 83.85% -5.17% 1 2 60.06% -24.80% 75.29% -12.29% 79.75% -9.80% 12 1 57.40% -28.13% 73.24% -14.68% 78.56% -11.15% 9 1 57.51% -27.99% 73.24% -14.68% 77.87% -11.94% 6 1 57.26% -28.30% 73.52% -14.35% 78.34% -11.40% 3 1 57.04% -28.58% 73.93% -13.87% 78.39% -11.34% 2 1 56.57% -29.17% 73.71% -14.13% 77.98% -11.81% 1 1 54.96% -31.18% 71.88% -16.26% 76.73% -13.22%
layers size compressed size method QPS Speedup 12 418 387 GPU 105.852 1.00 9 337 212 GPU 139.494 1.32 6 256 236 GPU 172.338 1.63 3 175 161 GPU 299.45 2.83 2 148 136 GPU 441.422 4.17 1 121 111 GPU 660.64 6.24 12 418 387 CPU 47.278 1.00 9 337 212 CPU 63.24 1.34 6 256 236 CPU 90.386 1.91 3 175 161 CPU 166.012 3.51 2 148 136 CPU 229.666 4.86 1 121 111 CPU 378.534 8.01
Table 3: Variation in model throughput according to the serving method and the number of transformer layers. Structural pruning can lead to a 6 and 8-layer performance increase on GPU and CPU and pruning a model to 3 layers allows a CPU to offer better inference performance than the GPU.

heading: 4 KL Alignment of Embeddings
text: While training asymmetric models can improve latency, it requires novel training regimes and experimentation, and existing workloads need to regenerate their entire index to take advantage of any inference speedups. Generation of the passage index can take longer than model training (Karpukhin et al., 2020), which makes regenerating a new index and retraining a model to meet changing latency requirements an inefficient experimentation pathway. Moreover, coupling asymmetry into training makes generating query encoder variants more difficult, as each encoder requires its own index and document encoder. Motivated by this bottleneck, we introduce Kullback-Leibler Allingment of Embeddings (KALE), a simple method of improving bi-encoder latency by aligning the embeddings of compressed models. KALE is applied after model training and leverages large batch sizes to make compression computationally inexpensive and independent of training. A single V100 GPU KALE can produce a compressed query encoder in less than 5 minutes. First, a bi-encoder model trains with separate query and document encoders. When training is complete, the document encoder, edocument, is frozen, and using the query encoder, eq, a structurally pruned copy, eq′ , is made. Then, using a sample of queries, the eq′ model is fine-tuned to minimize the KL divergence of their query representations as shown in equation 2. While the KL divergence is a measure of differences in probability distributions it has been applied successfully for representation alignment (Kim et al., 2023). To leverage it, we treat each of the representation vectors as a probability over a set of logits.
DKL(eq′ ∥ eq) = ∑
x∈X eq′(x) log
( eq′(x)
eq(x)
) . (2)
We explored the use of various distance functions such as cosine similarity, Manhattan distance, and the KL divergence but found little sensitivity in any metric besides KL divergence. We believe this is due to us freezing the document representations,
1https://onnx.ai/
and as a result, cosine distance allows the query embeddings to drift more than probability distribution matching methods. To explore this further, we experiment with tuning the temperature for the KL divergence and add a loss scaling factor but find a temperature of one and a scaling factor of ten to be most optimal. Additionally, we explored using a contrastive loss with random negative and hard negatives mined from the trained encoder but found no positive impact for either method. We leave further exploration of training objective improvement for future work.

heading: 4.1 Experimental Results
text: We evaluate the effectiveness of KALE by taking uncompressed BERTBASE models and pruning them with and without KALE on a variety of wellestablished passage retrieval benchmarks. First, models are trained, and indexes are generated using un-optimized BERTBASE models. Next, the document encoders are frozen, and the query encoders are structurally pruned to have 9,6,3,2 or 1 transformer layer. Finally, query encoders are aligned using KALE, and we compare the performance of compressed models by comparing the impact on retrieval accuracy at 20,100, and 200. To aid reproducibility, each model is trained using the Tevatron (Gao et al., 2022) 2 library, which makes use of hugginface’s transformers to provide a simple interface for exploring neural ranking models. Our experiments focus on the plain BERTBASE-uncased 12-layer transformer model. While never more capable models exist, the unaltered BERT model is widely used in production workloads, which our experiments seek to emulate. Our work aims not to produce the highest possible retrieval accuracy for a dense encoder. Instead, our goal is to find the role of asymmetry in biencoder models. As a result, we leverage the wellestablished parameters in all of our experiments without using an advanced methodology like contrastive or curriculum learning. There are fewer parameters for using KALE, and we deliberately do not optimize on anything but the loss between eq and eq′ . In general, higher degrees of pruning require longer training with smaller batches.
2https://github.com/texttron/tevatron
6
7
Datasets We use a wide variety of standard dense retrieval benchmarks, including MSMARCO V1.1 3 (Campos et al., 2016), NQ Passage Ranking 4 (Kwiatkowski et al., 2019), SciFact Passage Ranking 5 (Wadden et al., 2020), TriviaQA passage Ranking 6 (Joshi et al., 2017), and SQUAD Passage Ranking 7 (Rajpurkar et al., 2016). For each dataset, we evaluate performance by measuring the recall accuracy with retrieval depths of 20,100, and 200. Additionally, for the MSMARCO dataset, we also report MRR@10; for Scifact, we also report NDCG @10 and RR@10. Computational Experiments Our experimentation on fine-tuning our compressed models uses a 16 GB V100 GPU. Experiments in bi-encoder model training leverage 1 V100 for the MSMARCO and 4 for each other experiment. Due to the vast number of models and datasets we train on, each experiment happens with the same fixed seed.

heading: 4.2 Evaluating KALE
text: We compare the performance of using KALE for post-training compression in figure 3 across the five datasets and see a fairly consistent trend. When the recall set is small and the query encoders are pruned to a high degree, the impact of KALE is most visible, often driving over 50 improvements in retrieval accuracy. Additionally, using KALE allows the models to have a steady and gradual drop in recall accuracy relative to speedup instead of the sharp drop shown by the regular usage of structural pruning. Without KALE, post-training compression causes a 20-50% loss in retrieval accuracy. With the use of KALE, these losses are cut to 1-10%. In practice, this allows using one or 2-layer encoder models running with CPU-based inference with minor impacts on accuracy. We also notice a surprising performance improvement between 3 and 2-layer query encoders with and without KALE. We believe this shows the phenomena studied elsewhere: the first and last layers
3https://huggingface.co/datasets/Tevatron/msmarcopassage
4https://huggingface.co/datasets/Tevatron/wikipedia-nq 5https://huggingface.co/datasets/Tevatron/scifact 6https://huggingface.co/datasets/Tevatron/wikipedia-
trivia 7https://huggingface.co/datasets/Tevatron/wikipediasquad
do most of the work (Oh et al., 2022).

heading: 4.3 Aiding Asymmetry with KALE
text: Seeking to optimize compression further, we combine KALE with asymmetrical finetuning and evaluate the results similarly to our earlier experiments. Results on the impact of KALE and asymmetry on the five datasets on the recall accuracy at 100 can be found in table 5 where 3kd − 6kd denotes a three-layer query encoder and six-layer document encoder, 3kd − 3kd denotes dual three layer encoders. Full results and metrics for each task can be found in the appendix section .4.
First, it is immediately observable that posttraining compression via KALE performs worse than models natively designed for that size. We believe this is due to the convergence of the KALE models to have some distance from the uncompressed model because of dropout. We experimented with not using dropout in KALE, but model performance quickly suffered. Looking at the best retrieval accuracy vs. the model speedups shown in figure 4, we can see a substantial variation in the impact of compression across datasets. In tasks like SCIfacts, it is possible to get over 4x speedup while improving accuracy, while on tasks like SQuAD, even minor speedups lead to major losses in accuracy. We believe this variation is driven by the relative difficulty of each dataset, where easier tasks are more compressible than harder tasks. We believe these variations in results highlight the utility of post-training compression methods like KALE. Given the task variability in the impact of
8
compression, iteration speed and cost are essential to effectively tuning model inference speed and accuracy.

heading: 5 Limitations
text: While our work makes a broad study on how to improve model efficiency our scope is limited. Our work is limited to the usage of BERT-base and it is not clear how our compression approaches scale to more varied architectures like the sequence-tosequence models used by DocT5 (Lee et al., 2022) or more optimized models like RoBERTa (Liu et al., 2019) or compressed models like MiniLM (Wang et al., 2020).

heading: 6 Conclusion and Future Work
text: In this work, we have demonstrated how the use of asymmetry between the query and document encoders in bi-encoder models can be leveraged for improved inference efficiencies across CPUs and GPUs. Using our post-training compression framework, KALE, we can compress models up to 6x with little loss in accuracy. Compressing models without regenerating the document index or the document encoder makes it practical to have many query encoders tailored to each use case’s latency needs.
In the future, we wish to study how asymmetry in retrieval can be implemented with models which are widely different and may have different hidden sizes, such as using MiniLM for the query model and RoBERTA-Large for the document model.

heading: Future Work
text: In the future, we wish to study how asymmetry in retrieval can be implemented with models which are widely different and may have different hidden sizes, such as using MiniLM for the query model and RoBERTA-Large for the document model