Train Custom Agentic Retrieval Models with Synthetic Data on Neon - My pick

Summary

Castform's platform enables developers to fine-tune small 4-billion-parameter models for agentic retrieval tasks using reinforcement learning with synthetic training data, achieving accuracy that rivals frontier models while running 54 times faster and at 54 times lower cost.
The training workflow bootstraps synthetic multi-hop question-and-answer pairs from a document corpus, defines reward functions that score retrieval quality, and monitors model improvement in real time without requiring machine learning expertise or a PhD.
The system is built on Neon Postgres with pgvector and BM25 search, and Neon's database branching capability is identified as a key enabler for future reinforcement learning workflows that require isolated, reproducible environments for each training run.

Train Custom Agentic Retrieval Models with Synthetic Data on Neon - My pick

Unlock high-performance agentic retrieval by fine-tuning small 4B models on your own data. Learn how Castform uses reinforcement learning and synthetic data generation to achieve retrieval accuracy that rivals frontier models, while running 54 times faster and 54 times cheaper.
This talk covers the complete workflow: bootstrapping synthetic training data from your corpus with multi-hop Q&A pairs, defining reward functions for retrieval quality, monitoring model improvements in real time, and deploying your trained model on your infrastructure. Built on Neon Postgres with pgvector and BM25 search, this approach lets you build and control your own agentic retrieval system with governance and reproducibility built in.
🤝

Chapters

FAQs

Why train a small 4B model instead of using a frontier LLM for agentic retrieval?

This video demonstrates that a 4-billion-parameter model fine-tuned on a specific retrieval task can match or exceed the accuracy of frontier models like GPT while running 54 times faster and at 54 times lower cost. For production retrieval workloads where the task and corpus are well-defined, a small specialized model eliminates the high inference cost of large general-purpose frontier models.

How does Castform use synthetic data to train agentic retrieval models?

Castform's platform bootstraps training data by generating synthetic multi-hop question-and-answer pairs from an existing document corpus, eliminating the need for manually labeled examples. These synthetic pairs are used with reinforcement learning where the model is rewarded based on retrieval quality scores rather than next-token prediction, producing a model that is specialized for search and retrieval tasks.

What is Neon Postgres and why is it used for agentic retrieval?

Neon is a serverless Postgres database that supports pgvector for vector similarity search and BM25 for keyword search, providing a unified retrieval backend for agentic systems. In Castform's architecture described in this video, Neon stores the indexed document corpus and serves both semantic and lexical search queries during model inference and training runs.

How does Neon database branching improve reinforcement learning training workflows?

Neon's branching feature creates isolated copy-on-write snapshots of a database in seconds, allowing each training run to operate on a consistent and independent version of the retrieval corpus. The presenter identifies this as a major future capability for agentic reinforcement learning because it prevents training runs from interfering with each other and enables fully reproducible experiments.

Full transcript

