Efficient, Property-Aligned Fan-Out Retrieval via RL-Compiled Diffusion
Pengcheng Jiang, Judith Li, Moonkyung Ryu, Lily Hu, Kun Su, Zhong Yi Wan, Liam Hebert, Hao Peng, Jiawei Han, Dima Kuzmin, Craig Boutilier
Abstract
Many modern retrieval problems are set-valued: given a broad intent, the system must return a collection of results that optimizes higher-order properties (e.g., diversity, coverage, complementarity, coherence) while remaining grounded with respect to a fixed database. Set-valued objectives are typically nondecomposable and are not captured by existing supervised (query, content) datasets which only prioritize top-1 retrieval. Consequently, fan-out retrieval is often employed to generate diverse subqueries to retrieve item sets. While reinforcement learning (RL) can optimize set-level objectives via interaction, deploying an RL-tuned LLM for fan-out retrieval is prohibitively expensive at inference 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 large-scale fashion and music benchmarks consisting of curated item sets, we show that R4T improves retrieval quality relative to strong baselines while reducing query-time fan-out latency by an order of magnitude. In parallel, diffusion-based generative retrieval *Work done as a Student Researcher at Google Research.
Ask about this paper
Your agent reads all of it.
Lune indexed this paper to the last equation, along with the top-tier papers that cite it. Ask a question and the answer quotes them.
Your agent calls
Luneget_paper_fulltext
Free to start. No credit card required.
Terminal
Install the CLIlune papers fulltext 33c27ec3-2dd9-4f41-b3f3-86368e7a82c7Builds on18
- Learning Transferable Visual Models From Natural Language SupervisionAlec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh et al.ICML 2021 · 47,906 citations
- Denoising Diffusion Probabilistic ModelsJonathan Ho, Ajay Jain, Pieter AbbeelNeurIPS 2020 · 35,902 citations
- Scalable Diffusion Models with TransformersWilliam Peebles, Saining XieICCV 2023 · 5,568 citations
- Elucidating the Design Space of Diffusion-Based Generative ModelsTero Karras, Miika Aittala, Timo Aila, Samuli LaineNeurIPS 2022 · 3,959 citations
- Score-Based Generative Modeling through Stochastic Differential EquationsYang Song, Jascha Sohl-Dickstein, Diederik P. Kingma, Abhishek Kumar et al.ICLR 2021 · 1,270 citations
Related papers
- Optimizing Retrieval for RAG via Reinforcement LearningJiawei Zhou, Lei ChenNeurIPS 2025 · 1 citation
- Diffusion Augmented Retrieval: A Training-Free Approach to Interactive Text-to-Image RetrievalZijun Long, Kangheng Liang, Gerardo Aragon-Camarasa, Richard McCreadie et al.SIGIR 2025 · 6 citations
- Retrv-R1: A Reasoning-Driven MLLM Framework for Universal and Efficient Multimodal RetrievalLanyun Zhu, Deyi Ji, Tianrun Chen, Haiyang Wu et al.NeurIPS 2025 · 12 citations
- s3: You Don't Need That Much Data to Train a Search Agent via RLPengcheng Jiang, Xueqiang Xu, Jiacheng Lin, Jinfeng Xiao et al.EMNLP 2025
- Training Diffusion Models Towards Diverse Image Generation with Reinforcement LearningZichen Miao, Jiang Wang, Ze Wang, Zhengyuan Yang et al.CVPR 2024 · 12 citations
