告别昂贵的推理时思考,Retrieve-for-Train框架仅需一次强化学习训练一个轻量级扩散模型,就能绕过繁重的自回归"思考预算",即时生成连贯的专家级AI搜索结果集。

快速链接

现代搜索或推荐应用,越来越需要返回一组连贯的结果,而不仅仅是单一的最佳匹配。比如,当用户搜索"露营装备"时,他们不想要十个略有差异的四人帐篷,而是想要一套连贯、互补的装备清单,包括帐篷、睡袋、便携炉具和头灯等必需品。

为实现这一目标,系统通常采用查询扇出(query fan-out)技术,将一个宽泛的提示拆解成多个相关的子查询,以覆盖用户可能的兴趣点。然而,让大语言模型动态执行数据库感知的查询拆解,会大量消耗思考预算。从设计上看,零样本大语言模型本质上是通用的自回归文本预测器,它们并未针对目标语料库的特定几何结构进行优化。因此,为了在保持与固定数据库对齐的前提下,优化多样性、覆盖度、互补性、连贯性等高阶集合级属性,这些模型需要在推理时进行大量额外计算。

在我们发表于ICML 2026的论文《Efficient, Property-Aligned Fan-Out Retrieval via RL-Compiled Diffusion》中,我们通过一种"奖励转数据"的编译框架来解决这一拆解瓶颈。我们的Retrieve-for-Train框架并不强迫模型在推理阶段耗费大量思考预算,而是利用离线强化学习发现与奖励对齐的扇出方案,并将其编译为监督信号。通过将这些优化过的探索行为蒸馏到一个轻量级扩散检索器中,我们实现了推理阶段高效的单次查询扇出,从而在不产生测试时思考Token开销的情况下,达成了数学意义上明确定义的集合级属性。

为什么日常AI并非搜索专家

当需要构思一组复杂的搜索词时,人们很容易想到直接在推理阶段调用一个标准的现成大语言模型来完成任务。然而,依赖通用模型进行数据库感知的查询拆解,会带来两个关键挑战:

Retrieve-for-Train框架

Retrieve-for-Train将AI的训练过程视为一次离线练习,而不是用户等待时必须临场应对的实时考试。系统不会强迫AI每次用户输入查询时都要重新摸索良好搜索的规则,从而消耗巨大的处理预算,而是只运行一次离线强化学习训练程序。

该程序利用严格的奖励机制,将"确保结果多样且确实有库存"这类抽象目标,转化为精确的分步操作手册。一旦手册构建完成,AI在真实搜索场景中便可即时执行,无需延迟。

整个流程分为三个明确步骤:

Retrieve-for-Train框架概览。第一步:使用强化学习训练一个扇出语言模型(FOLM),生成与属性对齐的子查询。第二步:利用训练好的FOLM合成监督数据。第三步:训练一个基于扩散模型的扇出检索器,直接从查询嵌入采样出内容嵌入。

Retrieve-for-Train框架的成功,完全取决于我们如何定义"良好"的搜索行为。传统的监督训练通过学习排序来评估单点相关性,即孤立地为每个检索到的项目打分。然而,真正专业的搜索结果集,是由不可分解的集合级属性来定义的。你无法衡量单个项目的多样性或互补性,这些属性只有在评估整个检索结果集合时才具有数学意义。

Retrieve-for-Train没有依赖含糊的自然语言指令来强制实现这些扇出属性,而是通过强化学习,使用严格的数学复合奖励,对40亿参数的开源语言模型(Gemma3-4B和Qwen3-4B)进行微调。对于我们的开放式抽象检索任务,这一复合奖励是三个相互制衡因素的加权平衡。

在训练过程中,我们使用带软近端策略优化(soft PPO)的组相对策略优化(GRPO)算法,针对这些几何现实来优化扇出语言模型。

这组特定的奖励组合之所以关键,是因为它们相互构成制衡锚点。如果模型仅针对"扎根性"进行优化,它会通过生成一些恰好能在数学上映射到特定数据库坐标、但毫无意义的退化字符串来"钻空子"。如果为了修正这种无意义字符串而加入对齐性奖励,策略又会简单地通过对用户提示进行重复性改写来作弊。

通过引入Vendi分数作为制衡锚点,Retrieve-for-Train有效地堵住了这些捷径解法。为达到高奖励状态,策略被迫进入嵌入空间中的一个平衡区域,在这里它必须发现既严格扎根于数据库、又在语义上明显不同的原始意图变体。

实验

