The Magic of LLM Distillation — Rishabh Agarwal, Google DeepMind

Hey everyone, welcome back to a Latent Space Lightning Pod. This is Alessio, partner and CTO at Decibel, and I'm joined by my co-host Swyx, founder of Small AI.

Hey, and today we are very happy to have Rishi Agarwal from GDM. What is your, what is your role at GDM?

I'm a staff research scientist.

Research... Yeah, staff research scientist.

And...

Occasionally you do these talks, uh, it sounds like you do them internally, but then also you publish your slides. Uh, and so it caught my eye on a few of them, and, and I think th- this most recent one, you said that you, you did a talk on distillation, and it wasn't recorded.

So I was like, you know, why not let's jump on and talk about all this sort of post-training, uh, distillation for LLMs, because I think it's, uh, just the general field of post-training has just exploded this year.

Yeah, that's true.

Cool. Yeah, like, uh, maybe, maybe you can give us a little bit of context on, like, how or, like, why this topic was relevant enough for you to do a presentation on it.

Yeah. So I guess that's a good question. I usually am an, uh, like, an RL researcher generally, and distillation is something like a side quest I've done at, like, Google. But this time around, I guess there has been a lot of progress.

Uh, so if you think about distillation, I think what comes to mind is, like, the twenty fifteen paper from Geoff Hinton and others, which is, like, the seminal work that people refer to. But then since then it's been, like, a decade, right?

And there must be some progress that has happened since then. And especially now we are talking about LLMs, back then we were talking about classifiers, like, which will, let's say, train on ImageNet and would say. So I thought it would be a good idea to give a tutorial, at least internally at Google.

And after the tutorial, what I did was I also published the part of the slides which were, like, at least based on publicly available information. So-

Yeah

... so I thought that would be worth sharing, and that's it. And I think the reason is... So I mentioned I usually work as an RL researcher, and some of the new methods which are coming out are RL inspired, which look very different than what you typically think of when you think of distillation.

So I thought that would be a nice way to just talk about what people ty-typically think of as distillation and what are the new latest kids on the block, and maybe some of them make sense, maybe not, but still worthwhile knowing, like, what is possible.

Yeah. There, there's a purist definition, which is the Hinton, uh, definition. But I would say, like, the, the current usage of the term distillation has kind of expanded a little bit.

Yeah. I think also how you do it, right? So usually, like, let's say the definition is very broad. It says you want to transfer knowledge from a expensive teacher model to a smaller model, and sometimes the models could be of the same sizes, but then, like, this is very broad.

It doesn't say how do you transfer the knowledge. And usually people think that, oh, we'll collect some data from the teacher, and then we'll just fine-tune the student on that data. But that's one way of doing distillation, and there are other ways too, right?

And I thought maybe we should talk about that. And-

Okay

... and I guess the motivation for distillation is kind of easy to grasp. Like, you would think bigger LLMs are usually better, and that was the trend a few years ago, even now. But in practice, what you care about is, let's say, the graph, which may, uh, one of the, the hos-hosts here created, which is the graph of cost versus performance.

So you want something really good but at a really good low cost. Like, even if something is really good, but it costs, like, I don't know, thousand times more than something else, maybe we'll still not use that because it's just too expensive for the use cases people care about.

And then there are other use cases. Let's say you want to deploy a model in your smartphone. Uh, it's just not possible to fit, like, a really, really large model there. So that's another challenge, and I think that's why maybe we still care about distillation per se.

Yeah. I think one thing I'll, I'll comment about this is, like, I actually, um, used to think about this as distillation as, like, kind of like a, oh, it's like a fun way to, um, um, you know, to, to, to sort of, uh, fine-tune a model or whatever.

But now I always think about it, think about it as, like, a deployment technique, um, in the sense that, like, we're starting to get models that are too big to use in practice, right? Like the, the Llama four five Bs, the, uh, even GPT four point five is too expensive.

Nobody's really gonna use it in, in any reasonable quantity. Like, you know, Claude three point five Opus, like, if it, if it does exist, still not... Like, you know, the, the thing that we actually use is Sonnet, right?

So, like, it's almost like a deployment strategy. Like, you, you train the large model, the, the teacher model in order to distill. Like, you don't actually expect to use the, the teacher model for inference anymore.

Yeah, I think it, it is really the, the marginal benefit versus cost argument, right? Like, if the marginal benefit and capability is not worth the additional cost to you, then people don't wanna use that bigger model. And it, it is very, very task dependent.

Maybe you really want that capability, but sometimes, especially when you're doing things at bulk, maybe it's okay that you're slightly worse off, but it's much cheaper. And that's maybe one of the reasons it released this graph. That's why you wanna be at the Pareto frontier.

And I think the other thing is the capability which we have right now, maybe next year it'll be much cheaper to have that same thing. And that likely is the result of distillation, right? Like that's, that's... it's not just be- because we are doing or figured out something magical.

It is because distillation can help you bring the cost of any existing capability.

Yeah. There's a little bit of sawtooth. Like, you go up in capability, then you go down in efficiency, yeah, to, to make it more efficient and that kind of thing. The other thing I noticed as well, which I kind of pointed out, was the, the slope of the d-- um, the slope of the Pareto frontier is basically the state-of-the-art of distillation, and the gentler the slope, the better distillation is, which I thought is kind of cool.

Oh, see, that's a very good argument. Like, I didn't think of it that way, but you're, you're right.

You see it improve over time in the sense that it used to be this kind of very steep slope for-

Yeah, yeah

... intelligence, and now it's a lot gentler.

