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

Hugging Face Daily Papers Papers
on-policy-distillation control-variate reinforcement-learning large-language-models variance-reduction knowledge-distillation reasoning

Summary

Proposes vOPD, which stabilizes on-policy distillation for LLMs by introducing a control variate baseline from reinforcement learning, achieving performance comparable to expensive full-vocabulary methods at lower computational cost.

On-Policy Distillation (OPD) has emerged as a dominant post-training paradigm for large language models, especially for reasoning domains. However, OPD remains unstable in practice due to the high gradient variance of its single-sample Monte Carlo estimator, and recipes for stable training are still immature. We propose vOPD (On-Policy Distillation with a control variate baseline), which casts OPD as policy-gradient RL and stabilizes it by introducing a control variate baseline-canonically a value function -- from the RL literature. We show that the OPD value function admits a closed form as the per-token negative reverse KL divergence between the student and the teacher, available directly from the already-computed forward pass with no additional critic or inference. Existing stabilization methods either compute the full token-level reverse KL over the entire vocabulary, adding significant overhead, or restrict it to a top-k support, biasing the objective. vOPD instead preserves the lightweight single-sample estimator, subtracting the value function as a detached baseline to keep the gradient unbiased while reducing variance. Furthermore, we show that a top-k approximation of the baseline further lowers cost without compromising performance. Across mathematical and scientific reasoning benchmarks, vOPD consistently outperforms vanilla OPD and matches the most expensive full-vocabulary baseline, offering an efficient stabilization of On-Policy Distillation through principled RL variance reduction.
Original Article
View Cached Full Text

Cached at: 05/15/26, 04:24 AM

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

Source: https://huggingface.co/papers/2605.07865

Abstract

On-Policy Distillation with control variate baseline stabilizes training through policy-gradient reinforcement learning techniques while maintaining efficiency and performance.

On-Policy Distillation(OPD) has emerged as a dominant post-training paradigm for large language models, especially for reasoning domains. However, OPD remains unstable in practice due to the highgradient varianceof its single-sampleMonte Carlo estimator, and recipes for stable training are still immature. We propose vOPD (On-Policy Distillationwith acontrol variate baseline), which casts OPD aspolicy-gradient RLand stabilizes it by introducing acontrol variate baseline-canonically avalue function-- from the RL literature. We show that the OPDvalue functionadmits a closed form as the per-token negativereverse KL divergencebetween the student and the teacher, available directly from the already-computed forward pass with no additional critic or inference. Existing stabilization methods either compute the fulltoken-level reverse KLover the entire vocabulary, adding significant overhead, or restrict it to a top-k support, biasing the objective. vOPD instead preserves the lightweight single-sample estimator, subtracting thevalue functionas a detached baseline to keep the gradient unbiased while reducing variance. Furthermore, we show that atop-k approximationof the baseline further lowers cost without compromising performance. Across mathematical and scientific reasoning benchmarks, vOPD consistently outperforms vanilla OPD and matches the most expensive full-vocabulary baseline, offering an efficient stabilization ofOn-Policy Distillationthrough principled RLvariance reduction.

View arXiv pageView PDFAdd to collection

Get this paper in your agent:

hf papers read 2605\.07865

Don’t have the latest CLI?curl \-LsSf https://hf\.co/cli/install\.sh \| bash

Models citing this paper0

No model linking this paper

Cite arxiv.org/abs/2605.07865 in a model README.md to link it from this page.

Datasets citing this paper0

No dataset linking this paper

Cite arxiv.org/abs/2605.07865 in a dataset README.md to link it from this page.

Spaces citing this paper0

No Space linking this paper

Cite arxiv.org/abs/2605.07865 in a Space README.md to link it from this page.

Collections including this paper0

No Collection including this paper

Add this paper to acollectionto link it from this page.

Similar Articles

On the Geometry of On-Policy Distillation

Hugging Face Daily Papers

This paper characterizes the unique parameter space dynamics of on-policy distillation (OPD) for large language models, showing that it exhibits relaxed off-principal updates and subspace locking, distinguishing it from supervised fine-tuning and reinforcement learning with verifiable rewards.

DOPD: Dual On-policy Distillation

Hugging Face Daily Papers

DOPD proposes a dual on-policy distillation paradigm that dynamically routes token-level supervision between privileged teacher and student policies based on advantage gaps and probabilities, addressing privilege illusion and improving capability transfer in LLMs and VLMs.

OPRD: On-Policy Representation Distillation

Hugging Face Daily Papers

OPRD proposes a new knowledge distillation method that aligns student and teacher hidden states across layers during on-policy rollouts, eliminating sampling variance from token-space KL estimation. Empirically, OPRD outperforms output-space baselines on math reasoning benchmarks (AIME 2024/2025, AIMO) while being 1.44x faster and using 54% less memory.

ShortOPD: Recovering Pruned LLMs with Short-to-Long On-Policy Distillation

arXiv cs.LG

ShortOPD proposes a short-to-long on-policy distillation schedule that recovers pruned LLMs for free-form generation by focusing training on effective prefixes, achieving up to 9x improvement over unrecovered models and matching long-horizon distillation with a quarter of the training time.