Aligning Dense Retrievers with LLM Utility via DistillationAligning Dense Retrievers with LLM Utility via Distillation
For practitioners of retrieval-augmented generation, UAE offers a practical method to achieve LLM-level retrieval quality without test-time LLM inference, addressing the computational bottleneck of utility-based approaches.
The paper proposes Utility-Aligned Embeddings (UAE), a framework that trains a bi-encoder to imitate utility distributions from LLM perplexity reduction, improving retrieval Recall@1 by 30.59%, MAP by 30.16%, and Token F1 by 17.3% over BGE-Base on QASPER, while being 180x faster than LLM re-ranking methods.
Dense vector retrieval is the practical backbone of Retrieval- Augmented Generation (RAG), but similarity search can suffer from precision limitations. Conversely, utility-based approaches leveraging LLM re-ranking often achieve superior performance but are computationally prohibitive and prone to noise inherent in perplexity estimation. We propose Utility-Aligned Embeddings (UAE), a framework designed to merge these advantages into a practical, high-performance retrieval method. We formulate retrieval as a distribution matching problem, training a bi-encoder to imitate a utility distribution derived from perplexity reduction using a Utility-Modulated InfoNCE objective. This approach injects graded utility signals directly into the embedding space without requiring test-time LLM inference. On the QASPER benchmark, UAE improves retrieval Recall@1 by 30.59%, MAP by 30.16% and Token F1 by 17.3% over the strong semantic baseline BGE-Base. Crucially, UAE is over 180x faster than the efficient LLM re-ranking methods preserving competitive performance, demonstrating that aligning retrieval with generative utility yields reliable contexts at scale.