Play Around with Distillation, Part 1: Forward vs. Reverse KL, On vs. Off Policy, and Hard vs. Soft!

1Introduction#

First of all, as mentioned above, the language in this blog will be casual, but the information I give you will be legit and informative, hopefully! Expect the content to be a mix of my personal ramblings, random unfunny jokes, and my personal experiences :), fused with academic writing. It's been a long time since I last wrote a blog (those were just some fun experiments, like training a model to write recipes), and I'm going to do a fun thing here too.a^a Also, some psychology/neuro stuff that I read says that sharing what you learn by writing it down tidily helps you learn better. So, I encourage people to write blogs too :D.

Motivation. Alright, let's start with the motivation. I see lots of people either downplaying knowledge distillation (KD) or just taking it for granted. The first group is people who don't really keep up with the distillation world, while the second group is people jumping on LLM-world techniques without caring how they work. However, some people do want to know the basic idea behind KD. Since I am in this field and I still keep updating my knowledgeb^b fyi, this is a pain in the current agentic era. People publish too fast.., I would like to share my insights about it. I will especially focus on the intuition behind distillation, such as its behavior, its theoretical foundation, and so on, while also giving you some practical problems. Hopefully, those who want to understand the KD problem better or want to enter this field will love it.

Structure. This blog will cover KD, on/off-policy, and soft and hard distillation. I will start with the basics of each, and then move on to a toy problem for vanilla (non-sequence) distillation. The sequence version will come in a separate blog post (Part 2), so this one doesn't get too long. For the toy problem, I will not use any modern language models (LSTM included in this list). Instead, I will use a very simple statistical model (a single bell curve over the classes), as this will help you understand better. Sure, there are lots of works with more advanced techniques, but as far as I know, they mostly tinker around with the policy and the objective function. The foundation should be the one I explain in this blog.

Despite being in the KD field, I am, honestly, still learning about it (it keeps evolving!!). I've been here since 2018, back when distillation was just a way to compress models (e.g., Sanh et al., 2019). Nowadays, with the advent of LLMs, people want to make small models work, and distillation is the main key (Gemma Team, 2024; Guo et al., 2025). RL (Reinforcement Learning) for language models also uses KD as a reward (Agarwal et al., 2024; Lu and Thinking Machines Lab, 2025), since the common rewards (e.g., a reward model (Ouyang et al., 2022) or a verifiable reward (Guo et al., 2025)) are sparse, hence not as informative as KD. I will not dive further into it here, but maybe I will in the future, hopefully. Anyway, what I was going to say is: KD is in an exciting era, though there are definitely some problems still unsolved. Mainly, the data used for distillation, and the right objective function that makes the knowledge transfer go the way we want it to.c^c Sometimes, KD is also used to change a model's behavior, not only to make it better.

2What Is Knowledge Distillation?#

Core KD Explanation. Let's start with the definition. The term KD was coined by Hinton et al. (2015) (the idea of making a small model copy a big one is a bit older, see Bucilua et al., 2006), who call the extra information from the teacher ‘dark knowledge' (sounds cool, huh?). Basically, KD is a method to transfer knowledge from a larger, better model (the teacher) to a smaller model (the student).d^d Nowadays, the student can be as big as the teacher, or even bigger! E.g., a same-size student can beat its teacher (Furlanello et al., 2018), and KD can work as a regularizer that stops a big student from overfitting (Yuan et al., 2020). The dark knowledge itself comes from the output probabilities, which give more information than the hard label (see Figure 2). Because of this, the student gets a more precise target, and the training signal has lower variance than training with hard labels only.

KD Objective Function. To understand KD, we need to look at its objective function, since every training in Machine Learning is about optimizing one (people sometimes call it the loss function). Let's start from the intuition: as mentioned above, we want the student's distribution to be anchored to, or mimic, the teacher's. To keep things standard, we follow Hinton et al. (2015) and start by defining the distributions we get from a neural network. Let \(z^t\) and \(z^s\) denote the logits of the teacher and the student, respectively. Logits are converted into probabilities with a softmax at temperature \(T\):

\(T\) is useful for controlling how much knowledge is transferred to the student. If you want a sharp distribution, use a lower \(T\) (close to 0 means it becomes one-hot). A higher \(T\) flattens the distribution, so classes with close logits end up with similar probabilities, and the tiny probabilities (the dark knowledge!) become more visible. But if \(T\) is too high, everything looks almost the same, and you are mostly transferring noise. See Figure 3 for an illustration. The choice of \(T\) depends on your use case. Now, let's define the distributions:

Here \(p\) is the teacher's ‘softened' distribution (the soft targets), \(q_1^s\) is the student's distribution softened with the same temperature, and \(q_2^s\) is the student's ordinary distribution at \(T=1\). Both \(q_1^s\) and \(q_2^s\) come from the same student logits; they differ only in temperature.

The student parameters \(\theta\) are trained to minimize a weighted combination of two terms:

where

and \(y^{gt}\) is the one-hot ground-truth label. The KL term pulls the student's softened distribution toward the teacher's; the cross-entropy term keeps the student anchored to the true labels. \(\alpha \in [0,1]\) balances the two. The factor \(T^2\) keeps the two terms on a comparable scale.e^e The gradient of the KL term with respect to the student logits scales as \(1/T^2\), since the softmax divides the logits by \(T\) and the difference between the softened distributions shrinks by another factor of \(T\). Multiplying by \(T^2\) cancels this, so changing \(T\) does not require retuning \(\alpha\). In modern LLM distillation, \(T=1\) is the common choice, in which case \(q_1^s = q_2^s\) and the \(T^2\) disappears.

Note that in some cases, we don't have any ground truth and can only rely on the teacher, so the second term is dropped (\(\alpha = 1\)). This is what people do when they use an AI model to ‘distill' another one: there is no ground truth, so the focus is only on the teacher.

