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.
Figures & tables
Figure 1: Overview of the R4T. Step 1 (§ 2.3 ) trains a fan-out language model (FOLM) using RL to produce property-aligned sub-queries. Step 2 (§ 2.4 ) uses the trained FOLM to synthesize (q,c) supervision data. Step 3 (§ 2.5 ) trains a diffusion-based fan-out retriever that samples content embeddings directly from query embeddings.
Figure 2: Illustration of rewards used in OAR.
Polyvore
Music
Groundedness
Diversity
Alignment
Average
Groundedness
Diversity
Alignment
Average
No Fan-out
22.4
34.4
21.4
26.1
48.8
20.0
41.8
36.9
Gemini-2.5-Flash
Zero-shot
24.0
47.0
23.6
31.5
45.8
45.2
44.4
45.1
Best-of-N
26.1 ± 1.6
52.2 ± 2.3
25.2 ± 1.9
34.5 ± 1.8
48.2 ± 2.1
48.4 ± 2.6
49.0 ± 2.0
48.5 ± 2.2
Gemma3-4B
Table 1: Performance on Task 1: Open-Ended Abstract Retrieval (OAR). Results are reported as mean ± standard deviation. Bold denotes the best performance within each model family. R4T consistently outperforms fan-out baselines across datasets and metrics. Note that Groundedness is not applicable to R4T-Diffusion (indicated by “\”) due to the absence of intermediate sub-query.
Recall@5K
Hit@5K
VS
LLM-Based Retrieval
Gemini-2.5-Flash
15.7
52.1
33.4
Gemma3-4B
6.0
25.9
44.2
Qwen3-4B
10.1
33.9
46.4
R4T-FOLM (Gemma)
16.9
54.4
40.5
R4T-FOLM (Qwen)
20.9
64.6
27.5
Table 2: Performance on Task 2: Weakly Supervised Compositional Retrieval (WSCR) . We use an LLM (Gemini-2.5-Pro) to generate a broad query for each outfit in Polyvore, and set the items in each as the ground truths. For diversity evaluation with normalized Vendi Score, we set temperature=0.9 for LLM-based retrieval, and set #samples=5 (i.e., 5 × 10=50 sub-queries in total) for all the methods.
Figure 3: Qualitative comparison on the Polyvore dataset for the open-ended abstract retrieval task given the broad query “Bohemian festival style” . R4T generates semantically distinct, on-topic sub-queries that retrieve diverse outfit collections, while the Qwen3-4B zero-shot baseline produces largely paraphrastic sub-queries, leading to more homogeneous results.
Figure 4: Reward & Scores over Time (FOLM Training) for Open-Ended Abstract Retrieval.
Figure 5: Efficiency comparison between Autoregressive LLM and Diffusion Model for query fan-out ( k=10 ) generation. Note that the x-axis follows a logarithmic scale (doubling at each step). The Diffusion Model demonstrates superior scalability and lower latency, maintaining sub-second performance for small batches and achieving an order-of-magnitude speedup at larger batch sizes.
Appendix figures & tables5 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 8Figure 9Figure 10
Parameter
Value
Fan-Out LM (RL Training)
Optimizer
AdamW ( β1=0.9,β2=0.95 )
Learning Rate
1×10−7
Global Batch Size
512
Micro-batch Size
64
Grad. Accum. Steps
8
Appendix
Table 3: Implementation Hyperparameters for R4T.
Figure 6: Examples of generated subqueries and retrieved images from Polyvore dataset given a broad query “Labor day picnic outfit” for the open-ended abstract retrieval task.
Unlike traditional fact-based retrieval, rationale-based retrieval typically necessitates cross-encoding of query-document pairs using large language models, incurring substantial computational costs. To address this limitation, we propose Rabtriever, which independently encodes queries and documents, while providing comparable cross query-document comprehension capabilities to rerankers. We start from training a LLM-based generative reranker, which puts the document prior to the query and prompts the LLM to generate the relevance score by log probabilities. We then employ it as the teacher of an on-policy distillation framework, with Rabtriever as the student to reconstruct the teacher's contextual-aware query embedding. To achieve this effect, Rabtriever is first initialized from the teacher, with parameters frozen. The Joint-Embedding Predictive Architecture (JEPA) paradigm is then adopted, which integrates a lightweight, trainable predictor between LLM layers and heads, projecting the query embedding into a new hidden space, with the document embedding as the latent vector. JEPA then minimizes the distribution difference between this projected embedding and the teacher embedding. To strengthen the sampling efficiency of on-policy distillation, we also add an auxiliary loss on the reverse KL of LLM logits, to reshape the student's logit distribution. Rabtriever optimizes the teacher's quadratic complexity on the document length to linear, verified both theoretically and empirically. Experiments show that Rabtriever outperforms different retriever baselines across diverse rationale-based tasks, including empathetic conversations and robotic manipulations, with minor accuracy degradation from the reranker. Rabtriever also generalizes well on traditional retrieval benchmarks such as MS MARCO and BEIR, with comparable performance to the best retriever baseline.
Retrieval-augmented generation (RAG) systems typically rely on a single retriever and a single set of hyperparameters, despite facing highly heterogeneous queries that range from simple factoid questions to complex multi-hop reasoning. We propose a method that automatically selects a small, diverse subset of retrievers (a portfolio) from a large pool of candidates, to cover different regions of the target query distribution. We formalize this setting via an expected best-of-k objective over the query distribution and show that it admits an efficient portfolio construction algorithm with near-optimal guarantees. Across multiple QA benchmarks, our learned portfolios and router pipeline consistently outperform single-retriever and naive multi-retriever baselines on both retrieval metrics and answer quality. In addition, compared to inference-time hyperparameter tuning approaches, fixed portfolios enable parallel retrieval and LLM calls, achieving comparable (and sometimes better) accuracy with substantially lower latency and token cost.
Miltiadis Stouras, Vincent Cohen-Addad, Silvio Lattanzi +1
Dense vector retrieval is the practical backbone of Retrieval- Augmented Generation (RAG), but similarity search can suffer from precision limitations. Conversely, utility-based approaches leveraging LLM re-ranking often achieve superior performance but are computationally prohibitive and prone to noise inherent in perplexity estimation. We propose Utility-Aligned Embeddings (UAE), a framework designed to merge these advantages into a practical, high-performance retrieval method. We formulate retrieval as a distribution matching problem, training a bi-encoder to imitate a utility distribution derived from perplexity reduction using a Utility-Modulated InfoNCE objective. This approach injects graded utility signals directly into the embedding space without requiring test-time LLM inference. On the QASPER benchmark, UAE improves retrieval Recall@1 by 30.59%, MAP by 30.16% and Token F1 by 17.3% over the strong semantic baseline BGE-Base. Crucially, UAE is over 180x faster than the efficient LLM re-ranking methods preserving competitive performance, demonstrating that aligning retrieval with generative utility yields reliable contexts at scale.
Rajinder Sandhu, Di Mu, Cheng Chang +4
Layer 6 AI · Toronto, ON, Canada · Dalhousie University +1