Google — Making AI Search Ask Better Questions, 20x Faster
Episode notes
AI search often splits a broad question into several sub-searches (“query fan-out”), and that step can take nearly 50 seconds under load. Google's Retrieve-for-Train distills it into a tiny diffusion model that's 12–20x faster.
Key takeaways:
- Two problems. Zero-shot LLMs suffer “paraphrastic collapse,” producing near-duplicate sub-queries, and their token-by-token generation is too slow for a search bar.
- An RL teacher. Gemma3-4B and Qwen3-4B are fine-tuned with RL on a composite reward: stay grounded in real database items, cover the user's intent, and stay diverse (measured by the Vendi Score).
- Diversity blocks shortcuts. Without the Vendi term, the policy cheats by repeating paraphrases of the prompt.
- A tiny diffusion student. A 53.9M-parameter diffusion model maps a query embedding to the whole set of sub-query embeddings in one parallel pass, with no text generation.
- Results. 12–20x faster than autoregressive fan-out. Where autoregressive latency grows to nearly 50 seconds, the diffusion model stays between sub-second and a few seconds.
Chapters:
00:00 Waiting 50 seconds for search
01:32 Why complex AI search is slow
03:06 Paraphrastic collapse
04:36 Training the teacher with RL
06:56 Rewarding grounding and intent
08:29 The Vendi Score for diversity
10:02 The latency problem remains
11:34 Distilling into diffusion
13:07 The rubber-stamp analogy
13:53 Results
16:13 Our analysis
17:46 How Meta and others handle it
20:04 Recap
Sources:
- Bypassing inference bottlenecks: Accelerating complex AI search with Retrieve-for-Train (Google Research blog, Sep 2026 (ICML 2026 paper))
Related episode: E28: Databricks — Teaching a Search Agent When to Stop
Enjoyed it? Follow Data Science in the Wild (one real production ML/AI system or paper per episode, hosted by Peter Liu) and rate it on Apple Podcasts or Spotify.
Also from Peter: The GenAI Evolution Atlas, Wednesday AI deep dives and Friday news roundups.
Made with AI-generated voices; curated by Peter Liu from the sources above.