Hard Distillation. The setting above (Eq. (3)), where you have access to the teacher's logits or probability distribution, is called soft distillation (mimicking soft targets). In reality, though, we often don't get them. When people publish data generated by an LLM, they usually don't include the probability distribution of each token due to storage issues.f^f Imagine you have a vocabulary of 40,000 tokens and each sequence has 1,024 tokens on average. That's around 41 million numbers (\({\sim}\)160MB in float32!) for one sequence. The file gets heavy really fast. The same goes for using an API, which normally doesn't give them to you (this is a current hot topic in the AI world). Training on this kind of data is still called distillation, since the LLM that generated it acts as the teacher, even though we only have hard targets. This is what we call hard distillation (Kim and Rush, 2016). A famous example is Alpaca (Taori et al., 2023), which was trained on outputs generated through the OpenAI API.

Note that hard distillation is not necessarily the same as using a low \(T\), since it depends on the decoding strategy used to generate the data. With an LLM API, you can still adjust the temperature, but you only get hard targets back. Though, perhaps, these hard targets can approximate the teacher's output distribution: if you sample one token from the teacher, the one-hot vector of that sample is, on average, exactly the teacher's distribution. So with enough samples, they do approximate it (see Figures 4c and 4d). Whether this works nicely for whole sequences is another story. See Figure 4 for a comparison between hard and soft distillation. Later, I will show you how hard and soft distillation behave differently in the toy problems.

3Forward and Reverse KL#

So, the vanilla KD proposed by Hinton et al. (2015) is what people call forward KL, where they use the Kullback–Leibler (KL) divergenceg^g In reality, there are lots of objective functions that can be used for KD (e.g., MSE, JS divergence, and so on). However, from this distribution point of view, KL divergence or similar divergences feel like the natural choice, since they measure how far apart two distributions are. Some work (Kim et al., 2021) says that MSE (on the logits) may be equal to, or even outperform, KL divergence. Actually, my prior work (Nityasya et al., 2022) also found this. to make the student match the teacher's output distribution. Look at the KL term in Eq. (3): this is the forward KL. People say it has a characteristic called ‘mean-seeking' (some also call it mode-covering), where the student tries to cover the teacher's distribution as much as possible.

Note that if the student has a lower capacity (i.e., fewer parameters) than the teacher, it needs to compromise on what it can store. This is a problem: we want to make sure the student does store acceptable information from the teacher to perform well, but it doesn't need to cover every piece of information from the teacher. MiniLLM (Gu et al., 2024) tries the OPPOSITE of forward KL: it uses reverse KL with on-policy data (we will get to this in §4), and finds that it performs better than forward KL with off-policy data. Reverse KL has a ‘mode-seeking' characteristic, where the student focuses on fitting one (or a few) modes really well instead of trying to cover all modes of the teacher. For those of you who are still confused about these seeking things, please refer to the visualization in Figure 5.

The reverse KL is just the KL term in Eq. (3), but flipped (with \(T=1\) here, so \(q^s = q_1^s = q_2^s\)):

Notice that the average is now taken over the student's distribution \(q^s\), not the teacher's.h^h Fun fact: this is why reverse KL and on-policy often come together. To estimate an average over \(q^s\) for long sequences, you naturally sample the sequences from the student itself, which is exactly on-policy (§4).

Also, note that only the student gets updated! In both forward and reverse KL, the teacher is frozen: it just gives \(p\) as the target, and the gradient only flows into the student's \(q^s\) (i.e., only \(\theta\) in Eq. (3) changes). So when people say “reverse KL makes the student mode-seeking'', it is the student that moves to fit the teacher, never the other way around.

Note that, as I mentioned earlier, there are lots of objective functions, but in this blog I will only focus on these two. I will cover the others in the future.

Understanding the Forward and Reverse KL. Now, let's dissect the forward and reverse KL to understand their characteristics more. Look at a single term inside the sums of Eqs. (4) and (5), i.e., one class (or one token) \(i\). Since both \(p_i\) and \(q^s_i\) can be anything, there are three common cases: \(p_i\) is much higher than \(q^s_i\), \(q^s_i\) is much higher than \(p_i\), and \(p_i\) is close or equal to \(q^s_i\). Let's do this! See Table 1.

So, forward KL hates it when the student misses something the teacher likes, hence the student spreads out to cover everything, even if that means putting probability in places the teacher doesn't like (mean-seeking). Reverse KL hates it when the student puts probability where the teacher says nope, hence the student plays safe and sits on the modes it is sure about (mode-seeking). This is exactly what you see in Figure 5: the forward KL student puts a lot of probability in the middle (the second case is not punished), while the reverse KL student ignores the right bump (the first case is not punished).

In the toy problems, I will show you this seeking behavior in action and observe its characteristics more. While at it, I will also show you a mathematical derivation later (§7.1, with the full proofs in App. A).

4Policy#

So here is the last piece before we play with the toy problems: policy. The term comes from RL (Sutton and Barto, 2018), where a policy is the thing that decides which action to take. For an LLM, the policy is just the model itself: given the text so far, it picks the next token. So, when people say on-policy or off-policy in distillation, they are basically answering one question: who writes the sequences used for training? See Figure 6.

Off-Policy. The sequences come from someone other than the student: the ground-truth dataset, or text generated by the teacher (like hard distillation above, Figure 4; Kim and Rush, 2016). The student then learns on these sequences, either from the teacher's distribution at each token (soft) or just from the tokens (hard). This is cheap and easy, since you can prepare the data once and reuse it. The catch: at inference time, the student continues from its own outputs, not from the dataset. If it makes a small mistake early, it ends up in a situation it never saw during training, and the mistakes can pile up. People call this the train–inference mismatch (or exposure bias) (Ross et al., 2011; Agarwal et al., 2024). For instance, suppose the training data often contains “I like an apple'', but the student generates “I love ..'' instead. If “I love'' never appears in the training data, the student has never learned what should come next, so it may make a mistake or hallucinate (see Figure 7).

On-Policy. The student generates the sequences itself, and the teacher gives its distribution for every token the student wrote. So, the teacher is basically grading the student's own homework and telling it what it should have done at each step (you can think of it as a reward), including how to recover from its own mistakes (Agarwal et al., 2024; Gu et al., 2024). The problem: it is more expensive, since you need to sample from the student and run the teacher during training. Note that you can also mix the two; for example, GKD (Agarwal et al., 2024) uses on-policy data some of the time and off-policy data the rest of the time. See Figure 7 for what the teacher and the student actually see, token by token, in both settings.

