Google’s Retrieve-for-Train Turns AI Search Fan-Out Into a One-Time Training Bill

Ask a generic large language model to brainstorm search terms for a broad shopping prompt and it tends to hand back several near-identical phrasings of the same idea rather than genuinely different options. Google Research calls this “paraphrastic collapse,” and in a September 15 blog post the team gives a concrete example: prompted with “Bohemian festival style,” an off-the-shelf model without careful engineering tends to loop between “bohemian festival fashion” and “bohemian festival clothes” instead of branching out toward fringe jackets, crochet dresses, or suede boots. The results are not wrong. They are just redundant, and redundant results are the opposite of what a search or recommendation system is supposed to return when a query is broad enough to cover many valid answers at once.

That gap between what an autoregressive model naturally produces and what a genuinely diverse, database-grounded set of results requires is the subject of a new Google Research paper being presented at ICML 2026, “Efficient, Property-Aligned Fan-Out Retrieval via RL-Compiled Diffusion.” Its framework, called R4T (Retrieve-for-Train), does not try to make a language model reason its way to a better set of sub-queries at the moment a user types something. Instead, it spends reinforcement learning exactly once, offline, then compiles what that training taught into a small, fast diffusion model, so a production system doesn’t have to run the expensive RL-tuned model live to get the benefit of what it learned.

Why a Search Bar Can’t Afford to Think Every Time

The technique at the center of this problem is called query fan-out: breaking one broad prompt into several narrower sub-queries, each sent independently to a retrieval system, so the combined results form a coherent, complementary set rather than ten variations on a single idea. The paper motivates this with a search interface expanding a broad query into subtopics, a recommender assembling a diverse slate, and a bundling system retrieving complementary items, three examples of a system that has to return a whole set of results rather than one best match, all of which face the same underlying difficulty: the objective is set-valued. You cannot judge whether one retrieved item is “diverse” or “complementary” in isolation. Those properties only exist once you look at the whole collection, which makes them awkward to learn from the ordinary (query, single correct item) datasets that most retrieval systems are trained on.

Paraphrastic collapse is only half the problem. The other half is latency. To reliably decompose a complex query into complementary parts, a general-purpose model typically needs a large “thinking budget,” generating hundreds of intermediate chain-of-thought tokens before it ever outputs a usable search term. That kind of deliberate, token-by-token reasoning is exactly what autoregressive language models are built for, but it is also fundamentally at odds with the sub-second response times a production search bar has to hit. A model can be taught, through reinforcement learning, to produce better fan-outs. Actually running that RL-tuned model live, on every query, is a separate and much more expensive problem.

R4T Pays the Reinforcement Learning Bill Exactly Once

R4T’s answer is to separate “getting the behavior right” from “running the behavior cheaply.” The paper does this in three steps. First, the team trains a fan-out language model (FOLM), starting from a 4-billion-parameter open model, either Gemma3-4B or Qwen3-4B, using reinforcement learning to produce exactly 10 sub-queries for a given broad prompt. Second, that trained FOLM is used to generate its own supervision: a large batch of broad-query-to-sub-query pairs that already reflect what the reward function wants, without needing a single human label. Third, Google Research trains a much smaller, non-autoregressive diffusion model on that synthesized data, so the expensive RL-trained behavior gets compiled into a model that can generate an entire fan-out in one fast pass instead of one token at a time.

Three Rewards That Have to Check Each Other

Training the fan-out language model in step one uses Soft-GRPO, a variant of group relative policy optimization combined with soft proximal policy optimization regularization. The reward it optimizes is a weighted composite of three separate signals that Google’s post frames as “mutual counter-anchors,” because each one exists specifically to stop the others from being gamed:

• Groundedness: penalizes distance to the target database’s embedding manifold, so every generated sub-query has to correspond to something real and retrievable, not an invented phrase that merely sounds plausible.
• Diversity: measured with the Vendi Score, a reference-free diversity metric from Dan Friedman and Adji Bousso Dieng’s 2022 paper that scores an entire set of items on a scale running from 1, if every item is identical, up to the number of items in the set, if every item is maximally distinct. It is computed from the eigenvalues of a similarity matrix rather than from any fixed list of labels, which is what lets it apply to open-ended sets with no single correct answer.
• Alignment: anchors every candidate sub-query back to the original broad prompt, so the fan-out doesn’t drift into territory the user never asked about.

Google’s post is unusually candid about why all three have to work together. In ablation testing, training without the diversity term in that composite reward produced a textbook case of reward hacking: the model collapsed into repeating degenerate, meaningless strings, such as “line ending line ending,” because that gibberish still happened to land on a real coordinate inside the database’s embedding space, satisfying groundedness even though it meant nothing. Adding the diversity term was not just a nice-to-have for variety’s sake. It was the mechanism that forced the model out of that collapsed, nonsense-generating region of embedding space and into a stable region where the only way left to keep collecting reward was to behave like an actual search expert.

From a Trained Policy to a One-Pass Diffusion Retriever

The diffusion retriever trained in step three is, per the paper, a denoising transformer built on the timestep-preconditioning scheme from Karras et al.’s 2022 diffusion work. The broad query’s own embedding is injected into the denoiser through cross-attention, and the model is trained with classifier-free guidance by randomly dropping that query embedding during training. At inference, rather than reasoning token by token, the retriever solves a single probability-flow differential equation to produce a block of embeddings in one pass, then maps each resulting embedding to real content through nearest-neighbor lookup. There is no intermediate, human-readable sub-query text in this variant at all. The model goes directly from a broad prompt to a finished slate of results.

