Google Research's Retrieve-for-Train speeds up AI search fan-out 12 to 20 times

Google Research published a blog post on September 15, 2026 by Pengcheng Jiang, a student researcher, and Judith Yue Li, a senior research engineer, both at Google Research, describing Retrieve-for-Train, a technique from their ICML 2026 paper 'Efficient, Property-Aligned Fan-Out Retrieval via RL-Compiled Diffusion.' The problem it addresses: modern search and recommendation systems increasingly need to return a coherent set of complementary results, not a list of near-duplicates. Their example is a search for 'camping gear': a good result set includes a tent, a sleeping bag, a portable stove and a headlamp, not ten slightly different four-person tents. Systems get there through 'query fan-out,' splitting one broad query into several sub-queries that together cover the space. Doing that well, in a way that understands the specific structure of the target database, is the hard part: general-purpose LLMs are not built for that, so making them fan out queries at inference time burns a large 'thinking budget' of extended test-time computation.

Retrieve-for-Train moves that cost offline. Instead of asking a model to reason its way to a good fan-out every time someone searches, the authors run a reinforcement learning (RL) training program once, turning set-level goals like diversity and grounding into what they describe as a step-by-step instruction manual that a model can execute instantly at query time. Concretely, they fine-tune two 4-billion-parameter open-source language models, Gemma3-4B and Qwen3-4B, with RL to generate exactly 10 sub-queries per search prompt, using an approach they call Soft-GRPO: group relative policy optimization combined with soft PPO regularization. The reward is a composite balancing pressures the authors describe as 'mutual counter-anchors': optimizing purely for staying grounded in the database lets the model reward-hack by generating nonsensical strings, such as the degenerate 'line ending line ending,' that just happen to map to a database coordinate; adding alignment to the user's query then lets it cheat a different way, by collapsing into repetitive paraphrases of that same prompt. Injecting a diversity metric, the Vendi Score, closes off both shortcuts, forcing the policy into a region of the embedding space where it can only score well by producing sub-queries that are valid, grounded and genuinely distinct from each other.

The authors tested the approach on two kinds of retrieval built on multimodal embedding backbones: text-to-image search over a large fashion dataset of user-curated outfits, scored with a CLIP-based retriever, and text-to-music search over a proprietary, expert-curated set of music playlists, scored with MuLan. Across both tasks, Retrieve-for-Train beat single-query search, zero-shot query expansion and a heavily optimized Best-of-N baseline; the post does not give a numeric margin for that improvement. Qualitatively, the zero-shot LLM baselines tended to produce near-synonymous sub-queries, such as 'bohemian festival style' next to 'bohemian festival fashion,' which just duplicated results. Retrieve-for-Train instead branched into genuinely different but still grounded sub-queries, such as 'boots' or 'lace' for a fashion search.

Running the RL-tuned language model directly already produced strong results, but it still carried the slow, costly autoregressive reasoning the whole project was trying to avoid. The authors then distilled that learned behavior into a much smaller, 53.9-million-parameter diffusion model, Retrieve-for-Train-Diffusion, which generates every sub-query direction at once, in a single non-autoregressive pass through a continuous embedding space, instead of token by token. That distillation is what produces the headline number: a 12-to-20 times speedup over the autoregressive fan-out approach. At scale, the authors report autoregressive fan-out latency growing linearly to nearly 50 seconds under large context batches, while Retrieve-for-Train-Diffusion stays between sub-second and a few seconds.

The authors frame the broader lesson as evidence that reinforcement learning works well as a one-time 'objective transducer': used once, offline, to compress a hard-to-specify goal (a diverse, complementary, on-target set of results) into a fast model, rather than run online inside every inference call. They present Retrieve-for-Train as a data-efficient recipe for building set-level retrieval systems in specialized or multimodal domains where hand-labeled, property-aligned training examples are scarce or costly to obtain. The post does not say whether Retrieve-for-Train has shipped in any Google product.

Key facts

