A KL View of Post-Training Losses

This post is a note on Kullback–Leibler (KL) divergence in post-training, written after reading Thinking Machines Lab’s post on On-Policy Distillation.

I mainly want to understand one thing:

How does KL show up in different post-training losses, and what intuition does each KL direction give us?

The useful framing for me is to ask three questions:

1. Which distribution do we sample from?
2. Which distribution do we want to move toward?
3. Is the loss closer to forward KL or reverse KL?

Notation#

Let $x$ be the prefix and $y$ be the response or next token. For token-level formulas, I use $x$ as the current prefix.

I use:

• $\pi_\theta$ for the current student model.
• $\pi_T$ for the teacher model.
• $\pi_{\text{ref}}$ for the reference model.
• $p_{\text{data}}$ for the data distribution.

Forward KL

$$ D_{\mathrm{KL}}(\pi_{\text{ref}} \,\|\, \pi_\theta) = \mathbb{E}_{y\sim \pi_{\text{ref}}}\left[\log \pi_{\text{ref}}(y|x) - \log \pi_\theta(y|x)\right] $$

Reverse KL

$$ D_{\mathrm{KL}}(\pi_\theta \,\|\, \pi_{\text{ref}}) = \mathbb{E}_{y\sim \pi_\theta}\left[\log \pi_\theta(y|x) - \log \pi_{\text{ref}}(y|x)\right] $$

Forward KL: SFT#

SFT is like next-token prediction:

$$ L_{\mathrm{SFT}} = \mathbb{E}_{(x,y)\sim \mathcal D}\left[-\log \pi_\theta(y|x)\right] $$

We can also write this loss as:

$$ L_{\mathrm{SFT}} = D_{\mathrm{KL}}(p_{\text{data}} \,\|\, \pi_\theta) + \text{const} $$

For the relationship between cross-entropy and forward KL, see Understanding Cross-Entropy as a Loss Function.

So the intuition is: SFT asks the student to cover the data distribution. If the data contains a behavior, the student should assign probability to it. This can also help explain catastrophic forgetting: the model shifts probability mass away from its original distribution to match the target distribution.

SFT gives dense token-level supervision, but only on prefixes in the training data. During generation, the student conditions on its own outputs. If it makes an early mistake and reaches a prefix not covered by the dataset, it may not know how to proceed. Exposure bias arises because the model learns from dataset prefixes but must continue from its own generated prefixes at test time [1].

On-Policy Distillation: Reverse KL to the Teacher#

Instead of sampling trajectories from the teacher, we sample trajectories from the student. Then we ask the teacher to score the student’s tokens.

The objective is:

$$ L_{\mathrm{OPD}} = D_{\mathrm{KL}}(\pi_\theta \,\|\, \pi_T) $$

Expanding the KL gives a per-token log-prob loss:

$$ L_{\mathrm{OPD}} = \mathbb{E}_{y\sim \pi_\theta}\left[\log \pi_\theta(y|x) - \log \pi_T(y|x)\right] $$

It is on-policy like RL because the prefixes come from the student. But it is dense like distillation because every token can get a teacher log-prob signal.

SFT directly trains on demonstrated responses, including ones the student rarely generates. Sampled on-policy distillation provides feedback on the responses the student actually generates, so teacher behaviors that rarely appear in those rollouts receive less direct supervision.

Also, teacher supervision can be a biased proxy for our ultimate goal of maximizing expected reward on downstream tasks. Matching the teacher’s distribution may encourage its stylistic preferences or reproduce its mistakes, even when those changes do not improve task performance.

KL-Regularized RL: Reverse KL to an Implicit Target#

A common post-training RL objective has two parts: expected reward and a penalty for moving away from the reference policy.

$$ J(\pi_\theta) =\underbrace{\mathbb{E}_{y\sim\pi_\theta}[r(x,y)]}_{\text{expected reward}} -\underbrace{\beta D_{\mathrm{KL}}(\pi_\theta\|\pi_{\text{ref}})}_{\text{reference penalty}}. $$

For a fixed prompt $x$, reward, and reference policy, this can be rewritten as [2]:

So, our objective can be seen as optimizing our distribution to approach $p^*$, the reference distribution reweighted by reward. Higher-reward responses gain probability relative to lower-reward ones; $\beta$ controls the strength of this adjustment.

Maximizing this objective is therefore equivalent to minimizing reverse KL toward $p^*$. When PPO or GRPO is used to optimize this objective, sampled updates move the model toward this implicit target, which the model may not be able to represent exactly.

Summary#

References#

1. Samy Bengio, Oriol Vinyals, Navdeep Jaitly, and Noam Shazeer. Scheduled Sampling for Sequence Prediction with Recurrent Neural Networks. NeurIPS, 2015.
2. Rafael Rafailov, Archit Sharma, Eric Mitchell, Stefano Ermon, Christopher D. Manning, and Chelsea Finn. Direct Preference Optimization: Your Language Model is Secretly a Reward Model. NeurIPS, 2023. Eq. (4) and Appendix A.1.