为评估Retrieve-for-Train框架,我们结合使用了针对特定数据集冻结的多模态嵌入骨干网络,以及针对查询扩展优化的开源语言模型。我们在两种不同的集合值检索场景下进行了评估。

在多模态嵌入骨干网络方面,我们在两个领域进行了实验:一个是用于文本转图像实验的大规模时尚数据集(收录用户精选搭配,使用基于CLIP的检索器评估),另一个是用于文本转音乐评估的专有工业数据集(收录专家生成的播放列表,使用MuLan评估)。

在语言模型方面,查询扇出过程由40亿参数的开源模型驱动,具体为Gemma3-4B和Qwen3-4B,它们的任务是为每一个处理的主搜索提示生成恰好10个子查询。我们通过Soft-GRPO方法为这些扇出模型实现了强化学习训练,该方法采用带软PPO正则化的组相对策略优化。

结果

在两项检索任务中,Retrieve-for-Train的表现均优于传统单查询搜索、零样本扩展方案,甚至优于经过大量优化的Best-of-N基线方法。

从质量上看,零样本大语言模型基线方案往往生成近乎同义的改写(例如"波西米亚音乐节风格"与"波西米亚音乐节时尚"),导致结果冗余。而Retrieve-for-Train生成的子查询高度多样且各不相同(例如分支出"靴子"或"蕾丝"等方向),同时严格保持与数据库结构的一致性。

直接部署我们经强化学习调优的语言模型,能获得极佳的搜索质量,但它继承了标准自回归模型的延迟限制,并需要较高的计算思考预算。

通过将这一习得行为蒸馏到仅有5390万参数的Retrieve-for-Train扩散模型中,我们成功打破了延迟瓶颈。由于扩散模型在连续嵌入空间中以单次非自回归并行方式同时生成所有目标方向,其速度相比自回归方法实现了12到20倍的大幅提升。

在规模化场景下,自回归扇出方案的延迟会随大批量上下文线性增长至近50秒,而Retrieve-for-Train扩散模型的延迟则始终保持在亚秒级到几秒之间,以极低的计算成本实现了生产级、专家水准的搜索效果。

Retrieve-for-Train框架(包括FOLM和扩散模型)在开放式抽象检索(OAR)和弱监督组合检索(WSCR)两类任务上,均持续超越标准搜索和零样本基线方案,在多样性、对齐性和召回率方面均实现显著提升。

在奖励优化过程中,我们发现了一个关于训练搜索扇出语言模型的根本性问题。如果没有多样性约束项,模型会很快退化为生成诸如"行末 行末"之类毫无意义的字符串,以此在数学上钻数据库向量坐标的空子。而引入几何多样性指标(Vendi分数)能起到关键的制衡锚点作用,迫使模型进入嵌入空间中的稳定区域,只有真正表现出搜索专家的行为,才能获得最大化的奖励。

结论

我们证明了强化学习作为一次性的"目标转换器"使用时,比作为在线推理引擎使用效果更好。通过将奖励驱动行为探索的繁重计算,与最终部署的模型解耦,我们的框架成功绕开了在线大语言模型部署中典型的高推理延迟和高计算开销问题。

将这些复杂的集合级行为蒸馏为一个轻量级扩散先验模型,使生产环境中的检索系统能够有效地针对多样性、对齐性等高阶属性进行优化。归根结底,Retrieve-for-Train为专业或多模态领域的集合检索建立了一条高度可扩展、数据高效的流程,这些领域通常缺乏人工标注、属性对齐的训练数据对,获取成本高昂。更多细节请参阅论文原文。

快速链接

Q&A

Q1:Retrieve-for-Train框架解决了什么问题?

A:它解决了AI搜索中查询拆解耗费大量推理时计算资源的问题。传统方法需要大语言模型在推理阶段动态思考如何拆解查询,而Retrieve-for-Train通过一次性离线强化学习训练,将这些行为提前编译进轻量级模型,推理时无需再耗费思考预算。

Q2:Retrieve-for-Train的速度比传统方法快多少?

A:通过将学习到的行为蒸馏到5390万参数的扩散模型中,Retrieve-for-Train相比自回归方法实现了12到20倍的速度提升。在大批量场景下,自回归方案延迟可达近50秒,而该框架的延迟始终保持在亚秒级到几秒之间。

Q3:为什么搜索结果的多样性对AI模型训练很重要?

A:因为如果没有多样性约束,模型很容易通过生成毫无意义但能匹配数据库坐标的字符串来"钻空子"。研究团队引入Vendi分数作为制衡指标,迫使模型在保持结果与数据库高度相关的同时,生成语义上真正不同的多样化结果。

Google