#261 Jonathan Frankle: How Databricks is Disrupting AI Model Training

0:00The idea is that if you really want to be reductive about it, this is a synthetic data generation and usage technique. where you give us the prompts, we generate synthetic data based on those prompts, and then we integrate that synthetic data back into the model. And then you can repeat this process many times with more prompts as you collect them over time. And the nice thing about synthetic data generation at training time is like, I usually make recommendations to customers to be kind of as conservative as possible when they're doing something with AI. They're already doing something with AI and that's an aggressive choice to make.

0:30We're still learning how to do that. So, you know, don't push your luck in some sense. So what I typically tell people is don't deploy something that you haven't tested. And so my usual recommendation is deploy the model, you know, collect the data, and then redo TAU on all the data you've collected so far. Test that model. And then, you know, even consider A-B testing it to make sure that, you know, for whatever metrics you're looking at in production, this model isn't going to make things worse, and then deploy it. In business, they say you can have better, cheaper, or faster. You only get to pick two.

1:00What if you could have all three at the same time? That's exactly what Coher, Thompson Reuters, and Specialized Bikes have, since they upgraded to the next generation of the cloud, Oracle Cloud Infrastructure. OCI is the blazing fast platform for your infrastructure, database, application development, and AI needs, where you can run any workload in a high availability, consistently high performance environment, and spend less than you would with other clouds. How is it faster? OCI's block storage gives you more operations per second. Cheaper? OCI costs up to 50 % less for compute, 70 % less for storage, and 80 % less for networking.

1:48Better? In test after test, OCI customers report lower latency and higher bandwidth versus other clouds. This is a cloud built for AI and all your biggest workloads. Right now, with zero commitment, try OCI for free. Head to oracle.com slash IonAI. IonAI, all run together, E-Y-E-O-N-A-I. That's oracle.com slash IonAI. But why don't you begin by introducing yourself? Go ahead. Definitely. So I'm Jonathan Frankel. I'm chief AI scientist here at Databricks. I oversee a lot of our research efforts on trying to make AI more useful for our customers and understanding how AI and data interact in a productive way.

2:39I got here. Databricks acquired my startup Mosaic ML, I guess, about a year and a half ago now. We were kind of dedicated to helping everybody train their own AI models from scratch. Founded kind of out of my PhD along with, you know, co-founders Naveen Rao and Hanlon Tang and Michael Carbon. And before that, I was a PhD student studying how to make neural network training more efficient. So it's been a wild journey over the past few years. But, you know, if you're looking for one through line, I care a lot about making sure everybody can use and control this technology and do with it what they want.

3:14I think a lot about kind of the internet where everybody could have a website. And that led to just such an explosion of creativity, a lot of really bad ideas and a few absolutely brilliant ideas. And, you know, if I told you from the beginning, oh, why don't we create an encyclopedia online that anyone can edit and contribute to? You'd say definitely one of the bad ideas. And it turns out to change the world. So the idea that AI is similar and everybody kind of has the chance to do creative things with it to solve their problems, it'll change the world again. Yeah, absolutely. I agree, which is why I have this podcast.

3:52So I was reading, and at Databricks, what's your, I mean, describe for listeners what Databricks does, first of all. I mean, I've talked, I think, before on the podcast about the data lakes and or data. What's it called? Yeah, data lakes. Isn't that the term? We're the data lake house. Right. Exactly. Data lake house. That's right. And what that means and how what you're doing on tuning models fits in that. Yeah, so I think the Databricks journey actually from the very beginning, you know, the stories have lived in legend for me as someone who's watched Databricks from afar. And now I get to be a part of it and hear all the inside stories.

4:42But the story began kind of where it is today. People have lots of data. They need a way to manage it and they need a way to understand it. And kind of data and AI are intrinsically linked to each other. Because what do you do with a bunch of data? You understand it. And Databricks is, as far as I can tell, looking through the history books, always been about using whatever the latest AI was to understand your data. Now, that AI wasn't always, you know, large language models or neural networks or anything like that. But and machine learning have been fundamental because what do you do with your data?