To sum it up, distillation roughly has three components: the models (the teacher(s) and the student(s)), the objective function, and the data used for training (i.e., the policy). For the objective, this blog only covers forward KL and reverse KL (Eqs. (4) and (5)). See Figure 8 for the whole picture.

5Toy Problems#

Now, we come to the exciting part: observing how the components above (§§2 and 4) behave! First of all, this blog focuses only on the KL term in Eq. (3), as that's the part we are interested in (so \(\alpha = 1\), no ground-truth anchor, and \(T = 1\)). Here, we look at non-sequence KD (Toy Problem 1, §6) and change different variables: the policy (off vs. on), the objective (forward vs. reverse KL), hard vs. soft targets, and the teacher/student models. (Sequence KD will be in Part 2.)

As mentioned in §1, instead of using advanced neural models (well, I will do that in the future), I will use simple models, since it's much easier to see what's going on. The question we want to answer is:

6Toy Problem 1: Non-Sequence KD#

First of all, let me remind you that KD essentially wants the student to mimic the teacher's distribution. This DOES NOT NECESSARILY mean the teacher has to be a neural network. It can also be a statistical model with a pre-defined output distribution. For instance, suppose the teacher has 3 labels; we can just assign the probabilities directly (e.g., label 0: 0.2, label 1: 0.5, label 2: 0.3).

The Teacher. In this toy problem, we let the teacher be a BI-MODAL GAUSSIAN distribution, i.e., it has \(M = 2\) hills, each with its own predefined MEAN, VARIANCE, and weight. To make it look like an LLM's next-token distribution, we put it over \(K\) classes \(x \in \{0, 1, \dots, K-1\}\) (think of them as tokens) and normalize it so it sums to 1:

where \(\mu^t_m\), \(\sigma^t_m\), and \(w_m\) are the mean (center), standard deviation (width), and weight of hill \(m\). The superscript \(t\) means these belong to the teacher (same as \(z^t\) in Eq. (2)): they are fixed and never learned.

Basically, if you paid attention in your statistics class, you will get it at a glance: it's just a Gaussian, but double, and we tell the model the mean and variance of each hill. If you didn't, now you learn it lol. Quick reminder in case you forgot: in a normal/Gaussian distribution, the mean is where the peak is, while sigma is how wide it is.

So, the teacher is fully described by \(K\) and the \((\mu^t_m, \sigma^t_m, w_m)\) of each hill, and each class gets its own probability (see Figure 9a). The exact values I use are in Table 2. Note that the classes here are ordinal (e.g., class 5 is ‘more' than class 4), though I believe this does not mean much here, since the KL still treats each class independently.

The Student. For the student, we use a simple model that can only make ONE GAUSSIAN HILL, so there are only TWO parameters to learn: the mean \(\mu^s\) and the width \(\sigma^s\) (we actually learn \(\log\sigma^s\), so \(\sigma^s\) always stays positive). The superscript \(s\) means student: unlike the teacher's \(\mu^t_m\) and \(\sigma^t_m\), these two are the ones being learned. Here is the formula of the student:

Fun fact: this is just a softmax (Eq. (1)) whose logits are forced into the shape \(-(x - \mu^s)^2 / 2(\sigma^s)^2\), so you can think of it as a tiny ‘neural network' with only two weights.

So, the capacity of the student is LIMITED: it can never match the teacher's two hills exactly. This is actually what we want! If the student could match the teacher perfectly, every method would end up in the same place. Here, the objective and the policy decide how the student is wrong. So, how does distillation behave given this?

The Policy. In distillation, we need both the teacher's and the student's distributions, and the policy decides who generates the samples we train on: the teacher (off-policy) or the student (on-policy). Note that in this toy problem, each sample is only ONE TOKEN (one class \(x\)), not a sequence.

Wait, why do we need samples at all? In this toy, we could compute the KL exactly by summing over all \(K\) classes. But in real training (especially for sequences), there are way too many possible outputs to sum over, so people estimate the loss with samples instead. To keep things realistic, we do the same here.

What Is Sampling?. If you are not familiar with statistics, think of sampling as rolling a weighted die with \(K\) sides. Each side is a class, and the weight of each side is its probability. For off-policy, the die is the teacher's distribution \(p\) (Eq. (6)), and we write it as

Since \(p\) has its hills around \(\mu^t_1\) and \(\mu^t_2\), most rolls land near them (e.g., 15, 15, 43, 12, 38, …), while classes where \(p\) is almost zero (e.g., the valley around 28, or 55) basically never show up. The same class can also show up more than once (15 shows up twice here!). Each roll gives ONE token, and we roll \(N\) times to get one batch. Then, for every sampled token, we ask both models how likely it is, \(p(x_i)\) and \(q_\theta(x_i)\), and average the gap. For forward KL, this is

where the hat means it's an estimate of the real KL (Eq. (4)), and it gets closer to the real one as \(N\) grows. The whole pipeline is in Figure 10.

For on-policy, the pipeline is exactly the same, except that we roll the student's die instead: \(x_i \sim q_\theta\) (Eq. (7)). And since the student keeps changing during training, the die (and so the samples) keeps changing too. With reverse KL, the estimate is just Eq. (9) flipped: \(\frac{1}{N} \sum_i [\log q_\theta(x_i) - \log p(x_i)]\) with \(x_i \sim q_\theta\).iⁱ These are the natural pairs: forward KL is an average over \(p\), so it goes well with teacher samples, and reverse KL is an average over \(q_\theta\), so it goes well with student samples. You can also do the other two pairs (e.g., forward KL with student samples), but then you need importance weighting (Lemma 6.1).

What If the Pair Doesn't Match? Importance Weighting. Say we want forward KL (an average over the teacher's \(p\)), but our samples come from the student's \(q_\theta\) (on-policy). The trick is to give each sample a weight that corrects for how much more (or less) often \(q_\theta\) rolls it than \(p\) would:

