绕过推理瓶颈：用 Retrieve-for-Train 加速复杂 AI 搜索

Pengcheng Jiang, Student Researcher, and Judith Yue Li, Senior Research Engineer, Google Research Instead of relying on expensive inference-time reasoning, the Retrieve-for-Train framework uses reinforcement learning once to train a lightweight diffusion model. This bypasses the heavy autoregressive "thinking budget" to instantly generate a cohesive, expert-level slate of AI search results. 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: 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. 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 reason