5:15You don't just, you know, stick it in a data lake house and just let it sit there collecting dust. You want to understand it. It's only useful. It's only worth having there if you can do something useful with it. So this is kind of this concept we like to refer to these days as data intelligence, the idea that, you know, you want to be able to talk to your data, understand it, figure out what's there, interact with it, learn what it's trying to tell you, kind of make effective use of it. So, you know, we offer ways to manage both unstructured data like documents and structured data. You can think of that as like, you know, SQL or kind of individual entries into a database.

5:50and a lot of the work these days is on the question of how do we use all the cool modern AI that's out there to give people new ways to understand it and to give new people ways of understanding it there are a lot of people who in yesteryear would have had to like write a notebook and write Python and do a SQL query or write in Scala or Spark to be able to get insights from their data and today I think AI is going to make it possible for pretty much everybody to do this I have no SQL, but I've been able to do this for my data at Databricks. Yeah, and the data lake, lake house concept, it really started during the supervised learning phase of AI when everyone was labeling data and training models with labeled data, was largely sort of text reading models or image reading models.

6:50And you came in after the generative AI revolution where suddenly really those silos are breaking down. I mean, at the time during the supervised learning phase, I used to talk to people about, you know, you train your data, but there are all these pools or lakes, some data's around, and sooner or later they're going to merge and we're going to have an ocean of data, a lot of it labeled. But the science kind of overtook that. You don't need to do that any longer. I wanted to talk to you about tau, which I've got to call it up.

7:44Let me see here. Tau, tau. So test time adoptive optimization, which I found fascinating. I read your paper. before I start asking about it, what really fascinates me is it looks like it's a way to create an ever-improving model, at least in narrow domains, which is something that I haven't seen before. Maybe it exists and I just haven't stumbled on it. But But first of all, can you talk about training and tuning? How you, is that what, I mean, I guess that's what Mosaic was doing, right? And that's what Databricks bought you for. But is, what's, talk a little bit about your role in that at Databricks.

8:43Yeah. So I think it's worth taking a step back and just asking, you know, yeah, why am I even here? Like, what is my purpose and what do I contribute? You know, I'm an AI researcher. every day yeah yeah we're trying to solve real problems for customers like why do you need a guy like me around databricks and you know my best answer is there are a lot of cool things happening in the research world that are clearly very valuable somehow and there are a lot of important problems our customers have that they need to get solved and they really wish they had the magic to do it and my job is to connect the dots between these and tau is one of those examples so you know customers want to get their data and their models to interact with each other.

9:26They want to be able to bring all that data to their models and have the models be able to speak fluently or answer questions about their model or reason in whatever domain they care about. And there are lots of ways of doing that. And those ways are things like you can prompt or do rag to kind of bring your data in. You could fine tune a model on your data. There are lots of ways to do this. You could even have a model write calls to a SQL database. And fine tuning is probably the most powerful of these techniques. It lets you literally change the model and teach it more about your data directly and intrinsically, as opposed to having the model kind of start fresh every time you interact with it.

10:04And fine tuning is a bit of a tricky thing to do for a lot of reasons. But one of the reasons we saw again and again was that fine tuning is really demanding in terms of data. You basically, to fine tune a model to do something, you need examples of doing that exact thing with your data. So if you want a model to serve as a customer service bot, you need to have examples of user asking question and what the customer service agent should say in return. And you need thousands of those examples. And what we found was that people have lots of incredible data. It's very rare that they magically have the perfect fine tuning set.

10:42That's just not the kind of data that gets created in nature. In nature, you have lots of documents that accumulate and you have, you know, some telemetry data and you have some user inputs and you have, you know, some manuals on how to respond or some guides or some decision trees or whatever else you have floating around, plus a SQL database of relevant information. But in nature, you don't have this perfect, pristine data showing what an interaction should look like with an LLM. It's just a very, it's a hard, it's a lot to ask of someone. And, you know, the correct answer when you come to a customer and they say, I really want an LLM that's fine-tuned for my task is not, well, go annotate 10 ,000 examples and come back to me when you're done.