Think of dice: you want the average face of a fair die, but you can only roll a loaded one that gives a 6 half the time. The 6s show up too often, so each 6 gets a weight of \(\frac{1/6}{1/2} = 0.33\), while the other faces show up too rarely and get \(\frac{1/6}{1/10} = 1.67\), and the weighted average comes back to 3.5. So, forward KL with student samples becomes \(\frac{1}{N}\sum_i \frac{p(x_i)}{q_\theta(x_i)} [\log p(x_i) - \log q_\theta(x_i)]\) with \(x_i \sim q_\theta\), and reverse KL with teacher samples uses the weight \(q_\theta / p\) instead. The catch: if \(q\) almost never rolls a place where \(p\) is big, then those rare samples get a huge weight, and the estimate jumps around a lot. Keep this in mind for the results!

In our runs, each training step uses \(N\) samples, and the student starts from \(\theta_0 = (\mu^s_0, \log\sigma^s_0)\), on the left side: near the first hill, but far away from the second one. Also, note that an LLM outputs a token given an input (the prompt), i.e., \(P_\theta(\text{token} \mid \text{prompt})\), while our toy has no input at all: it's just \(P_\theta(x)\), which is our \(q_\theta(x)\). Doing the experiment with inputs would be interesting, but that'll be for a future blog, I guess XD.

The setup can be seen in Figure 9. Note that I will also try what happens if the teacher has \(M = 3\) hills instead of two, while the student still has only one (Figure 9b).

Try It Yourself!. Before we look at the results, you can play with this toy problem right here. Pick the objective, the policy, and soft or hard targets, then press Step to move forward 20 training steps at a time (roll, compare, update), or Play to watch the student learn to mimic the teacher. (All the runs were trained beforehand with the same code and settings as above; the page just plays them back.) Try reverse KL with the student starting on the right side, or the close hills, and see what happens!

7Result of Toy Problem 1#

Here are the results for the two-hill teacher (Figure 9a). The soft settings compute the KL exactly over all \(K\) classes, so there is no sampling and no randomness at all. That also means soft has no policy here: with only ONE token and no input, there is nothing for the teacher or the student to generate (this changes with sequences, where the prefixes are sampled; more on that in Part 2). The hard (sampled) settings use \(N\) sampled tokens per step, and each is run with 3 seeds. Everything is summarized in Table 3. Note that I trained for 4000 steps, since with forward KL, the student was still moving at 1500 steps (see Figure 13).

Wait, Why Doesn't Soft Have a Policy?. You might think: teacher samples and student samples are different data, so soft off-policy and soft on-policy should be different, right? Yes, the data is different. But whether that matters depends on whether the sample actually shows up inside the loss. Let's write the loss for one sampled token \(x_i\):

Look closely: \(\ell_{\text{soft}}\) sums over all classes \(k\), so the sampled \(x_i\) doesn't appear in it at all. \(\ell_{\text{hard}}\) only looks at the one token we sampled, so \(x_i\) is right there.

Why Does \(\ell_{\text{hard}}\) Look Like That?. Hard targets mean that, other than the target label, all the other classes get 0. For example, if the sampled token is class 3 (out of 4), the hard target is \([0, 0, 0, 1]\), while the student might say \(q_\theta = [0.3, 0.1, 0.1, 0.5]\). If we plug this one-hot target into the KL (Eq. (4)), every class with target 0 disappears (since \(0 \cdot \log 0 = 0\)), and only the sampled class is left:

which is just the plain cross-entropy on the sampled token (basic log arithmetic, in case you forgot: \(\log\frac{a}{b} = \log a - \log b\)). Our \(\ell_{\text{hard}}\) in Eq. (10) only adds \(\log p(x_i)\) on top. That extra term doesn't depend on the student at all, so both give exactly the same gradient; the \(\log p(x_i)\) is only there so that averaging \(\ell_{\text{hard}}\) over many teacher samples gives back the real KL (Eq. (9)).

A Tiny Example. Let's take only 3 classes, with teacher \(p = [0.2,\ 0.5,\ 0.3]\) and student \(q_\theta = [0.4,\ 0.4,\ 0.2]\). The soft loss compares the whole vectors at once:

while the hard loss depends on which token we rolled: \(\ln\tfrac{0.2}{0.4} = -0.69\) for class 0, \(\ln\tfrac{0.5}{0.4} = 0.22\) for class 1, and \(\ln\tfrac{0.3}{0.2} = 0.41\) for class 2. Now, let the teacher and the student each roll 4 tokens (Table 4).

So, the teacher and the student still disagree (\(p \neq q_\theta\), the soft loss is 0.095, not 0), but who rolled the tokens doesn't change the soft loss at all. That's why soft has no policy in Toy 1. The hard loss, on the other hand, gives different numbers for the teacher's batch and the student's batch, so off- and on-policy really differ. Notice that the student's batch even gives a negative hard loss, which can't be right for a KL! That's because forward KL should be weighted by the teacher's \(p(k)\), but the student rolls with its own \(q_\theta(k)\). This is exactly why forward KL with student samples needs importance weights (\(p / q_\theta\)) to fix the weighting.

In sequences, it's different again: the sampled prefix \(y_{<t}\) becomes the context, so it shows up inside the soft loss too, through \(p(k \mid y_{<t})\) and \(q_\theta(k \mid y_{<t})\). So there, soft off- and on-policy really give different losses (remember “I like ___'' vs. “I love ___'' in Figure 7).

7.1Forward vs. Reverse KL#

Mean-Seeking vs. Mode-Seeking. Figure 11 shows where the student ends up. As expected, with forward KL, the student covers both teacher hills (\(\mu^s \approx 26.9\), \(\sigma^s \approx 16.9\)): its center \(\mu^s\) sits between the two hills, and its width \(\sigma^s\) gets much bigger to reach both. So, it puts a lot of probability in the valley (around class 28), where the teacher has almost none. This is what being mean-seeking means: the student goes for the overall mean of the teacher's distribution, even if the teacher has almost nothing there.

