How we taught agents to use good retrieval - Hanna Lichtenberg, Mixedbread AI

Talk transcript

1. 0:00

   Hey. Today we're gonna talk about how we taught agents to use code retrieval or how we call it internally, closing the oracle gap with knowledge agents.

2. 0:10

   I'm Hanna. I'm an AI engineer at Mixedbread, and I'm leading the agentic search at Mixedbread.

3. 0:17

   Yeah, I'm Aamir. I'm the co-founder of Mixedbread, and let's jump into it. So we have seen over the past years with the LLMs just getting better and better, that the reasoning capabilities of the models is growing exponential.

4. 0:31

   So just think back how good GPT-3.5 was and how good GPT-5.5 is, right? It's like a clearly exponential curve. But if you look at search over the last 20 years, basically, we see that it's improving, but it's improving very, very slowly.

5. 0:49

   So there's a huge gap between how basically LLMs and the reasoning is evolving and how retrieval is evolving. But these retrieval tools are the main access pattern for this reasoning layer 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.

6. 1:13

   Internally, we call this gap happening right now between reasoning and search the knowledge gap. And it's not just one obscure theory of ours. We see it actually with real benchmarks and real tasks that this gap exists in real world as well.

7. 1:28

   Let's take two benchmarks here, BrowseComp+ and OfficeQA Pro. BrowseComp+ is like a browsing benchmark where we have pretty complex queries and a corpus of one hundred thousand documents, and we just try to answer these questions.

8. 1:44

   It's like real world deep search task, basically. And OfficeQA Pro, which has the treasuries of the U.S. Treasury of the past, um, hundred years, basically, and we ask really complex question over this.

9. 2:01

   OfficeQA Pro was created by Databricks, and BrowseComp was created by OpenAI, and BrowseComp+ is a version of it with a fixed size corpus, while BrowseComp was always an open web.

10. 2:11

    And 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 with the question, with the question.

11. 2:26

    We see for Office, uh, for BrowseComp, it's ninety-three percent, and for OfficeQA Pro it's sixty-four percent. And now we take something like Codex with its default tools and want to see, hey, what is the performance these tools get right now?

12. 2:42

    And you see there's a sharp drop in the quality of the answers Codex produces. For BrowseComp it's nine points, and for OfficeQA Pro it's eight points. So we see that the models are extremely capable if they would get the right documents.

13. 2:58

    But if we put them into the noisy corpus, the performance drops sharply. Meaning that actually the bottleneck here is not the reasoning, it's actually the access to the right knowledge it needs to answer this question.

14. 3:16

    So we see this knowledge gap in real world.

15. 3:21

    And if you would just drop in way better search using Mixed Bread, which is a search tool using late interaction, we can recover most of the performance. So for BrowseComp, the difference between the Oracle and the, uh, GPT-5 with Mixed Bread is just three point.

16. 3:38

    And for OfficeQA Pro, we even almost completely close the gap. With giving the model better search tool, we can recover most of the knowledge gap. But looking at the queries the model asks or writes right now to the search system, it gets super, super interesting.

17. 3:55

    So 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." It's basically gibberish, right?

18. 4:08

    If you have a search system which, uh, wants to have neural questions or semantic question and just gets these keywords over keywords, it gets confused. And the reason why the models write this type of queries is they're mostly trained for coding tasks, like for coding agents, which are then optimized for code-based exploration using tools like grep.

19. 4:31

    And grep just tries to find regular expressions, right? So the model tries to just write as many expression it thinks i- in the document to find the right thing.

20. 4:42

    The second thing is the models are trained to use the web pretty efficiently. And to use the web tools properly, which are highly optimized for humans, they try to mimic human-like query patterns.

21. 4:55

    So just keyword on keyword on keyword. And number three, obviously, is the benchmark bias. Most benchmarks we have right now like BEIR, NanoBEIR use caveman-style queries, which are entity-based queries that structurally favor heavily BM25.

22. 5:11

    So 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.

23. 5:23

    This 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, um, and to work beyond code search for knowledge work.

24. 5:40

    And most importantly, the agent should be precise, fast, and cheap. The first step for building this agent was to define a very powerful harness. Um, our harness is fully built on the Mixed Bread platform.

25. 5:54

    Our agent has four main search tools, which are the overview search. This is used as a very wide semantic search where the agent receives up to 50 retrieved chunks, um, and it sees only summaries of the chunk contents to really have just like an overview of the corpus, so we, uh, what exists and to not fill up

26. 6:18

    the context too much. Then there is the main semantic search tool where the agent gets the full payload of the top 10 retrieved chunks. Um, it has a filter chunks tool where it can sort by-- sort and find chunks based on metadata facets.

