2026年9月15日
彭程江，学生研究员；李悦（Judith Yue Li），高级研究工程师，谷歌研究院
Retrieve-for-Train 框架不依赖昂贵的推理时推理，而是使用强化学习一次性训练一个轻量级扩散模型。这绕过了繁重的自回归“思考预算”，即时生成一组连贯、专家级的 AI 搜索结果。
现代搜索或推荐应用越来越被期望返回一组连贯的结果，而不是单一的最佳匹配。例如，当用户搜索“露营装备”时，他们不想要十个略有不同的四人帐篷。他们想要一个连贯、互补的组合，包括必要的露营装备，如帐篷、睡袋、便携式炉灶和头灯。
为此，系统使用查询扇出技术，将单个宽泛的提示分解为多个相关的子查询，以覆盖潜在的用户兴趣。然而，教 LLM 执行数据库感知的查询分解会动态消耗大量的思考预算。按照设计，零样本 LLM 是通用的自回归文本预测器；它们并未针对导航目标语料库特定的几何流形进行优化。因此，它们需要扩展的测试时计算，以返回一组优化高阶集合级属性（例如多样性、覆盖度、互补性、连贯性）的结果，同时保持相对于固定数据库的接地。
在我们 ICML 2026 的论文“通过 RL 编译扩散实现高效、属性对齐的扇出检索”中，我们通过奖励到数据的编译框架解决了这一分解瓶颈。我们的 Retrieve-for-Train 框架不是强迫模型在推理时花费大量思考预算，而是使用离线强化学习 (RL) 来发现奖励对齐的扇出并将其编译为监督。通过将这些优化的探索行为蒸馏到一个轻量级扩散检索器中，我们能够在推理时实现高效的单次查询扇出。这实现了数学公式化的集合级属性，而无需测试时思考令牌的开销。
当任务是为一个复杂的搜索词组进行头脑风暴时，很容易在推理时直接部署一个标准的、现成的 LLM 来处理这项工作。然而，依赖通用模型进行数据库感知的查询分解会带来两个关键挑战：
Retrieve-for-Train 将 AI 的训练视为一次离线练习，而不是在用户等待时必须当场进行的测试。Retrieve-for-Train 不是强迫 AI 在每次有人输入查询时缓慢地找出良好搜索的规则并消耗大量的处理预算，而是运行一次离线 RL 训练程序。
该程序使用严格的奖励系统，将“确保结果多样且确实有库存”这类抽象目标转化为精确的分步指令手册。一旦该手册构建完成，AI 便能在真实搜索过程中即时执行，无需延迟。
该流程分为三个不同的步骤：
Retrieve-for-Train 框架概览。 步骤 1：* 使用强化学习训练一个扇出语言模型（FOLM），以生成与属性对齐的子查询。* 步骤 2：* 使用训练好的 FOLM 合成监督数据。* 步骤 3：* 训练一个基于扩散的扇出检索器，直接从查询嵌入中采样内容嵌入。*
Retrieve-for-Train 框架的成功完全取决于我们如何定义“好的”搜索行为。传统的监督训练通过学习排序来评估逐点相关性，孤立地对每个检索项进行评分。然而，真正专家级的搜索列表是由不可分解的、集合级别的属性定义的。你无法衡量单个项目的多样性或互补性；这些属性只有在评估整个检索结果集合时，才在数学上存在。
Retrieve-for-Train 不依赖模糊的自然语言指令来强制实现这些扇出属性，而是通过强化学习，使用严格的数学复合奖励，对 4B 开源语言模型（Gemma3-4B 和 Qwen3-4B）进行微调。对于我们的开放式抽象检索任务，该复合奖励是三个相互竞争支柱的加权平衡：
在训练过程中，我们使用组相对策略优化（GRPO）和软近端策略优化（PPO），针对这些几何现实来优化扇出语言模型。
这一特定的三重奖励至关重要，因为它们充当相互的反向锚点。如果模型仅针对接地性进行优化，它会通过生成退化的、无意义的字符串来奖励黑客攻击系统，而这些字符串恰好能在数学上映射到特定的数据库坐标。如果加入对齐来修复无意义的问题，策略就会简单地通过坍缩为用户提示的重复性复述来作弊。
通过注入 Vendi 分数作为反向锚点，Retrieve-for-Train 有效地封堵了这些捷径解决方案。为了达到高奖励状态，策略被迫进入嵌入空间的一个平衡区域，在那里它必须发现原始意图的有效、严格接地但语义上不同的变体。
为了评估 Retrieve-for-Train 框架，我们结合使用了冻结的、数据集特定的多模态嵌入骨干网络和针对查询扩展优化的开源语言模型。我们在两种不同的集合值检索机制下评估了该设置：
对于多模态嵌入骨干网络，我们在两个领域进行了实验：一个用于文本到图像实验的大规模时尚数据集，由用户策划的服装搭配组成（使用基于 CLIP 的检索器进行评估），以及一个专有的工业数据集，由专家生成的音乐播放列表组成，用于文本到音乐评估（使用 MuLan 进行评估）。
对于语言模型，查询扇出过程由 4B 开源模型驱动，具体是 Gemma3-4B 和 Qwen3-4B，它们被要求为处理的每个主搜索提示生成恰好 10 个子查询。我们通过 Soft-GRPO 为这些扇出模型实现了 RL 训练，该方法使用组相对策略优化与软 PPO 正则化。
在两个检索任务中，Retrieve-for-Train 均优于传统的单查询搜索、零样本扩展，甚至优于经过高度优化的 Best-of-N 基线。
定性上，零样本 LLM 基线倾向于生成近乎同义的改写（例如，“波西米亚节日风格”与“波西米亚节日时尚”），导致冗余结果。Retrieve-for-Train 生成了高度多样化、独特的子查询（例如，分支到“靴子”或“蕾丝”），这些子查询严格地扎根于数据库流形内。
直接部署我们经过 RL 调优的语言模型产生了卓越的搜索质量，但它继承了标准的自回归延迟约束，并需要高计算思考预算。
通过将学习到的行为蒸馏到 53.9M 参数的 Retrieve-for-Train 扩散模型中，我们成功打破了延迟瓶颈。由于扩散模型在连续嵌入空间中通过单次非自回归并行传递同时生成所有目标方向，它比自回归方法实现了 12 到 20 倍的巨大加速。
在大规模下，虽然自回归扇出延迟在大上下文批次下线性扩展到近 50 秒，但 Retrieve-for-Train-Diffusion 保持在亚秒到几秒之间，以一小部分计算成本提供生产就绪的专家级搜索。
Retrieve-for-Train 框架（FOLM 和 Diffusion）在开放式抽象检索（OAR）和弱监督组合检索（WSCR）任务中始终优于标准搜索和零样本基线，在多样性、对齐和召回方面带来显著提升。
在我们的奖励优化过程中，我们发现了关于训练搜索扇出语言模型的一些基本问题。如果没有多样性项，模型会迅速崩溃，生成退化的、无意义的字符串（如 “行尾 行尾”），以数学方式利用数据库的向量坐标。注入几何多样性度量（Vendi Score）作为重要的反锚点，迫使模型进入嵌入空间的稳定区域，在那里它只能通过像真正的搜索专家一样行动来最大化其奖励。
我们证明了，当强化学习被用作一次性的“目标转换器”而非在线推理引擎时，可以非常高效。通过将奖励驱动行为探索的繁重计算与最终部署的模型解耦，我们的框架成功绕开了在线大语言模型部署中典型的严重推理延迟和高计算开销。
将这些复杂的集合级行为蒸馏为轻量级扩散先验，使生产检索系统能够有效优化多样性和对齐等更高阶属性。最终，Retrieve-for-Train 为专业或多模态领域的集合检索建立了一条高度可扩展、数据高效的流水线，而在这些领域中，人工标注的、属性对齐的训练对原本稀缺或获取成本高昂。更多细节请参见论文。