So you're saying basically we're getting better at... At least this, this is costs. We don't know underlying what's going on, but you're getting better at serving the same intelligence at much cheaper and the slope itself saying the methods themselves might be improving.

That's awesome.

Yeah, it's just, it's just, like, easing. Yeah. Okay.

Yeah. The other fun reason which I think, uh, if you go to this next slide, I think this is, like, one of the coolest reason, and this is something that people have lots and lots of theories, and you can go to this thread.

I think there's like hundreds of reasons people have given, but this is Something that people have found again and again, that basically you can train a model on some data, or you can train a bigger model on that data and distill that model to another model-

Yeah

... and that distilled model is better. Like th-this has been true. People-- There's a lot of theories, some I know, and there's a lot more in this, uh, you can see based on the popularity of this thread. And that, that's another reason that at the end, if you just want to train a good small model, maybe it's still better to train a bigger model.

Like that's what you were mentioning earlier also, that maybe that's why we are training teacher models. But, uh, this kind of finding has still stayed relevant, I think, and, uh, at least I would say this is still true.

And, and this is another good motivation for actually thinking about distillation.

Yeah. Um, is, is this really called dark knowledge or what, what is this phenomenon called?

So the dark knowledge actually is funny. If you go to the slides back, the dark knowledge comes from Jeff Hinton's slide themselves. Like he gave-- he used to give talks about distillation where his sl-slides were just titled Dark Knowledge.

So...

Okay, interesting. Interesting.

So those are also very fun slide, slide to see from Jeff himself, like what's the motivation of all this thing?

And when you, when you think about using this model for distillation, is there any changes you would make to how you approach RLHF and kind of preference tuning? You know, a lot of time the models kind of get narrowed down to be good for user in facing things, but if you're just using them to then distill things, maybe you can skip some of that work.

Depends, right? For example, your distilled model could itself be RLHF, right? Maybe you want to skip the RLHF stage, let's say. So what if you have the final, final model you want to serve, and then you just distill it?

Maybe you don't need to do anything else after distillation. So think of it as like, what do you want out of this? Let's say you say, "Oh no, I just want distillation so that I distill some base capabilities in my model, and then I will still do some sort of specific post-training afterwards."

Then it makes sense to have a different teacher. So I think it depends on what you want out of distillation, because distillation is just saying, "I can transfer capabilities from this teacher model to the student one," but that's not perfect, right?

There's always some gap. So maybe you still wanna self-improve that student, and that's what RL is for. So, so I would think maybe an optimal pipeline would look like you do distillation heavily, but then you still do some RL afterwards, because maybe there's still something you can get out of your reward functions or whatever your post-training stack is.

Uh, at least in, in my thinking. Unless your distillation gets you really, really close to the teacher, that at that point you'll be like, "Okay, there's no point doing anything more. Uh, let's just, like, just do distillation and move on."

That could happen, right? If a distillation was that good. So let's say you take like the, the best model out there and you can just distill it, and for the thing you care about, let's say if it's coding, and you distill it so well that on coding benchmarks and all the ki- kind of coding tasks you're doing, it's just really nearly matching the teacher, then why would you do anything else afterwards?

So it depends on how good distillation is, and, and that's more of an empirical thing, I would say.

Yeah. Cool. Um, I can keep going. Uh-

Yeah, yeah

... I guess.

So-

Yeah, section. Yeah

... so I guess, yeah. So this is like the, the thing which was popularized by the, the original seminal paper where we basically take the student, have the teacher, we run some input through it and get logits and match them.

But this is not actually what, what is used in LLM, so maybe we can skip and go like a bunch of slides further. One thing I'll say is, so if you go up, it's just a bunch of math, but basically it's minimizing some sort of KL divergence or objective by matching the distribution of the student and the teacher.

Just keep that in mind. Like this is what the Hinton style distillation is doing, and it's using logits. It's basically rather than... The, the intuitive way to think about it is it's doing next token prediction for LLMs, except you are using soft tokens.

So if you go down a couple slides, I think that will come out. So the objective is the same, except rather than having a single hard token that you're saying, "I want to predict this token next," I am actually gonna-- I have soft labels that I have basically some probabilities for each of these tokens, and I'm predicting all of them together.

So it's basically a generalized version of next token prediction. But at the end of the day, loss is exactly the same. The only thing changes is whether we have one hot labels, which is only one token gets a value of one, everything else zero, or do we have soft labels.

And this can... Uh, now you can see that this can be easily used for pre-training. And one notable example of this is Gemma 2. So Gemma 2 used distillation for pre-training, where they used logits or these soft labels, which is rather than having hard zero one tokens, which is I want to predict this next token, they have like soft labels for all possible tokens.

And then they just run the next token prediction loss on this.

Is there an estimate of how much better... Uh, because obviously you're, you're just us-utilizing all the inf- all the available information, right?

So yeah. So this is like basically saying you observed some next token in the data, and you assume that's ground truth. Here you're saying, "Oh, I'm actually relying on my teacher to tell me what is the possible next token after this sentence or this part of this," and then you assume that as ground truth.

Right? The objective is the same. The only thing is what information you assumed access to. And here we have more information because, well, I guess the teacher was trained on that human data, but we're hoping that it generalized.

That's the dark knowledge bit, that maybe it actually knows more, that maybe there are multiple tokens which are possible. It's just that I only observed one of them in the data, but the teacher is telling me that, "Oh, there are these few tokens that are very likely after this sentence, uh, so maybe I should be trying to predict all of these."

