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

Google Research researchers Pengcheng Jiang and Judith Yue Li published an ICML 2026 paper describing Retrieve-for-Train, a reward-to-data compilation framework that uses offline reinforcement learning to train a lightweight diffusion retriever for single-pass query fan-out. The framework replaces expensive inference-time reasoning, avoiding the large "thinking budget" that autoregressive LLMs require to produce set-level search results optimized for diversity, coverage, complementarity, and coherence against a fixed database. The work targets database-aware query decomposition in search and recommendation systems, where a single broad prompt such as "camping gear" must return a coherent slate of complementary results.

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 https://blog.google/products-and-platforms/products/search/ai-mode-search/ 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 https://www.emergentmind.com/topics/query-decomposition dynamically drains a massive thinking budget. By design, zero-shot https://www.promptingguide.ai/techniques/zeroshot LLMs are general autoregressive https://aws.amazon.com/what-is/autoregressive-models/ text predictors; they aren’t optimized to navigate the specific, geometric manifold https://medium.com/@adnan.mazraeh1993/manifold-learning-and-geometry-based-approaches-a-comprehensive-explanation-7bc33d29cc04 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 https://icml.cc/virtual/2026/poster/66354 paper, “ Efficient, Property-Aligned Fan-Out Retrieval via RL-Compiled Diffusion https://arxiv.org/abs/2603.06397 ”, 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 https://en.wikipedia.org/wiki/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 https://huggingface.co/google/gemma-3-4b-it and Qwen3-4B https://huggingface.co/Qwen/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 https://arxiv.org/abs/2505.22257 GRPO with soft proximal policy optimization https://medium.com/@kdk199604/ppo-efficient-stable-and-scalable-policy-optimization-15b5b9c74a88%5C 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 https://deepmind.google/models/gemma/gemma-3/ 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 https://openai.com/index/clip/ -based retriever , and a proprietary industrial dataset of expert-generated music playlists used for text-to-music evaluations evaluated using MuLan https://research.google/pubs/mulan-a-joint-embedding-of-music-audio-and-natural-language/ . For the language models, the query fan-out process was driven by 4B open-source models, specifically Gemma3-4B https://huggingface.co/google/gemma-3-4b-it and Qwen3-4B https://huggingface.co/Qwen/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 https://arxiv.org/abs/2511.06411 , an approach that uses group relative policy optimization https://cameronrwolfe.substack.com/p/grpo with soft PPO regularization https://medium.com/@kdk199604/ppo-efficient-stable-and-scalable-policy-optimization-15b5b9c74a88 . Across both retrieval tasks, Retrieve-for-Train outperformed traditional single-query search, zero-shot expansion, and even the heavily optimized Best-of-N baseline https://openai.com/index/measuring-goodharts-law/ . 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 https://arxiv.org/abs/2603.06397 for more details.