Forward and Reverse KL Divergence

I recently read this awesome blog post by Wei Hao about SFT, RL and On-Policy Distillation. It really brought to life to me how these methods differ in terms of (1) where the training signal is coming from and (2) the strength and nature of parameter updates. The article offers different explanations for why SFT is much more prone to catastrophic forgetting than RL and I recommend reading it for that alone.

One of the explanations has to do with the fact that the loss in SFT is using forward KL divergence, while the loss in RL is using backward KL divergence. Turns out this is not actually the author's favourite explanation, but I noticed that I didn't fully understand the distinction between forward and reverse KL divergence. So that's the main focus of this article.

We're going to look at the following post-training methods:

Supervised Fine Tuning (SFT): take starting distribution and make it similar to target distribution using negative log likelihood (NLL)
Reinforcement Learning (RL), with 2 flavours: Reinforcement Learning with Verifiable Rewards (RLVR) and Reinforcement Learning from Human Feedback (RLHF). There is no target distribution. sample from starting distribution and score by reward function.
On-Policy Distillation (OPD): sample from starting distribution like in RL, but score using target distribution like SFT or teacher model

tldr: the fact that SFT minimises forward KL divergence forces a model to match the target distribution regardless of its starting distribution. This gives rise to its mode covering behaviour and potentially bigger risk of catastrophic forgetting. RL/OPD on the other hand minimise reverse KL divergence, which lead to mode seeking behaviour. It's not forced to cover all of the target distribution, just that the states that it assigns high probabilities to should also have high probabilities under the target/teacher.

Chen et al. 2025

SFT/NLL minimizes forward KL divergence

In SFT, we have a labelled dataset representing the target distribution $p(y)$ and try to make the model's distribution $\pi_{\theta}(y)$ as similar to it as possible. The model is trained by minimizing the negative log likelihood:

L_{SFT} =

-\sum_{y} p(y) \log \pi_{\theta}(y)

which is the same as minimizing the forward KL divergence $D_{\mathrm{KL}}(p \| \pi_{\theta})$ :

D_{\mathrm{KL}}(p \| \pi_{\theta})

= \sum_{y} p(y) \log \frac{p(y)}{\pi_{\theta}(y)}

= \sum_{y} p(y) \log p(y) -

\sum_{y} p(y) \log \pi_{\theta}(y)

where the first term $\sum_{x} p(y) \log p(y)$ is the negative entropy of $p$ which doesn't depend on model parameters $\theta$ , and so the cross entropy term $-\sum_{y} p(y) \log \pi_{\theta}(y)$ is what matters, and which is the same as the SFT loss.

It's called forward KL divergence by convention because the target distribution $p$ comes first and $\pi_{\theta}$ second. It asks:

When samples come from the target distribution $p$ how surprised is the model $\pi$ ?

You can see why forward KL/ SFT training is quite a violent technique in the sense that, no matter what the initial model distribution is, you will pull the model towards the target distribution. That's because any training example that has a high probability under the target distribution $p(y)$ but a low probability under the model $\pi(y)$ makes the SFT loss explode.

KL-regularized RL minimizes reverse KL divergence

That's different from the case of reverse KL divergence $D_{\mathrm{KL}}(\pi_{\theta} \| p)$ . By a similar derivation as above, minimizing the reverse KL divergence essentially means minimizing the term $- \sum_{y} \pi_{\theta}(y) \log p(y)$

We just swapped the order of the arguments, model distribution first, then target distribution, but the resulting loss function is quite different. This says:

Whatever you generate, make sure it lies in a high-probability region of p.

Note that training examples which the model never visits, i.e. where $\pi_{\theta}(y)$ is low, don't contribute much to the overall loss. So the model is only judged on samples it generates, not all of what the target distribution prescribes.

Except in the case of RL, target distribution doesn't mean a fixed training data set like in SFT, but instead the current/previous/frozen version of the model, which we're going to refer to as the reference policy $\pi_{ref}$ .

Next let's look at why the (KL-regularized) RL loss is equivalent to reverse KL divergence. Let's start with the KL-regularized RL objective which is to maximize the expected reward subject to staying close to the reference policy:

\operatorname*{max}_{\theta} \sum_{y} \pi_{\theta}(y) \left[r(y) - \beta \mathrm{log} \frac{\pi_{\theta}(y)}{\pi_{ref}(y)}\right]