Google Research's Retrieve-for-Train, from the ICML 2026 paper 'Efficient, Property-Aligned Fan-Out Retrieval via RL-Compiled Diffusion,' replaces live, inference-time LLM reasoning for AI-search query fan-out with a one-time offline reinforcement-learning training stage.
That trained behavior is distilled into a 53.9-million-parameter diffusion model, Retrieve-for-Train-Diffusion, which generates a full set of sub-queries in a single non-autoregressive pass.
The diffusion model delivers a 12-to-20 times speedup over autoregressive fan-out: at scale, autoregressive latency grows to nearly 50 seconds under large context batches, while the diffusion model stays between sub-second and a few seconds.
The RL stage fine-tunes two 4-billion-parameter open models, Gemma3-4B and Qwen3-4B, with a composite reward balancing groundedness, alignment to the query and diversity (via the Vendi Score) to generate exactly 10 sub-queries per search prompt.
Across text-to-image (fashion outfits, CLIP-based retriever) and text-to-music (proprietary playlists, MuLan) retrieval, Retrieve-for-Train beat single-query search, zero-shot expansion and a Best-of-N baseline, though the post gives no numeric margin for the gain.

Why it matters

Search and recommendation systems increasingly need to hand back a coherent set of results, not a ranked list of near-duplicates. Google Research's own example: someone searching 'camping gear' wants a tent, a sleeping bag, a portable stove and a headlamp, not ten slightly different four-person tents. Getting there requires 'query fan-out,' splitting one broad request into several sub-queries that together cover the space, and doing it well requires the model to understand the specific structure of the target database, not just language in general. Normally that means burning a large 'thinking budget' of test-time reasoning on every single query. Retrieve-for-Train instead pays that cost once, offline, and trains a small model to do the whole fan-out in one fast pass at query time.

Who it affects

The work targets teams building search, shopping or recommendation systems that need complementary sets of results rather than single best matches, especially in specialized or multimodal domains, such as fashion or music, where hand-labeled, property-aligned training examples are scarce or expensive to collect. It also speaks to anyone running LLM-based query expansion in production today and paying the live inference cost for it, since the paper's premise is that the cost can be moved offline instead. The post is framed strictly as research behind an ICML 2026 paper; it does not say whether Retrieve-for-Train has shipped inside any Google product.

How to use it

There is no product or download here: the post describes a training recipe from the paper 'Efficient, Property-Aligned Fan-Out Retrieval via RL-Compiled Diffusion,' not a released tool. The recipe is concrete enough to reproduce in shape: fine-tune an open-source language model (the authors used Gemma3-4B and Qwen3-4B, 4 billion parameters each) with reinforcement learning against a composite reward that balances staying grounded in the target database, staying aligned with the user's query and staying diverse across results, measured with the Vendi Score, using what the authors call Soft-GRPO, group relative policy optimization combined with soft PPO regularization. Once that model reliably produces good fan-outs, exactly 10 sub-queries per prompt in the paper's setup, its behavior is distilled into a much smaller diffusion model, 53.9 million parameters, that performs the same job in one pass instead of token by token.

How solid is it

The method comes from a peer-reviewed ICML 2026 paper, described here on Google Research's own blog; the post names the paper but gives no month, date or location for the conference itself. Evaluation covered two domains: text-to-image search over a large fashion dataset of user-curated outfits, scored with a CLIP-based retriever, and text-to-music search over an industrial playlist dataset the post explicitly calls proprietary, scored with MuLan. Across both, Retrieve-for-Train is reported to beat single-query search, zero-shot query expansion and a heavily optimized Best-of-N baseline, but the post gives no numeric margin, score or percentage for that gain, and no detail on training data size, compute budget or hardware. Neither author's affiliation is given beyond 'Google Research,' and no university is named for Pengcheng Jiang despite the 'Student Researcher' title. The qualitative result, better and faster, rests on the authors' own account rather than on independently checkable numbers.

Risks and caveats

The headline figure, a 12-to-20 times speedup, is quoted from a source that gives the number without attaching a multiplier symbol or unit to it; read it as an order-of-magnitude claim from the authors rather than an audited benchmark result. One of the two evaluation datasets is proprietary, so outside researchers cannot rerun that half of the comparison, and neither dataset is named as a public benchmark. The 10-sub-query setup and the 4-billion-parameter base models define the tested scope; the post does not say how the approach behaves at a different fan-out width or model size. And since Retrieve-for-Train is presented as research rather than as a shipped feature, none of these numbers yet describe a product anyone can actually use.