blog: the intuition behind KL divergence in post-training
For the longest time I understood KL as a pretty term in all RL papers and a "how different two probability distributions are" measurement with the mental image of my friend pinching their fingers at the vertical gap between two deformed curves scribbled on a whiteboard in Fine Hall, our math building. This year after reading some more papers on RL specifically on learning capabilities and temporal credit assignment, I found that picture to be inadequate.
i want to spend some time writing out how to gain intuition beyond this pinching 🤏, which could help readers understand its relation to phenomena in post training like forgetting, how things are learned and probability mass on weights moved, and its relation to the type of samples.
as we see neolabs and research show more and more interest in agents and llm systems that can self-update, self-improve, and self-learn, we must pay attention to the fundamental trade-off between placistiy and stability, between learning and forgetting. sitting at the the fundamental layer, among many other things, is KL (as a cousine of cross entropy), and understanding it will provide a good baseline for choosing the minimizers or guarding principals to control the learning and forgetting of these models.
Notation:
- $\pi_\theta$ is the policy being trained, $\pi_{\theta_0}$ is the initial policy before post-training, $\pi^*$ is the target policy
- $r(y)$ is a reward
- $\beta$ is the KL-regularization coefficient
- $\nabla_\theta \log \pi_\theta(y)$ is the score function = the direction in parameter space that most increases the log-probability of the specific outcome $y$.
entropy, cross entropy, and KL
for those unfamiliar, here's some intuition on entropy, cross entropy, and KL that I found useful: read the appendix background.
What KL measures and punishes as loss
Formally,
here's a great youtube video https://youtu.be/SxGYPqCgJWM?si=0Fj0CafUAq3ZeJ7a that explains KL intuitively your surprise if reality was p and you believed in q. in more formal language, this is the expected number of extra nats you pay to encode samples drawn from $p$ using a codebook optimized for $q$.
key points to internalize:
- The first slot is "reality" and generates the samples. the second slot is the one being judged against reality p.
- Look at the term $\log(p(y)/q(y))$ at a single outcome $y$. If $p(y)$ is large and $q(y) \to 0$, the term approaches infinity: an unbounded penalty for calling something impossible when it actually happens. on the other hand, if instead $p(y) \approx 0$, the term is multiplied by $p(y) \approx 0$ in the expectation, so whatever $q$ does there is nearly free -- if something is rare in reality, however we predict it in the model gets no penalties under this model.
$\mathrm{KL}[p \Vert q]$ only punishes $q$ for missing regions where $p$ has mass. It never punishes $q$ for placing mass where $p$ has none.
SFT is forward KL, RL is reverse KL: different punishments
a large part of this section and the next two will draw on a paper by he Princeton Language and Intelligence(Chen et al. 2025). they found that, when post training across Llama and Qwen models, on instruction following, general knowledge, and arithmetic reasoning, they find a consistent pattern: SFT forgets while learning a new task, RL doesn’t forget while learning decently well. The paper introduced the notion of SFT minimizing the forward KL divergence, while RL minimizes the reverse KL. when i read it i said i didn’t know the direction in which you pinched mattered, especially in post training. so let’s dig in!
SFT. The loss is negative log-likelihood of samples $y$ drawn from the target distribution $\pi^*$, with $\pi_\theta$ being the distribution based on the weights we are training:
i'm not going to delve into the math here, but basically what we see here is that this loss is equivalent to cross-entropy.
cross-entropy is entropy plus KL (shown in the appendix background). the entropy of $\pi^*$ is a constant in $\theta$, so minimizing the SFT loss is exactly minimizing $\mathrm{KL}[\pi^* \Vert \pi_{\theta}]$, which we call forward KL: the tunable model sits in q, the judged slot.
this says that SFT tunes the model in a way such that anywhere the data has mass, the model must place mass or get a huge penalty. If there is some information in the target, real-life data, like “the sky is purple”, the model must learn it or get punished.
RL. in RL the notion of reverse-KL loss arises from our analysis of the standard KL-regularized objective that we are attempting to maximize:
Korbak et al. (2022) showed that the policy maximizing this objective has a closed form:
intuitively, this means: take the initial policy, multiply up the outputs that earn reward, renormalize. we see that the RL target $\pi^*$ literally contains the initial policy $\pi_{\theta_0}$ as a factor; RL is not chasing some external ideal but rather a reweighted version of what the model already was. intuitively, think about when the model gets a binary reward = 0 on tasks using old knowledge, then $e^0 = 1$, and those portions of $\pi_{\theta_0}$ pass into $\pi^*$ unchanged apart from renormalization. so the model explicitly retains "memory".
now, with some manipulation purely arithmetically we get
now write out the reverse KL from the current policy to the destination and substitute:
The sum splits into three pieces:
- $\sum_y \pi_\theta (\log \pi_\theta-\log \pi_{\theta_0})$ = $\mathrm{KL}[\pi_\theta \Vert \pi_{\theta_0}]$ = the regularizer.
- $-\frac{1}{\beta}\ \mathbb{E}_{y \sim \pi_\theta}[r(y)]$ = the expected reward just scaled.
- $\log Z \cdot \sum_y \pi_\theta(y) = \log Z$ = constant
Thus, all together,
Since $\log Z$ is a constant in $\theta$, the two sides are equal up to an additive constant:
maximizing the RL objective is minimizing the reverse KL to $\pi^*$. (Korbak et al. 2022; Chen et al. 2025, §3.1).
what does this mean? the model is only penalized where it places probability mass and $\pi^*$ does not (equivalently, where the base model does not: since $\pi^* \propto \pi_{\theta_0}e^{r/\beta}$, the two share the same support for bounded rewards); it gets penalized for learning things drastically out of distribution , i.e. drifting into new territory and learning very wild things. it doesn't get penalized for ignoring parts of the target model - for not learning. so it can be lazy!
The data (p) is what determines the KL direction
The previous section showed that SFT minimizes forward KL and RL minimizes reverse KL. This section explains how and why it is determined by which distribution produces the training samples.
p is the sampler. what is a sampler? the sampler is whoever hands you the examples you train on this step.
Your loss is an average over some distribution: $\mathbb{E}_{y \sim D}[\text{stuff}(y)]$. That average is over every possible output $y$. For a language model that is every possible token sequence, so you can never actually compute it. The only thing you can do is grab a few concrete $y$'s, compute $\text{stuff}(y)$ on those, and average. The distribution $D$ you grab from is the sampler. It is the answer to the question "where did this batch physically come from?"
There are two samplers in post-training:
- SFT = the sampler is the dataset
- Someone (humans, a bigger model) wrote demonstrations, they are sitting in a file, and each batch you read some. They are drawn from $\pi^*$ regardless of what the model currently believes.
- RL = the sampler is the model itself
- for each batch, you run the current $\pi_\theta$, it generates outputs, and those generations are the training examples, which reflect what the model currently believes. Why this matters mechanically: the gradient is computed at the samples, meaning $\nabla_\theta \log \pi_\theta(y)$ needs a specific $y$ to be evaluated at, and it only gets evaluated at the $y$'s in your batch. So the sampler decides where in output space the model receives any training signal at all. If no sample lands in a region, there is no loss about that region this step.
an example to gain intuition: say the model weights include some weights that can do poetry and you are fine-tuning on math.
- SFT batch = math problems from a dataset. Every gradient is evaluated at a math solution and poetry never appears in a batch, so no gradient term ever mentions the poetry region. It is unmonitored (we will see what happens to this unmonitored region mechanistically, but spoiler: it gets forgotten).
- RL batch = whatever the model generates. If the model still writes poetry-adjacent outputs sometimes, those show up in batches, get scored, and get gradient terms attached. The poetry region stays under observation.
Read the technical explanation of sampling and KL in the appendix.
To reiterate:
SFT computes gradients at locations chosen by the data. RL computes gradients at locations chosen by the model.
The next section shows what each pattern of coverage does to the model's probability mass.
How the gradients move probability mass and governs learning/unlearning
this part is again inspired by chen’s paper on the forgetting difference in a simple setting: the target is a mixture of two univariate Gaussians, an "old" mode standing in for prior knowledge and a "new" mode for the target task, and the training policy is a Gaussian whose parameters are trained by sampled gradients on each objective. they found that the forward KL first stretches the policy toward the new mode and then transfers mass out of the old mode, while reverse KL slides its new component over and leaves the old component intact. these verbs, stretch and slide, fall directly out of the score functions, and walking through why is the most instructive calculation in this post.
The gradient analysis here is my own unpacking of their setup.
For a Gaussian, the score functions with respect to the mean and the width are
The $\sigma$ derivative says, any sample landing more than one standard deviation away makes the numerator positive, i.e. pushes the distribution to widen.
Why forward KL stretches.
SFT evaluates the model's score at the data's locations, and the data oftentimes sit far from the model's mass. the $\mu$-gradient scales like $d/\sigma^2$ where $d=y-\mu$ while the $\sigma$-gradient scales like $d^2/\sigma^3$; their ratio is $d/\sigma$ (this ratio argument depends on parameterizing in $(\mu, \sigma)$ specifically; other parameterizations change the ratio, but the qualitative conclusion is what Chen et al.'s simulations show). thus, when the data is many standard deviations out, the steepest descent points primarily along "get wider." and notice KL only ever perceives distance in units of the model's own $\sigma$. It has no ruler of its own. Inflating $\sigma$ rescales the model's unit of distance and attacks the $d^2/2\sigma^2$ penalty multiplicatively, which is why widening is so effective. The stretch is also fully legal, because laying density across the empty region between modes costs forward KL nothing - recall that adding mass to where the data has none is unpenalized.
Where the forgetting actually comes from: normalization
nothing in the SFT gradient ever says "reduce probability on old capabilities." The erosion happens through a side channel: normalization. Probabilities sum to one, so probability is a fixed budget, and every step that deposits density at the data's location implicitly withdraws it from somewhere. because the forward KL objective specifies where to deposit but is silent about where the mass is moved from, the model can just secretly forget.
now an intuitive math explanation with softmax:
Three tokens with logits $(2, 1, 0)$ have probabilities about $(0.665, 0.245, 0.090)$. Raise the third logit by one, touching nothing else: the probabilities become about $(0.576, 0.212, 0.212)$. The first two tokens each lost roughly thirteen percent of their mass with no gradient ever pointing at them. this is all because of normalization. In the Gaussian mixture the same mechanism runs through the mixture weight: every training sample is better explained by the new component, so weight drains from the old one. The old mode is just defunded. This also explains the learning-rate tradeoff Chen et al. observe for SFT (their Figure 3): a lower learning rate slows forgetting and learning together, because they are the same transaction.
Why reverse KL slides.
recall the setup: the training policy $\pi_\theta$ is a two-component Gaussian mixture, with an old component initialized on the region of prior knowledge and a new component that must move to the target task. the target $\pi^*$ is also a two-component mixture, and its old component is essentially identical to the model's, because zero-reward regions of $\pi_{\theta_0}$ pass into $\pi^*$ unchanged (reminder again that $\pi^*$ is the closed-form optimizer of the KL-regularized objective, $\pi^*(y) = \frac{1}{Z}\pi_{\theta_0}(y)e^{r(y)/\beta}$. not the base model! not the "real data"!). The on-policy gradient draws outputs $y \sim \pi_\theta$ and weights the gradient term at each draw by the log-ratio $\log\big(\pi_\theta(y)/\pi^*(y)\big)$. Every draw comes from one of the model's two components; there are 3 cases to check:
- Draws y from the old component, then the model's density comes from its old component and the target's density comes from its old component, and the two are nearly equal. The ratio $\pi_\theta(y)/\pi^*(y)$ is close to one, its log is close to zero, and the gradient contribution is negligible.
- Draws y from the new component, then the new component sits partway between its initialization and the target's new component, so draws on the side nearer the target land where $\pi^*(y)$ is comparatively large (log-ratio negative, probability pushed up), and draws on the far side land where $\pi^*(y)$ is small (log-ratio positive, probability pushed down). The net effect on the component's mean is a step toward the target's new component. The update repeats from the new position, so the component translates across output space step by step.
- the components don’t widen for because widening would place the model's own density in the low-density region between the target's two components, and that is exactly the configuration reverse KL penalizes: $\pi_\theta(y)$ large where $\pi^*(y)$ is near zero makes the log-ratio weight large and positive.
my takes
what i personally find magical is that all of this behavior of forgetting or lack of forgetting or stretching and sliding which is explained via KL divergence is inherent in the data, not in the presence of absence of the KL term in the objective function: Chen et al. (§4)'s ablation ran 1) GRPO with $\beta = 0$ 2) GRPO with the KL penalty 3) plain REINFORCE with no advantage estimator at all, and all three have similar levels of forgetting. the forgetting asymmetry is downstream of where the score function gets evaluated, which is downstream of who generates the samples, which in this case is on-policy data. some try to explain this with the parameter space: RL updates sparse subnetworks while SFT's updates are denser (Mukherjee et al. 2025), and concurrent work attributes RL's robustness to an implicit bias toward small KL from initialization (Shenfeld et al. 2025, "RL's Razor"). These are complementary - note that no claim above required examining which weights moved. but i would like to examine it in the next section.
KL vs geometry
KL is a vertical comparison: we stack the densities at each point, take the log-ratio, and average them.
there is a second tradition for measuring distance between distributions.
there's a measurement called Wasserstein distance used in optimal transport which is horizontal: the Wasserstein distance treats $p$ and $q$ as piles of sand and asks for the minimum effort (mass times distance carried) to reshape one pile into the other.
KL is a coding cost; Wasserstein is a moving cost. KL grades like a multiple-choice test, where a nearly-right answer in the wrong location earns zero; Wasserstein grades like a map-reading test, where one centimeter off is nearly perfect and one kilometer off is bad.
the difference is that KL is blind to the geometry of the underlying space. for example, 2 narrow Gaussians with negligible overlap have an astronomically large KL whether their means are five units apart or five million. this means that KL saturates and stops encoding distance at all once the two distributions don't touch. Their Wasserstein distance is simply $|\mu_1 - \mu_2|$, informative at every range. This blindness has a famous gradient-side symptom: KL-family objectives give essentially no useful gradient between disjoint distributions, the motivating pathology behind the Wasserstein GAN (Arjovsky et al. 2017). A geometry-blind judge cannot say "the target is over there," because "there" is not a concept it has. the forward KL's stretching makes sense here: since the Gaussian penalty $d^2/2\sigma^2$ measures distance only in the model's own units of width, inflating $\sigma$ is the model's only way to shrink a distance under a judge that owns no ruler.
i want to also briefly talk about parameter space. When gradient descent nudges a location parameter like a Gaussian mean, moving the parameter is transporting the mass: shifting $\mu$ by $\delta$ carries the whole mode by $\delta$. And the reverse KL behavior of sliding (instead of widening) looks very much like this sand-moving and shaping action in optimal transport. I find this little smuggling of Wasserstein transport extremely interesting.
parting note
hopefully by now you have some intuition behind KL and its magic. i find these 3 ideas to be the most important to chew on:
- The direction of the KL is not a design choice; it is a consequence of who generates the training samples.
- SFT forgetting is a normalization side effect where the loss deposits probability at the data's location and the conserved budget withdraws it from regions unchecked by the gradient; RL's old capabilities instead sit at a checked, zero-error fixed point of the model's own inspection.
- KL has no knowledge of the distance of drift but has interesting interactions with optimal transport which are up for more digging.
some open research questions is this interaction between KL and other means of measurement from information theory like the geometry of the subspace the agent learns from or updates into, and empirical ways to utilize such findings of forward and backwards KL.
References
[1] Chen, Razin, Narasimhan, Chen. "Retaining by Doing: The Role of On-Policy Data in Mitigating Forgetting." arXiv:2510.18874, 2025.
[2] Korbak, Elsahar, Kruszewski, Dymetman. "On Reinforcement Learning and Distribution Matching for Fine-Tuning Language Models with no Catastrophic Forgetting." NeurIPS 2022.
[3] Williams. "Simple Statistical Gradient-Following Algorithms for Connectionist Reinforcement Learning." Machine Learning, 8:229–256, 1992.
[4] Minka. "Divergence Measures and Message Passing." Microsoft Research Technical Report, 2005.
[5] Bishop. Pattern Recognition and Machine Learning. Springer, 2006.
[6] Tajwar, Singh, Sharma, Rafailov, Schneider, Xie, Ermon, Finn, Kumar. "Preference Fine-Tuning of LLMs Should Leverage Suboptimal, On-Policy Data." ICML 2024.
[7] Shao et al. "DeepSeekMath: Pushing the Limits of Mathematical Reasoning in Open Language Models." arXiv:2402.03300, 2024.
[8] Shenfeld, Pari, Agrawal. "RL's Razor: Why Online Reinforcement Learning Forgets Less." arXiv:2509.04259, 2025.
[9] Lai et al. "Reinforcement Fine-Tuning Naturally Mitigates Forgetting in Continual Post-Training." arXiv:2507.05386, 2025.
[10] Mukherjee, Yuan, Hakkani-Tur, Peng. "Reinforcement Learning Finetunes Small Subnetworks in Large Language Models." arXiv:2505.11711, 2025.
[11] Arjovsky, Chintala, Bottou. "Wasserstein GAN." ICML 2017.
[12] Jordan, Kinderlehrer, Otto. "The Variational Formulation of the Fokker–Planck Equation." SIAM Journal on Mathematical Analysis, 29(1):1–17, 1998.
[13] Talagrand. "Transportation Cost for Gaussian and Other Product Measures." Geometric and Functional Analysis, 6:587–600, 1996.
[14] Cuturi. "Sinkhorn Distances: Lightspeed Computation of Optimal Transport." NeurIPS 2013.
[15] Peyré, Cuturi. "Computational Optimal Transport." Foundations and Trends in Machine Learning, 2019.
appendix
background: entropy, cross-entropy, and KL: one picture
we have two distributions:
- reference distribution $P$ = produces the irl outcomes
- model distribution $Q$ = predicts the outcomes
with respect to these distributions,
- Entropy $H(P)$ is how unpredictable $P$ inherently is.
- Cross-entropy $H(P,Q)$ is how surprised you are when predicting $P$ using $Q$.
- KL divergence $D_{\mathrm{KL}}(P \Vert Q)$ is how much extra surprise comes specifically from $Q$ being wrong.
The relationship:
Concretely, take a real, fair coin, $P(H) = P(T) = 0.5$. your model is overconfident, $Q(H) = 0.9$, $Q(T) = 0.1$.
- A fair coin inherently costs $H(P) = -\log_2(1/2) = 1$ bit to predict.
- the cost, or surprisal, of predicting this reality with our overconfident model is the cross entropy:
the KL divergence is the difference, $0.737$ bits, which can be seen as the penalty for the wrong probabilities.
i like to think of it as: entropy = unavoidable surprise, KL = avoidable surprise, cross-entropy = all the surprise.
surprise is $-\log p$ which is the code length of encoding the event that has probability p of occurring, and it grows without bound as the assigned probability shrinks. to gain intuition: if you believe an event happens with probability $1/2$, the surprise or the cost is 1 bit when it happens; for $1/1024$ it costs 10. thus, if your model calls something rare but in reality it happens often, the cost is huge. this is what papers call "exploding tail" and what "causes asymmetry". hold on to this!
technical explanation of sampling and KL
SFT: Here $D = \pi^*$, which does not depend on $\theta$ and is fixed, so the expectation gradient
read the right-hand side. Each sample $y$ comes from the dataset, and the gradient term $\nabla_\theta \log \pi_\theta(y)$ is computed at that sample. So every gradient contribution is computed at a location chosen by the data. If the data lies in a region where the model currently assigns low probability, the gradient is still computed there. The model has no say in where it receives training signal.
RL Here $D = \pi_\theta$. Now the sampling distribution itself depends on $\theta$, so the gradient cannot simply move inside the expectation: changing $\theta$ changes both the integrand and which samples you see. The standard fix is the log-derivative trick. Start from the identity
which is just the chain rule applied to $\log$. Applying it to the expected reward:
This is the REINFORCE estimator (Williams 1992). The procedure it describes: generate outputs from the current model, score each one with the reward, and weight each output's gradient by its score. Again read where the gradient term is computed: at the model's own samples. The model only receives training signal in regions it actually visits.
another piece of intuition: the expected gradient of the log-probability, taken under the model's own distribution, is always zero:
Any parameter change that raises probability on some outputs must lower it on others, and averaged over the model's own distribution, these changes cancel. this has 2 consequences. First, adding a constant to the reward does not change the expected gradient, which is why baselines and group-relative advantages (as in GRPO) are valid. Second, constant terms like $\log Z$ drop out of any gradient expression of this form.
The two cases produce the two KL directions. Differentiate the reverse KL, $\mathrm{KL}[\pi_\theta \Vert \pi^*]$, using the log-derivative trick and the lemma. Substituting $\log \pi^* = \log \pi_{\theta_0} + r/\beta - \log Z$ and dropping the constant via the lemma:
Compare this to the policy gradient of the KL-regularized RL objective from the previous section: they are the same!!
what we are basically saying is that the update an RL trainer computes is a Monte Carlo estimate of the reverse-KL gradient. Similarly, the SFT update is a Monte Carlo estimate of the forward-KL gradient, since the SFT loss equals the forward KL up to a constant.
Could either method use the other direction? No, because each is missing the required sampler. Estimating the forward-KL gradient requires samples from $\pi^*$. An RL setup has a reward function, but no dataset of samples from the optimal policy, so it cannot form that estimate. Estimating the reverse-KL gradient requires scoring the model's own samples against $\pi^*$. An SFT setup has samples from $\pi^*$, but no reward function to score model generations, so it cannot form that estimate either. Each method uses the only KL direction its data source can support.