And that, that is where maybe the magic of distillation is from. So you can see in these plots, right? If you can... So again, this is the kind of dark knowledge we are talking about, that, oh, if you train the student model from scratch, that's like always worse.

So higher is... Like I guess its perplexity, so higher perplexity is worse. But similarly on the previous plot in terms of performance also, you can see from scratch is the top bar and distilled is the bottom bar, so there's no change.

I have the same data and everything. The only thing is whether I'm just training this model from scratch or I'm distilling using a bigger model on that data, and you see that effect of, oh, distillation actually doing better.

So that's the reason maybe people do distillation. It's just not cost savings. Maybe you're getting a better student model at, at the end of the day. You couldn't-

Yeah

... have just done it if you were to do scratch training for that.

Would there, would there not be more, quote-unquote, "flops" spent on training on the distilled model because there's the, there's the flop spent on the teacher, then there's the flop spent on the student?

Yeah. That, that's actually a very good thing. I think there's definitely more compute. I This is something I don't know if the Gemma two folks reported in the report. Was it compute matched or was it data matched? But that's a very good comparison.

Like, you should be ideally comparing on compute, like whether if you do more pre-end updates on the same data which matches, let's say, the amount of compute you used. Okay, so this is actually the kind of distillation which is very, very common, right?

Which is, well, this is like the one we commonly use for LLMs. We call it distillation using synthetic data. It was actually introduced back in twenty sixteen by-- from this paper, the sequence level knowledge distillation. But what they do is they have some prompts.

You just take your teacher model, generate some outputs, and then you just do supervised fine-tuning on this data. That's like the simplest thing possible you can do, and it is also a very common method. The nice thing about this method is it only needs API access, right?

Like, so you can take any model out there, and you can just prompt that model and generate outputs using that, and then you can just use that data. Now, this was also in the news recently, if you were following it.

The nice thing is, like, I guess if you go down, next slide, this is exactly the same loss that we were optimizing earlier with the principal distillation, like the KL divergence. It's just a different way of doing it.

So it is a very principal objective. It is minimizing some sort of divergence or KL divergence between my teacher distribution Q and the student, and it's just a different way of doing it, where I'm just basically sampling something from the teacher and fine-tuning the student on this.

So the simple thing which people do is very principled and is optimizing the same loss. It's just assuming access to different kind of information. So here we don't assume access to any sort of logits or next token probability distributions.

Here we are just sampling from a teacher, and that's the only access we assume, and then we just fine-tune the student on that data. So same loss, different way of optimizing it, and it just turns out to be simpler, so that's why everyone uses it, right?

Yeah. Yeah.

And I guess a better version of this is you can do something like best of n, which is, let's say this was something which was used in DeepSeek. So, uh, if you saw that paper, they trained the DeepSeek reasoner model, and then they distilled the, the, basically that model to other models.

What they did was they took the, the best model they had. They generated a bunch of samples. They filtered them based on correctness, because these were on tasks like coding, mathematical problem-solving, and a bunch of those things. So they saw whatever-- what are the outputs which actually get the right answer, or they saw they get past the test cases or whatnot.

They pick those, and then they just fine-tune the model on that data.

Yeah.

So if you go down, I think that's what the results from Deep... Yeah. So you can see these are taken from the, the DeepSeek report itself, and people can read that, uh, for more details. The one nice thing about this kind of distillation is it doesn't matter if you have a vocabulary mismatch, because we're not using the next token distribution or probability labels.

You can distill from one model which uses some random tokenizer to another model. That's like a very good feature. Like, that's the cool thing about it. That's why they were able to distill from DeepSeek model to Llama or Qwen, and you don't have to even think about that tokenizer.

And that's why this is common, right? Like, it's very easy to do. You can just need an API and just collect data, fine-tune that on that data with the student model, and it's done.

Right. But-

And you can already get pretty close.

It's still a lot less efficient than the logic matching distillation, right?

It depends. I guess I'll-- that's a complicated, uh, question. It depends on the thing. But yeah, logits are definitely assuming access to more information. So it, it can be... But then again, you don't have access to logits if you want to distill, let's say, a black box model.

So... But then again, no one want-- it's not legal, I guess, to distill a black box model. Or at least the terms and conditions don't allow you to do that.

Yeah.

But for open models, yeah, you can do better because you have access to logits, so why not use them? At least that's my take. So, um, and I guess this maybe we can skip, but this is just saying when you're doing best of n why not keep the top K rather than just the top?

Because if you're already generating, let's say, a lot of samples, why, why throw away, let's say, if you have some sort of scoring mechanism of how good something is, and this thing is called MBR in the translation literature.

But anyways, this is another way of best of n. So the key idea is very simple. Just, like, use some of the, the da-- Yeah, just don't throw away everything. Just use the top K samples rather than just the top one, and this might be better, and that's true.

So th-that's like a worth, a simple thing worth trying if people are doing synthetic data distillation.

Okay.

But yeah, I'll maybe... Yeah, so there's small gains you can see. If MBR is, let's say, using just the top one, and then if you increase from five, ten, like top ten, top twenty, top forty, you can see slight gains as you use more samples.

I think the one below is much more interesting, and this got a lot of popularity on Twitter, and it keep-- like, I finally love every few months, this keeps getting retweeted and getting viral. But maybe I'll clarify what's going on here.

But basically, the simple idea here is just that you have-- if you have two models, right? There is a sampling cost you're paying. So you can sam-- generate data from your really big model, or you can generate data from a smaller model, but you can generate a lot more of it because both of that use the same compute.