Interestingly, with reverse KL, the student sits exactly on the first hill of the teacher (\(\mu^s = 15\), \(\sigma^s = 3\), which are exactly that hill's own mean and width) and ignores the second one. This is what we call mode-seeking: the student picks one mode and fits it really well.

Reverse KL Depends on the Starting Point. Which hill reverse KL picks depends on where the student starts (Figure 11b). Starting from \(\mu^s_0 = 10\), it picks the hill at 15, but starting from \(\mu^s_0 = 50\), it picks the hill at 42 (\(\mu^s = 42\), \(\sigma^s = 4\)). Forward KL, on the other hand, always ends at the same place from both starts.

This makes sense if we look at the gradients. Our student is a softmax (Eq. (7)): \(q_\theta = \operatorname{softmax}(z)\) with logits \(z_k = -(k - \mu^s)^2 / 2(\sigma^s)^2\). So, it's enough to look at how each KL reacts when we nudge one logit \(z_k\) (the gradient for \(\mu^s\) and \(\sigma^s\) then just follows from the chain rule).

Proof in App. A

Proof in App. A

Comparing the two gradients gives the “bounded vs. not bounded'' story:

Proof in App. A

In words: with forward KL (i), the gradient stays between \(-1\) and \(1\), and every class the teacher likes (\(p_k > 0\)) pulls the student toward it, no matter how far that class is from where the student currently is. So, the far hill always gets covered.

With reverse KL, two things happen instead. First, (iii): look at Lemma 7.2 again, and notice that the whole signal for class \(k\) is multiplied by the student's own probability \(q_k\). So, if the student gives a class almost zero probability, the teacher's opinion about that class barely reaches the student, even if the teacher loves that class. Let's check it on our two-hill teacher, with the student sitting on the first hill (\(\mu^s = 15\), \(\sigma^s = 3\)), and look at class 42, the peak of the far hill. The teacher likes it (\(p_{42} \approx 0.047\)), but the student almost ignores it (\(q_{42} \approx 3 \times 10^{-19}\)). With reverse KL, the gradient for class 42 is about \(-1 \times 10^{-17}\), basically nothing, so the student never hears “the teacher wants class 42''. With forward KL, it is \(q_{42} - p_{42} \approx -0.047\), a clear pull toward class 42. Think of reverse KL as checking your answers with a flashlight: you only see the teacher's opinion where your flashlight (your own probability) is shining. The far hill is in the dark, so as far as the student knows, it isn't there, and the student just rolls into the nearest hill (that's why the starting point matters). Forward KL, on the other hand, turns on the lights in the whole room.

Second, (ii): where the teacher has almost nothing (the valley) but the student still puts some probability, the gradient blows up, so the student gets pushed out of the valley really hard (mode-seeking).

Proof in App. A

This is mean-seeking, written as math: the student's mean lands on the teacher's overall mean, which, with two far-apart hills, is right in the valley. Checking it on our two-hill teacher: the student's mean is \(27.70\) vs. the teacher's \(27.71\), and its \(\mathbb{E}[x^2]\) is \(961.2\) vs. \(961.5\). Reverse KL has no such guarantee (because of Lemma 7.2), so every hill can be its own resting point.

So, Which One Is Better?. Neither, really: each one wins on its own KL. The forward KL student gets \(\mathrm{KL}(p\,\|\,q_\theta) = 0.66\), while the reverse KL student gets 18.3 on that same metric. Flip the metric, and the reverse KL student gets \(\mathrm{KL}(q_\theta\,\|\,p) = 0.64\), while the forward KL student gets 1.55. Again, which one you want depends on your use case: covering everything the teacher can do, or doing one thing really well.

What If the Hills Are Close to Each Other?. In Figure 12, I move the hills closer (centers at 25 and 35), so the valley between them is shallow. Now, forward and reverse KL give almost the same wide hill (\(\mu^s \approx 29.7\) vs. 30.1, \(\sigma^s \approx 6.1\) vs. 5.8), so the mode-seeking behavior doesn't really show. Perhaps, the mode-seeking behavior only shows up when the valley is deep?

7.2Off- vs. On-Policy#

Here, the “natural'' pairs are forward KL with teacher samples (off-policy) and reverse KL with student samples (on-policy). The other two pairs need importance weights (Lemma 6.1).

Final Destination. With the natural pairs, hard training lands at the same place as soft training (\(\mu^s = 26.7\) vs. 26.9 for forward KL; 15.0 vs. 15.0 for reverse KL, Table 3).

Forward KL + On-Policy Is the Least Reliable. Across the 3 runs, it ends up in different places (\(\mu^s = 23.5 \pm 2.8\), vs. only \(\pm 0.07\) for off-policy), and it wobbles the most (0.64). From what I see, the student rarely samples the far hill, so it rarely sees what the teacher wants there. And when it finally does, the importance weight for that sample becomes huge, which makes the gradient jump around (see the bumpy on-policy curve in Figure 13a).

Reverse KL + Off-Policy. This one was okay (all 3 seeds reach the first hill), but it is a bit less stable than reverse KL + on-policy (wobble 0.02 vs. 0.01), and at 1500 steps, one of the three seeds had not made it yet.

Conclusion. For forward KL, off-policy is better, whereas for reverse KL, on-policy is better. Basically, use the policy that matches the KL.

7.3Hard vs. Soft#

Figure 14 zooms in on what happens with hard targets.

Final Destination. Basically, both end at the same place (Table 3). On average, the hard gradient also points the same way as the soft one: at the start, the true (soft) gradient for \(\mu^s\) is \(-0.70\), and the average of 2000 hard estimates is also \(-0.70\).

Noisiness. However, each single hard estimate is a noisy guess: those 2000 estimates spread out with a standard deviation of \(0.07\), which is about 10% of the true value (Figure 14a). Because of this noise, the student never fully stops. Soft training stays still once it arrives (wobble \(= 0\)), while hard training keeps wobbling around the answer (wobble \(\approx 0.27\) for forward KL, off-policy, Figure 14b).

Reducing the Noise. Remember that the hard gradient is just an average over the \(N\) sampled tokens: each token \(x_i\) gives its own gradient \(g_i\), and we use \(\hat{g} = \frac{1}{N}\sum_{i} g_i\). Two basic facts from statistics tell us how good this average is:

