Google Research distills query fan-out into a 53.9-million-parameter retriever

Google Research introduced Retrieve-for-Train, a framework for optimizing groups of search results offline and distilling that behavior into a 53.9-million-parameter diffusion retriever. At deployment, the model produces 10 retrieval directions together in one non-autoregressive pass, without running the reinforcement-learning-tuned language model and its repeated retrieval calls for every live query.

In the researchers’ wall-clock tests, the diffusion model generated 10 directions in 0.07 seconds for a batch of eight and 4.21 seconds for a batch of 1,024. The autoregressive comparison took about 1.46 seconds and nearly 50 seconds, respectively. The authors described those measurements as a 12-to-20-times speedup; the results have not been independently reproduced and do not establish production performance.

Retrieve-for-Train targets “fan-out” retrieval: a broad request must return a set whose members work well together, not simply rank the closest individual match. Depending on the task, the set-level objective can reward properties including groundedness, diversity, alignment, coverage, complementarity and coherence.

The paper describes three stages. Reinforcement learning first tunes a fan-out language model against set-level rewards. The tuned model then creates query-to-result-set examples offline. Finally, a diffusion model learns the resulting distribution and generates continuous embedding-space directions at serving time; nearest-neighbor lookup maps those directions to database items.

The shift from live inference to offline supervision parallels a separate data-generation approach in Google Research’s ToolGrad system, which builds verified tool-use workflows before writing paired prompts. ToolGrad addresses tool-use training rather than set-valued retrieval.

The Retrieve-for-Train experiments used Gemma3-4B and Qwen3-4B as the 4-billion-parameter fan-out models. They were trained with Soft-GRPO to produce 10 subqueries for each main query, and the authors generated 128 fan-out samples per training query at a temperature of 0.9. The paper says this synthetic supervision removes the need for human-labeled, property-aligned training pairs.

The authors evaluated the method on Polyvore fashion data and a proprietary dataset of expert-generated music playlists. In open-ended retrieval, the Gemma-distilled diffusion model scored 74.3 for diversity and 37.6 for alignment on Polyvore, compared with 56.0 and 31.2 for zero-shot Gemma3-4B. On the music dataset, the respective scores were 46.7 and 59.6, versus 42.6 and 51.8. These are author-reported benchmark results, and the proprietary music data limits independent inspection.

Results were mixed on one weakly supervised Polyvore task. The Qwen-distilled model reached Recall@5K of 16.5 and Hit@5K of 57.5, above the zero-shot Qwen3-4B results of 10.1 and 33.9, while its Vendi diversity score was lower at 34.7 versus 46.4.

The researchers say the framework amortizes computation rather than eliminating it. Repeated interactions with a frozen retriever and reward calculations can impose substantial upfront training costs on extremely large or frequently changing databases. The paper also says subjective preferences can be difficult to encode as scalar rewards and notes that its open-ended evaluation partly relies on an LLM judge; it inconsistently identifies that judge as Gemini 2.5 Pro in one section and Gemini 2.5 Flash in an appendix.