Quick Dense Retrievers Consume KALE: Post Training Kullback Leibler Alignment of Embeddings for Asymmetrical dual encoders
KALE uses post-training KL alignment and asymmetric pruning to accelerate dense retrieval models, outperforming DistilBERT with 3x faster inference.
Key Findings
Methodology
This work employs a bi-encoder framework, exploring structural pruning combined with asymmetric training. KALE aligns compressed query embeddings via KL divergence after training, avoiding index regeneration. Experiments on MSMARCO, NQ, TriviaQA validate that models compressed to 2 layers with KALE maintain near-original accuracy while achieving 4.5x inference speedup. The study analyzes the impact of encoder size asymmetry, revealing document encoder size as more influential. Results demonstrate that asymmetric models with KALE outperform symmetric counterparts and traditional distillation, offering a practical approach for efficient retrieval.
Key Results
- On MSMARCO, compressing to 2 layers with KALE yields less than 2% accuracy loss, while boosting QPS by 4.5 times. Similar trends observed across datasets, with models maintaining high accuracy (>88%) at significant speedups.
- Asymmetric configurations, such as 3 query layers and 12 document layers, outperform symmetric models, confirming the benefit of non-uniform compression. KALE effectively preserves performance at extreme compression levels.
- Across multiple benchmarks, models with KALE exhibit robustness, with speedups up to 6x and accuracy loss below 10%, demonstrating broad applicability.
Significance
This research addresses the critical challenge of balancing retrieval accuracy and inference efficiency in large-scale systems. By enabling post-training compression without index regeneration, KALE significantly reduces deployment costs and latency. Its ability to leverage asymmetric encoder sizes offers flexible model design, crucial for real-time applications like search engines and question answering. The approach paves the way for scalable, high-performance retrieval systems that can adapt dynamically to resource constraints, thus impacting both academia and industry profoundly.
Technical Contribution
The core innovation lies in the KL divergence-based post-training alignment mechanism, which allows effective compression of query encoders without retraining the entire index. Combining structural pruning with this alignment, the method simplifies deployment workflows and enhances inference speed. The study systematically analyzes the role of encoder asymmetry, providing theoretical and empirical evidence that non-uniform models outperform symmetric ones in retrieval tasks. This work extends existing knowledge distillation and pruning techniques, offering a novel, practical solution for large-scale dense retrieval systems.
Novelty
This is the first comprehensive study to quantify the impact of query-document encoder asymmetry in dense retrieval, introducing a KL divergence-based post-training alignment method (KALE). Unlike traditional distillation or pruning, KALE does not require retraining indices, making it highly practical. Its ability to maintain high accuracy at extreme compression levels and outperform symmetric models demonstrates a significant advancement in model efficiency. The integration of structural pruning with a simple yet effective alignment mechanism marks a new direction in model compression research.
Limitations
- Extreme compression (e.g., 1-layer query encoder) still results in notable accuracy degradation, limiting applicability in highly resource-constrained environments.
- The current approach is validated mainly on BERT-base; its effectiveness on larger models or different architectures remains to be tested.
- KL alignment relies on sufficient query samples; in low-data scenarios, alignment quality may decline, affecting performance.
Future Work
Future research will explore adaptive multi-task and multi-modal compression strategies, incorporating dynamic structure adjustment and self-supervised alignment. Extending KALE to larger models like T5 or GPT series, and multi-lingual settings, will broaden its applicability. Investigating alternative divergence measures and training objectives could further enhance compression quality. Additionally, integrating KALE with multi-task learning frameworks may improve robustness and generalization, enabling deployment in diverse real-world scenarios.
AI Executive Summary
Dense retrieval systems based on bi-encoders have revolutionized information access, offering high efficiency and scalability. However, their large model sizes pose significant challenges for deployment, especially in latency-sensitive environments. Traditional compression techniques like knowledge distillation and structured pruning can reduce parameters but often at the cost of performance loss and require costly index regeneration. This paper introduces KALE, a novel post-training KL divergence alignment method that enables effective compression of query encoders without retraining indices. By structurally pruning query models and aligning their embeddings with the original via KL divergence, KALE achieves remarkable speedups—up to 4.5 times—while maintaining near-original accuracy across multiple benchmarks such as MSMARCO, NQ, and TriviaQA. The key insight is that asymmetric models—smaller query encoders paired with larger document encoders—perform better than symmetric counterparts, a finding validated through extensive experiments. This approach not only simplifies deployment workflows but also offers a flexible framework adaptable to various hardware constraints. The results demonstrate that even heavily compressed models can outperform traditional distilled models like DistilBERT, marking a significant step toward scalable, efficient dense retrieval. While promising, the method faces limitations at extreme compression levels, and further work is needed to generalize across architectures and modalities. Nonetheless, KALE’s simplicity and effectiveness position it as a transformative tool for real-world search and question-answering systems, promising faster responses, lower costs, and broader accessibility in the era of large-scale NLP applications.
Deep Dive
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.