(See, e.g., Wasserman, 2004; and Mohamed et al., 2020 for this in the context of gradient estimators in machine learning.) In words: the spread shrinks like \(1/\sqrt{N}\), so to cut the noise in half, you need \(4\times\) more samples.j^j For on-policy reverse KL, the samples share a baseline (the average reward of the batch), so they are not perfectly independent, and Eq. (12) is only approximately true. It still follows the rule nicely in Figure 14c. And indeed, every time \(N\) goes up \(4\times\), the spread of the gradient is cut in half (\(0.14 \rightarrow 0.07 \rightarrow 0.036 \rightarrow 0.018\) for \(N = 16, 64, 256, 1024\), Figure 14c), and the wobble at the end also goes down (0.47, 0.27, 0.18, 0.08). So, more samples give more stability, but of course, more compute per step.

7.4Three Hills: What If the Student Has One or Two Hills?#

Now, let's make the teacher harder: three hills at classes 10, 30, and 50, all with the same width (\(\sigma^t_m = 3\)) and weights \(0.40\), \(0.35\), and \(0.25\) (Figure 9b). And this time, I also give the student a second hill. With two hills, the student is a mix of two copies of Eq. (7):

where each \(q_j\) has its own \(\mu^s_j\) and \(\sigma^s_j\), and \(\pi_j\) says how much weight each student hill gets (also learned).

Setup. To keep things clean, I only use soft (exact) training here, so there is no sampling noise at all. The only randomness is where the student starts: for each of 10 seeds, the student's means \(\mu^s_j\) start at random spots between 0 and 59, and its widths start at random values between 2 and 8. Each run uses Adam (learning rate 0.05) for 8000 steps, which is long enough that every seed stops moving.

A Handy Rule for Reverse KL. Before looking at the results, here's a rule that helps a lot.

So, copying only the hill at 10 costs \(-\log 0.40 = 0.92\), only the hill at 30 costs \(1.05\), and only the hill at 50 costs \(1.39\). Copying two hills costs \(0.29\) (10 and 30), \(0.43\) (10 and 50), or \(0.51\) (30 and 50). Fun fact: this also explains the two-hill result above (§7.1), where the first hill holds about 53% of the teacher, so copying it costs \(-\log 0.53 = 0.64\), exactly the reverse KL in Table 3.

One Student Hill. With forward KL, all 10 seeds end up at the exact same place: one very wide, almost flat hill (\(\mu^s \approx 20.6\), \(\sigma^s \approx 30.5\)) that covers everything (Figure 15a). Again, forward KL has only one best answer, no matter where the student starts.

The surprise is reverse KL: 8 out of 10 seeds also end up as a wide hill (\(\mu^s \approx 25\), \(\sigma^s \approx 28.4\), loss \(0.84\)), and only 2 seeds sit on a single teacher hill, at 30 (loss \(1.05\)) and at 50 (loss \(1.39\)) (Figure 15b). So, reverse KL is not mode-seeking here! Lemma 7.6 tells us why: copying any single hill costs at least \(0.92\), while spreading out costs only \(0.84\), so spreading out is simply cheaper. In the two-hill case, it's the other way around (copying the first hill costs \(0.64\), spreading costs \(1.55\)), so the student picks a hill. Together with the close-hills result (Figure 12), this tells us that whether reverse KL is mode-seeking depends on the teacher's shape: it picks one hill only when that hill holds enough of the teacher's probability.

Two Student Hills. Now, the seeds start to disagree (Table 5). With forward KL, they end up in three different places (Figure 15c). The best one (4 seeds, loss \(0.29\)) puts a narrow hill on 10 and stretches a wide hill over 30 and 50; 5 seeds do it the other way around (wide over 10 and 30, narrow on 50, loss \(0.37\)); and 1 seed basically gives up on its second hill (weight \(\approx 0\)) and ends as a single wide hill (loss \(0.48\)). So even forward KL can get stuck now: with a mix of hills, the “only one best answer'' property is gone. Note that each student hill still covers, though: a wide student hill happily stretches over two teacher hills.

With reverse KL, it gets even messier: the seeds end up in four different places, and none of them is the best one (Figure 15d). By Lemma 7.6, the best possible answer is to copy the hills at 10 and 30 exactly (loss \(0.29\)), and I checked that it really is a resting point by starting the student right there. But 0 out of 10 random seeds found it. Instead, 3 seeds put a narrow hill on 10 and a wide one over 30 and 50 (loss \(0.42\)), 2 seeds copy the hills at 30 and 50 exactly and ignore the hill at 10 (loss \(0.51 = -\log 0.60\), the classic mode-seeking picture), and the rest end up with one wide and one narrow hill (loss \(0.55\) and \(0.57\)).

Why Seeds Matter. With a simple student (one hill) or a simple teacher (two hills), forward KL always lands at the same place, and the starting point only matters for reverse KL. But once the student gets more flexible (two hills), there are many more places to get stuck, for both objectives, and the result really depends on the seed. So, if you try this yourself, run several seeds (or start the student somewhere smart), and don't trust a single run!

8Closing Remarks with More Intuition#

In this blog, I shared the results of me playing around with distillation, especially the policies (on/off), reverse and forward KL, and hard and soft distillation. The goal was to dive deep into their characteristics through toy problems and mathematical proofs.

How Does This Relate to KD in NLP Now?. As you can see, people are leaning toward on-policy distillation, often together with reverse KL (Gu et al., 2024; Agarwal et al., 2024; Lu and Thinking Machines Lab, 2025). As the proofs show, reverse KL has no bound on its gradient, unlike forward KL (Corollary 7.3), so it can push the student very hard; on the other hand, the distribution it gives is quite aligned with what we usually want (high-quality outputs). But since the result depends on where the student starts (Lemma 7.2), and the starting point is the pre-trained model, it would be wise to give the student a proper prior first (for instance, by doing SFT), since otherwise it may get stuck in a bad local minimum (remember the seeds in §7.4). This is also what MiniLLM and GKD do: they start from a student that was already fine-tuned (Gu et al., 2024; Agarwal et al., 2024). How to do this well is still open research. Some papers look at related problems: for instance, Xin et al. (2026) show that during on-policy distillation, the student can drift into a bad prefix (e.g., a reasoning mistake or repeated tokens) where the teacher locally agrees with it, so the reverse KL is low but gives almost no corrective signal (they call it the “low-KL agreement trap''), and they propose stopping such rollouts early. This is very much in the spirit of Corollary 7.3 (iii): with reverse KL, the teacher only gives feedback on the states the student actually visits. Again, it depends on your use case: I think the intuition is to figure out which distribution you want your student model to have.

