Bypassing inference bottlenecks: Accelerating complex AI search with Retrieve-for-Train

Modern search or recommendation applications are increasingly expected to return a coherent set of results rather than a single best match. For example, when a user searches for "camping gear", they don’t want ten slight variations of four-person tents. They want a coherent, complementary slate that includes essential camping gear, such as a tent, sleeping bag, portable stove, and headlamp.

To do this, systems use a query fan-out technique that breaks a single broad prompt into several related sub-queries to cover potential user interests. However, teaching an LLM to perform database-aware query decomposition dynamically drains a massive thinking budget. By design, zero-shot LLMs are general autoregressive text predictors; they aren’t optimized to navigate the specific, geometric manifold of a target corpus. Consequently, they need extended test-time computation to return a collection of results that optimizes higher-order set-level properties (e.g., diversity, coverage, complementarity, coherence) while remaining grounded with respect to a fixed database.

In our ICML 2026 paper, “Efficient, Property-Aligned Fan-Out Retrieval via RL-Compiled Diffusion”, we address this decomposition bottleneck via a reward-to-data compilation framework. Instead of forcing the model to expend a large thinking budget at inference, our Retrieve-for-Train framework uses offline reinforcement learning (RL) to discover reward-aligned fan-outs and compile them into supervision. By distilling these optimized exploration behaviors into a lightweight diffusion retriever, we enable highly efficient, single-pass query fan-out at inference time. This achieves mathematically formulated, set-level properties without the overhead of test-time thinking tokens.

When tasked with brainstorming a complex group of search terms, it’s tempting to simply deploy a standard, off-the-shelf LLM at inference time to handle the job. However, relying on generic models for database-aware query decomposition introduces two critical challenges:

1. Paraphrastic collapse: Without database-aware optimization, zero-shot LLMs frequently suffer from paraphrastic collapse. Rather than exploring complementary facets of a topic, they tend to generate redundant, near-synonymous queries. For example, given the broad prompt "Bohemian festival style”, a standard LLM without careful prompt engineering might lazily generate "bohemian festival fashion" and "bohemian festival clothes”. This semantic looping produces a homogeneous slate of results, entirely missing the distinct, helpful semantic directions a fashion expert would identify, such as fringe jackets, crochet dresses, or suede boots.
2. Autoregressive latency bottlenecks: Standard LLMs are fundamentally constrained by sequential, autoregressive generation. To successfully decompose a complex query into complementary facets, modern models typically require a substantial thinking budget, generating hundreds of intermediate chain-of-thought (CoT) reasoning tokens (i,e., the intermediate steps or internal processing units an AI model generates before answering a complex question) to plan their expansion before outputting the actual search terms. While this deliberate reasoning is acceptable for conversational AI, it introduces a severe structural bottleneck for set-valued search (e.g., retrieving a complementary slate of results, such as fringe jackets or crochet dresses mentioned above). When a system must brainstorm a large slate of sub-queries simultaneously, the combined overhead of continuous context processing and generating extended reasoning tokens scales poorly. Even with advanced serving optimizations, this token-by-token architecture creates a latency floor that is fundamentally at odds with the sub-second response times required by a production search bar.

The Retrieve-for-Train treats the AI's training like an offline practice session rather than a test it has to take on the spot while a user is waiting. Instead of forcing the AI to slowly figure out the rules of a good search and drain a massive processing budget every single time someone types a query, Retrieve-for-Train runs an offline RL training program once.

This program uses a rigorous reward system to turn abstract goals like "ensure the results are diverse and actually in stock" into an exact step-by-step instruction manual. Once that manual is built, the AI can execute it instantly during a real search without delay.

The pipeline operates in three distinct steps:

• Fan-out language model training: RL trains a fan-out language model to emit property-aligned sub-queries scored by a set-level property-check reward. This evaluates the entire group of results as a whole, rather than scoring each result in isolation.
• Supervision synthesis: The frozen fan-out language model synthesizes (query → target-set) pairs entirely offline for supervised learning, requiring no human labels.
• Diffusive retriever training: A compact, 53.9M-parameter diffusion model learns to map a query embedding directly to a complete set of target embeddings in one non-autoregressive pass, officially bypassing the need for text-based CoT reasoning tokens.

The success of the Retrieve-for-Train framework hinges entirely on how we define "good" search behavior. Traditional supervised training evaluates pointwise relevance via learning to rank, scoring each retrieved item in isolation. However, a truly expert search slate is defined by non-decomposable, set-level properties. You can’t measure the diversity or complementarity of a single item; these properties only exist mathematically when evaluating the entire collection of retrieved results.

