Efficient, Property-Aligned Fan-Out Retrieval via RL-Compiled Diffusion
Abstract
Many modern retrieval problems are set-valued: given a broad intent, the system must return a collection of results that optimizes higherorder properties (e.g., diversity, coverage, complementarity, coherence) while staying grounded to a fixed database. Set-valued objectives are inherently non-decomposable and are not captured by existing supervised (query, content) datasets which only prioritize top-1 retrieval. While reinforcement learning (RL) can optimize set-level objectives via interaction, deploying an RL-tuned LLM for fan-out retrieval is prohibitively expensive at query time. Conversely, diffusion-based generative retrieval enables efficient single-pass fan-out in embedding space, but requires objective-aligned training targets. To address these issues, we propose R4T (Retrieve-for-Train), which uses RL once as an objective transducer in a three step process: (i) train a fan-out LLM with composite set-level rewards, (ii) synthesize objective-consistent training pairs, and (iii) train a lightweight diffusion retriever to model the conditional distribution of set-valued outputs. Across Polyvore and a music playlist dataset, R4T improves retrieval quality over strong baselines while reducing query-time fan-out latency by an order of magnitude.
Lay Summary
When you ask a search engine, music app, or shopping site for something broad, like songs for a rainy Sunday, outfits for a winter wedding, or books similar to one you just finished, there isn't a single right answer. A good response is a set of items that together feel diverse, on-topic, and well-matched to what's actually in the catalog. Training computer models to produce sets like this is hard, because the usual approach needs huge collections of human-labeled "perfect" examples that don't exist for tasks with many valid answers. We developed a method called R4T (Retrieve-for-Train). It first uses a powerful but slow language model that learns, through trial and error, how to break a broad request into several focused sub-requests that collectively give a good result. We then use that model as a teacher to train a much smaller, faster model that produces the whole set in one shot. The result is search and recommendation systems that return richer, more thoughtful sets of suggestions while running over ten times faster, practical enough for real products.