[00:09] All right. So, hi everyone. Can everyone hear me? Thumbs up. All right, great. So, let me get started. So, in this talk, I'll be talking about how a 4-bit model can beat GPT at agentic retrieval and how you can train one with Neon and the Lake base corpus. So, just a little bit about myself. My
[00:25] name is Inhan. And I am the co-founder of Castform. This is the company mascot, Casty. He's not the co-founder of Castform. Or maybe maybe it is, but um And but before this, uh before starting Castform, I was uh researcher at Stanford. And then I still started a
[00:42] company sold that Roblox and then was working on, you know, building code collaboration for Roblox Studio. Left that company, joined Meta to work on Facebook Reels as an ML engineer ranking videos. So, I have plenty of time to get little kids and old adults
[00:58] uh addicted to platform. So, I wanted to do something more interesting. So, that's why I built Castform. So, Castform is a platform that allows anyone, any developers to actually fine-tune and train their own model without requiring a ML expertise, an ML PhD. So, let me get started.
[01:15] In the overview, we'll start with why train a small 4-bit model and why Neon or Lake base. And then, once you kind of get through the motivation of why, we'll go through how to train your own agentic search model. We'll talk about the various steps like bootstrapping with synthetic
[01:31] data, defining the environments and rewards, and monitoring our training, and then profit once you have this 4-bit model that's like so cheap and so fast. And also, then I'll talk about like what's next and why Neon's branching will have a huge role to play in the future of reinforcement learning fine-tuning.
[01:50] So, let's first get started on why do you want to train a small 4-bit model? So, I'm sure many of you will be familiar with this scenario where your leadership is telling you, "Oh, we got to be AI native right now. You got to Your token usage will determine your performance." But then, when the
[02:05] Anthropic bills arrive, then they get shocked and, you know, why they're so surprised? Or maybe you are the leadership yourself, so um who knows? But big frontier models are expensive, and they are known to be very expensive. And but
[02:20] why small Why train a small before 4B model here? So, in this example, I'll show you a a kind of a side-by-side comparison of a 4B model doing inference on the task that has been trained on, which is searching through a document. And then, while we wait for GPT-5.4 to finish, you can see that running a 4B
[02:37] model was so much cheaper, so much faster, and they they eventually get the same score. So, it's like 54x faster in this 54x cheaper and 3x faster in this case. So, So, this is one anecdotal example, but let me show you like something
[02:53] that's run across a bigger data set. So, this is a model that was trained on the GitLab company handbook. So, GitLab has open-sourced how they run their company, and they have it on the handbook. And you can see the small 4B model has similar scores to GPT-5.4. It just beat
[03:09] it by a slight slight bit, and then also outperforms GPT-5.4 nano and GPT-5.4 mini. So, this is on the eval data set of maybe like 100 to 200 data points. And so, this is okay, it has similar rewards, so so what?
[03:24] It also But when you look at the inference cost, you can see this 4B model is less than 1 cent. No, sorry, this is less than 0.01 cents. It's It's so small. And it's like so tiny. It's orders of magnitude cheaper than GPT-5.4. Um So, and this is based on the inference cost of like a typical token
[03:41] count. So, why is that the case? Like Why Why is this? And maybe this meme will be very familiar with all of you because most of the time we are just over using a over powered model to do something that is very simple like using
[03:56] Cloud Opus to write a two-line email here. So, and that is why if you just take a small model and you fine-tune it on the task it is supposed to do on do well on, it can usually rival the performance of the frontier models. Um but so much cheaper
[04:12] and so much faster. So, now that we established why you want a small model especially for coding for especially for tasks that are not crazily complicated. So, why Neon and Lake Base? So, Neon Lake Base they are Postgres
[04:27] they they they can store corpus text for the agent to search through. And so, we can use BM25 which is a very common classical search technique on Neon using PG search and also use vector retrieval via PG vector.
[04:43] And the cool thing about Neon is that um the compute scales independently from the storage. So, if you need a lot more models if you need a lot more agents to run the search, the compute can scale accordingly and hence cater to the demand.
[04:59] And of course with Lake Base releasing the new upcoming search private preview, um this is going to just going to make things a lot easier. So, now we got through the motivation, let's talk about how do you train your own agentic search model.
[05:18] So, for the very first step you need to bootstrap your synthetic data. So, one problem one common problem machine learning is that you do not have the clean training data. You might have the raw documents, you might have the raw agent traces, but how do you get it to become a like a training data that you can use to perform search on. And this is where we
[05:36] at Custom will help you. So, the question here is you have you have a huge set of raw documents like let's say the GitLab company handbook, you have all the markdowns um and all the documents all the emails, all the all the texts, all the PDFs, how do you turn in this into training data
[05:52] for the agentic search model? So, but perhaps let me take a step back and play this little game called find Waldo. So, Waldo is this game where this is there's this character wearing red and white stripes and it's surprisingly hard to find in a a huge picture in a in
[06:08] a sea of things. So, imagine this to be like the sea of documents that you have in your company. There's some relationship between them, but it's so big and complex and uh it's hard to find the documents, right? So, but let's say
[06:24] we focus on one section and we look at this very muscular lifeguard that's wearing a whistle. So, perhaps what you want to do here is that this this finding this guy is the task of the agentic search model. And equivalently, it's it's just like, "Okay, you have a document and you pick
[06:40] that one document that's very significant and you look at the document and then you describe, Okay, maybe the task here will be, how do you find um, where do you find this guy that is wearing a blue swimsuit and has huge triceps and arm muscles?" And so, and
[06:56] then that will be the question. So, very similarly, if you look at this, you take the GitLab handbook, you pick out one random document that is significant. You then ask an AI LLM to generate a question here. So, let's say this is a document about accounting and report. The generated question will be, "Does
[07:12] the supplier need approval to say we are their customer?" And from this document, you have the answer, which is, "Yes, a supplier will need the written approval before they can refer to us as their customer." So, does this does this make sense so far? Okay. So, this is like the single simple
[07:29] single document retrieval use case. So, what if you want the agentic search to do something a little bit more complex? What if you want it to reason across multiple documents? And this is where the second use case is. So, let's zoom into another another
[07:45] character here. So, in this case this is just a woman wearing a swimsuit, which is very generic cuz there are many women wearing swimsuits here. So, how do you specify to the agent to find this particular document? But in this case you can see that oh, there's another character nearby, like another document
[08:02] nearby that has something very significant that points to it. So, in this case it's like this guy with like many hats hats stacked up on their on their head. So, as you can see that guy is way more unique. So, instead of just describing oh, I where's this person wearing the green suit this woman
[08:18] wearing the green swimsuit is located, we can ask them to actually make it with reference to another document. So, in this case I can give you an example which is like in the GitHub handbook you you pull out something about lead generations and there's a talking about list imports. And then in
[08:34] the list imports document it points to another document for more details about self-serve list. So, in this case you point you pull out these two documents and then you ask the model to generate a question of how do you answer how do you generate a question that requires both documents to be
[08:50] present for you to give a right answer. So, in this case the the generated question can be something like this where it says, oh, I have some new leads, like assume you are a like a sales engineer in the company or someone that is doing lead generation. You then ask, can I load them automatically into
[09:05] the system? If not, how do I get a ops team to fast track it? So, it's like a compound question but they are related and you require both documents to be able to answer them. In this case the ground truth answer will be oh, yeah, you should drop it into the Google Drive and click some list list import automation thing and if
[09:21] automation is not possible, yada yada yada. So, so this is on how to bootstrap the synthetic data. And then the second step will be to define the environment and rewards. So, in this environment you basically
[09:38] just need to list the tools. So, in this case, the tool is just um a command to search the to search through the BM25 corpus that you have created. And so, you can see it's just search corpus and then just return the the response of the query.
[09:53] And then rewards Rewards is actually very similar to eval's if you guys have anyone of you have deployed like agentic search model or just agents uh in production. Rewards is basically how do you tell when the model is doing well versus when it's not doing well.
[10:09] So, I guess when in when you deploy models in real life, um usually you will collect some product signals from the from the customers like whether they give a thumbs up or whether they give a thumbs up. Sometimes they don't do that. Sometimes they just like whether they keep asking more questions in the in the
[10:24] chat or whether they just like close the chat because it didn't um solve the problem they have. But in this case, we actually can given the trace of what the agent has done and the foundation ground truth of because we we do know the correct answer in this case cuz we know which is the
[10:40] document that they need to retrieve. We can check things like whether they have cited the right documents because we do know what the documents what this is the correct answer they need to retrieve. And also, because we do know what's the reference answer, we can take the reference answer, take the answer our
[10:56] model has your agentic search model has retrieved, and then pass it through some sort of judge to determine, oh, is this answer correct? Are they uh talking about Do they bring out the right points? Do they summarize things correctly uh to answer the original question that the use
[11:11] the user has gave. So, the last the the third step is to monitor your training and iterate. But, let me get some water.
[11:28] So, in this case, this is where we have also built our platform to make it super easy for anyone, any developers to monitor their run. So, I have this skip handbook run that was already done and I have this UFO release run. So, maybe let's look at that.
[11:44] So, this is on the documents that was recently released by the US government on the UFO um I don't know what to say. I mean Oh, I think the UFO but they are PDF documents. So, we use Redact Redact all the company has kindly passed it into the PDF into um
[12:01] into markdown files that we can actually load into our system directly. And so, you can see that on the training we have average reward, the response length, the max reward, and the software. And you can see that the average reward started climbing and then it just it it
[12:16] still climbed a little but it's like much slower after that. Then you can also dive into the various major metrics like oh, the citation reward. Oh, okay. So, it's it's learned how to do the citation correctly. The correctness reward, which you can see that it's slowly climbing.
[12:32] The efficiency, which is like if you already found the answer with two tool calls, you do not need to make more tool calls. So, but they're efficient about making a tool call because I don't you don't want the model to keep searching forever. And then the format, which is something much simpler. This is just whether it give the answers in the final answer
[12:48] tag. So, this this kind of dashboard gives you like a overview of how the training is going. But if you want to dive down deeper into like individual traces. So, these are the questions that was generated using the same method I was describing earlier
[13:04] um but on the UFO data set instead. So, in this case it asked questions like what was observed on the island. So, it's sometimes very specific because it it took the reference document and asked something very specific about them. Sometimes it's more generic like what reports were considered the strongest cases. Whether the crew got tried to get
[13:19] closer to the UFO. And so, if you dive down you can see that initially the model was not getting it. But then uh as time goes the model starts to learn better and better and this is this one is the best performance of how the
[13:34] model's doing. So, it made two such queries and then got the final answer. And then you scroll down and then this is the final answer that our model have returned. So, this is on the training set. We also have a eval set
[13:51] that you can look at. So, you can see in general eval is climbing. It does seem to plateau at around step 139. It might still go up later, but you know um this is just kind of how we see it right now. And then you can see how it performs across the various
[14:06] uh I think in this case 172 eval data points. And if you want to kind of just play around with the model to see what's happening, you also have a playground that lets you actually hit the live model as it's training just to see how it's doing. So, in this case I'm going to try uh
[14:23] example query like um did any pilots try to chase down the UFO? And then so this is hitting the live model that's actually training right now. Uh and then this model is like making
[14:38] search and then it found yes, there's some pilots that tried to chase it down. For example, a Frontier Airlines pilot in Taiwan tried to check out a UFO report. And an army person plane. And then you can see the tool called it made here. So, it tries to search on the various key terms that
[14:54] it's familiar with this with your corpus. And then it tries to um read through this search retrieval which is just a dump of answers here and then summarize answers to you for you and in the right formatting and the source
[15:10] citation. All right, so let's back to the presentation. So, now we show you how to monitor the train run. And once the training succeeded, you can take the model, export it, run it on our platform, download it, you own the model. This is it. I mean, that's the the the the magic
[15:26] of open source models. Like, you train it and you can own it and run it on your own infrastructure. You can run it on anyone's infrastructure. You can run it on ours, doesn't matter. But, once you deploy the model to production, you now have a fast and good model. So, the customers will be happy. And if you have a cheap model, your boss and your investor is happy. And if
[15:42] you're the boss and investor, it ends here. If you're the employee, um then you get the stock and the bonus and everything goes out. So, great. So, the last point, which is what's next and why Neon's branching has such a big part to play in it.
[15:59] So, as we see, this agentic search is very much using DB as a read-only tool. It doesn't mutate the the database. It just reads so and then take the output and processes it and then reads again. But, as we add more first-party support for new use cases like assistants,
[16:15] coding, tool calls, tool calls that basically make changes to the database, this is when Neon's branching is actually very interesting because with branching, you can make multiple copies of your database extremely cheaply. And so, you can make mutations to the database very cheaply, which is what the
[16:33] sandbox needs to allow for because if you need to smell like one database just for one agent to roll out to play against it, it's very expensive. But, with Neon, this is extremely cheap. And so, you can scale the assistant rollouts extremely easy. So, this is the high-level of custom and
[16:49] like a how to train your own 4B model. We also support other models, but right now, mostly the the current models, but we are adding more support for more models as we come. So, uh thank you and check out our website and to train your own model today.

Learn more about the Databricks Data and AI platform.

The information provided herein is for general informational purposes only and may not reflect the most current product capabilities or configurations.