Retrieved article excerpt
Open article · Retrieved 2026-09-18T12:22:40.038773+00:00
A conceptual diagram illustrating a system retrieving various camping gear items based on a user's text query.
# Bypassing inference bottlenecks: Accelerating complex AI search with 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.
## Quick links
- [Paper](https://arxiv.org/abs/2603.06397)
- Share
- Copy link
×
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.
## Why everyday AI isn't a search expert
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:
1. *Paraphrastic collapse****:*** Without database-aware optimization, zero-shot LLMs frequently suffer from [paraphrastic collapse](https://arxiv.org/html/2605.04665v2). Rather than exploring complementary facets of a topic, they tend to generate redundant, near-synonymous queries. For example, given the broad prompt "Bohemian festival style”, a standard LLM without careful prompt engineering might lazily generate "bohemian festival fashion" and "bohemian festival clothes”. This semantic looping produces a homogeneous slate of results, entirely missing the distinct, helpful semantic directions a fashion expert would identify, such as fringe jackets, crochet dresses, or suede boots.
2. *Autoregressive latency bottlenecks****:*** Standard LLMs are fundamentally constrained by sequential, autoregressive generation. To successfully decompose a complex query into complementary facets, modern models typically require a substantial thinking budget, generating hundreds of intermediate [chain-of-thought](https://blog.bluedot.org/p/faithful-chain-of-thought?utm_source=google&utm_medium=pmax&utm_campaign=FoAI_&utm_term=&utm_content=&gad_source=1&gad_campaignid=22554833691&gbraid=0AAAAA_kXnYMdypNLRSDdkci909e4FZdzE&gclid=Cj0KCQjwnbrUBhDOARIsAKKhPpe_x5PmJtrYhjcdmCSdnrM6xrWWnbjdI_8RAn6WxkMZ55NPyB399VMaAv0WEALw_wcB) (CoT) reasoning tokens (i,e., the intermediate steps or internal processing units an AI model generates before answering a complex question) to plan their expansion before outputting the actual search terms. While this deliberate reasoning is acceptable for conversational AI, it introduces a severe structural bottleneck for set-valued search (e.g., retrieving a complementary slate of results, such as fringe jackets or crochet dresses mentioned above). When a system must brainstorm a large slate of sub-queries simultaneously, the combined overhead of continuous context processing and generating extended reasoning tokens scales poorly. Even with advanced serving optimizations, this token-by-token architecture creates a latency floor that is fundamentally at odds with the sub-second response times required by a production search bar.
## The Retrieve-for-Train framework
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:
- *Fan-out language model training:* RL trains a fan-out language model to emit property-aligned sub-queries scored by a set-level property-check reward. This evaluates the entire group of results as a whole, rather than scoring each result in isolation.
- *Supervision synthesis:* The frozen fan-out language model synthesizes (query → target-set) pairs entirely offline for supervised learning, requiring no human labels.
- *Diffusive retriever training:* A compact, 53.9M-parameter diffusion model learns to map a query embedding directly to a complete set of target embeddings in one non-autoregressive pass, officially bypassing the need for text-based CoT reasoning tokens.
A three-step framework diagram for a machine learning model, detailing Fan-Out LM Training, Supervision Synthesis, and Diffusive Retriever Training.
*Overview of the Retrieve-for-Train framework.* ***Step 1:*** *trains a fan-out language model (FOLM) using RL to produce property-aligned sub-queries.* ***Step 2:*** *uses the trained FOLM to synthesize supervision data.* ***Step 3:*** *trains a diffusion-based fan-out retriever that samples content embeddings directly from query embeddings.*
### Designing for the set: The power of composite rewards
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:
- *Groundedness:* Penalizes distance to the database manifold, ensuring every generated sub-query corresponds to a real, retrievable item in the database.
- *Diversity:* Measured using the [Vendi Score](https://arxiv.org/abs/2210.02410) over the entire set of sub-queries, forcing the model to explore broad semantic breadth.
- *Alignment:* Anchors candidate sub-queries to the original broad prompt to prevent semantic drift.
### Mutual counter-anchors and soft-GRPO training
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.
## Experiments
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:
- *Open-ended abstract retrieval:* A setting where no unique ground truth exists and quality is exclusively measured by set-level properties, including diversity, query alignment, and database groundedness.
- *Weakly supervised compositional retrieval:* A setting where queries are paired with a weak reference set that serves as just one plausible realization of the query intent.
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).
## Results
### Retrieval quality and accuracy
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:/