Because the sampling cost or FLOPS is proportional to the number of parameters in the model times the number of inference tokens. And there are some assumptions here, and I'm assuming we are using KV cache, but assume this mental model, that the cost is proportional to number of parameters time number the number of tokens.

And for a fixed number of tokens, if you have a much smaller model, so let's say Llama seventy B versus Llama seven B, you can generate ten times more data in principle, right? So now you can ask this question.

Well, is it actually better to generate data from Llama seventy B, or is it better to just generate a lot more data from Llama seven B and assume we have-- we can do something like best of n, which is we have access to some sort of filtering function or reward function.

Then it becomes even more interesting, right? Because now you can filter this data too. So if you go down...

So one thing we found consistently, so here what we-- we had two models, nine Gemma nine B and Gemma twenty-seven B, and we found consistently that actually generating data from nine B in a compute match setting is always better, even better for distilling or actually improving the twenty-seven B model itself.

So which is we are fine-tuning the bigger twenty-seven B model. There are two things we can do. We can generate data from the twenty-seven B model And filter that and fine-tune it, so that's self-improvement using some sort of reward function.

Or we can generate data from the smaller nine B model. It's just that we can generate a lot more of it and fine-tune the twenty-seven B. We found that actually it's-- that's better to do.

Wow.

So, so this is again, the magic is because we can generate a lot more samples, and we have a way to actually select what is a good sample because we have-- we are doing this on domains like reasoning, like math coding, so we can filter that.

So if you go down to the slide, uh-

Quality based on this.

It can be-- Yeah, if you do cost match. So I was-- so far, I was talking about compute match, but if you think about API pricing, like let's say we can-- we have two models, Gemini Pro and Gemini Flash.

The pricing of Flash is thirty-five times cheaper than Pro. And now actually it's even better to generate data from the Flash for the, the trade-off you get, 'cause you can generate thirty-five times more data for the same cost.

Let's say if you have some dollar budget of hundred dollars, you can generate samples from these two models. And it turns out that actually just generating a lot more data from Flash and filtering that would be much better.

These are, are all results on some math reasoning benchmarks. I think it's Hendrix Math.

Yeah, I mean-

So-

The intuition is really that you're injecting more data through the filtering process, right?

Yes. So it's not just more data. It's more data, but also going through this filtering. If you just do more data, well, again, you have a weaker model, right? So more data from a weaker model cannot just be better than less data from a stronger model.

Maybe true, but that's unclear. But because of the filtering function, which means for domains like reasoning, this could actually very well work. So that's another reason to use a small model. Okay, now I think this is gonna be switching gears.

So this is what I was talking earlier, that RL inspired distillations. If you think about distillations or generally just next token prediction, there is a mismatch at training time and test time or inference time, which is at training time, we are basically training on a fixed data set of outputs, which is I already give you the output, and you fine-tune on that.

Let's say that's what-- that, that, that was being done in the synthetic data distillation. I already gave you some input and output, and you just supervise fine-tune on that data. But at test time, the student itself or the student model itself is generating tokens one by one autoregressively.

This is also the same difference which happens in SFT versus RLHF. In SFT, you're teaching it to copy something. In RLHF, it is learning what to do using a reward function. So there's a big difference, and there is a mismatch basically at training time and inference time, and just think-- imagine the autoregressive nature of models and what we are doing at training versus what happens at test time.

It's called exposure bias?

This is like something... Yeah, this-- there's-- this is like an-another term for this mismatch, but there's a bunch of papers you can find. But yeah, this is also known as that in the NL-NLP community. But really all that there is to it is this is-- this thing is train inference mismatch, right?

And, uh, and then I guess really, like, think about it this way. This is like also very well known. This in-inference mismatch is very common in the RL community. It's like this work is done back in the... So when people are trying to copy agents, which might be relevant now again, you can do behavior cloning or supervised fine-tuning, which is very common.

But the RL folks figured out a much better way. They call it DAgger. I'll explain what that is very quickly, but just imagine what's going on. So let's say you're training a student model on some fixed data distribution.

Now, at test time, the student model makes a mistake. So here, let's say the student is learning how to drive a car, and all the student was trained on how the teacher learns or how the teacher is driving the car, right?

But there are scenarios which the teacher would never do because the teacher is a good driver, so maybe he'll never run into really bad situations. But the student might run into such situations. Let's say I see some tree, and I'm gonna crash into that, right?

The teacher is never gonna show you such a scenario because it's-- generally, teacher is a good driver, and if you collect data from that teacher, they'll just show you what does good driving looks like.

Yeah.

But that means that you cannot recover from these things. Like, if the student is gonna, like, crash into something, it has not seen this kind of data before. And this is the mismatch I was talking about, that it hasn't seen, "Oh, if I made this kind of mistake, what do I do?"

And I guess pictorially, pic-- uh, yeah, if you see the left picture here, that you can see the learner is falling off the cliff, and the teacher just, just never goes there. Like, it's just like you can-- this is, like, from-- this is adapted from someone else's slide.

But, uh, this is, like, a common problem that also happens in RL or any sort of agentic tasks. That's-- so this is, like, the de facto approach people use when they are distilling in self-driving style situations or other agentic situations.

This is a method from two thousand and ten or something, I think. But this was designed for RL agents, and we thought, well, we can close our eyes and think of LLMs as agents, and eventually LLMs will become agents.

But LLMs are also producing token by token, right? And maybe they are do-- think of these as actions or they are doing things, and then we want feedback on, "Oh, well, this is-- is this good enough or not?"

So you can assume the similar methodology can be applied. So if you go down, what's the fix? The fix is very simple. The main change is, rather than sampling the output sequences from the teacher, we just sample that from the, the, basically the student itself.

