内容来源：https://research.google/blog/bypassing-inference-bottlenecks-accelerating-complex-ai-search-with-retrieve-for-train/

内容总结：

谷歌研究团队提出“先检索后训练”框架，通过离线强化学习一次性训练轻量级扩散模型，绕开昂贵的推理时推理过程，实现高效AI搜索。该框架将抽象搜索目标转化为具体训练信号，使模型在单次并行生成中产出多样化、互补且贴合数据库的搜索结果，速度较传统自回归方法提升12至20倍。实验表明，该方法在时尚和音乐等多模态检索任务中均优于传统单查询搜索和零样本扩展基线，为专业领域集合检索提供了可扩展、数据高效的解决方案。

中文翻译：

2026年9月15日
彭城江，学生研究员；李悦琪，高级研究工程师，谷歌研究

Retrieve-for-Train框架不依赖昂贵的推理时推理，而是使用强化学习一次性训练一个轻量级扩散模型。这绕过了沉重的自回归“思考预算”，即时生成一组连贯的、专家级的AI搜索结果。

现代搜索或推荐应用越来越被期望返回一组连贯的结果，而非单一的最佳匹配。例如，当用户搜索“露营装备”时，他们不想要十个略有差异的四四人帐篷。他们想要的是一组连贯、互补的搭配，包括帐篷、睡袋、便携炉具和头灯等必备露营装备。

为此，系统采用查询扇出技术，将单个宽泛的提示词拆解为若干相关子查询，以覆盖用户潜在的兴趣点。然而，教会大语言模型动态执行数据库感知的查询分解会消耗大量思考预算。从设计上讲，零样本大语言模型是通用的自回归文本预测器；它们并未针对目标语料库特定的几何流形进行优化。因此，它们需要大量的测试时计算才能返回一组在固定数据库约束下优化高阶集合级属性（如多样性、覆盖度、互补性、连贯性）的结果。

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

当面对头脑风暴一组复杂搜索词的任务时，人们很容易想到在推理时直接部署一个标准的现成大语言模型来处理。然而，依赖通用模型进行数据库感知的查询分解会带来两个关键挑战：

Retrieve-for-Train将AI的训练视为一次离线练习，而非在用户等待时必须当场完成的考试。它不是强迫AI在每次有人输入查询时缓慢摸索好的搜索规则并消耗大量处理预算，而是离线运行一次强化学习训练程序。

该程序使用严格的奖励系统，将“确保结果多样且确实有库存”等抽象目标转化为精确的分步操作手册。一旦手册构建完成，AI便可在真实搜索中即时执行，毫无延迟。

该流程分为三个 distinct 步骤：

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为这些扇出模型实现了强化学习训练，这是一种使用组相对策略优化配合软PPO正则化的方法。

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

从定性上看，零样本大语言模型基线倾向于生成近乎同义的改写（如“波西米亚节日风格”与“波西米亚节日时尚”），导致结果冗余。Retrieve-for-Train则生成了高度多样、彼此不同的子查询（如分支到“靴子”或“蕾丝”），同时严格保持在数据库流形范围内。

直接部署我们经强化学习调优的语言模型产生了卓越的搜索质量，但它继承了标准自回归的延迟约束，并需要高昂的计算思考预算。

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

在大规模场景下，当自回归扇出延迟在大上下文批次下线性扩展至近50秒时，Retrieve-for-Train-Diffusion保持在亚秒级到数秒之间，以极低的计算成本提供生产就绪的专家级搜索。

在奖励优化过程中，我们发现了关于训练搜索扇出语言模型的一个根本性规律。如果没有多样性项，模型会迅速坍缩为生成退化的、无意义的字符串（如“行尾 行尾”），以在数学上利用数据库的向量坐标。注入几何多样性指标（Vendi分数）充当了关键的反向锚点，迫使模型进入嵌入空间中的一个稳定区域，在那里它只有像真正的搜索专家那样行动才能最大化奖励。

