arXiv: 2604.22722 · PDF

作者: Rajinder Sandhu, Di Mu, Cheng Chang, Md Shahriar Tasjid, Himanshu Rai, Maksims Volkovs, Ga Wu

单位: Layer 6 AI, Dalhousie University

主分类: cs.IR · 全部: cs.AI, cs.IR, cs.LG

命中关键词: llm, retrieval, rag, inference, serving


TL;DR

UAE 把 dense retriever 通过 distillation 对齐到 LLM 的 utility 分布,用 Utility-Modulated InfoNCE 让 bi-encoder 模仿 perplexity reduction 信号,在 QASPER 上大幅超过 BGE-Base,且比 LLM re-ranking 快 180 倍。

核心观点

  • Similarity-based dense retrieval 精度有限,utility-based LLM re-ranking 性能好但开销大、易受 perplexity 噪声影响。
  • 把 retrieval 形式化为 distribution matching 问题,直接把 graded utility 信号注入 embedding 空间。
  • 无需 test-time LLM 推理即可获得 utility-aligned 检索,兼顾高性能和可部署性。

方法

提出 Utility-Aligned Embeddings (UAE) 框架:由 perplexity reduction 构造一个 utility 分布,训练 bi-encoder 去模仿该分布,使用 Utility-Modulated InfoNCE 目标函数作为 distillation 损失。通过这种方式,把原本依赖 LLM re-ranking 的 graded utility 信号蒸馏进稠密向量空间,检索时只需普通 bi-encoder 相似度查询。

实验

基准数据集:QASPER。基线:强语义检索模型 BGE-Base,以及 efficient LLM re-ranking 方法。指标:Recall@1、MAP、Token F1,以及推理效率对比。

结果

在 QASPER 上相对 BGE-Base:Recall@1 +30.59%、MAP +30.16%、Token F1 +17.3%。相较高效 LLM re-ranking 方法,速度快 180 倍以上,同时保持有竞争力的性能。

为什么重要

对 RAG 从业者而言,UAE 提供了把生成式 utility 对齐到检索阶段的实用方案,避免线上 LLM re-ranking 的延迟和成本,又能显著改善上下文质量,适合大规模部署的 agent / RAG pipeline。

与已有工作的关系

延续 dense retrieval(如 BGE 系列)与 LLM-based re-ranking / utility retrieval 研究,借鉴 InfoNCE 对比学习与 distillation 思路,把 perplexity-based utility 信号(类似 RankGPT、UPR 风格)压缩进 embedding。

尚未回答的问题

  • 是否在 QASPER 之外(多领域、多语言、长文档)同样稳健?
  • 对不同规模 teacher LLM 以及 perplexity 噪声的敏感度如何?
  • 与 cross-encoder re-ranker 级联能否进一步提升上限?
  • utility 分布的在线更新和持续学习策略。

原始摘要(中文翻译)

稠密向量检索是 Retrieval-Augmented Generation (RAG) 的实际骨干,但相似度搜索可能存在精度上的局限。相反,基于效用(utility-based)的方法借助 LLM 重排序通常能取得更优的性能,但计算代价高昂,且容易受到 perplexity 估计中固有噪声的影响。我们提出 Utility-Aligned Embeddings (UAE),一个旨在将上述两方面优势融合为一种实用且高性能检索方法的框架。我们将检索形式化为一个分布匹配问题,训练一个 bi-encoder 去模仿由 perplexity reduction 得到的效用分布,所用目标为 Utility-Modulated InfoNCE。该方法将分级的效用信号直接注入到嵌入空间中,而在测试时不需要进行 LLM 推理。在 QASPER 基准上,UAE 相较强语义基线 BGE-Base,将检索的 Recall@1 提升 30.59%、MAP 提升 30.16%、Token F1 提升 17.3%。更关键的是,UAE 比高效的 LLM 重排序方法快 180 倍以上,同时保持具有竞争力的性能,表明将检索与生成式效用对齐可以在大规模场景下产出可靠的上下文。