The student is sampling these tokens or sequence of words or text or whatever you can think of that as. Now, we run inference using the teacher on these tokens, and this can be parallelized because once you sample the sequence from the student, you can just send it the entire sequence to the teacher.

Yeah.

Yeah. And then you can get logits for all of this. This is just running inference, right? So this is not autoregressive. You just send it in parallel, and you get logits, and then you minimize the logits. So the only main difference is you're sampling data from the student rather than teacher.

And what it looks like is basically at every point in time, I'm asking the teacher that, "Hey, what do you think about this? Do you think is this good enough? Like, would you have done this, or would you have done something else?"

So at every point in time, the student is asking teacher for feedback, and the feedback is the stu-- teacher is going like, "I will not do this. I'll do this instead." And the student is like, "Okay, I guess I'll, I'll keep that in mind."

But it's getting this side kind of correction feedback. So imagine the scenario I was saying earlier that, oh, if you run into a scenario where basically you're really close to colliding to a tree or something, the teacher would say, "Oh, actually do this thing so that you avoid this collision."

Or, or it will-- before it's actually even going to the collision, teacher would be like, "Oh, I'll never do this kind of thing. Do something else." So it's, it's-- this is like, uh, something... And the-- if you go up, that it seems, yeah, I guess this is like clearly.

The nice thing is this, basically, there's no train test mismatch. The, the, the way we are generating the data is also the, the same thing done at test time, which is at training time, the student is generating the data autoregressively.

At test time, the student is again generating the data autoregressively, so there is no train test mismatch of that sort, and that's like the main kind of fix this kind of on-policy thing approaches. Now, I'll, I'll maybe mention one interesting detail afterwards, but, uh, yeah, if you go down the slides, I guess it's pictorially mentioning what's going on.

So we sample data, we get logits, we minimize some sort of loss. If you go one more slide down-

Yeah. It's comfortable this, yeah.

So the nice thing is you can derive this actually. So all you need to do, at least in math, if you just flip the direction of your KL divergence. So earlier we were di-- minimizing KL between a teacher and student model.

If we just swap that order, 'cause KL is not symmetric, you will actually get this kind of distillation.

Oh, that's cool.

So it's like that's, that's like nice cute math you can write to say, "Oh, actually, I can derive this," and this is very principle, it's not a hack. The other nice thing is you don't have to do full RL.

Basically, you can sample from the student, and I said, "Oh, you just optimize some loss." You don't have to backdrop through the sampling distribution, 'cause that also depends on student parameters. That becomes full RL, and there's another people on that.

But basically, if you still wanna do supervised training and you're only changing which data distribution you care about, and we are saying we care about the data distribution with student samples, 'cause that's what we an-anyways want to use at this time.

So that's the other trick which people have done in the literature of-- RL literature, and we're saying, "Well, what if we just do that again with language models?" So this was paper... this is a paper we wrote in twenty twenty-four, but basically it's saying there's a math way of doing distillation which RL people use all the time.

What if we try that for language models? And that works really well. So if you go to the next slide, you can see the results. And it also answered the, the question. So th-this was done on very small scale tasks, right?

This is like very specific tasks like summarization, translation, or reasoning with chain of thought. The supervised fine-tuning or red line is the synthetic data distillation we were discussing earlier. So not-- sorry, that's synthetic da-- not synthetic data, it's distillation on human-written outputs.

CKD, if you look at the light sky blue line, that's synthetic data distillation, which is if I just take the teacher model and just generate data from that and distill. Supervised KD is using logits, and this green guy is basically sampling from the student and using logits.

Yeah.

But you can see, uh, I guess one cool thing you'll see is it's not always the case that synthetic data distillation beats... oh, sorry, it's worse than logits. Like, sometimes logits can be worse off. So if you look at the T5 base two fifty million scenario on GSM8K on the last plot, you can see that synthetic data distillation can actually be better than using logits, right?

Okay.

So that can also happen. Like it's, it's, it's basically... but usually I would say if you have access to logits, it's probably worth using them, and the green curve, which is doing on-policy data, is using logits anyway. So, so that, that, that does help the efficiency argument.

So this kind of thing was used in Gemma du post-training. It's, uh, mentioned briefly, like there's no details there, but it was used there.

Excellent. But there's no way to know ahead of time. Like, you don't know ahead of time based on the task whether or not distillation logits are gonna be better. You just gotta do it.

Uh, I would generally say if-- I think if you have more information, usually it's probably better to use it.

Yeah.

It might be a capacity thing, right? Like w-look at the T5 base model. How many parameters are-- two fifty million. That's not even a language model at this point we consider. Like th-that's like micro size, uh, 'cause people call point five billion tiny, and this was a base model back in the day.

So I would say it, it depends heavily on the capacity to some extent. Uh, my opinion still is if you can use logits and it's not too inconvenient, probably worth it. Maybe your method is not as good. And again, the method which we are competing against is like a method developed in twenty fifteen, right?

Like the synthetic-- the logit distillation thingy, uh, the supervised KDs, and better methods like the green one can easily make logits much more useful. So I would say that's maybe the blocker with logits. Uh, and more research would be really helpful there.

Anyways, uh, so I would say my answer is use logits if you have them.

Yeah.

The, the other fun thing, and I think this is a very nice plot. This is worth internalizing, actually. So there is this notion of well, okay, how do I say this? So when you're-- let's say you're doing distillation, there's always a model capacity mismatch that the student model is generally much smaller than the teacher model, right?

