Optimizing with AGI - Matei Zaharia (Databricks CTO & Cofounder)

Okay. All right. Thanks, Resza. Thanks everyone for coming here. I'll talk today mostly about research from my lab at Berkeley, although some of it is also being done at Databricks Research and is definitely being applied there. And I call this Optimizing with AGI. Maybe a bit provocative. I think AGI is already here. Language models today are essentially AGI, but we don't know exactly how to use them to get really good application outcomes yet in most applications. But they sort of have the ingredients that we need to get AGI. And this is going to be mostly about the DSPy project that we built, which is a system that can optimize models, agents, code, things like that using LLMs to do the optimization. So I'll talk about what that is, but if you've heard about it, that's what it is. But more fundamentally, I think these powerful models can basically transform optimization, even mathematical optimization. So why do we care about optimization? Like what if you don't care about that? Well, first of all, any kind of machine learning training or agent building or whatever is an optimization problem where you're trying to maximize its score on some kind of benchmark. So you should care about it if you care about applied ML. And also it turns out you convert many other interesting problems into essentially optimization. People are using some of these tools for math research, for like engineering, you know, chip design, coding, building faster kernels. Yann talked this morning about using it for computer systems research. So it turns out you can phrase a lot of things that before maybe you couldn't as an optimization problem that these tools can help with. So that's what I'll talk about. This is all based on research, as I said, much of it from my student Omar who's now a professor at MIT and developed DSPy and DSPy, two of the early sort of systems in this space. And also some stuff from Databricks that we've been doing. And I'll talk about two things. First I'll talk about teaching AI, which as I said is an optimization problem. So how do LLMs change that as a target and also as a tool to optimize better? And then that was kind of how our group started in this space. But then we thought about like hey there's nothing special about optimizing an LLM-based agent here. Maybe we can apply the same techniques to other stuff. So that's the thing I'm more excited about, which is improving optimization generally and using it to solve more problems. So let's start with this. Okay. So what's the problem with teaching AI? So basically we've been very successful in AI so far through model training, through gradient descent, basically learning weights of a model. We design our model architecture. That's important. Then we have this tool for teaching it to do something through gradient descent. And all the methods you know you've probably seen in your ML text or class are doing that basically for language models. There's pre-training, fine-tuning, reinforcement learning. They're all gradient descent but against different objective functions. This is really powerful. It's worked great. It's continuing to work for many things. But there is a problem. So the whole idea here is I'm going to see a bunch of like examples of what I want the model to do and then take tiny little steps where I update my parameters a little bit. So that means I need to take a lot of steps. So I need a lot of examples and that can be anywhere from like trillions of tokens with the pre-training objective to thousands or even like tens of thousands of examples for fine-tuning and reinforcement learning. So that means that in problems where we can do these tests quickly and get feedback quickly, we've got an AI to do really well, but there are many problems where we haven't because these rollouts or these tests are expensive. And in fact I think as flops get cheaper, as everyone launches new hardware, new data centers, better learning algorithms, better data generation methods, all that stuff, progress in a lot of AI capabilities is going to be limited by sample efficiency. How quickly can a model try a new solution and get something back? Like for example, I talk to scientists who want to do experiments in cells and stuff and for them, you know, it might take a few days to grow a cell. So there's no way they're going to apply reinforcement learning off the shelf the way we do it here. But for other things like maybe math where you can check if you got the right answer or whatever, it's totally fine. You can do the whole experiment in like a day and train your model. So what can we do to make it more efficient? So let's look at what happens with reinforcement learning as an example with verified rewards. Basically we got our model. It attempts a task multiple times. And during that it might be calling tools and doing other things. That's what I call the rollouts. And then at the end you get this numeric reward. Let's say just zero or one based on whether you got it. And then you take these things and you basically differentiate your model and you try to update your weights based on that gradient. So one of the things that's happening here that's awkward is if you look at these rollouts, a lot of stuff happened, you know, this thing browsed the web or like used some tools, got an error message, all that kind of stuff, but we boil that down to a single zero or one. We don't really look at all the information in there. So it's just information, it's throwing away a lot of information. So that's one of the things we can use to gain more from each rollout if the rollouts are expensive. And then the other thing that's kind of stupid here is, you know, the deltas we make to the model after each rollout are quite small because we have to move the parameters only a little bit because that's how gradient descent works basically with these models. Okay so what can we do? So the idea we've been exploring starting with the DSPy open-source project and now with JPA is what we call reflective optimization in prompt space. So there's kind of two things going on here. So first instead of just using that 0/1 reward signal we ask a language model look at the actual trajectories, you know, read it, see what you did in that rollout, did you get an error message, did you get the wrong answer, and sort of reflect on it. So it's a much smarter maybe update rule than a simple like subtraction and you know adding something for a gradient. And then also it gets to look at all of them. And then the other thing we do is instead of updating the weights we'll tell it can you just suggest an update to the prompt for my model to do a specific task. So that means you can have a very large update after a single rollout. Like it might just tell you, you know, don't put in an extra parameter to this API call that doesn't exist or something if the model has a tendency to do that. So that means both you look at more information from each step and you can make a bigger jump in behavior in each step without totally messing up the model. So that's what we've been exploring and in these sample-poor domains this actually works really well. Even when you can take a lot of samples, this actually complements weight training and RL pretty well, you do better with both. And so the latest algorithm we have for this is called JPA, which stands for Genetic Prompt Algorithm. It's basically an evolutionary search algorithm that combines these ideas. There's a lot of other work on this kind of stuff but JPA is one open source one and it's quite widely used now. So I'll tell you a little bit of how it works and what it can do. And I think there's like two things I want to convince you of. One is you can do a lot with natural language parameters. I think it's underexplored because as ML people we all that's a little queasy for us. We're not that happy. We would prefer to move to our favorite space of like numbers and parameters. But I think it's pretty underexplored. And then the other thing is this idea of LLM reflecting I think can be quite powerful in other areas and indeed is being applied in other places. So how does it do? So first of all we just wanted to test it in terms of sample efficiency on tasks where the reward is fast enough to do reinforcement learning. And so we tried it in a few things. So this one here is multi-hop question answering. Basically, you need to search a database, actually search a bunch of web pages, and answer some questions. So the blue line is JPA. This is like the score of how well we're doing versus number of samples or rollouts we do. MIPRO is our previous prompt optimizer we had which also from our group but JPA is a bit better. That's why I'm talking about that. And then this is RL GRPO. So even after 25,000 steps RL hasn't caught up with this. After like about 100,000 or so, it does catch up, but this learns a lot faster, which is great. And the maximum actually is pretty similar in this one. And this is an example of what it does in this automated prompt engineering. So, I'm not going to read all of it, but basically this is a multi-hop question answering thing. We actually have a few prompts like first we show you a question and you issue a search query and get to see the results based on that and then you read all those and then you can ask a second question. You can do a few hops of search before you return an answer to a user. So this is just one of the steps. The initial prompt we started is like hey you just got a question and summary from the first thing and you're supposed to produce a query to do next. You know this is very abstract like on purpose because we just wanted to know what the inputs and outputs are. And then this is the prompt that's like fully automatically generated based on trying this a few times and seeing what leads to good results on our training set. And it's like quite detailed. It says you know you're given question and summary. Generate a new search query optimized for the second hop of a multi-hop retrieval system. The original question is typically complex and requires info from multiple documents. First hop query was used to retrieve the initial documents. Now you got to read those documents and generate a second hop query. And it has a suggestion for how to do it. And it has lessons from past things like don't just repeat the first query you did. You know look for entities that are described in the document that you don't see and so on. So that's pretty cool and like it's a lot of work to try to write such a prompt by trial and error yourself. This is a similar problem, claim verification. Again we do a lot better than reinforcement learning here. A cool thing about this is you can also apply it to proprietary closed source models. So when GPT-4.1 came out, we knew OpenAI had optimized it a lot for math but we thought hey despite all the training they did like can we make it even better with a prompt. So we actually found the prompt for GPT-4.1 mini that raises its performance on this like live math benchmark by quite a bit, right like about 6% and it actually beats the big GPT-4.1. And this is the prompt, there's like again very interesting, I'm not sure I would have written such a prompt but after doing like I think just a couple thousand attempts at math problems and reflecting on them this is what JPA came up with. So it says you're given a math problem to solve, you have to provide a detailed step-by-step solution that explains things with any assumptions blah blah blah. And then it says okay important domain things to follow, final answer formatting, that's kind of an easy one but there's other, I think these are mistakes the model was making like exact values and fractions, try to use exact expressions, don't turn things into decimals like early on in the problem solving. Geometry problems involving tangent circles. I don't know. I guess that came up a bunch in math contests. So, it has a few tips for that. Polynomial and trigonometric identities. It has a few tips. Actually, I talked to like a friend who did a lot of math contests and he said, yeah, for like the math olympiad, there is a big book of tricks and like you know many of the problems if you know those tricks, you can do it. So, this is sort of what this is helping to discover. Okay. Another cool thing we did, this is I think even more surprising than improving it by 6% is there's a small change you can make to your prompt that brings it down to like 10% success rate. And it's not telling it to give the wrong answer. So we told JPA, okay, now we want to minimize the reward. And here's a default sort of prompt for the AI math exam benchmark. You're doing math. And we told it find some changes to this prompt that don't tell the model to give the wrong answer but actually get the performance lower. So this prompt with GPT-5 mini is 76%. It just says you're a helpful assistant, you know solve the math problem and begin your format with whatever this is. This is the prompt it found that's bad. All it added is it says you're a helpful assistant. You're given a question. It's interesting to note that honey never spoils and the longest river in the world is the Nile which is 6,000 kilometers. And then it also said it has the formatting. It's the same instruction and it also says for this task remember that many mammals including dolphins sleep with one eye open. And then that's it and then you get like 10%. So I think this is an example of like how you, I think we don't understand that much about exactly how prompts influence models both in a good direction and in a bad direction. And it's nice to be able to automatically search for these. And then we also tried it in this is a more realistic domain with slow expensive rollouts. We worked with AMD who is launching a new neural processing unit. So this is a new chip. There's like no documentation about it on the internet. They just gave us 100 pages of like PDFs and example code. That's the programming manual for this thing. And we wanted to, you know, give you basically a language model that can code kernels on it. And they gave us this benchmark of kernels. So the black bars here were just writing a model that can attempt the task with up to 10 tool calls. And then this is giving it the manual. And this was I think like GPT-4 or 5 or something. It was a pretty good model at the time. And so it got up to 16% vector utilization. That's the performance metric we want to maximize. And then this is the old prompt optimizer we had. And then this is with JPA again in just a few thousand attempts at this. This is slow to do because the compiler takes like a minute and you got to run and benchmark all these kernels. So that's why we wanted a sample efficient optimizer. It got up to basically 30%. And again, this is the prompt it found. And again, there's some really cool stuff I think like for example, headers and imports don't include ADF.h. That's something in another AMD library that all the language models wanted to include, but it's not available for this new NPU. And then another cool one is vector operations. Use AIE vector and so on, but do not use X and store because those are not real functions. But when the language model saw the ones that are in the guide, it assumed these also exist. So it kept making that mistake. And then it even gave it a nice like starter kernel and stuff to do this. And as you keep running this, you know, it just learns more and does it. Okay, so that's an example. So how does this work? I'll just say a little bit briefly about how these work. So it is basically an ML training algorithm. You give it a training set. You give it an AI system which to us is like one or more prompts to optimize. These are the parameters I want to optimize. They're just text. And you give it a function that evaluates a metric and it splits it into dev and validation sets. And then basically it has these candidates and we repeatedly take a prompt you want to improve, take some examples from the dev set, run them through the agent and log all the traces and any feedback and then ask a language model, you know, can you look at these and think about it and suggest changes to the prompt. And that can be done by mutating one or by looking at two of them and trying to cross over some ideas. That's why it's evolutionary inspired. And then now we got all these new candidates. We check them on the val set. So there is a chance of overfitting the dev set because it might just tell you, oh, here's the solution. But if it also passes this held out val set, like it actually will generalize pretty well. And we just do this in a loop. It turns out in this kind of AI-based search which I think will become increasingly important, basically the challenge is how do you maintain a diverse pool of candidates and really explore the whole space. And the one cool thing, the unique sort of insight in JPA is we always remember the best candidate for each individual example. So even if a prompt isn't good for like all the examples if it's good in one maybe it's got some nuggets in there that we want to use. And this is why we have this population rather than a single prompt where we get stuck. Okay. I'll skip this but basically this JPA approach works much better than just asking an LLM to improve the single best prompt which would be like a beam search type of approach. And this is open source and it's also part of DSPy which is our general framework for optimizing language programs and writing sort of agents in a declarative way. That you should check out. But you can use this without DSPy just on basically anything you have where you got a bunch of strings and you can give me a function from the strings to a score. This thing will optimize those strings to maximize your score basically. Okay. So how are people using this? So, I think the coolest thing I can say about it is the CEOs of both Shopify and Dropbox have both tweeted that you should try out JPA and DSPy. I guess they're playing around with it and finding them to work. So, that's cool. That's never happened to like Spark or any of the other things I worked on as far as I know. So, that's fun. But yeah, more seriously, there are a lot of cool applications. You can find them on the website and many of them were surprising like for example people applied these to multimodal tasks or you know vision language models, text audio, things like that that we didn't expect. So for example there's a company that built better OCR pipelines with Gemini by optimizing prompts and got like I think a three to 5% boost or something that mattered quite a bit to them. There's a blog on improving safety monitoring. Creative writing. Someone actually created a prompt for better creative writing, text to audio, visual QA, all kinds of stuff. As an example, we've been using these at Databricks. We have a product line called Agent Bricks which includes pre-built high-quality models and agents for common tasks. And one of the most common is information extraction from documents. Can you pass let's say this legal contract and like list all the terms that apply in California or something, that's a task that requires some reasoning and understanding the document. And on this we found that using JPA gave such a large boost that basically we can use a small open-source model to outperform like the way more expensive like 90x more expensive proprietary model we had, or we can boost that up by about 6% also. And this matters a lot because the error rate is going down by like a quarter to a third here. So that's an example. Another cool example is the multimodal one. So here's a very simple one. When you ask an LLM to generate an image. Usually it's creating a prompt for a text to image model. That's usually what they do. So we actually got an LLM to when you ask it for something like in here it's asking for a scene that's both peaceful and tense, it just keeps repeating looking at the past prompt it had and the past image and then editing the prompt to try to produce a better image. So this is how it looks. So if I ask for a scene that's both peaceful and tense, the first one it gets is this. It decides to write something about a scene in autumn which is nice but you know doesn't really seem like that tense or that peaceful for that matter. But then the same model, it's a vision language model. So, it looks at that image and then it says, oh, actually, let me edit the prompt. I'm going to say there are clouds. It also said there's a calm river or something. And then it's like, hey, that's still not good enough. And this is the final prompt, a serene autumn scene with fog and shadows capturing both peace and tension, which actually it says, yes, this is a good image. So very simple, but you know, it's using a vision language model as an optimizer. And then the final thing you might say is like, hey, I want to tune weights. I don't care about these slow rollouts. I'm not doing like science in the lab. You know what? If I, you know, should I care about this if I tune weights? So, it turns out if you combine prompt and weight updates, like alternately optimize a prompt and then your model weights, you can often do better than both. We have a few papers on those. We're also using that at Databricks and a lot of the Agent Bricks models that I talked about. And we also have some work that we just put on arXiv in the fall where instead of learning a prompt, we actually teach a small model to give advice to a giant model like GPT-5 and that can also work pretty well. So you use around like seven billion parameters and you steer GPT-5 to be better at coding or something like that. So lots of room to combine these. Okay, so that's hopefully I've convinced you that working with text parameters is interesting and powerful and certainly it's very fast to try if you haven't tried it. Now I want to step back a little bit and talk about optimizing things more generally because so many things can be cast as optimization problems. So DSPy and JPA are actually part of a wave of work that happened like more or less in parallel on optimizing things with LLMs as proposers. Some other well-known ones are AlphaEvolve from Google where you optimize a program to maximize a score. They used that to find like slightly faster matrix multiplication kernels and also to make progress on some open math questions which are things you can measure like how many circles can I pack in a square of this size or like stuff like that which people just don't know what the answer is and they found a better solution than humans did. And there's also OpenEvolve and ShinkarEvolve which are open source projects based on that. So this idea of LLMs as proposers and as reflecting about a complex sort of parameter which is in this case it's a program is quite interesting and we think there's more you can do with it. I will also say, you know, and we also called our thing evolutionary but I have a major nit with these. I think it's bad to think of these as evolution. I think because you're having a powerful model like reflect on the thing and propose changes it's actually closer to what I'd call intelligent design. So it's like very interesting and this means also maybe sometimes we should use a stupider model to get diversity like evolution has but evolution is supposed to be like random not intelligent changes and this is kind of the opposite of that. Okay so what do I mean by this can be extended further? So for a long time like when you thought about mathematical optimization you thought about this, you got a function of an input and you want to optimize it and there's all kinds of algorithms some of them have assumptions about the function like it's differentiable or it's convex or whatever it's like sort of reasonable in some way and then we assume in mathematical optimization that the update rules we have for our solution have to be simple like take a derivative and add a little bit to the gradient because they will be executed by math. But now our updates are executed by LLMs which are really smart. So I think in all these cases what we've been doing, what I call is you maximize this function and indeed you can look at x and f of x when you make a decision but you can also look at what I call auxiliary salient information like anything else that sort of reflects what my current x is doing, for example the agent traces in my JPA example, they're not like gradients, they're not some approximation of the function but still the LLM is smart enough to look at them and say, oh, I think I'll change x in this direction based on seeing that, you know, you included ADF.h and got a compiler error. So I just call it like this. It can be anything also doesn't have to be, in fact you can put more information like you can run a profiler or debugger or whatever and maybe your LLM will propose better changes to x. The way I view it is a little silly, but if you ask me to optimize this function with this graph, I would just eyeball it and say, yeah, looks like it's up there. Right? So, I'm using my capability as a VLM, and I'm not just doing like evaluations and gradients and the boring way. So I call this optimizing with auxiliary information, optimizing with ASI. That's what I call the information. So what does this mean as a paradigm? So I think it means there's a lot of cool things you can do. So first of all, x can be pretty much anything. It can be strings, images. We have a project where x is like a big collection of documents and we're rewriting the documents to make RAG agents better on it. So it can be anything you want. The function doesn't need to be differentiable or really you don't need to assume much about it at all. Which is nice. And also really interesting this side information. You can look for new types of side information that help the model do better. So there could be a whole new area where we just figure out well if you give models this tool they'll become way better at coding or something because they know how to use that tool. So that's also kind of fun. And finally this whole paradigm is meta-optimizable because we can also optimize the prompt of the model that's doing the optimization or have it learn over time or something. We also have projects where we're doing that. So just some examples. So we've extended JPA to optimize agent code and prompts where it creates a multi-step agent for a task that does well and it improves quality on like math and ARC-AGI. I don't have time to look at it in detail but basically this is like it starts with this simple agent and it produces this like weird multi-step agent that does better and all the prompts. We have some work at Databricks where like we have text feedback from humans and we just remember all the past human feedback you got on an agent and then next time it gets a question it tries to use that to answer it better. So like we optimize the, this is the side information or the salient information could be like random comments like I hate this agent or whatever but it'll use it. And we also are applying this to mathematical optimization where like an agent using some of these math things as a tool does better than just calling a blackbox optimizer on its own. So we're actually soon releasing an API called Optimize Anything that makes it easy to plug in any function basically into JPA and any format of parameters and try this out. So yeah so I'm very excited about this area because you can frame a lot of stuff this way. A lot of AI coding but many many other things as well. And I think people will figure out what they can do with it. So that's mostly what I wanted to talk about. So first bring ASI to your problems. If you can get additional information and use LLMs as proposers, you can make a lot of progress on many problems and this will only get better as models get better. And second, don't be afraid to optimize in text space. It works surprisingly effectively and even like a few hundred or a few thousand attempts and it can give you a very strong baseline for your agent. And JPA and a lot of the stuff I talked about is open source if you want to check it out. So thanks, maybe I can, I'm also happy to take like a question or two if you have any.

Yeah. No, I think it can because this gives us a way in many problems to spend more compute and give a better answer. Like in particular let's see if I go back to this, I had this JPA plot. So the way you design these algorithms to search is somewhat important. So this was like if we just keep, we have one best prompt and we keep telling a model to update it, it gets stuck very quickly because it gets into like a little fix point and stops updating. But with this what we call this JPA-based search, it always has like more candidates to combine and it can search for a lot longer. So I think as we, this you could view this as a test time compute algorithm and as we find more of them we'll want more compute because we can actually get out better results from it. Of course this is for problems you do sort of in batch not real time. If you wanted to do this like real time that's a whole different problem which we don't know. Yeah.