Bypassing inference bottlenecks: Accelerating complex AI search with Retrieve-for-Train

绕过推理瓶颈：通过检索以训练加速复杂AI搜索

2026年9月15日

姜鹏程，学生研究员，Judith Yue Li，高级研究工程师，Google Research

与依赖昂贵的推理时推理不同，检索以训练框架通过一次强化学习训练轻量级扩散模型。这绕过了耗时的自回归“思考预算”，可即时生成连贯、专家级的AI搜索结果组合。

快速链接

• 论文

• 分享 复制链接 ×

现代搜索或推荐应用越来越期望返回一组连贯的结果，而非单一的最佳匹配。例如，当用户搜索“露营装备”时，他们不希望看到四人帐篷的十种细微变体。他们希望得到包含帐篷、睡袋、便携炉具和头灯等基本露营装备的连贯互补组合。

为此，系统使用查询扩展技术，将单个宽泛提示分解为多个相关子查询以覆盖潜在用户兴趣。然而，训练大语言模型动态执行数据库感知的查询分解会消耗巨大的思考预算。按设计，零样本大语言模型是通用的自回归文本预测器；它们并未优化以导航目标语料库特定的几何流形。因此，它们需要延长测试时计算时间，才能返回在固定数据库中保持一致性的结果集合，同时优化更高阶的集合级属性（如多样性、覆盖率、互补性、连贯性）。

在我们的ICML 2026论文《通过强化学习编译的扩散模型实现高效、属性对齐的查询扩展检索》中，我们通过奖励到数据的编译框架解决这种分解瓶颈。我们检索以训练框架不强制模型在推理时消耗大量思考预算，而是使用离线强化学习（RL）发现奖励对齐的扩展，并将其编译为监督信号。通过将这些优化的探索行为蒸馏到轻量级扩散检索器中，我们在推理时实现了高度高效的单次查询扩展。这在无需测试时思考令牌开销的情况下实现了数学公式化定义的集合级属性。

• 自回归延迟瓶颈：标准的大型语言模型（LLMs）本质上受到顺序自回归生成的限制。要成功将复杂查询分解为互补的多个方面，现代模型通常需要大量的思考预算，生成数百个中间思维链（CoT）推理标记（即AI模型在回答复杂问题前生成的中间步骤或内部处理单元），以规划扩展路径后再输出实际的搜索关键词。虽然这种刻意推理对于对话式AI是可以接受的，但会为集合值搜索（例如检索上述提到的边缘夹克或钩针连衣裙等互补结果集）引入严重的结构性瓶颈。当系统需要同时头脑风暴大量子查询时，连续上下文处理和生成扩展推理标记的综合开销会显著增加。即使采用先进的服务优化，这种逐标记的架构仍会产生与生产搜索栏所需亚秒级响应时间相矛盾的根本性延迟下限。

Retrieve-for-Train 框架

Retrieve-for-Train 将AI的训练视为离线练习而非用户等待时必须实时进行的测试。与其强制AI每次用户输入查询时都缓慢推导出良好搜索的规则并消耗大量处理预算，Retrieve-for-Train 只需运行一次离线强化学习（RL）训练程序。

该程序通过严格的奖励系统，将"确保结果多样且实际有货"等抽象目标转化为精确的分步操作手册。一旦构建完成该手册，AI便可在真实搜索中立即执行而无需延迟。

该流程分为三个明确阶段：

• 扇出语言模型训练：通过强化学习训练扇出语言模型，使其生成由集合级属性检查奖励评分的属性对齐子查询。该奖励系统将对整个结果组进行整体评估，而非单独对每个结果进行评分。

• 监督合成：冻结后的扇出语言模型完全离线合成（查询→目标集）配对，无需人工标注。

• 扩散检索器训练：一个紧凑的53.9M参数扩散模型通过单次非自回归传递，直接将查询嵌入映射到完整的目标嵌入集合，正式绕过了基于文本的CoT推理标记需求。

根据image_width确定适当宽度

对于移动图像，使用默认宽度

Retrieve-for-Train框架概述。步骤1：使用强化学习训练扇出语言模型（FOLM）生成属性对齐子查询。步骤2：使用训练好的FOLM合成监督数据。步骤3：训练基于扩散的扇出检索器，直接从查询嵌入中采样内容嵌入。

为集合设计：复合奖励的力量

检索-训练框架的成功完全取决于我们如何定义“优质”搜索行为。传统监督训练通过学习排序（learning to rank）评估单点相关性，对每个检索结果独立打分。然而，真正专业的搜索结果集合应由不可分解的集合级属性定义。你无法单独衡量某个项目的多样性或互补性，这些属性只有在数学上评估整个检索结果集合时才存在。

与其依赖模糊的自然语言指令来强制这些发散属性，检索-训练框架通过强化学习，使用严格的数学复合奖励对4B开源语言模型（Gemma3-4B和Qwen3-4B）进行微调。对于我们的开放式抽象检索任务，该复合奖励是三个竞争性支柱的加权平衡：

