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

Forum topic · 小凯 · 2026-04-28

Summary

This paper introduces Utility-Aligned Embeddings (UAE), a framework that distills LLM-based utility signals into dense retrievers for Retrieval-Augmented Generation (RAG). While dense vector retrieval is the practical backbone of RAG, similarity search can suffer from precision limitations, and LLM re-ranking approaches achieve better performance but are computationally prohibitive. UAE formulates 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 injects 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 while maintaining competitive performance. Authors: Rajinder Sandhu, Di Mu, Cheng Chang (arXiv 2504.19766).

Paper Overview

Field: Machine Learning Authors: Rajinder Sandhu, Di Mu, Cheng Chang Published: 2025-04-28 arXiv: 2504.19766

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. The authors propose Utility-Aligned Embeddings (UAE), a framework designed to merge these advantages into a practical, high-performance retrieval method.

UAE formulates 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.

Key Results

On the QASPER benchmark, UAE improves over BGE-Base by:
+30.59% Recall@1
+30.16% MAP
+17.3% Token F1
UAE is over 180x faster than efficient LLM re-ranking methods while maintaining competitive performance.