How we taught agents to use good retrieval - Hanna Lichtenberg, Mixedbread AI
AI Engineer · 14 min · 116 sentences · from YouTube's caption track
Each timecode opens YouTube at the start of that sentence. Line anchors (#s42) are the cue ids in the WebVTT, and every line carries its start and end seconds. All transcripts has every talk, and the whole corpus as one file.
- 00:00Hey, today we're going to talk about how we taught agents to use good retrieval or how we call it internally closing the Oracle gap with knowledge agents.
- 00:10I'm Hannah.
- 00:11I'm an AI engineer at Mix bread and I'm leading the agentic search at Mix bread.
- 00:17Yeah, I'm Amir.
- 00:18I'm the co-founder of Mix bread and let's jump into it.
- 00:21So, we have seen over the past years with the LLMs just getting better and better that the reasoning capabilities of the models is growing exponentially.
- 00:31So, just think back how good GPT-3.5 was and how good GPT-5.5 was, right?
- 00:37So, I can clearly exponential curve.
- 00:39But, if you look at search over the last 20 years basically, we see that it's improving, but it's improving very very slowly.
- 00:49So, there's a huge gap between how basically LLMs and the reasoning is evolving and how retrieval is evolving.
- 00:59But, the retrieval tools are the main access pattern for this reasoning there to get the right knowledge and to be truly useful beyond code like in legal work, in finance work and so on and so forth.
- 01:13Internally, we call this gap happening right now between reasoning and search the knowledge gap.
- 01:18And it's not just one obscure theory of ours.
- 01:22We see that actually with real benchmarks and real tasks that this gap exists in real world as well.
- 01:28Let's take two benchmarks here, BroSque Plus and Office QA Pro.
- 01:33BroSque Plus is like a browsing benchmark where we have pre-complex queries and a corpus of 100,000 documents and we just try to answer these questions.
- 01:44It's like real world deep search task basically and Office QA Pro which has the treasuries of the US treasuries of the past um 100 years basically and we ask really complex question over this.
- 02:00Office QA Pro was created by Databricks and BrowseComp was created by OpenAI and BrowseComp Plus is a version of it with a fixed size corpus while BrowseComp was over the open web.
- 02:11And we see here's a Oracle performance in dotted lines and Oracle means what is the maximum theoretical performance of the models if you would put in the right documents
- 02:23with the question with with the question.
- 02:25We see for Office for BrowseComp it's 93% and for Office QA Pro it's 64%.
- 02:34And now we take something like Codex with its default tools and want to see hey, what is the performance this tools get right now?
- 02:42And you see there's a sharp drop in the quality of the answers Codex produces.
- 02:48For BrowseComp it's nine points and for Office QA Pro it's eight points.
- 02:53So we see that the models are extremely capable if they would get the right documents but if you put them into the noisy corpus the performance drops sharply.
- 03:05Meaning that actually the bottleneck here is not the reasoning.
- 03:11It's actually the access to the right knowledge it needs to answer this question.
- 03:16So we see this knowledge gap in real world.
- 03:21And if you would just drop in way better search using Mixbread with just a search tool using latent interaction we can recover most of the performance.
- 03:30So for BrowseComp the difference between the Oracle and the uh GPT-5 with Mixbread is just three point and for Office QA Pro we even almost completely closed the gap.
- 03:41With giving the model better search tool we can recover most of the knowledge gap.
- 03:47But looking at the queries the model asks or writes right now to the search system, it gets super super interesting.
- 03:55So, here's an example query we found uh during some benchmarking, which is senator woman questions billionaires not a company then okay thank you staff will check hearing.
- 04:07It's basically gibberish, right?
- 04:08If you have a search system which wants to have neural questions or semantic question and just gets these keywords over keywords, it gets confused.
- 04:18And the reason why the models write these type of queries is they're mostly trained for coding task, like for coding agents, which are then optimized for code base exploration
- 04:29using tools like grep.
- 04:31And grep just tries to find regular expressions, right?
- 04:36So, the model tries to just write as many expression it thinks are in the document to find the right thing.
- 04:41The second thing is the models are trained to use the web pretty efficiently.
- 04:47And to use the web tools properly, which are highly optimized for humans, they try to mimic human-like query patterns.
- 04:55So, just keyword on keyword on keyword.
- 04:57And number three, obviously, is the benchmark bias.
- 05:00Most benchmarks we have right now, like Beer, Nano Beer, use Caveman's style queries, which are entity-based queries that structurally favor heavily BM25.
- 05:11So, right now, the agent guesses the keywords to actually increase the overlap between the query and documents and can't really use powerful search tools properly.
- 05:23This motivated us to make our own search agent, to teach it to use powerful search properly, especially to use the right search tools for different use cases.
- 05:35Um and to work beyond code search for knowledge work.
- 05:40And most importantly, the agent should be precise, fast, and cheap.
- 05:44The first step for building this agent was to define a very powerful harness.
- 05:50Um [clears throat] our harness is fully built on a mixed back platform.
- 05:53Our agent has four main search tools, which are overview search.
- 05:59This is used as a very wide semantic search where the agent receives up to 50 retrieve chunks.
- 06:07Um and it sees only summaries of the chunk contents to really have just like an overview of the corpus every uh what exists and to not fill up the context too much.
- 06:19Then there is the main semantic search tool where the agent gets the full payload of the top 10 retrieve chunks.
- 06:27Um it has a filter chunks tool where it can sort sort and find chunks based on metadata facets.
- 06:36And of course we also have a grep tool for the keyword match searches.
- 06:41Our um agent loop, the harness itself is very simple and short cuz the agent should be fast.
- 06:48But we still wanted it to uh have a lot of exploration possibilities.
- 06:54That's why we decided to um define that it has a maximum four search rounds, but within each search round it can have parallel searches.
- 07:06So it can start several search at that searches at once.
- 07:10Um initially our agent sees the user query, but also already um results from an initial semantic search on the user query plus um hints on which metadata facets are available.
- 07:25Um so that based on this preview of the corpus the agent can start its search planning.
- 07:31It splits its search intent into maximum four queries um which cover separate aspects.
- 07:41And then for each query and each search intent, it will pick the best tool.
- 07:49Um as you can see like an example here on the right.
- 07:53Then when um the results of the tools are returned to the agent via deduplicating chunks, so that it's never seeing several chunks um several times to not fill up the context.
- 08:06[snorts]
- 08:07Um and yeah, the next search one can start, and as soon as the agent has enough evidence to answer the user query, it submits its ranking.
- 08:17Um yeah, we are just outputs all the chunks that are relevant to the query in a plausible ranking.
- 08:25So, why does this harness encourages better semantics queries and like better use of search tools?
- 08:32Um there are five more main points for this, which is first the goal framing.
- 08:38The agent has to articulate articulate what evidence it needs before writing the query.
- 08:45Um then the different tools are very important, so it can really use semantic search only if it needs aspects, and also only use a scrap when it needs exact keyword matching.
- 08:58Um then the way in which we frame the task of query uh writing is an important point.
- 09:07We kind of trick the um model into not thinking it has to write the typical BM25 base query by just instructing it to write one concise sentence describing what it wants to find,
- 09:21instead of directly instructing write a search query, so it cannot fall into this old pattern.
- 09:27Um then we put also provide examples in a prompt just to show the a few good queries and also how to divide an input query into different aspects to explore
- 09:41um the corpus.
- 09:43Yeah.
- 09:44What also helps is that we provide the original original query semantic search results.
- 09:51So, it can see what the corpus is about and really define um see the language, define where to dig deeper in which aspect.
- 10:01The second important step into building our own search agent was of course the training itself.
- 10:08So, here for training, we just decided to have a very small LLM to make the agent even faster.
- 10:16So, we are training the small agent um to optimize the search strategy itself to have better tool choice, um higher quality of semantic queries, more exploration ranking, and to be more efficient.
- 10:31So, the first step in training is um supervised fine-tuning with a larger teacher LLM.
- 10:37And then we uh did on-policy reinforcement learning with an own search reward.
- 10:45Our search reward is a combination of both of a retrieval reward and a trajectory reward.
- 10:54The retrieval reward um is based on like the metric result, typical retrieval metric results like we call it NDCG, that's the agent's final ranked list achieves.
- 11:08Plus also um an LLM retrieval judging.
- 11:13We had thought like it gets several rubrics where it decides um if the rubrics is hit or not, the rubric is hit or not, which are if the agent result relevant for this query,
- 11:24are all chunks relevant, and is the ranking itself plausible.
- 11:31For the trajectory reward, we also using an LLM judge and here um is the point to like really improve the tool queries, make it more efficient and to um
- 11:42improve the quality of the queries.
- 11:46Of course, like we have rubrics that are where the judge is deciding if the query is really a natural sentence.
- 11:53Also, if the amount of exploration is sufficient, is it too much or too less.
- 11:59Here you can see an example trajectory.
- 12:03You see that in the beginning we have the initial search and the initial metadata hints.
- 12:09Um then the agent did four parallel searches in the first round and then a second round just a simple grep tool.
- 12:16And here's an extract of um a trajectory.
- 12:21Here we see the agent queries.
- 12:24Um the input query is on the left.
- 12:27This is a very long query typical rambling style query from the Oblique Congress benchmark.
- 12:35And yeah, we see the agent queries here for the first semantic search tool.
- 12:40It's really a sentence describing what it wants to find and not a weird keywordy added behind each other formulation.
- 12:52The grep tools of course are like typical keyword patterns, which is exactly as intended.
- 12:58Our trained agent is not released yet, unfortunately.
- 13:02However, we have some intermediate results.
- 13:05Here on the left side for the Oblique Congress benchmark, we see that we achieved an NDCG of 10 at 10 of 0.4, which is a huge jump um towards the model that performed best on the paper of this benchmark, which is the GPT multi-hop agent,
- 13:24and which achieves 0. 18.
- 13:28On the right side, you can see the result of our beta version of the search agent, which we have right now in production.
- 13:36It's the mixed better genetic search.
- 13:39And this agent is top one on the snowflakes match QA benchmark, achieving an accuracy of 93.4 when we give the Gemini 3.5 flash model our genetic search as search tool.
- 13:55And this uh this performance was also achieved while having way less effort than comparable LLMs with other search agents.
- 14:06You can check out the match QA leaderboard.
- 14:09These results show that there's still a lot of room for improvement when it comes to huge language models and their search tools.
- 14:18And if you want to be part of pushing the boundaries of retrieval even further, um we are happy awaiting your application.