And now actually it depends on which direction of KL you optimize. It changes what the student would learn. So here it's a very idealized scenario where I have a bimodal Gaussian distribution, that's your purple guy, and then the green curve is the student model.

It's a unimodal Gaussian, which means I can only fit one of the peaks, right? And now depending on which direction of KL I optimize, I fit-- I have very different behavior. So these are called mode-covering and mode-seeking KL.

The-- typically what we use is this mode-covering KL, the one on the left. That is the standard distillation loss that everyone uses. But you can see the behavior, and it-- you can already see what's weird about it. It's putting a lot of mass on places where there's no mass put by the teacher.

Like, the teacher would not sample anything in that regime where the peak of the green curve is. But that is exactly the behavior you're getting when you optimize for this direction of KL. While if you look at the other guy, you can see that what it's doing is it's covering one of the peaks, but it's basically putting no mass all together on other parts of the distribution.

So this is called mode seeking.

So that is basically the, the difference in these two directions, and supervised KD minimizes the direction of KL between teacher and student, that's similar to the one on the left, while on-policy distillation minimizes the other direction. So you can see the behaviors.

It's basically saying, "I'm happy fitting one of the modes," or, "I'm happy basically getting one way of getting this capability right, rather than getting all these different ways of getting the capability right." So if capacity is a constraint, then you can see the behavior, or at least why there is a basically the difference in divergence use.

Yeah. I would say this is like a very exaggerated example. Obviously, most, uh, actual text distribution doesn't look like that

Yeah. Also, this is Gaussians, right? This is like saying if world was perfect, then pictures look like Gaussian. So if you go down, actually that's good, you mentioned it's... So in practice, this is what things look like. You get a trade-off.

There's a trade-off between diversity and performance. So on the Y-axis, I'm showing performance. On the X-axis, I'm showing similarity or basically how, like one minus diversity, so more similar things are less diverse. And you can see, depending on the KL, the divergence you use, you get like a trade-off.

So if you use the one direction of KL, you get like pretty high diversity or low similarity, but low performance. On the other end, you have high performance but low diversity or high similarity. So if you look at just like one of the curves, as you go along the X-axis, you're actually losing diversity because things are getting more and more similar.

Mm-hmm.

So there is that trade-off you're facing. And now for any task you care about, I would say like basically you'll have this kind of trade-off. The general recommendation I would give people is that maybe use a mixture of half and half, like that's what some people have used.

Oh, right.

Right. That's like saying-

mixture of KLs

... yeah, basically saying, "I don't know what I want. I just want something to work well enough. I'll just do half and half and hope that it works out." And generally it would, I think. So there are different ways of mixing half and half.

This is one way. JSD, the Jensen-Shannon divergence is another, but this is just like our way to say, "I don't know what I want. Let me just, like do this and hope it works," and usually it does. Okay, now the most fun slide.

I think I really like this slide. There are three easy steps to implement this distillation. You go to an RLHF framework. Oh, I think my airport's died. Can you hear me still?

Yeah, can hear you still.

Yeah. Okay, that's good. So there are basically three easy steps to do this thing. First step is you go to an RL-RLHF, RLX, or whatever framework you have. You turn off the reward term, so you delete the reward part of it.

All you're left is some KL term, and now you swap your KL, the anchor policy, to a bigger teacher policy, and there you go. You have implemented the distillation we are talking about. So RLHF or RLXF is already doing two things.

It's doing reward maximization subject to some KL constraint. The KL constraint is usually to the same policy as the student or some-

Right

... SFT policy, but it doesn't have to be. It could be a bigger model. And yeah, basically that's it. So that's like the, the cool thing that-

I mean, it's, it's a, I mean, it's, it's a, it's very interesting, interesting. Why are we turning off reward maximization when we're just talking about, uh, KL divergence minimization, right? Like what, like why is this so profound to you?

Sorry. So saying we are basically just interested in distillation.

Why is this so profound to you?

I'm just saying the reason it's interesting is because, I mean, banking on any existing RL framework, this is how the Hugging Face people implemented this, but I'm saying that usually infra is a big bottleneck for these kind of approaches.

But what I'm trying to say is that the infra already exists. All you have to do is delete a bunch of pieces in that infra, and you're good to go. And I guess that's why it's already-- You also asked a good question, "Why do I have to turn the reward term?"

You don't. That's the next slide. You can actually combine RLHF and distillation together, 'cause now you're doing two things at the same time.

Basically, I'm saying you already get this free infra for free. Like you have it. You just have to make sure that you can fit in their teacher model there in that infra, and that's the main change you have to do, and you already have implemented this sort of distillation.

'Cause we are anyways sampling from the student when we are running RL. The main difference is rather than using some logits from a teacher, we are usually maximizing rewards and logits are coming from the SFT policy. While here we are saying, "Well, I'd actually want to use a bigger model as my teacher and do this."

Got it. Okay. Nice, ni-nice qualification.

Um, yeah, I think the rest of the slide is just some application. I'll actually-- C-can you skip this one? I think I'll go for one more. Yeah, this one we can skip. I think there's a more fun application which is much more relevant to people.

So if you go down-

Sure

... to the next slide.

Sure.

Yeah. If, if you go further.

I think, uh, the share might have frozen. No, no, no. My, my configuration, my configuration acting up a little bit.

Yeah. Okay, this one.

Okay. Speculative decoding. Speculative decoding.

Yeah. So speculative decoding, I think people-- I'm assuming people know about it, but basically the way to understand it, we have a smaller model which is proposing some tokens, and the bigger model is accepting those tokens. So you can propose multiple tokens at a time, and this bigger model can say, "I like them.