Another Fun Fact. If you've realized it by now, the teacher that I showed you is basically like training data, which has its own distribution. This is actually how you should see normal training: the data usually comes as hard labels, and we do off-policy training on it. In reality, the data itself was produced by some process (or by another model), and your model maps it into a distribution (e.g., a transformer outputs a probability distribution), which may not cover every possible pattern in the data, again due to capacity. That's why crafting appropriate training data is IMPORTANT, so that you get the distribution you want. And since you train with the cross-entropy loss, which is just forward KL with hard targets (see the \(\ell_{\text{hard}}\) paragraph in §7), you get mean-seeking behavior toward the whole population of the data. Then, when that model is distilled into a smaller model, the behavior again depends on the loss used. So, based on the experiments, you can pick whichever suits your use case, or mix them like GKD (Agarwal et al., 2024) and DistiLLM (Ko et al., 2024).

Feedback Please!. I am still eager to learn about this field, even though I have been in it for a while. Here, I used LLMs to assist me with the proofs, writing (grammar and better phrases), and the tedious stuff (especially making this blog nicer to read and play with). The topics I talk about here are things I was already familiar with, but I had never played around with them before this write-up. I just took it for granted that reverse and forward KL are mode-seeking and mean-seeking, respectively, but I never checked it with my own eyes. So here I am! I revisited distillation and wrote up the conclusions of these fun experiments. I'm happy to get any feedback, as it will help me get better!

What's Next?. As I said, I have actually also played with KD on sequences, but I felt it would make this blog way too long, so I will put it in a separate part. I would like to show you how sequence distillation behaves, since it has some interesting findings, especially for soft on- vs. off-policy. There, I will use a simple Markov chain (think of it as a tiny language model that only looks at the last token), and besides that, I will also try a more language-like sequence problem (spoiler: I will use that old tree-based statistical model, a Probabilistic Context-Free Grammar (PCFG), and check how distillation behaves there). So, stay tuned for Part 2!

The GitHub repository will be uploaded soon, as the code is still messy… please wait for it. Thank you for understanding!

References#

1. Agarwal et al., 2024Rishabh Agarwal, Nino Vieillard, Yongchao Zhou, Piotr Stanczyk, Sabela Ramos Garea, Matthieu Geist, and Olivier Bachem. On-Policy Distillation of Language Models: Learning from Self-Generated Mistakes. International Conference on Learning Representations, 2024.LinkScholarCopy BibTeXCited in §1, §4, §8
2. Bucilua et al., 2006Cristian Bucilua, Rich Caruana, and Alexandru Niculescu-Mizil. Model Compression. Proceedings of the 12th ACM SIGKDD International Conference on Knowledge Discovery and Data Mining, pages 535–541, 2006.DOIScholarCopy BibTeXCited in §2
3. Furlanello et al., 2018Tommaso Furlanello, Zachary Lipton, Michael Tschannen, Laurent Itti, and Anima Anandkumar. Born Again Neural Networks. Proceedings of the 35th International Conference on Machine Learning, pages 1607–1616, 2018.LinkScholarCopy BibTeXCited in §2
4. Gemma Team, 2024Gemma Team. Gemma 2: Improving Open Language Models at a Practical Size. arXiv preprint arXiv:2408.00118, 2024.LinkarXivScholarCopy BibTeXCited in §1
5. Gu et al., 2024Yuxian Gu, Li Dong, Furu Wei, and Minlie Huang. MiniLLM: Knowledge Distillation of Large Language Models. International Conference on Learning Representations, 2024.LinkScholarCopy BibTeXCited in §3, §4, §8
6. Guo et al., 2025Daya Guo, Dejian Yang, Haowei Zhang et al.. DeepSeek-R1 Incentivizes Reasoning in LLMs through Reinforcement Learning. Nature, 645(8081):633–638, 2025.DOIScholarCopy BibTeXCited in §1
7. Hinton et al., 2015Geoffrey Hinton, Oriol Vinyals, and Jeff Dean. Distilling the Knowledge in a Neural Network. arXiv preprint arXiv:1503.02531, 2015.LinkarXivScholarCopy BibTeXCited in §2, §3
8. Kim and Rush, 2016Yoon Kim and Alexander M. Rush. Sequence-Level Knowledge Distillation. Proceedings of the 2016 Conference on Empirical Methods in Natural Language Processing, pages 1317–1327, 2016.DOILinkScholarCopy BibTeXCited in §2, §4
9. Kim et al., 2021Taehyeon Kim, Jaehoon Oh, NakYil Kim, Sangwook Cho, and Se-Young Yun. Comparing Kullback-Leibler Divergence and Mean Squared Error Loss in Knowledge Distillation. Proceedings of the Thirtieth International Joint Conference on Artificial Intelligence, pages 2628–2635, 2021.DOIScholarCopy BibTeXCited in §3
10. Ko et al., 2024Jongwoo Ko, Sungnyun Kim, Tianyi Chen, and Se-Young Yun. DistiLLM: Towards Streamlined Distillation for Large Language Models. Proceedings of the 41st International Conference on Machine Learning, 235:24872–24895, 2024.LinkScholarCopy BibTeXCited in §8
11. Lu and Thinking Machines Lab, 2025Kevin Lu and Thinking Machines Lab. On-Policy Distillation. Thinking Machines Lab: Connectionism, 2025.DOILinkScholarCopy BibTeXCited in §1, §8
12. Mohamed et al., 2020Shakir Mohamed, Mihaela Rosca, Michael Figurnov, and Andriy Mnih. Monte Carlo Gradient Estimation in Machine Learning. Journal of Machine Learning Research, 21(132):1–62, 2020.LinkScholarCopy BibTeXCited in §7.3
13. Nityasya et al., 2022Made Nindyatama Nityasya, Haryo Akbarianto Wibowo, Rendi Chevi, Radityo Eko Prasojo, and Alham Fikri Aji. Which Student is Best? A Comprehensive Knowledge Distillation Exam for Task-Specific BERT Models. 2022.LinkarXivScholarCopy BibTeXCited in §3
14. Ouyang et al., 2022Long Ouyang, Jeffrey Wu, Xu Jiang, Diogo Almeida, Carroll Wainwright, Pamela Mishkin, Chong Zhang, Sandhini Agarwal, Katarina Slama, Alex Ray et al.. Training Language Models to Follow Instructions with Human Feedback. Advances in Neural Information Processing Systems, 35:27730–27744, 2022.LinkScholarCopy BibTeXCited in §1
15. Ross et al., 2011Stéphane Ross, Geoffrey Gordon, and Drew Bagnell. A Reduction of Imitation Learning and Structured Prediction to No-Regret Online Learning. Proceedings of the Fourteenth International Conference on Artificial Intelligence and Statistics, pages 627–635, 2011.LinkScholarCopy BibTeXCited in §4
16. Sanh et al., 2019Victor Sanh, Lysandre Debut, Julien Chaumond, and Thomas Wolf. DistilBERT, a Distilled Version of BERT: Smaller, Faster, Cheaper and Lighter. arXiv preprint arXiv:1910.01108, 2019.LinkarXivScholarCopy BibTeXCited in §1
17. Sutton and Barto, 2018Richard S. Sutton and Andrew G. Barto. Reinforcement Learning: An Introduction. MIT Press, 2018.LinkScholarCopy BibTeXCited in §4
18. Taori et al., 2023Rohan Taori, Ishaan Gulrajani, Tianyi Zhang, Yann Dubois, Xuechen Li, Carlos Guestrin, Percy Liang, and Tatsunori B. Hashimoto. Stanford Alpaca: An Instruction-following LLaMA Model. GitHub repository, 2023.LinkScholarCopy BibTeXCited in §2
19. Wasserman, 2004Larry Wasserman. All of Statistics: A Concise Course in Statistical Inference. Springer, 2004.DOIScholarCopy BibTeXCited in §7.3
20. Xin et al., 2026Haoran Xin, Anhao Zhao, Ying Sun, Jin Li, Xiaoyu Shen, and Hui Xiong. Escaping the KL Agreement Trap in On-Policy Distillation. 2026.LinkarXivScholarCopy BibTeXCited in §8
21. Yuan et al., 2020Li Yuan, Francis E. H. Tay, Guilin Li, Tao Wang, and Jiashi Feng. Revisiting Knowledge Distillation via Label Smoothing Regularization. Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 3903–3911, 2020.DOIScholarCopy BibTeXCited in §2