11:22That's not an adequate answer. And I think to date, that's been the answer that the research community has had for all of our customers. Just, you know, it's your problem. You didn't give me the right data. And I think that's a bad answer. That's not an answer that's going to lead to people successfully deploying AI. Nobody has time for that. Nobody has resources for that. By the time they do that, the field will have moved on or their business will have moved on. And so the question we asked ourselves with Tao was kind of, on the one hand, what can we do to simplify this? Like, is there something we can take away or something we can, some constraint we can remove that will make this easier?

11:56And on the flip side, does the technology exist to actually do anything useful with that? If I have the customer give me just inputs, but no outputs for how the LLM should behave? Well, do I have science that can help me to actually tune a model on that? Or that it sounds kind of unreasonable, like you have no idea what someone actually wants the LLM to do. You know, there could be many good responses to one input, depending on what the user wants. And so my job, in a sense, is to move back and forth between these worlds. On the one hand, I'm always asking, like, what's making our customers' lives difficult, or what did they wish they could do?

12:32And if I had a magic wand and I could make their life a little easier, what would I do to make that easier? And on the other hand, okay, my magic wand is really not magic, it's science. What techniques exist? What techniques can I modify? What can I build that will bridge these gaps? So half my time, I feel like I'm a researcher. Half my time, I feel like I'm a product manager. I actually spend a huge amount of time with our customers, just talking to them about what they're trying to do and understanding their problems, not the ones in abstract that I think many of us think people have, but really where they're stuck.

13:03And if only I could remove this one thing. So with Tao, that's exactly what happened. We said, getting labeled data is really hard. How can we simplify this? And there are lots of ways you could simplify it. You could ask someone to write a detailed rubric of how something should behave, or write a bunch of judges or graders that tell you for any arbitrary output, how good is it? It turns out those things are also really hard. And it would be nice if we could just simplify it further. You don't even need any labels at all. That means that all you have to do is show up with a bunch of inputs reflecting what people might ask your LLM.

13:37And you just get an LLM that's good at those inputs. Now, the nice thing about that is it's really easy to get inputs. Just set up a really basic LLM. You can set up Llama, give it a really basic prompt, and ask your friends or your teammates or yourself to just interact with it and ask it a bunch of questions. It's probably going to do a bad job. But you don't care whether it gives you a good answer. You just care that you've asked it a good question. And so this is great. This means that it's really easy to get the training data because it's something that people can naturally provide the system.

14:06There's no real work involved. We're not asking someone to sit down and literally write 10 ,000 questions. Just put this in the wild. Beta test it. Prototype it. Ask your team to do an hour hackathon and just ask it all the questions that they might want to. Pretend to be the customer. Ask a customer to prototype it with you. and you can collect the data you need. Then you get to the really tricky part, which is, what the hell do you do with that data? Because it sounds, I mean, I hope I'm not like being overly dramatic, but it sounds kind of insane. Like all I know are the inputs, and you're going to tune a model that will give me the outputs I want.

14:38That's like, what's going on? And, you know, this is a place where we turn to techniques from reinforcement learning. And reinforcement learning is a set of techniques where rather than giving yourself inputs and good outputs, you give yourself inputs and just an environment where a model can explore. You can think of this like, you know, when, um, when, you know, people have built models for go, for example, deep mind built this model that was really good at go or game playing, let it explore the environment and figure out for itself. You need to have an environment it can explore and some kind of signal that tells you what outcomes are good and which outcomes aren't.

15:13This is why people love games so much because you just, you explore the environment, you know, when you're scoring points, you know, and you're winning and, you know, you get good signal. And here, what we did is the environment is in some sense, this special model we call DBRM or the Databricks reward model. This is a type of model that's very common in the reinforcement learning world. And all it's meant to do is tell you how good is a particular output. And, you know, it's, that sounds very generic, like just generically, how good is it? But the way you train this is you You show it lots of examples of, you know, pairs of outputs for a given input and which one a human liked more.

