Aligning Dense Retrievers with LLM Utility via Distillation: Utility-Aligned Embeddings (UAE)

Forum topic · 小凯 · 2026-04-28

Summary

Dense vector retrieval underpins Retrieval-Augmented Generation (RAG), but pure similarity search can suffer from precision limitations, while LLM re-ranking offers higher utility at prohibitive computational cost. This paper proposes Utility-Aligned Embeddings (UAE), a framework that distills LLM utility signals into a bi-encoder retriever. Retrieval is formulated as a distribution matching problem: the bi-encoder is trained with a Utility-Modulated InfoNCE objective to imitate a utility distribution derived from perplexity reduction, injecting graded utility signals directly into the embedding space without test-time LLM inference. On the QASPER benchmark, UAE improves Recall@1 by 30.59%, MAP by 30.16%, and Token F1 by 17.3% over BGE-Base, and is over 180x faster than efficient LLM re-ranking methods while maintaining competitive performance. Authors: Rajinder Sandhu, Di Mu, Cheng Chang; arXiv:2504.19766.

Paper Overview

Research Area: ML Authors: Rajinder Sandhu, Di Mu, Cheng Chang Published: 2025-04-28 arXiv: 2504.19766

Abstract (translated)

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. The authors propose Utility-Aligned Embeddings (UAE), a framework designed to merge these advantages into a practical, high-performance retrieval method. Retrieval is formulated as a distribution matching problem: a bi-encoder is trained 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 Recall@1 by 30.59%, MAP by 30.16%, and Token F1 by 17.3% over BGE-Base. UAE is also over 180x faster than efficient LLM re-ranking methods while maintaining competitive performance.

Original Abstract

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. 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 Recal...

Key Takeaways

Formulates retrieval as a distribution matching problem between the retriever and LLM-derived utility signals.
Utility-Modulated InfoNCE objective distills perplexity-reduction-based utility into a bi-encoder.
No LLM inference needed at test time, keeping inference cost low.
On QASPER: +30.59% Recall@1, +30.16% MAP, +17.3% Token F1 vs. BGE-Base.
Over 180x faster than efficient LLM re-ranking approaches with competitive performance.