• 扎根性：惩罚与数据库流形的距离，确保每个生成的子查询都对应数据库中的真实可检索项目。
• 多样性：通过所有子查询集的Vendi Score衡量，迫使模型探索广泛的语义广度。
• 对齐性：将候选子查询锚定到原始宽泛提示，防止语义漂移。

互为反锚点与软GRPO训练

在训练过程中，我们使用带软近端策略优化（PPO）的组相对策略优化（GRPO），将发散语言模型优化到这些几何现实。

这种三元奖励组合至关重要，因为它们互为反锚点。如果模型仅优化扎根性，它会通过生成退化、无意义的字符串来破解系统，这些字符串恰好数学上对应特定数据库坐标。如果添加对齐性来修正无意义内容，策略会通过重复用户提示的同义句来作弊。

通过注入Vendi Score作为反锚点，检索-训练框架有效关闭了这些捷径方案。要达到高奖励状态，策略被迫进入嵌入空间的平衡区域，必须发现原始意图的有效、严格扎根且语义不同的变体。

实验

为评估检索-训练框架，我们结合了冻结的、特定数据集的多模态嵌入主干网络和优化用于查询扩展的开源语言模型。我们在两种不同的集合值检索场景中评估了该设置：

• 开放式抽象检索：一种不存在唯一真实值的场景，质量仅通过集合级属性（包括多样性、查询对齐性和数据库扎根性）衡量。
• 弱监督组合检索：一种查询与弱参考集配对的场景，该参考集只是查询意图的一种可能实现。

在多模态嵌入主干网络方面，我们在两个领域进行了实验：一个用于文本到图像实验的大型时尚数据集（使用基于CLIP的检索器评估），以及一个用于文本到音乐评估的专有工业数据集（使用MuLan评估），包含专家生成的音乐播放列表。

对于语言模型而言，查询扇出过程由4B开源模型驱动，具体为Gemma3-4B和Qwen3-4B。这些模型被要求为每个主搜索提示生成恰好10个子查询。我们通过Soft-GRPO方法实现这些扇出模型的强化学习训练，该方法采用带软PPO正则化的组相对策略优化。

结果

检索质量和准确性

在两个检索任务中，Retrieve-for-Train的表现均优于传统单查询搜索、零样本扩展，甚至超越了经过高度优化的Best-of-N基线。

定性分析表明，零样本LLM基线倾向于生成近义改写（例如"波西米亚节庆风格" vs. "波西米亚节庆时尚"），导致结果冗余。而Retrieve-for-Train生成高度多样化且明确的子查询（例如分支到"靴子"或"蕾丝"），始终严格限定在数据库流形范围内。

数量级更快的推理速度

直接部署我们的RL调优语言模型可实现卓越的搜索质量，但继承了标准自回归模型的延迟约束，并需要高昂的计算资源。

通过将学到的行为蒸馏到53.9M参数的Retrieve-for-Train扩散模型中，我们成功突破了延迟瓶颈。由于扩散模型在连续嵌入空间中通过单次非自回归并行传递同时生成所有目标方向，其速度比自回归方法提升了12至20倍。

在规模化场景中，当自回归扇出延迟在大上下文批量处理时线性扩展至近50秒时，Retrieve-for-Train-Diffusion仍能保持在亚秒级到几秒之间，以不到1/10的计算成本实现生产级专家级搜索。

Retrieve-for-Train框架（FOLM和扩散模型）在开放端抽象检索（OAR）和弱监督组合检索（WSCR）任务中，始终优于标准搜索和零样本基线，在多样性、对齐性和召回率方面实现显著提升。

抗攻击锚点（消融实验见解）

在奖励优化过程中，我们发现了一个关于训练扇出语言模型的基本规律。如果没有多样性项，模型会迅速退化为生成退化、无意义的字符串（如"line ending line ending"），以数学方式利用数据库的向量坐标。注入几何多样性度量（Vendi Score）可作为关键的反向锚点，迫使模型进入嵌入空间的稳定区域，只有在此区域中，模型才能通过真正搜索专家的行为最大化奖励。

将这些复杂的集合级行为提炼为轻量级扩散先验，使生产级检索系统能够有效优化多样性、对齐性等高阶属性。最终，Retrieve-for-Train 为专业领域或多模态场景建立了高度可扩展、数据高效的集合检索流水线，这些场景中人工标注的属性对齐训练样本通常稀缺或获取成本高昂。详见论文了解更多详情。

• 标签：

• 算法与理论

• 数据挖掘与建模

• 生成式AI

其他相关文章

• 2026年9月18日 MilleMiglia：一个用于中段物流的现实实例生成器 算法与理论

• 2026年9月17日 教学的未来：赋能教师通过生成式UI创建学习互动 教育创新 · 生成式AI · 机器智能

• 2026年8月27日 行星预测引擎：通过地球AI自动化全球模型 地球AI · 生成式AI · 机器智能

×

❮

❯

两个柱状图对比了不同AI模型在任务1（OAR）和任务2（WSCR）上的评估指标，突出了Retrieve-for-Train方法的高性能表现。

一个机器学习模型的三步框架图解，详细说明了扇出式语言模型训练、监督合成和扩散检索器训练三个阶段。