15:50And you just show it tons and tons of this data. And it turns out it actually generalizes pretty well. So, you know, you combine that with a bunch of inputs you have, do some synthetic data generation and kind of filter it with this reward model, use the right techniques to train the model after that. You know, a lot of, and this sounds really easy, a lot of very careful work on the part of the scientists who worked on this to just get the details right. And it turns out that it actually, there are very reasonable choices you can make a lot of the time. If I showed you two outputs and I showed you the input, you'll probably have a pretty good intuition as to which one is better.

16:25You won't know which one is right. But as long as you can kind of keep steering things in the right direction, the model will get better and better. And that's what we saw. Not only did this actually lead to improvements, even though you had no examples of how the model should behave, but it even beat training on labeled data, which is kind of crazy. And I can say a bit more about that at some point if you'd like, but it's, you know, it's not as crazy as it sounds, I think, but it was just very surprising to us. And this worked across a pretty wide range of tasks, like document question answering and SQL, which again was very surprising to me that a reward model could kind of generalize in that way.

17:00It's, it's really exciting for the future. Maybe enterprise tasks are more similar than we think. yeah i mean but but supervise fine-tuning you need these question answer pairs in in uh for uh a language model uh and what you're talking about you you still need question answer pairs right but you're getting them through reinforcement learning with human feedback i i wouldn't call them question answer pairs here because I think it's it ends up being more complicated than that I would call it questions and kind of steering on which answers are better or worse because it's not really none of them are there's nothing in tau that ever checks is the answer right or even is the answer good there's just asking is the answer better and which answers are better and trying to continually move the model toward better answers, which is very, I don't know, I found that to be a little bit mind bending when I started this project.

18:03I have a bunch of reinforcement learning kind of enthusiasts on my team and specialists on my team who have been, you know, spouting the gospel of reinforcement learning for as long as they've been here. And I've just been politely ignoring them and patting them on the head. And, you know, I've learned actually maybe there's something to what they do and I should listen to them a little bit more. Yeah. So in this process, you do use RLHF? That's correct. Yeah. And then you, once you have that, that data, then you use pure reinforcement learning. That's correct. So we're kind of using, I would think of this as techniques from the RLHF literature and just kind of combining them in the right way.

18:42I won't claim that there's anything we invented that's mind bending from like a scientific perspective. I, you know, I think there were, when we released this, there were some folks on Twitter who were like, this doesn't sound very novel. And my response was, it works. And getting it to work is the hard part. Like that's the novelty is easy. Like getting stuff to work for real customers is hard. So the details of like, okay, we're going to take some inputs. We're going to generate in the right way. We're going to select in the right way, build the right reward model. All that stuff took time. This reinforcement learning step after you've gathered the data is happening at test time.

19:18Is that right? It's very similar to DeepSeq and this test time compute training that has become kind of a paradigm for refining models. but it's different in that you're only doing it during the tuning phase. You don't do it during the inference phase. Can you talk about why that's an advantage? And what really fascinates me is unlike test time compute training is temporary, It improves the inference accuracy. But once you're done with that inference, the model reverts to its training state and you lose that improvement. But in your case, it continually improves the model. Can you talk about that?

20:29Yeah, yeah. So I think there was actually enormous debate internally about whether we should use the word test time. because it's not really it's like using test time compute for training and so is it really test time there was kind of you know there there was some debate over this word so the way i think about it is even ignoring the word test time the idea is that you can if you really want to be reductive about it this is a synthetic data generation and usage technique where you give us the prompts we generate synthetic data based on those prompts and then we integrate that synthetic data back into the model and then you can repeat this process many times with more prompts as you collect them over time and the nice thing about synthetic data generation at training time is like well let's take a step back if you were to just use let's say o1 or r1 which are both excellent models you're going to use kind of an uncertain amount of compute at inference time you don't know how much is going to get used.

21:28The model may explore for a while. It kind of, you know, there's some uncertainty there. And, you know, that can get expensive as well. We saw kind of in practice, 01 and 03, we're using about five times as many tokens somewhere around then to solve problems, which is not unexpected. They're designed to do that. That's the whole point. That's why they're so powerful. And one other way of looking at this is what if you use all that extra compute to generate the synthetic data. And then you create a standard model that when you actually go to use it at inference time, doesn't use any extra compute, just behaves normally.

