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

FastOPD: On-Policy Distillation for Lightweight VLA Deployment

Hugging Face Daily Papers

FastOPD proposes an on-policy distillation framework that compresses large Vision-Language-Action (VLA) foundation models into lightweight students with few inference steps, retaining 84% of π₀.₅'s LIBERO performance at a 78.1% latency reduction and deploying a compact student distilled from MolmoAct2 on a real robot.

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.