Mixedbread AI on Teaching Agents Better Retrieval

In a recent presentation, Hanna Lichtenberg, an AI Engineer at Mixedbread AI, detailed the company's approach to improving how AI agents utilize retrieval. The core challenge, as highlighted by Lichtenberg, is the growing "knowledge gap" between the advanced reasoning capabilities of large language models (LLMs) and their ability to effectively retrieve relevant information. This gap, she explained, is becoming increasingly pronounced as LLMs evolve, necessitating better tools for information retrieval.

Bridging the Knowledge Gap

Lichtenberg illustrated this challenge with a graph showing that while LLM reasoning capabilities are experiencing exponential growth, retrieval capabilities are advancing much more slowly. This disparity means that even powerful LLMs can struggle to access and utilize the precise information needed for complex tasks, particularly in domains like legal or financial work. The team at Mixedbread AI recognized this limitation and set out to build a solution.

The Mixedbread AI Search Agent Harness

To address this, Mixedbread AI developed a "search agent harness" designed to teach agents to use powerful search tools more effectively. The harness is built upon the Mixedbread AI platform and incorporates several key search tools:

• Overview Search: This tool provides a broad overview by retrieving summaries of the top 50 relevant chunks of information.
• Semantic Search: This tool focuses on more nuanced retrieval, executing semantic searches and returning the top 10 retrieved chunks.
• Filter Chunks: This tool refines the search by filtering both raw and chunked data based on metadata facets.
• Grep: A traditional keyword-based search tool is also included for specific use cases.

The harness operates by first planning the search, then executing multiple queries using these tools, and finally filtering the results to identify relevant information. The agent is trained to guess keywords that increase the overlap between queries and documents, aiming to improve the efficiency and accuracy of the retrieval process.

Targeted Rewards for Better Search Behavior

A critical component of Mixedbread AI's approach is the use of "targeted rewards" to guide the agent's learning. The total reward (R) is a combination of retrieval quality and trajectory quality. This means the agent is rewarded not only for finding the right information but also for the efficiency and effectiveness of its search process.

The framework employs two types of judges:

• Retrieval Judge: This judge assesses the quality of the retrieved information based on recall and NDCG (Normalized Discounted Cumulative Gain), ensuring that relevant documents are identified and ranked appropriately.
• LLM Query and Explanation Judge: This judge evaluates the LLM's ability to generate natural language queries and explanations, ensuring clarity and relevance in the agent's output.

By combining these reward signals, the agent learns to optimize its search strategy, leading to more precise and efficient information retrieval.

Benchmarking and Results

The presentation included benchmarking results comparing the Mixedbread AI search agent against other models. On the Obliqua-congress benchmark, the Mixedbread AI search agent demonstrated superior performance, achieving a higher score in both precision and recall compared to other models like the OPP 0.2 Multi-hop Agent and Gemini embedding-based approaches. The team also shared preliminary results on the MTEB benchmark, showing competitive accuracy and effort scores for their models.

The data presented indicates that while LLMs are becoming increasingly powerful, their ability to perform complex search tasks is still limited by the quality of their retrieval mechanisms. Mixedbread AI's harness and reward system represent a significant step towards closing this gap, enabling agents to become more effective and reliable in knowledge-intensive applications.