22:04And the idea here is we expect that for our customers who are using this technique, if they have a bunch of inputs already, they've probably already got an application deployed. They at least have some amount of confidence that it's worth doing this training. And so So they're willing to invest a little more in training to make inference a lot cheaper. So the idea is, OK, let's spend some more compute at training time to generate better synthetic data. And then let's generate a normal model that when we use it at inference time has exactly the same inference cost as, you know, Lama had we not done Tau.

22:37So it's a bit of kind of a it's taking this test time compute paradigm and flipping it on its head a little bit. The test time compute is to help you build the training set. But at the end of the day, you're training a model that you use normally. And we think this is kind of the right mix for customers who are deploying applications at high volume. If you're doing something like Tau, you're trying to cut down on the cost of inference time anyway. Often you're trying to train a custom model to replace something else that you might have. So let's really lean into the fact that we can help you reduce your inference cost and a little more upfront investment, probably worth it to you.

23:11Yeah. And then once you're in inference mode, you're collecting more queries, more prompts that then you can use in tuning mode, just the tuning time compute. uh does do you periodically go back and retune the model on the data that's built up or is there a way of doing that in the background even if the model is being used to uh for inference or or do you split it into a uh two instances of the model and one you're improving through this tune time compute, the others doing inference out in the wild, and then you swap the models periodically. I mean, how does that work? It's a great question. It's something we're still kind of experimenting with customers on just to kind of get it in exactly the right shape.

24:17I usually make recommendations to customers to be kind of as conservative as possible when they're doing something with AI. They're already doing something with AI, and that's an aggressive choice to make. We're still learning how to do that. So, you know, don't push your luck in some sense. So what I typically tell people is don't deploy something that you haven't tested. And so my usual recommendation is deploy the model, you know, collect the data, and then redo tau on all the data you've collected so far, test that model, and then, you know, even consider A-B testing it to make sure that, you know, for whatever metrics you're looking at in production, this model isn't going to make things worse, and then deploy it.

24:53And the idea there is, you know, first of all, it's just playing it safe. And, you know, again, you're making a very aggressive choice to use AI at all right now. It's exciting. But, you know, be careful. And the other piece is that you don't necessarily want to just want to train a model that you've already done Tau on, on more recent data. You can end up in this interesting place where the model will kind of get extra good at the latest stuff and may forget things because you're emphasizing the latest data. Sometimes that makes a ton of sense. If you're in a chatbot application, let's say for retail, you may actually want to use the latest data because, you know, you don't want it to be talking about Valentine's Day when it's Thanksgiving.

25:31Sometimes you do want to mix all the data together just to, you know, make sure that you're fully representing the earlier stuff. So there are, this kind of gets into more application specific considerations where like, I'm guessing the big recommender systems powering like ad placement at Google and Facebook. I'm guessing those are actually trained mostly on most recent data because human behavior changes so quickly. But for a lot of our enterprise applications, taking things one step at a time and kind of combining everything together, again, it's the safest possible choice to make. And then you can experiment and try lots of different things.

26:01But I wouldn't recommend kind of trying to update in real time, especially without testing the models first, just to be extra safe. Yeah. And so you mentioned sort of overriding previous parameters, which has been the problem that restricts continuous learning, is catastrophic forgetting. So yours doesn't get around the catastrophic forgetting problem. No, I wouldn't make that claim. I would say very narrowly, this gets around the need for labeled data, or at least assists you by reducing the need for labeled data. When it comes to catastrophic forgetting, it still puts you in the same trade-off that you're always in of like, do you train on the latest data and hope that the model doesn't forget?

26:52Do you combine all the data together? Do you upweight more recent data? I don't think we have an amazing answer for catastrophic forgetting. I will say forgetting is a lot less catastrophic than it used to be. Back in my day, five whole years ago, catastrophic forgetting was a very severe problem in general when you did fine tuning. And these days, models are a little more robust to that just because they're so much bigger and they're able to store a lot more knowledge and also because we have better techniques for managing catastrophic forgetting. To give you one example, there was a paper we worked on last year as a lab.