where $r(y)$ is the reward coming from a human in the case of RLHF or a verifier function in the case of RLVR. The term $\frac{\pi_{\theta}(y)}{\pi_{ref}(y)}$ is what penalizes the model distribution for drifting too far away from the reference distribution (that's the KL-regularization).

Maximising the above is the same as minimising the negative objective:

L_{\mathrm{RL}}(\theta) = \sum_{y} \pi_{\theta}(y) \left[\beta \mathrm{log} \frac{\pi_{\theta}(y)}{\pi_{ref}(y)} - r(y)\right]

And minimizing this loss is equivalent (up to a constant) to minimizing the reverse KL divergence $D_{\mathrm{KL}}(\pi_{\theta} \| p^{*})$ :

\begin{aligned} D_{\mathrm{KL}}(\pi_{\theta} \| p^{*}) &= \sum_{y} \pi_{\theta}(y) \log \frac{\pi_{\theta}(y)}{p^{*}(y)} \\ &= \sum_{y} \pi_{\theta}(y) \left[ \log \frac{\pi_{\theta}(y)}{\pi_{\mathrm{ref}}(y)} - \frac{r(y)}{\beta} \right] + \log Z \end{aligned}

where $p^{*}$ is a reward-tilted target distribution we can define as

p^{*}(y) = \frac{\pi_{ref}(y)\mathrm{exp}(r(y)/\beta)}{Z}

with constant normalizer

Z = \sum_{y'}\pi_{ref}(y')\mathrm{exp}(r(y')/\beta)

So multiplying out the above, here is what we're left with in the end:

\begin{aligned} D_{\mathrm{KL}}(\pi_{\theta} \| p^{*}) &= \sum_{y} \pi_{\theta}(y) \mathrm{log}\pi_{\theta}(y) - \sum_{y} \pi_{\theta}(y) \mathrm{log}\pi_{ref}(y) \\ &\quad - \frac{1}{\beta}\sum_{y}\pi_{\theta}(y)r(y) + \mathrm{log}Z \\ &= \underbrace{-H(\pi_{\theta})}_{\text{negative entropy}} + \underbrace{ \left[ -\sum_{y}\pi_{\theta}(y) \log \pi_{\mathrm{ref}}(y) \right] }_{\text{cross-entropy against reference}} \\ &\quad - \frac{1}{\beta} \underbrace{ \mathbb{E}_{\pi_{\theta}}\left[r(y)\right] }_{\text{reward term}} + \log Z \end{aligned}

To sum up

The KL-regularized RL objective gives rise to the objective of minimizing reverse KL divergence (in addition to maximizing the expected reward)
The entropy term can be thought of as a regularizer
The cross-entropy term $-\sum_{y}\pi_{\theta}(y) \log \pi_{\mathrm{ref}}(y)$ contributes to making the model less prone to catastrophic forgetting under RL than under SFT

OPD also minimizes reverse KL divergence

On-policy distillation is similar to SFT in that the model receives token-level feedback, rather than an overall training signal for the entire generation like in RL. In that sense parameter updates are more dense in OPD/SFT than in RL.

But it's also similar to RL in that we don't train on what the target distribution gives us but on what the model itself generates. In OPD, we have a student policy, the model getting trained, $\pi_{\theta}(y)$ and a teacher policy $\pi_{T}(y)$ . We sample from the student $y \sim \pi_{\theta}$ and try to move the student closer to the teacher on those student-sampled outputs by minimising the loss:

L_{\mathrm{OPD}} = \sum_{y} \pi_{\theta}(y) \left[\log \pi_{\theta}(y) - \log \pi_{T}(y) \right]

which is exactly the definition of reverse KL:

D_{\mathrm{KL}}(\pi_{\theta}||\pi_{T}) = \sum_{y} \pi_{\theta}(y) \log \frac{ \pi_{\theta}(y)}{\pi_{T}(y)}

Notice that outputs are weighted by whether the student actually attaches any weight to them, $\pi_{\theta}$ , rather than what the teacher says, in which case we would use $\pi_{T}$ . So similar to the case of RL, we ask:

among the things the model generates, what should get reinforced/suppressed according to the teacher?

We do this by comparing every token the student generates and compare it against the likelihood that the teacher would have generated the same token. If not ( $\pi_{\theta}$ is high and $\pi_{T}$ is low) this causes a large loss and therefore a training signal to reduce those tokens getting generated.

Wrap up

In this article we reviewed how the training objectives differ in SFT, RL and OPD, under the lense of forward and reverse KL divergence.

In SFT we're using forward KL divergence, which means that the expectation is under the target distribution $p$ . This forces the model to become more like the target distribution everywhere the target distribution assigns positive probability. In other words, we force the model to become the target distribution as much as possible. This makes SFT a very effective training method, but also risks catastrophic forgetting: by becoming more like the target distribution, the model forgets it's own priors.

In RL/OPD, we're using reverse KL divergence which means the expectation is under the model's distribution $\pi$ . This means the loss is weighted by the policy, i.e. by the model's probability of generating a certain token. States that are never visited by the model don't contribute to the loss. The model is judged based on whether the target/teacher would have assigned high probability on those tokens too, but only if the model itself generates them. This allows the model to keep more of it's prior knowledge when learning something new than in SFT, thereby reducing the risk of catastrophic forgetting.

Caveat: as Wei Hao points out, RL is actually less prone to catastrophic forgetting even if we omit or scale down the KL-regularization. So while this all makes intuitive sense, maybe it's only part of the story ¯\_(ツ)_/¯