Rather than relying on ambiguous natural language instructions to enforce these fan-out properties, Retrieve-for-Train fine-tunes the 4B open-source language models (Gemma3-4B and Qwen3-4B) via reinforcement learning using a strict mathematical composite reward. For our open-ended abstract retrieval tasks, this composite reward is a weighted balance of three competing pillars:

• Groundedness: Penalizes distance to the database manifold, ensuring every generated sub-query corresponds to a real, retrievable item in the database.
• Diversity: Measured using the Vendi Score over the entire set of sub-queries, forcing the model to explore broad semantic breadth.
• Alignment: Anchors candidate sub-queries to the original broad prompt to prevent semantic drift.

During training, we optimize the fan-out language model against these geometric realities using group relative policy optimization (GRPO) with soft proximal policy optimization (PPO).

This specific triad of rewards is critical because they act as mutual counter-anchors. If a model is optimized purely for groundedness, it will reward-hack the system by generating degenerate, nonsensical strings that happen to mathematically map to a specific database coordinate. If alignment is added to fix the nonsense, the policy simply cheats by collapsing into repetitive paraphrases of the user's prompt.

By injecting the Vendi Score as a counter-anchor, Retrieve-for-Train effectively closes off these shortcut solutions. To achieve a high-reward state, the policy is forced into a balanced region of the embedding space where it must discover valid, strictly grounded, yet semantically distinct variations of the original intent.

To evaluate the Retrieve-for-Train framework, we used a combination of frozen, dataset-specific multimodal embedding backbones and open-source language models optimized for query expansion. We evaluated this setup across two distinct set-valued retrieval regimes:

• Open-ended abstract retrieval: A setting where no unique ground truth exists and quality is exclusively measured by set-level properties, including diversity, query alignment, and database groundedness.
• Weakly supervised compositional retrieval: A setting where queries are paired with a weak reference set that serves as just one plausible realization of the query intent.

For the multimodal embedding backbones, we conducted experiments across two domains: A large-scale fashion dataset of user-curated outfits used for text-to-image experiments (evaluated using a CLIP-based retriever), and a proprietary industrial dataset of expert-generated music playlists used for text-to-music evaluations (evaluated using MuLan).

For the language models, the query fan-out process was driven by 4B open-source models, specifically Gemma3-4B and Qwen3-4B, which were tasked with generating exactly 10 sub-queries for every single main search prompt they processed. We implemented the RL training for these fan-out models via Soft-GRPO, an approach that uses group relative policy optimization with soft PPO regularization.

Across both retrieval tasks, Retrieve-for-Train outperformed traditional single-query search, zero-shot expansion, and even the heavily optimized Best-of-N baseline.

Qualitatively, zero-shot LLM baselines tended to generate near-synonymous paraphrases (e.g., "bohemian festival style" vs. "bohemian festival fashion"), causing redundant results. Retrieve-for-Train generated highly diverse, distinct sub-queries (e.g., branching into "boots" or "lace") that remained strictly grounded within the database manifold.

Directly deploying our RL-tuned language model yielded exceptional search quality, but it inherited standard autoregressive latency constraints and demanded a high computational thinking budget.

By distilling that learned behavior into the 53.9M-parameter Retrieve-for-Train diffusion model, we successfully smashed the latency bottleneck. Because the diffusion model generates all target directions simultaneously in a single, non-autoregressive parallel pass in continuous embedding space, it delivers a massive 12 to 20 speedup over autoregressive approaches.

At scale, while autoregressive fan-out latency expands linearly to nearly 50 seconds under large context batches, Retrieve-for-Train-Diffusion stays between sub-second to a few seconds, delivering production-ready, expert-level search at a fraction of the computational cost.

During our reward optimization process, we discovered something fundamental about training a fan-out language model for search. Without a diversity term, the model quickly collapses into generating degenerate, nonsensical strings (like "line ending line ending") to mathematically exploit the vector coordinates of the database. Injecting a geometric diversity metric (the Vendi Score) acts as a vital counter-anchor, forcing the model into a stable region of the embedding space where it can only maximize its reward by acting like a true search expert.

We demonstrated that RL can be highly effective when used as a one-time "objective transducer" rather than an online inference engine. By decoupling the heavy computation of reward-driven behavior exploration from the final deployed model, our framework successfully bypasses the steep inference latency and high computational overhead typical of online LLM deployment.

Distilling these complex, set-level behaviors into a lightweight diffusion prior allows production retrieval systems to optimize for higher-order properties like diversity and alignment effectively. Ultimately, Retrieve-for-Train establishes a highly scalable, data-efficient pipeline for set retrieval in specialized or multimodal domains where human-labeled, property-aligned training pairs are otherwise scarce or costly to obtain. See the paper for more details.