27. 6:36

    And of course, we also have a GREP tool for the keyword match searches. Our, um, agent loop, the harness itself is very simple and short because the agent should be fast, but we still wanted it to, uh, have a lot of exploration possibilities.

28. 6:54

    That's why we decided to, um, define that it has at maximum four search rounds, but within each search rounds, it can have parallel searches. So it can start several search at...

29. 7:09

    searches at once. Um, 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.

30. 7:25

    Um, so that based on this preview of the corpus, the agent can start its search planning. It splits its search intent into a maximum four queries,

31. 7:37

    um, which cover separate aspects. And then for each query and each search intent, it will se- uh, pick the best tool.

32. 7:49

    Um, as you can see like an example here on the right.

33. 7:53

    Then when, um, the results of the tools are returned to the agent, we are de-duplicating chunks so that it's never seeing several chunks, um, several times to not fill up the context.

34. 8:07

    Um, and yeah, the next search run can start, and as soon as the agent has enough evidence to answer the user query, it submits its ranking. Um, yeah, where it just outputs all the chunks that are relevant to the query in a plausible ranking.

35. 8:25

    So why does this harness encourages better semantic s- queries and like better use of search tools? Um, there are five more main points for this, which is first the goal framing.

36. 8:38

    The agent has to articulate, articulate what evidence it needs before writing the query. Um, then the different tools are very important, so it can really use semantic search only if it needs aspects, and also only uses GREP when it needs exact keyword matching.

37. 8:58

    Um, then the way in which we frame the task of query, uh, writing is an important point. We kind of trick the, um, model into not thinking it has to write the typical BM25 based query by just instructing it to write one concise sentence describing what it wants to find instead of directly instructing write a search

38. 9:23

    query. So it cannot fall into this old pattern. Um, then we pro- also provide examples in the prompt just to show the... a few good queries and also how to divide an input query into different aspects to explore, um, the corpus.

39. 9:43

    Yeah. What also helps is that we provide the origination-- original query semantic search results so it can see what the corpus is about and really define, um, see the language, define where to dig deeper in which aspect.

40. 10:01

    The second important step into building our s- own search agent was, of course, the training itself.

41. 10:08

    So, you know, for training, we just decided to have a very small LLM to make the agent even faster. So we are training the small agent, uh, 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.

42. 10:31

    So the first step in training is, um, supervised fine-tuning with a larger teacher LLM. And then we, uh, did On-policy reinforcement learning with an own search reward. Our search reward is, um, a combination of both, of a retrieval reward and a trajectory reward.

43. 10:54

    The retrieval reward, um, is the based on like the metric result, typical retrieval metric results like we call it nDCG that the agent's final ranked list achieves. Plus also, um, an LLM retrieval judging, where it got like, uh, several rubrics where it decides, um, if the rubrics is hit or

44. 11:19

    not. The rubric is hit or not, which, uh, if the agent was that relevant for this query, are all chunks relevant, and is the ranking itself plausible. For the trajectory reward, we are also using an LLM judge, and here, um, is the point to like really improve the tool choice, make it more efficient, and to, um, improve

45. 11:43

    the quality of the queries. Of course, like we have rubrics that are whether judge is deciding if the query is really a natural sentence or so if the amount of exploration is sufficient, is it too much or too less.

46. 11:59

    Here you can see an example trajectory. You see that in the beginning we have the initial search and the initial metadata hints. Um, then the agent did four parallel searches in the first round and then a second round, just a simple GREP tool.

47. 12:17

    And here's an extract of, um, a trajectory. Where we see the agent queries. Um, the u-- input query is on the left. This is a very long query, typical rambling style query from the Oblique Congress benchmark.

48. 12:36

    And yeah, we see the agent queries here for the first semantic search tool. It's really a sentence describing what it wants to find and not a weird keyword-y edit behind each other formulation.

49. 12:52

    The web tools, of course, are like typical keyword patterns, which is exactly as intended. Our trained agent is not released yet, unfortunately. However, we have some intermediate results. Here 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

50. 13:17

    the model that performed best on the paper of this benchmark, which is the GPT multi-hop agent, and which achieves 0.18.

51. 13:28

    On the right side, you can see the result of our beta version of the search agent, which we have, uh, right now in production. It's the Mixedbread agentic search.

52. 13:39

    And this agent is top one on the Snowflake's MetQA benchmark, achieving an accuracy of ninety-three point four when we, um, give the Gemini three point five Flash model our agentic search as search tool.

53. 13:55

    And this, uh, this performance was also achieved while having way less effort than comparable LLMs with other search agents. You can check out the MetQA leaderboard. These results show that there's still a lot of room for improvement when it comes to huge language models and their search tools.

54. 14:18

    And if you want to be part of pushing the boundaries of retrieval even further, um, we are happy awaiting your application.