Let's just do the-- go with that." And basically, you do not change the distribution or sampling samples from the bigger model, but you can make it much faster. 'Cause now rather than producing token one by one, you can produce a bunch of tokens at the same time, 'cause you can use the small model to produce those, which is cheaper, and then the bigger model can just accept those.

Anyways, read about it. It's very cool as a technique. Now the thing is, the effectiveness of this method depends on how close the, the sampler, the small model is to the bigger model that we want to speed up.

And actually, distillation exactly fixes that, which is by distillation, you can make things closer to each other. And we found that actually if you do that, you can make speculative decoding faster, which means you distill. So you're a student, you have a better student model, but then you can also speed up your teacher model using this distilled student-

Sure

... with speculative decoding. So you get two things for free.

So every large model should always have-

Like basically something which-

Every large model should always have a student model for-

Yeah. Basically, even if, let's say, you don't care about the student, the large model itself would benefit from having someone close enough to itself so that we can speed up serving of that large model. And this was actually used-- So I guess the cool application of this I can mention is the next slide, which is AI overviews at Google.

I don't know if people have seen this or this kind of thing. Like there is this thing that comes up and that actually uses this technique of speculative decoding, but then you can all obviously imagine that, oh, they have some small model which is distilled from the big one, which makes this even faster.

Oh, actually last slide. I think this is like hot off the press, I, I guess. So now you can merge all the ideas we have discussed so far. This is like i-- this paper that I cleared this year, where basically we are saying, "Well, we are sampling from, from the student."

Oh, and then we're getting feed, feedback from the teacher. But then there are scenarios where maybe the student just does something really terrible. The teacher is like, "Come on, man. Like, I can't tell you anything here. Like, this is just too bad.

I don't know what to do." And that's exactly what this is for. This is saying basically you get inspiration from speculative decoding, so the student is sampling, but it's making sure that the teacher is like, "Yeah, okay, this is reasonable."

If something is too far off, the teacher will take over and the teacher will say, "Okay, I don't think this should be sampled." I would sample one token from the teacher and then continue with the student. So basically, it's, it's like think-- imagine the driving scenario again.

It's basically saying the teacher is like, "I'm never gonna crash, but okay, you are in this situation. Let me actually fix this trajectory and let's continue rather than you crashing, 'cause after that point I cannot give you any meaningful feedback."

So it's much more expensive 'cause now your sampling involves speculative sampling itself, so now you're interleaving sampling with the student and teacher.

Right.

So the infra for this is even more complex. But you can see that this is more efficient. So if you go down, I guess that's the, the-

Yeah. Um, so like the, the, the, the comp is hidden in the actual inference. Yes. Am I still-

If you go to the last slide, I think that's the most interesting one in the current context. So this I think is maybe the main trade-offs when you're thinking about what distillation to use. So let's say if you're thinking of offline distillation, and I guess I've mentioned soft, but it could d- soft meaning you're using logits or next token probability distributions or not.

One thing is when you're doing online distillation, you don't need to have any sort of annotations. You don't need to save these to disk. You can just like press a button and basically it'll happen, 'cause you assume access to a teacher available online.

But then that, that ends up being more compute expensive, 'cause you're running a student, but you're also having access to a live teacher for running inference on that teacher during training itself, which is obviously more expensive. You're also not relying on just offline data.

You're sampling from the student, same as RL. Like you are running samples from that student and it's on policy, which is you're always sampling from your current model, so that is gonna be more expensive. That said, there are t- nice things about this, which is it manages the strain test mismatch.

The other thing is in the other literature there are like results which show that actually this kind of distillation is much more optimal for agentic tasks or really long horizon tasks. Think about like when you're outputting hundreds of thousands of tokens, if you're let's say doing that, it's very easy to imagine that there might be some places you'll make mistakes if you only learn from offline data, and basically you need this sort of online sampling to fix that.

Now this is TBD, whether how far we can go without it, but at least there are already a lot of results in the deeper literature not work on LLMs, but other tasks where they show that, okay, if you do just offline distillation, that's gonna be much more suboptimal than doing this on policy distillation thing.

So that's-- these are like the some of the trade-offs you'll face when picking on what distillation method to use. And I think-

I think it's gonna-

Oh, sorry.

Oh, sorry. I think that was an echo. I think that was an echo.

Oh, okay. Yeah. I, I think this is maybe the m- one of the most important things to consider when people are thinking about distillation, because they, they have to make some of these decisions, like what do they have?

Do they care about, let's say, trying things and training quite quickly, or do they want a process where they just have to do this once and that's it, they never have to come back to it again? So based on that, trade-offs are very different.

And also I guess more recently when you're thinking about agentic task or long horizon tasks, that's where maybe things get even more interesting, 'cause maybe it looks very different from what distillation has been done for let's say small, let's say a few tokens, thousands of tokens, a kind of task versus hundreds of thousands of tokens tasks.

Those I think might have very different behavior.

Cool. Uh, let me, let me pause this screen share. Uh, that was really comprehensive. Uh, actually I didn't realize this was accepted for ICLR, uh, twenty twenty-five. So, uh, it sounds like we'll-- you'll be presenting this in Singapore.

I think it's basically I told you about the frontier already, so people can now go and think about this and be like, "Well, this looks still dumb. Can I do better?" And I think that's generally a good mindset.

And also, well, I know about this now, maybe I can try it. 'Cause I think a lot of people don't know about some of these new things, partly 'cause they're very new and partly 'cause they have not propagated that widely.

But it's easy enough to try that it's worth trying given the fact that it's so closely tied-