That leaves the paper with two distinct ways to deploy what the RL training learned: R4T-FOLM, the RL-tuned language model itself, still autoregressive but far better trained than a zero-shot model; and R4T-Diffusion, the distilled model, which trades away readable intermediate sub-queries for a single, fast forward pass.

What the Benchmarks Actually Show

The team evaluated R4T on two domains, using a CLIP-based image-text encoder with matryoshka representation learning for Polyvore, a large public dataset of user-curated fashion outfits, and Google’s own MuLan joint music-text embedding model for a proprietary set of expert-curated music playlists. Two tasks were tested: Open-Ended Abstract Retrieval, scored on groundedness, diversity, and alignment against a broad prompt with no single correct answer, and Weakly Supervised Compositional Retrieval, in which the team used Gemini-2.5-Pro to generate a plausible broad query for each real Polyvore outfit and then treated that outfit’s actual items as ground truth, scoring with Recall@5K, Hit@5K, and the Vendi Score.

The compositional-retrieval numbers are a useful check on what “R4T wins” actually means in practice, because the two deployment variants do not win the same way. Starting from Qwen3-4B’s zero-shot baseline (10.1 recall, 33.9 hit rate, 46.4 diversity), training it into R4T-FOLM lifts recall to 20.9 and hit rate to 64.6, the strongest score any method reaches on either metric, but its diversity score drops to 27.5, well below the untrained baseline it started from. R4T-Diffusion, built on the smaller Gemma3-4B backbone, tells a different story: it more than doubles Gemma’s own zero-shot recall (6.0 to 15.0) and hit rate (25.9 to 54.1), while posting the highest diversity score anywhere in the table, 46.2, close to Gemma’s own untrained diversity (44.2) rather than sacrificing it. R4T-Diffusion (Gemma) even lands a nearly identical hit rate to R4T-FOLM (Gemma), 54.1 versus 54.4, while running as a single forward pass and scoring meaningfully higher on diversity, 46.2 versus 40.5. Which variant to deploy, on this evidence, depends on whether a given product cares more about raw hit rate or about how varied the results actually look to the person reading them.

The Real Payoff Is Latency

The efficiency comparison is far more one-sided. Benchmarking wall-clock latency for generating a ten-way fan-out across a range of batch sizes, the paper reports that the diffusion retriever, generating all ten directions in a single forward pass, was “maintaining sub-second performance for small batches and achieving an order-of-magnitude speedup at larger batch sizes” compared with the autoregressive baselines, which generate sub-queries sequentially and invoke retrieval calls along the way. That speed difference is the entire argument behind the framework’s design: instead of treating reinforcement learning as something that has to run every time a user types a query, R4T uses it once, as what the researchers call an “objective transducer,” to compile expensive, reward-driven exploration into a cheap, deployable model.

An Academic-Industry Pairing That’s Becoming Routine

The paper lists 11 authors. Two, lead author Pengcheng Jiang and Jiawei Han, carry a University of Illinois Urbana-Champaign affiliation in the paper itself, while Google Research’s blog post credits Jiang as a student researcher working alongside Judith Yue Li, a senior research engineer at Google Research. The remaining co-authors, including longtime Google Research scientist Craig Boutilier, round out the Google side of the collaboration. It is a pattern that has become close to standard for this kind of frontier retrieval and reinforcement learning work: a PhD student runs experiments inside a corporate lab’s compute budget, a professor supervises and co-authors, and the resulting paper lands at a peer-reviewed venue, ICML 2026 in this case, alongside a company blog post the same week.

Worth noting: none of the reported results come from a live, shipping product. Both benchmarks are curated, offline datasets, Polyvore’s public fashion catalog and an unnamed proprietary music-playlist set, and neither the paper nor the blog post claims R4T is running inside any deployed Google Search or Shopping feature. What the paper’s own tables show, read past the blog post’s summary claim that R4T “consistently outperforms” on every metric, is a more mixed but still useful result: compiling reinforcement learning into a one-time training cost, rather than an always-on inference step, made set-valued AI search meaningfully faster, and better on most axes, though not automatically better on every axis for every model choice.

sxz.io has covered adjacent parts of this problem from other angles. A tutorial on detecting AI query fan-out impressions in Google Search Console looks at the same “one prompt becomes several sub-queries” mechanic from a website owner’s side rather than a model builder’s. A separate tutorial on evaluating RAG retrieval quality with precision, recall, and MRR builds the same family of metrics R4T’s own Recall@5K and Hit@5K scores are drawn from. And diffusion models showing up outside image generation fits a broader shift sxz.io has tracked before in a look at how AI image decoders are moving beyond the VAE.

Sources: Google Research’s Bypassing inference bottlenecks: Accelerating complex AI search with Retrieve-for-Train, the underlying paper Efficient, Property-Aligned Fan-Out Retrieval via RL-Compiled Diffusion (arXiv:2603.06397, also read in full via its HTML version), and Craig Boutilier’s Google Research profile page.

Featured image: real photograph of an antique mother-of-pearl and Brussels lace folding fan, fully spread, from the Los Angeles County Museum of Art’s collection (accession M.83.189.70), released by LACMA under its Public Domain High Resolution Image Available mark, chosen as a literal visual match for fanning a single broad query out into many distinct directions, cropped to 16:9 and converted to WebP.