我们证明了，当强化学习被用作一次性的“目标转换器”而非在线推理引擎时，可以非常有效。通过将奖励驱动行为探索的重计算与最终部署模型解耦，我们的框架成功绕过了在线大语言模型部署典型的陡峭推理延迟和高计算开销。

将这些复杂的集合级行为蒸馏到轻量级扩散先验中，使生产检索系统能够有效优化多样性和对齐等高阶属性。最终，Retrieve-for-Train为专业或多模态领域的集合检索建立了一个高度可扩展、数据高效的流程，在这些领域中，人工标注的属性对齐训练对否则稀缺或获取成本高昂。更多详情请参阅论文。

英文来源：

September 15, 2026
Pengcheng Jiang, Student Researcher, and Judith Yue Li, Senior Research Engineer, Google Research
Instead of relying on expensive inference-time reasoning, the Retrieve-for-Train framework uses reinforcement learning once to train a lightweight diffusion model. This bypasses the heavy autoregressive "thinking budget" to instantly generate a cohesive, expert-level slate of AI search results.
Modern search or recommendation applications are increasingly expected to return a coherent set of results rather than a single best match. For example, when a user searches for "camping gear", they don’t want ten slight variations of four-person tents. They want a coherent, complementary slate that includes essential camping gear, such as a tent, sleeping bag, portable stove, and headlamp.
To do this, systems use a query fan-out technique that breaks a single broad prompt into several related sub-queries to cover potential user interests. However, teaching an LLM to perform database-aware query decomposition dynamically drains a massive thinking budget. By design, zero-shot LLMs are general autoregressive text predictors; they aren’t optimized to navigate the specific, geometric manifold of a target corpus. Consequently, they need extended test-time computation to return a collection of results that optimizes higher-order set-level properties (e.g., diversity, coverage, complementarity, coherence) while remaining grounded with respect to a fixed database.
In our ICML 2026 paper, “Efficient, Property-Aligned Fan-Out Retrieval via RL-Compiled Diffusion”, we address this decomposition bottleneck via a reward-to-data compilation framework. Instead of forcing the model to expend a large thinking budget at inference, our Retrieve-for-Train framework uses offline reinforcement learning (RL) to discover reward-aligned fan-outs and compile them into supervision. By distilling these optimized exploration behaviors into a lightweight diffusion retriever, we enable highly efficient, single-pass query fan-out at inference time. This achieves mathematically formulated, set-level properties without the overhead of test-time thinking tokens.
When tasked with brainstorming a complex group of search terms, it’s tempting to simply deploy a standard, off-the-shelf LLM at inference time to handle the job. However, relying on generic models for database-aware query decomposition introduces two critical challenges:
The Retrieve-for-Train treats the AI's training like an offline practice session rather than a test it has to take on the spot while a user is waiting. Instead of forcing the AI to slowly figure out the rules of a good search and drain a massive processing budget every single time someone types a query, Retrieve-for-Train runs an offline RL training program once.
This program uses a rigorous reward system to turn abstract goals like "ensure the results are diverse and actually in stock" into an exact step-by-step instruction manual. Once that manual is built, the AI can execute it instantly during a real search without delay.
The pipeline operates in three distinct steps:
The success of the Retrieve-for-Train framework hinges entirely on how we define "good" search behavior. Traditional supervised training evaluates pointwise relevance via learning to rank, scoring each retrieved item in isolation. However, a truly expert search slate is defined by non-decomposable, set-level properties. You can’t measure the diversity or complementarity of a single item; these properties only exist mathematically when evaluating the entire collection of retrieved results.
Rather than relying on ambiguous natural language instructions to enforce these fan-out properties, Retrieve-for-Train fine-tunes the 4B open-source language models (Gemma3-4B and Qwen3-4B) via reinforcement learning using a strict mathematical composite reward. For our open-ended abstract retrieval tasks, this composite reward is a weighted balance of three competing pillars:
During training, we optimize the fan-out language model against these geometric realities using group relative policy optimization (GRPO) with soft proximal policy optimization (PPO).
This specific triad of rewards is critical because they act as mutual counter-anchors. If a model is optimized purely for groundedness, it will reward-hack the system by generating degenerate, nonsensical strings that happen to mathematically map to a specific database coordinate. If alignment is added to fix the nonsense, the policy simply cheats by collapsing into repetitive paraphrases of the user's prompt.
By injecting the Vendi Score as a counter-anchor, Retrieve-for-Train effectively closes off these shortcut solutions. To achieve a high-reward state, the policy is forced into a balanced region of the embedding space where it must discover valid, strictly grounded, yet semantically distinct variations of the original intent.
To evaluate the Retrieve-for-Train framework, we used a combination of frozen, dataset-specific multimodal embedding backbones and open-source language models optimized for query expansion. We evaluated this setup across two distinct set-valued retrieval regimes:
For the multimodal embedding backbones, we conducted experiments across two domains: A large-scale fashion dataset of user-curated outfits used for text-to-image experiments (evaluated using a CLIP-based retriever), and a proprietary industrial dataset of expert-generated music playlists used for text-to-music evaluations (evaluated using MuLan).
For the language models, the query fan-out process was driven by 4B open-source models, specifically Gemma3-4B and Qwen3-4B, which were tasked with generating exactly 10 sub-queries for every single main search prompt they processed. We implemented the RL training for these fan-out models via Soft-GRPO, an approach that uses group relative policy optimization with soft PPO regularization.
Across both retrieval tasks, Retrieve-for-Train outperformed traditional single-query search, zero-shot expansion, and even the heavily optimized Best-of-N baseline.
Qualitatively, zero-shot LLM baselines tended to generate near-synonymous paraphrases (e.g., "bohemian festival style" vs. "bohemian festival fashion"), causing redundant results. Retrieve-for-Train generated highly diverse, distinct sub-queries (e.g., branching into "boots" or "lace") that remained strictly grounded within the database manifold.
Directly deploying our RL-tuned language model yielded exceptional search quality, but it inherited standard autoregressive latency constraints and demanded a high computational thinking budget.
By distilling that learned behavior into the 53.9M-parameter Retrieve-for-Train diffusion model, we successfully smashed the latency bottleneck. Because the diffusion model generates all target directions simultaneously in a single, non-autoregressive parallel pass in continuous embedding space, it delivers a massive 12 to 20 speedup over autoregressive approaches.
At scale, while autoregressive fan-out latency expands linearly to nearly 50 seconds under large context batches, Retrieve-for-Train-Diffusion stays between sub-second to a few seconds, delivering production-ready, expert-level search at a fraction of the computational cost.
During our reward optimization process, we discovered something fundamental about training a fan-out language model for search. Without a diversity term, the model quickly collapses into generating degenerate, nonsensical strings (like "line ending line ending") to mathematically exploit the vector coordinates of the database. Injecting a geometric diversity metric (the Vendi Score) acts as a vital counter-anchor, forcing the model into a stable region of the embedding space where it can only maximize its reward by acting like a true search expert.
We demonstrated that RL can be highly effective when used as a one-time "objective transducer" rather than an online inference engine. By decoupling the heavy computation of reward-driven behavior exploration from the final deployed model, our framework successfully bypasses the steep inference latency and high computational overhead typical of online LLM deployment.
Distilling these complex, set-level behaviors into a lightweight diffusion prior allows production retrieval systems to optimize for higher-order properties like diversity and alignment effectively. Ultimately, Retrieve-for-Train establishes a highly scalable, data-efficient pipeline for set retrieval in specialized or multimodal domains where human-labeled, property-aligned training pairs are otherwise scarce or costly to obtain. See the paper for more details.