Efficient, Property-Aligned Fan-Out Retrieval via RL-Compiled Diffusion

TL;DR

This paper introduces R4T, a method combining reinforcement learning and diffusion models to improve set-valued retrieval tasks, achieving higher quality results with faster inference in large-scale benchmarks.

Contribution

The paper presents a novel three-step approach that uses RL for training data synthesis and diffusion models for efficient retrieval, addressing the limitations of existing methods.

Findings

R4T outperforms strong baselines in retrieval quality.

Reduces fan-out latency by an order of magnitude.

Effective on large-scale fashion and music benchmarks.

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 non-decomposable 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.