How we taught agents to use good retrieval - Hanna Lichtenberg, Mixedbread AI
Mixedbread AI addresses the "Oracle Gap" – the disparity between LLM reasoning and retrieval capabilities – by developing agents trained to use advanced search tools effectively. They demonstrate how current LLMs generate poor queries due to training biases and introduce a sophisticated agent harness with diverse search tools and a unique training regimen, including supervised fine-tuning and reinforcement learning with custom rewards, to teach agents to form precise semantic queries. This approach significantly improves performance on benchmarks like Oblique Congress and Snowflake's Match QA, closing the gap between theoretical perfect retrieval and real-world agent performance.
The presentation addresses a critical issue in the field of Large Language Models (LLMs) and their application in knowledge-intensive tasks: the “Oracle Gap.” While LLM reasoning capabilities have advanced exponentially (e.g., GPT-3.5 to GPT-5.5), retrieval tools have improved much slower. This creates a significant disparity, dubbed the “knowledge gap,” where LLMs, despite advanced reasoning, are bottlenecked by their inability to access the right information effectively. The bottleneck is not the LLM’s reasoning ability but its access to relevant knowledge. Mixedbread AI’s enhanced search, leveraging latent interaction, significantly closes this gap, recovering most of the Oracle performance. The core issue lies in how LLMs generate search queries. Models often produce “gibberish” or keyword-stacked queries (e.g., “senator woman questions billionaires not a company then okay thank you staff will check hearing”) which are ineffective for semantic search systems. This behavior stems from three main factors:
- Training for Code Agents: Models are predominantly trained for coding tasks and code exploration, where tools like
grep(regular expression matching) are common. They learn to write queries that maximize keyword overlap. - Web Tool Mimicry: LLMs are trained to use web tools efficiently, which are optimized for human-like, keyword-based search patterns.
- Benchmark Bias: Many existing benchmarks (e.g., Beer, Nano Beer) favor “caveman-style,” entity-based queries that inherently benefit BM25-like lexical search, reinforcing poor query generation habits. To overcome these limitations, Mixedbread AI developed its own search agent, designed to teach LLMs to use powerful search tools properly, extending beyond code search into general knowledge work. Key design principles prioritize precision, speed, and cost-effectiveness. The agent operates within a powerful harness with four main search tools:
- Overview Search: Wide semantic search returning up to 50 summarized chunks for an initial corpus overview.
- Main Semantic Search: Returns full payload of the top 10 retrieved chunks.
- Filter Chunks: Sorts and finds chunks based on metadata facets.
- Grep Tool: For exact keyword matching. The agent loop is designed for efficiency and exploration:
- It allows for a maximum of four search rounds, with parallel searches within each round.
- Initial input includes the user query, results from an initial semantic search, and available metadata hints. This “corpus preview” informs the agent’s search planning.
- The agent splits its intent into up to four queries, selecting the best tool for each.
- Results are deduplicated to prevent context filling.
- The process continues until sufficient evidence is gathered, upon which the agent submits a ranked list of relevant chunks. Crucially, the harness encourages better semantic queries through:
- Goal Framing: The agent must articulate the evidence needed before writing a query.
- Tool Differentiation: Explicitly guiding the agent to use semantic search for conceptual aspects and
grepfor exact matches. - Task Framing: Instructing the model to write “one concise sentence describing what it wants to find,” rather than a generic “search query,” to break old BM25-style habits.
- Examples: Providing good query examples and demonstrating how to split input queries into aspects.
- Contextual Preview: Giving the agent initial semantic search results helps it understand the corpus and tailor its queries. Mixedbread’s agent employs a small LLM for faster operation and is trained to optimize its search strategy, tool choice, semantic query quality, exploration, and efficiency. The training involves two main stages:
- Supervised Fine-Tuning (SFT): Using a larger “teacher” LLM to guide the smaller agent.
- On-Policy Reinforcement Learning (RL): With a custom “search reward” function. The search reward is a composite metric:
- Retrieval Reward: Combines standard metrics like NDCG on the agent’s final ranked list with an LLM judge evaluating relevance, overall chunk relevance, and plausibility of ranking based on specific rubrics.
- Trajectory Reward: Uses an LLM judge to evaluate the quality of tool queries (e.g., natural sentence structure) and the sufficiency of exploration (not too much, not too little). Preliminary results demonstrate significant improvements:
- Oblique Congress Benchmark: The agent achieved an NDCG@10 of 0.4, a substantial leap from the GPT multi-hop agent’s 0.18.
- Snowflake’s Match QA Benchmark: Mixedbread’s beta agent (in production) ranks as top one, achieving 93.4% accuracy when paired with the Gemini 3.5 flash model, with less computational effort than comparable LLM-agent combinations. These results underscore the vast potential for improving LLM performance by focusing on sophisticated retrieval tools and teaching agents to leverage them effectively, thereby pushing the boundaries of retrieval in AI systems.