Yeah

... to the RL infrastructure people have these days. Or at least that's my feeling.

I would say one challenge I have, and maybe we can kind of close it off from there- I would say one challenge I have, and maybe we can kind of close it off from there is, is that a live, a live teacher kind of defeats the purpose of teacher.

Kind of defeats the purpose of teacher.

Not really. Depends on how valuable the task is. Like I guess it's a question of would a ten percent jump in capability... 'Cause I think this is the, this is something someone told me, and I kind of internalized it.

They said, "You distill a model once, but you serve it billions of times." And I, I think I realized that, okay, what they're trying to say is it's worth spending the compute, additional compute now, 'cause if this model is the model which you're gonna just serve to like lots and lots of people, maybe it's worth the extra capability improvement.

But then again, it really depends on what-- how much capability is the threshold. 'Cause I would say try the simplest thing first, which is synthetic data distillation. That already gets you to eighty percent of the job or ninety percent of the job.

And now the question to ask is, is this good enough or do we wanna do more? And I think that's where you decide whether you wanna do this more complex distillation things which have this infra needs. My answer is oftentimes we do, 'cause I guess if we can push that Pareto frontier really up, right, and be like, "Oh, we can serve you really, really good capabilities but at a really cheap cost," that's how maybe these newer, better distillation methods can come into the picture.

And again, hopefully it's only done once. Like you're not doing it again and again. So the point is, it depends on how valuable this model is and would like making this model twice as better but with the same cost.

Is that meaningful enough and worth the additional complexity? Or at least that's my answer. So I would say always try the simple thing first and then decide whether it's worthwhile doing the more complicated thing. Sometimes it is, sometimes it's not.

And, uh, yeah.

Any, any thoughts- Any, any thoughts

Reasoning

Reasoning

... distillation? Distillation?

Yeah. So I would say that's what I was trying to allude to earlier, that I guess synthetic distillation would get pretty far, and people are mostly working on that. But my hunch is that the logic-based distillation can go even further, and you might be able to even close the gap with the biggest of the teachers you have.

'Cause I don't think you need a huge number of parameters 'cause the reasoning process is very, very... Like it's, it's basically a really long sequence, right, at the end of the day, and the model is producing these thoughts.

And to reason really well, maybe it's not about capacity, it's about using this inference time compute. Well, my hunch is, like, synthetic data distillation would go very far, but a better distillation method can maybe even start getting to close the gap, and you might have a very strong small reasoner, and that will be like a big uplift.

People are not trying it because, again, we don't have access to logics. It's annoying to use, like a... Let's say if you were trying to distill from DeepSeek, you can't serve that huge model and have like a live teacher running 'cause it is gonna be expensive.

But then again, maybe some people can afford to do that, and it's worth trying to push that and see if we can push much further. 'Cause the results I showed you earlier, I basically, on math tasks, we're able to close the gap reasonably well.

On coding, there was a huge gap still after synthetic data distillation. Now, how do you bridge that gap? There are c- easy ways. Either you collect a lot more data, a lot more prompts, that's obviously a simple way to get about it, or you try to use your existing data and annotate this with logics, now we are at supervised distillation, or you go the extra mile and be like, "Let me just generate data from my student and ask my teacher model to see what happens."

And I don't know the answer, but I feel like that is where it depends on how, how good of a student you want. That will dictate what kind of distillation method you use. Likely most, for most people, maybe the simple method suffices, but then I'm guessing for the, the...

Like, let's say you want your model to be much better than others, and that is maybe the winning edge you have, then it's worthwhile pushing for that extra. Because it's one time complexity cost, I feel like, and we already are paying that when we are building, doing RL, so it's not that big of a cost, honestly.

Cool. Um, any call to actions to wrap up? Obviously, we are gonna be back here. We're gonna be there too.

Yeah. Call to action is mostly like assume things are suboptimal and we can do better. I think that's maybe my call to action in anything. Think about what method we are using and how old that is, and very likely you'll find that, oh, we have been using this kind of okay thing.

It's simple, that's why everyone's using it, but it's not really good. And usually that helps you. Even if you are like doing something really good, you might end up finding something better, and I think that's a good hunch for people.

Even when you're doing practical things, like just don't get satisfied easily. Like just ask these questions of, "Why am I happy with this, and can I do better?" And usually if you ar- start asking these questions, I think you move further along.

And the other thing is, I think use some of the tools we have to actually discover these... Like I think one problem I have in AI is there's too many things going on, a lot of cool stuff. It's very hard to find the cool stuff, and a lot gets lost in this.

And I think even finding some of these cool stuff i- is gonna be very important in the coming years to connect the dots and maybe try out the relevant things. That's, uh... So I guess maybe I, I'm saying maybe you...

I saw the deep research style podcast, but yeah, use those kind of tools a lot and maybe that will help you make the connections, uh, you need. Still have to do a lot of your own thinking, but, but like finding some of these things is...

Like for example, the kind of things I presented were mostly saying, "People used to do this thing in RL this way. We are doing LLMs now, but we are not using any of the things RL people used to do.

What if we just do that?" That's it, and suddenly you have a much better method. So th- that I think kind of helps, especially when we are again talking about things like agents and reasoning and whatnot. I think a lot of the things are probably done on a small scale in the RL literature, in games and whatnot, and a lot of those ideas would be relevant again now.

It's just worth thinking about which ones and I guess both in practice and also for research.

Awesome. Um, yeah. Thank you for the time, Rishabh. It's great.

Yeah. Cool. All right. Thanks for having me. It was really great.