Research map / Reinforcement learning / RL from feedback
KL for a KL: On-Policy Distillation with Control Variate Baseline
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes vOPD, which treats on-policy distillation as policy-gradient RL and subtracts a closed-form value baseline (negative per-token reverse KL, optionally top-k approximated) to reduce gradient variance while keeping the single-sample estimator unbiased.
Problem
On-policy distillation (OPD) with a single-sample Monte Carlo estimator of reverse KL has high gradient variance and unstable training; existing stabilizers either use full-vocabulary KL (costly) or top-k KL (biased).
Approach
OPD's per-token reward is a policy-gradient reward whose value function equals the negative per-token reverse KL between student and teacher, which is computable from the existing forward pass. Subtracting it as a detached baseline keeps the gradient unbiased while reducing variance; a top-k approximation of this baseline lowers cost further without biasing the gradient.
Key results
- vOPD top-k reaches 33.0 average on Qwen3-1.7B-Base vs 29.8 for OPD, with +6.2 on MATH500 for vOPD full-V over OPD (per Table, MINERVA/AIME included).
- On Qwen3-4B-Base, vOPD full-V averages 45.4 vs 42.9 for OPD and 45.1 for OPD full-V.
- vOPD keeps gradient norms 1-2 orders of magnitude lower than base OPD; vOPD top-k reduces wall-clock time up to 57.7% relative to full-V OPD as claimed in the introduction.
Surprising findings
- Plain top-k OPD (biased KL objective) yields only marginal gains over base OPD, whereas the same top-k approximation used as a baseline gives large gains.
- Subtracting the value baseline from full-vocabulary OPD gives no change in gradient, since its gradient already has zero variance at each context.
Limitations
- Experiments limited to models up to 7B parameters.
- Wall-clock results reflect the authors' implementation and are not definitive.
- Requires a stronger teacher; self-distillation not explored.
Experimental setup
Datasets: DAPO-Math-17K SciKnowEval chemistry subset
Benchmarks: MATH500 Minerva Math AMC23 AIME24/25 SciKnowEval GPQA-Diamond
Metrics: avg@n pass@n wall-clock time per step gradient norm
Baselines: OPD OPD top-k OPD full-V
Builds on: on-policy distillation MiniLLM GKD GRPO RLOO
Closest papers by meaning
- Learning to Foresee: Unveiling the Unlocking Efficiency of On-Policy Distillation (NeurIPS 2026 · similarity 0.93)
- Entropy-Aware On-Policy Distillation of Language Models (ICML 2026 · similarity 0.92)
- Self-Distilled Reasoner: On-Policy Self-Distillation for Large Language Models (ICML 2026 · similarity 0.92)
- The Many Faces of On-Policy Distillation: Pitfalls, Mechanisms, and Fixes (NeurIPS 2026 · similarity 0.92)
- Hybrid Policy Distillation for LLMs (ICML 2026 · similarity 0.92)
- AsyncOPD: How Stale Can On-Policy Distillation Be? (NeurIPS 2026 · similarity 0.90)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.
