Retrieve-for-Train explained: how Google Research is speeding up AI search fan-out

By Sean Chaudhary

On September 15, 2026, Google Research published a post by Pengcheng Jiang and Judith Yue Li on a method called Retrieve-for-Train. It tackles a practical problem in AI search: the smarter a system gets at breaking a question into sub-queries, the slower it becomes.

This post explains what the method does, what Google Research reports, where the limits are, and what we think it may mean for brands that want to be found and cited in AI search. If you want the fundamentals first, start with our explainer on how AI search works, from RAG to grounding.

The problem: fan-out is useful but slow

Many AI search systems do not run one search. They run several. Google documents this for AI Overviews and AI Mode as query fan-out: issuing multiple related searches across subtopics and data sources. We cover the practical side in our guide to query fan-out for GEO and AI SEO.

The Google Research team frames the harder version of this as set retrieval. Some requests do not have one right answer. They need a curated collection, such as an outfit or a playlist, where the items should be diverse, cover the request, and fit together. The paper calls these higher-order properties: diversity, coverage, complementarity and coherence.

Using a language model to plan that fan-out runs into two problems, according to the blog post:

How Retrieve-for-Train works

The core idea is to treat training as offline practice. The slow, careful reasoning happens before anyone searches. At query time, a much smaller model does the work in a single pass. The method has three steps.

1. Train a fan-out language model with reinforcement learning

   A language model learns to write sub-queries using a composite reward that scores the whole set of results, not each result alone. The reward has three parts: groundedness, which penalizes sub-queries that drift away from what actually exists in the database; diversity, measured with the Vendi Score; and alignment, which keeps sub-queries tied to the original request. The blog says training used group relative policy optimization (GRPO).

2. Synthesize training data offline

   The trained model then generates pairs of queries and target result sets on its own, with no human labels. This becomes the supervision for the next step.

3. Train a compact diffusion retriever

   A 53.9M-parameter diffusion model learns to map a query embedding straight to the full set of target embeddings in one non-autoregressive pass. It does not write sub-queries word by word, which is where the speed comes from.

The paper describing the method is titled “Efficient, Property-Aligned Fan-Out Retrieval via RL-Compiled Diffusion,” first posted to arXiv on March 6, 2026, with eleven listed authors. The blog says it appears at ICML 2026.

What Google Research reports

The team tested R4T on two tasks. The first was open-ended retrieval on a large fashion dataset of user-curated outfits, scored on set-level properties because there is no single right answer. The second was a proprietary dataset of expert-made music playlists, where queries had reference sets to compare against. The fan-out models were open-source 4B-parameter models, Gemma3-4B and Qwen3-4B, each producing 10 sub-queries per request.

The reported results, as stated by the authors:

• Quality. R4T consistently beat single-query search, zero-shot language-model expansion and a Best-of-N baseline on both tasks. Zero-shot baselines tended to produce near-synonymous paraphrases, while R4T produced distinct sub-queries that stayed grounded in the database.
• Speed. The blog reports a 12 to 20 times speedup over autoregressive approaches. The paper’s abstract describes it as reducing query-time fan-out latency by an order of magnitude.
• Latency in seconds. According to the paper, autoregressive fan-out took about 1.46 seconds even at a batch of 8 and grew to nearly 50 seconds at a batch of 1,024. The diffusion retriever took 0.07 seconds on the small batch and 4.21 seconds on the largest.

Limitations to keep in mind

This is research, and the paper is candid about its limits. Read the results with these in mind.

• Not web search. The benchmarks are fashion outfits and music playlists. The paper does not test web pages, and neither the blog nor the paper claims R4T is deployed in Google Search, AI Overviews, AI Mode or Gemini.
• Upfront cost. The paper notes that reinforcement learning needs repeated interaction with the retriever and reward computation, which adds substantial overhead for large or fast-changing databases.
• Rewards must be expressible as numbers. The method assumes preferences can be turned into a scalar reward. The authors note that qualities like creativity, novelty or cultural sensitivity may be hard to encode.
• Evaluation and generalization. Some evaluation used a language model as a judge, which can carry the judge’s biases, and results on other model backbones and media types are untested.

What it may mean for brands and AI search

What follows is AlchemyLeads’ interpretation, not a claim about how any Google product works today. Treat it as a reading of direction, not a checklist.

Fan-out is likely to get cheaper and broader. If methods like this reach production, latency stops being the main reason to keep fan-out small. More sub-queries per question means more distinct facets of a buyer’s need get their own retrieval. Brands should think in terms of the full set of questions a buyer’s prompt implies, not one head term.

Set-level selection favors distinct contributions. R4T rewards a result set for being diverse, aligned with the request and grounded in what exists. Ten near-identical pages answering the same angle add little to such a set. A page that answers one facet clearly, such as a spec comparison, a compliance question or an installation constraint, has a clearer role.

Groundedness points back to the index. The reward penalizes sub-queries that drift from what is actually in the database. In web terms, the content has to exist, be crawlable and be indexed to be retrieved at all. Google’s own guidance for AI features says the same thing in plainer terms: pages must be indexed and eligible to show with a snippet. Our guide to earning visibility across Google’s AI search surfaces covers those requirements.

Retrieval models that learn from feedback will keep changing. Systems trained offline against reward functions can be retrained as goals shift. That makes durable fundamentals, such as clear entities, original data and accurate corroboration, a safer bet than tactics tuned to one system’s current behavior. For how retrieval chains steps together, see our breakdown of agentic RAG and multi-step retrieval.

FAQ

What is Retrieve-for-Train?

Retrieve-for-Train (R4T) is a Google Research method that trains a language model with reinforcement learning to fan out diverse, grounded sub-queries, uses it to generate training data offline, and then trains a small diffusion model to reproduce that fan-out in a single fast pass at query time.

Is Retrieve-for-Train used in Google Search or AI Overviews?

Neither the Google Research blog post nor the paper says so. The experiments used fashion and music datasets. It is published research, so treat it as a signal of where retrieval research is heading rather than a description of a live product.

How much faster is Retrieve-for-Train?

Google Research reports a 12 to 20 times speedup over autoregressive fan-out. In the paper, autoregressive fan-out grew to nearly 50 seconds at the largest batch size, while the 53.9M-parameter diffusion retriever took 4.21 seconds on that batch and 0.07 seconds on the smallest.

What should brands do about it?

Nothing specific to R4T. Our view is that it reinforces existing fundamentals: make pages crawlable and indexed, answer distinct buyer questions clearly, and publish original data that other sources corroborate. Terms like fan-out, grounding and RAG are defined in our AI search glossary.

Sources

• Google Research Blog: Bypassing inference bottlenecks, accelerating complex AI search with Retrieve-for-Train (September 15, 2026)
• Jiang et al., Efficient, Property-Aligned Fan-Out Retrieval via RL-Compiled Diffusion, arXiv:2603.06397 (March 2026)
• Google Search Central: AI features and your website
• Google Search Central: Optimizing for generative AI features on Google Search