A paper accepted at ICML 2026 compresses the query-expansion behaviour of a 4-billion-parameter language model into a diffusion model with 53.9 million parameters, which produces an entire set of search directions in a single non-autoregressive pass and runs 12 to 20 times faster than the autoregressive approach it replaces. The method is called Retrieve-for-Train, and the paper is "Efficient, Property-Aligned Fan-Out Retrieval via RL-Compiled Diffusion." The claim is architectural rather than incremental: reinforcement learning runs once, offline, to work out what good retrieval looks like, and the answer is compiled into weights instead of being rediscovered at every user query.
The problem it addresses is set retrieval. Search and recommendation systems increasingly have to return a coherent group of results rather than one best match. A user shopping for camping gear does not want ten nearly identical four-person tents; they want a tent, a sleeping bag, a portable stove and a headlamp. Systems handle this with query expansion, breaking a broad query into several related sub-queries that cover the user's likely interests.
Doing that well against a particular catalogue is where general-purpose models struggle. A model working without additional examples is an autoregressive text predictor; it has not been optimised to navigate the geometry of any specific corpus, so it burns compute at test time trying to produce a set with the right properties — diversity, coverage, complementarity, coherence — while staying inside a fixed database.
The authors name two concrete failures. The first is semantic collapse. Given the broad query "bohemian festival style," a zero-shot model without careful prompt engineering is liable to return "bohemian festival fashion" and "bohemian festival clothing" — near-synonyms that pin the search to one meaning and produce a homogeneous result set. A fashion expert would have gone somewhere else entirely: fringe jackets, crochet dresses, suede boots.
The second is latency. To decompose a complex query into complementary aspects, current models spend a large reasoning budget and emit hundreds of intermediate chain-of-thought tokens before producing any actual search terms. For a conversational system that is tolerable. For set retrieval it is a structural ceiling: generating many sub-queries at once multiplies both context processing and sequential token generation, and serving optimisations do not remove a delay imposed by the architecture. A production search box generally needs an answer in under a second.
Retrieve-for-Train moves that work from query time to offline training, in three stages. First, a query-expansion model is tuned with reinforcement learning to produce sub-queries that satisfy properties of the set as a whole; the reward scores the group, not each result individually. Second, the frozen expansion model generates query-to-target-set pairs offline as supervised training data, with no human annotation. Third, a compact diffusion model with 53.9 million parameters learns to map a query embedding directly to the full set of target embeddings in one non-autoregressive pass, which removes the need for chain-of-thought tokens altogether.
Source: research.google
The reward is where the design lives. Ordinary supervised training scores relevance one result at a time, as in learning to rank. But an expert's set is defined by properties that do not decompose into individual items — diversity and complementarity only mean anything when you evaluate the whole collection. So instead of vague natural-language instructions, the authors fine-tune open 4-billion-parameter models, Gemma3-4B and Qwen3-4B, against a strict composite mathematical reward.
For open-ended abstract retrieval that reward combines three competing terms. Database grounding penalises distance from the database manifold, so every sub-query corresponds to something actually findable. Diversity is measured across the whole set of sub-queries with the Vendi Score, pushing the model to explore a wide spread of semantic directions. Query alignment keeps candidate sub-queries inside the original broad intent so they do not wander onto another topic. Training uses Soft-GRPO, which pairs group relative policy optimisation with soft PPO regularisation.
The three terms are needed simultaneously because they restrain each other. Optimise for database grounding alone and the model starts hunting for meaningless strings that happen to land on useful coordinates. Add query alignment and it evades the constraint by falling back into repeated paraphrases of the user's query. The Vendi Score closes that route: to score well the model has to find a balanced region of the embedding space and generate variants that are simultaneously valid, grounded in the catalogue and semantically distinct.
The experiments used fixed multimodal embedding models tuned to particular datasets, across two retrieval regimes. In open-ended abstract retrieval there is no single correct answer, so quality is judged purely on properties of the set: diversity, query alignment and database grounding. In compositional retrieval with weak supervision, a query is paired with a weak reference set that represents only one possible realisation of the user's intent. Two domains were tested: a large collection of user looks for text-to-image experiments, retrieved with CLIP, and a closed industrial collection of expert-built music playlists for text-to-music, retrieved with MuLan. In every case the expansion models — Gemma3-4B and Qwen3-4B — had to produce exactly ten sub-queries per main query.
On both task types Retrieve-for-Train beat single-query retrieval, zero-shot expansion and a carefully optimised Best-of-N baseline. Where zero-shot models produced "bohemian festival style" and "bohemian festival fashion," Retrieve-for-Train moved to "boots" or "lace" while staying on the database manifold.
Applying the RL-trained language model directly gave good retrieval quality but kept the usual autoregressive constraints and the large reasoning budget. Distilling that behaviour into the 53.9-million-parameter diffusion model, which generates all target directions at once in a single parallel pass through continuous embedding space, produced the 12-to-20-times speedup. Under scaling, the authors report that autoregressive expansion latency grows linearly and approaches 50 seconds at large context batches, while Retrieve-for-Train-Diffusion stays between fractions of a second and a few seconds.
Source: research.google
The most honest passage in the paper is the reward-hacking section. Stripped of its diversity term, the expansion model quickly degenerated into nonsense strings — "end of line end of line" — that mathematically hit the right database coordinates. The Vendi Score functions as the counterweight that keeps it in a stable region of embedding space, where the only way to maximise reward is to behave like an actual search expert.
That finding is also the warning label, and it complicates the pipeline's cleanest selling point. "No human annotation required" is true of the training pairs and false of the system: the human work moved into the reward function, which is three hand-chosen terms, weighted against each other, whose failure mode the authors had to discover empirically by watching the thing break. That is specification work, and it has to be redone for each new corpus and each new definition of a good set. The paper has not eliminated the labelling cost so much as converted it into a smaller quantity of much scarcer labour.
Two other things deserve a sceptical eye. The parameter ratio is roughly 74 to 1 — 4 billion down to 53.9 million — but the speedup is 12 to 20 times, which is the tell that the win comes from dropping autoregression rather than from shrinking the model; a small autoregressive expander would not have bought this. And both evaluation regimes are ones where nobody can produce a definitive wrong answer: open-ended retrieval is scored only against the set properties the system was trained to satisfy, and the compositional setting is explicitly weakly supervised. The music results, run on a closed industrial playlist collection, cannot be reproduced by anyone outside the company that owns it. The near-50-second autoregressive baseline is real, but it is also a configuration nobody would put in front of a user, so the headline multiple is measured against something that was already unshippable.
None of which undermines the central idea, which is genuinely worth taking seriously: reinforcement learning as a one-time objective compiler rather than a permanent inference-time tax. The heavy search for reward-satisfying behaviour happens once, and the shipped model is decoupled from expensive online reasoning. The trade is in the word "compiled." A prior that has been compiled from a reward inherits whatever its designer believed about good retrieval, frozen at training time. When the catalogue shifts or taste moves, a reasoning model can be re-prompted; a 53.9-million-parameter diffusion prior has to be recompiled from scratch.