AProofs#

Don't worry, these proofs only use a few basic rules. Let's collect them first, then build one helper fact that every proof below uses.

Toolbox. We only need these:

• Log of a fraction: \(\log \frac{a}{b} = \log a - \log b\).
• Log undoes exp: \(\log e^{z} = z\).
• Derivatives: \(\frac{d}{dz} e^{z} = e^{z}\), \(\frac{d}{dx} \log x = \frac{1}{x}\), and the derivative of a constant is \(0\).
• Chain rule: \(\frac{d}{dz} \log f(z) = \frac{1}{f(z)} \cdot \frac{d f(z)}{dz}\).
• Product rule: \(\frac{d}{dz} \big[f(z)\, g(z)\big] = \frac{d f}{dz}\, g + f\, \frac{d g}{dz}\).

A Tiny New Symbol. We will write \(\mathbb{1}[j = k]\) (“indicator'') for a number that is \(1\) if \(j = k\) and \(0\) if \(j \neq k\). For example, \(\mathbb{1}[3 = 3] = 1\) and \(\mathbb{1}[2 = 3] = 0\). It shows up naturally when we take the derivative of \(z_j\) with respect to \(z_k\): if \(j = k\), it's the derivative of \(z_k\) with respect to itself, which is \(1\); if \(j \neq k\), then \(z_j\) is a different variable (a constant as far as \(z_k\) is concerned), so the derivative is \(0\). In short,

It also has a handy property: in a sum over \(j\), it keeps only the \(j = k\) term, so \(\sum_j a_j\, \mathbb{1}[j = k] = a_k\).

Helper: How the Softmax Reacts to One Logit. The student is \(q = \operatorname{softmax}(z)\), i.e.,

is just the sum that makes everything add up to 1. Let's see how \(\log q_j\) changes when we nudge one logit \(z_k\), step by step.

1. Take the log (log of a fraction, then log undoes exp, see the toolbox):

   \[ \log q_j = \log e^{z_j} - \log S = z_j - \log S. \tag{16}\]
2. Derivative of the first part: \(\frac{\partial z_j}{\partial z_k} = \mathbb{1}[j = k]\), by Eq. (15).
3. Derivative of the second part: by the chain rule, \(\frac{\partial \log S}{\partial z_k} = \frac{1}{S} \cdot \frac{\partial S}{\partial z_k}\). Inside \(S = e^{z_1} + e^{z_2} + \dots\), only the term \(e^{z_k}\) depends on \(z_k\), and its derivative is \(e^{z_k}\). So \(\frac{\partial \log S}{\partial z_k} = \frac{e^{z_k}}{S} = q_k\).

4. Put them together:

   \[ \frac{\partial \log q_j}{\partial z_k} = \mathbb{1}[j = k] - q_k. \tag{17}\]

In words: if we push up logit \(k\), then \(\log q_k\) goes up by \(1 - q_k\) (its own class gains), and every other \(\log q_j\) goes down by \(q_k\) (they lose, since everything must still add up to 1). We will also need the version without the log. By the chain rule, \(\frac{\partial \log q_j}{\partial z_k} = \frac{1}{q_j} \frac{\partial q_j}{\partial z_k}\), so multiplying both sides by \(q_j\):

Back to Lemma 7.1

Back to Lemma 7.2

Back to Corollary 7.3

Back to Proposition 7.4

BCitation & Changelog#

@misc{wibowo2026playdistillation,
  title        = {Play Around with Distillation, Part 1: Forward vs. Reverse {KL},
                  On vs. Off Policy, and Hard vs. Soft!},
  author       = {Wibowo, Haryo Akbarianto},
  year         = {2026},
  month        = oct,
  howpublished = {Blog post}
}