KL for a KL: On-Policy Distillation with a Control Variate Baseline

Abstract

On-Policy Distillation (OPD) has become a dominant post-training paradigm for reasoning with large language models, training a student on its own rollouts against a teacher's dense token-level signal. Yet OPD remains unstable: its single-sample Monte-Carlo estimator has high gradient variance, especially at high-mismatch tokens where student and teacher disagree sharply. We recast OPD as policy-gradient reinforcement learning and stabilize it with a classical tool — a control-variate baseline. We show the OPD value function admits a closed form as the per-token negative reverse KL divergence between student and teacher, computable directly from the already-available forward pass with no extra critic or inference. The resulting method, vOPD, subtracts this value as a detached baseline: the gradient stays unbiased while its variance drops sharply. A top-\(k\) approximation of the baseline lowers cost further without hurting quality. Across math and science reasoning benchmarks, vOPD consistently outperforms vanilla OPD and matches the most expensive full-vocabulary baseline at substantially lower wall-clock cost.

Method

Vanilla OPD minimizes the token-level reverse KL to the teacher with a single sampled continuation, giving the policy-gradient estimator

\[ \nabla_\theta J_{\text{OPD}} = \mathbb{E}\!\left[\sum_t \big(\log \pi_T(y_t\mid c_t) - \log \pi_\theta(y_t\mid c_t)\big)\, \nabla_\theta \log \pi_\theta(y_t\mid c_t)\right]. \]

The bracketed per-token reward is heavy-tailed at high-mismatch tokens, which is the source of the instability. Reading OPD as RL, the value of a context is the expected reward under the student — and it collapses to a closed form:

\[ V^{\pi_\theta}(c_t) = -\,D_{\mathrm{KL}}\!\big(\pi_\theta(\cdot\mid c_t)\,\|\,\pi_T(\cdot\mid c_t)\big). \]

Subtracting it as a detached control-variate baseline turns the reward into an advantage and yields vOPD:

\[ \nabla_\theta J_{\text{vOPD}} = \mathbb{E}\!\left[\sum_t \Big(r_t(c_t,y_t) + D_{\mathrm{KL}}\!\big(\pi_\theta(\cdot\mid c_t)\,\|\,\pi_T(\cdot\mid c_t)\big)\Big)\, \nabla_\theta \log \pi_\theta(y_t\mid c_t)\right]. \]

• Unbiased & free. The baseline is action-independent, so the gradient is unchanged in expectation; it is read off the same forward pass — no critic, no extra inference.
• Single-sample & lightweight. vOPD keeps OPD's one-sample backward pass; the baseline adds only a forward-pass KL evaluation.
• Top-\(k\) baseline. Approximating the baseline KL over the student's top-\(k\) tokens (even \(k\!=\!5\)) preserves unbiasedness and nearly all of the gain, cutting cost further.

Results

Across student/teacher pairs from the Qwen3 (1.7B, 4B) and Olmo-3 (7B) families, distilled on DAPO-Math-17K and evaluated on math (MATH500, Minerva, AMC23, AIME24/25) and science (SciKnowEval, GPQA-Diamond) reasoning:

• Consistent gains over vanilla OPD on both math and science reasoning, with the largest improvements on MATH500.
• Matches full-vocabulary OPD — the most expensive baseline — while reducing wall-clock cost substantially (up to ~58%).
• Order-of-magnitude lower gradient variance: vOPD's gradient norms stay 1–2 orders of magnitude below vanilla OPD throughout training, despite higher final accuracy.
• Unbiasedness matters: the biased OPDtop-k barely helps, while vOPDtop-k retains the full gains.

Full quantitative tables and ablations are in the paper.

BibTeX

@article{oh2026klforakl,
  title   = {KL for a KL: On-Policy Distillation with Control Variate Baseline},
  author  = {Oh, Minjae and Song, Sangjun and Choi, Gyubin and Choi, Yunho and Jo, Yohan},
  journal = {arXiv preprint arXiv:2605.07865},
  year    = {2026}
}