Google moves query fan-out offline with Retrieve-for-Train

Google Research has introduced Retrieve-for-Train, a framework that shifts complex query decomposition from inference to training. It uses reinforcement learning to train a language model, synthesizes retrieval targets offline, and distills the resulting behavior into a 53.9 million-parameter diffusion model for serving.

Query fan-out splits a broad request into several targeted searches. A query such as “camping gear” should retrieve a complementary set containing a tent, sleeping bag, stove, and headlamp. Conventional LLM fan-out generates each sub-query token by token, which adds latency and often produces near-duplicates.

Retrieve-for-Train replaces that runtime generation with a model that maps the original query embedding, a numerical representation of its meaning, to a complete set of target embeddings. Those vectors can then retrieve matching items from an index without producing text sub-queries or chain-of-thought tokens.

Why zero-shot fan-out collapses

The associated ICML 2026 paper identifies paraphrastic collapse as a common failure mode. A general-purpose LLM asked to expand “Bohemian festival style” may return “bohemian festival fashion” and “bohemian festival clothes.” Useful coverage would span distinct products such as fringe jackets, crochet dresses, and suede boots.

Zero-shot models lack feedback from the underlying catalogue, so their expansions may also point toward concepts with few or no retrievable items. Database-aware prompting can improve the results, but the paper reports that strong decompositions may require hundreds of reasoning tokens before the model emits any sub-query.

Autoregressive generation imposes a serial serving cost because every token depends on the preceding tokens. In the paper’s large-context batch tests, fan-out latency approached 50 seconds as the generated sequence grew.

Training carries the reasoning

1. Train the teacher. A 4 billion-parameter language model learns to generate property-aligned sub-queries through reinforcement learning. A set-level reward scores the complete group of results.
2. Synthesize supervision. The frozen teacher generates query-to-target-set pairs offline. This stage requires no human-labelled target sets.
3. Distill the behavior. A compact diffusion retriever learns to map each query embedding directly to the corresponding set of target embeddings.

The teacher models are Gemma 3 4B and Qwen3 4B. The authors optimize them with group relative policy optimization and a soft PPO-style objective, then use their outputs as supervision for the smaller serving model.

Three rewards constrain the policy

The three objectives close different reward-hacking paths. Groundedness alone can favor malformed strings that happen to map near specific database coordinates. Adding alignment can push the policy toward repetitive paraphrases. The Vendi Score rewards variation across the set while groundedness and alignment keep that variation useful.

A 53.9M-parameter serving model

Serving the reinforcement-learned teacher would preserve its token-by-token latency. Distillation moves the final workload into continuous embedding space, where the 53.9 million-parameter diffusion retriever generates all target directions together through a parallel process.

The paper reports a 12× to 20× speedup over autoregressive fan-out. Under the tested batch and context settings, the diffusion system completed requests in less than a second to a few seconds, while the autoregressive baseline reached nearly 50 seconds at the largest scale.

Fashion and music test set quality

The researchers evaluated two retrieval regimes. Open-ended abstract retrieval had no single ground-truth answer and was scored through set-level properties. Weakly supervised compositional retrieval compared generated sets with imperfect reference sets.

Retrieve-for-Train outperformed single-query retrieval, zero-shot query expansion, and the optimized Best-of-N baseline across both regimes. Best-of-N spends additional inference compute generating several candidates and selecting the highest-scoring result, so the comparison tests whether offline distillation can retain quality without that runtime sampling cost.

The compute bill moves to training

The paper describes reinforcement learning as an “objective transducer” that converts goals such as diversity, relevance, and coverage into synthetic training targets. The expensive language model performs that conversion offline, and the serving model learns the resulting mapping.

This design applies to retrieval, recommendation, and agent systems that must return coherent sets. Developers evaluating the approach should account for several implementation requirements:

• Use set-level objectives. Pointwise ranking losses score items independently and do not directly optimize complementarity across a result set.
• Align the embedding stack. The corpus, teacher rewards, synthetic targets, and diffusion retriever must operate in compatible embedding spaces.
• Budget for offline generation. Human target labels are unnecessary, but teacher inference, reward computation, and diffusion training still require substantial offline compute.
• Test reward interactions. Groundedness, alignment, and diversity constrain different degenerate solutions, so ablations should measure both retrieval quality and set composition.
• Benchmark end to end. Serving tests should include embedding generation, diffusion sampling, nearest-neighbor lookup, batching, and tail latency.

Evidence and adoption limits

The published evidence covers fashion and music datasets, fixed outputs of 10 sub-queries, and specific embedding backbones. Production behavior may change with larger catalogues, frequently updated inventories, different set sizes, or domains whose complementary relationships are harder to encode.

Google Research describes the framework and experiments in the paper but has not provided an open-source implementation. Adoption therefore requires reproducing the reinforcement-learning teacher, synthetic-data pipeline, reward functions, and diffusion retriever against a project’s own corpus and embedding index.