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

• URL: http://arxiv.org/abs/2605.07865v1
• Date: Fri, 08 May 2026 15:24:51 GMT
• Title: KL for a KL: On-Policy Distillation with Control Variate Baseline
• Abstract summary: On-Policy Distillation (OPD) has emerged as a dominant post-training paradigm for large language models.<n>We propose vOPD, which casts OPD as policy-gradient RL.<n>We show that vOPD consistently outperforms vanilla OPD and matches the most expensive full-vocabulary baseline.
• Score: 13.281263788199219
• License: http://creativecommons.org/licenses/by/4.0/
• Abstract: 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.

Related papers

• Stabilizing On-Policy Distillation for MLLM Reasoning with Global Normalization [9.860371691501076]
  On-policy distillation (OPD) has recently emerged as an important post-training paradigm.<n>We propose Globally Normalized Distillation Policy Optimization (GNDPO), a practical method that stabilizes optimization by transforming raw KL scores into batch-level relative advantages.
  arXiv  Detail & Related papers  (2026-06-08T06:41:31Z)
• Self-Distilled Policy Gradient [51.95967256628261]
  On-policy self-distillation is a promising source of dense supervision for sparse-reward reinforcement learning.<n>We propose SDPG, a self-distilled policy-gradient framework that combines group-relative verifier advantages with normalized standard deviation.
  arXiv  Detail & Related papers  (2026-06-02T02:31:13Z)
• One-Way Policy Optimization for Self-Evolving LLMs [63.8638342097375]
  Reinforcement Learning with Verifiable Rewards (RLVR) has become a promising paradigm for scaling reasoning capabilities of Large Language Models (LLMs)<n>We propose One-Way Policy Optimization (OWPO), a method based on the principle of decoupling optimization direction from update magnitude.<n> Experimental results demonstrate that OWPO outperforms strong baselines, including DAPO, OPD, and MOPD.
  arXiv  Detail & Related papers  (2026-05-21T08:25:27Z)
• Rethinking Importance Sampling in LLM Policy Optimization: A Cumulative Token Perspective [22.848847562976633]
  Existing methods face a fundamental bias-variance dilemma: token-level IS ratios, as adopted by PPO and GRPO, introduce bias by ignoring prefix state distribution mismatch.<n>We propose CTPO, which combines the cumulative token IS ratio with position-adaptive clipping that scales log-space clip bounds according to the natural $sqrtt$ growth of the cumulative log-ratio.
  arXiv  Detail & Related papers  (2026-05-08T06:35:02Z)
• EVPO: Explained Variance Policy Optimization for Adaptive Critic Utilization in LLM Post-Training [69.32453275232662]
  We show that a learned critic can inject estimation noise that exceeds the state signal it captures, increasing rather than reducing advantage variance.<n>We propose Explained Variance Policy Optimization (EVPO), which monitors batch-level EV at each training step and adaptively switches between critic-based and batch-mean advantage estimation.
  arXiv  Detail & Related papers  (2026-04-21T14:07:39Z)
• Improving Search Agent with One Line of Code [68.58667107354253]
  Tool-based Agentic Reinforcement Learning (TARL) has emerged as a promising paradigm for training search agents to interact with external tools.<n>We propose textbfSearch textbfAgent textbfPolicy textbfOptimization (textbfSAPO), which stabilizes training via a conditional token-level KL constraint.
  arXiv  Detail & Related papers  (2026-03-10T04:07:39Z)
• Diffusion Controller: Framework, Algorithms and Parameterization [54.82539154511621]
  We present a unified control-theoretic view that casts reverse diffusion sampling as state-only control within (generalized) linearly-solvable Markov Decision Processes.<n>Under this framework, control acts by reweighting the pretrained reverse-time transition kernels, balancing terminal objectives against an $f$-divergence cost.<n>Experiments on Stable Diffusion v1.4 show consistent gains in preference-alignment win rates and improved quality-efficiency trade-offs.
  arXiv  Detail & Related papers  (2026-03-07T01:49:59Z)
• Stabilizing Reinforcement Learning for Diffusion Language Models [27.22260417844743]
  Group Relative Policy Optimization (GRPO) is highly effective for post-training autoregressive (AR) language models.<n>But its direct application to diffusion large language models (dLLMs) often triggers reward collapse.<n>We show these effects form a self-reinforcing instability loop that drives policy drift and further increases ratio variance.<n>We propose StableDRL, a reformulation of GRPO tailored for dLLMs.
  arXiv  Detail & Related papers  (2026-03-06T08:55:37Z)
• Rethinking the Trust Region in LLM Reinforcement Learning [72.25890308541334]
  Proximal Policy Optimization (PPO) serves as the de facto standard algorithm for Large Language Models (LLMs)<n>We propose Divergence Proximal Policy Optimization (DPPO), which substitutes clipping with a more principled constraint.<n>DPPO achieves superior training and efficiency compared to existing methods, offering a more robust foundation for RL-based fine-tuning.
  arXiv  Detail & Related papers  (2026-02-04T18:59:04Z)
• A Step Back: Prefix Importance Ratio Stabilizes Policy Optimization [58.116300485427764]
  Reinforcement learning post-training can elicit reasoning behaviors in large language models.<n> token-level correction often leads to unstable training dynamics when the degree of off-policyness is large.<n>We propose a simple yet effective objective, Minimum Prefix Ratio (MinPRO)
  arXiv  Detail & Related papers  (2026-01-30T08:47:19Z)
• Shrinking the Variance: Shrinkage Baselines for Reinforcement Learning with Verifiable Rewards [12.074691741125044]
  Reinforcement Learning with Verifiable Rewards (RLVR) has emerged as a powerful paradigm for post-training large reasoning models.<n>We propose using shrinkage estimators that combine per-prompt and across-prompt means to improve the overall per-prompt mean estimation accuracy.
  arXiv  Detail & Related papers  (2025-11-05T18:43:15Z)
• The Choice of Divergence: A Neglected Key to Mitigating Diversity Collapse in Reinforcement Learning with Verifiable Reward [57.56453588632619]
  A central paradox in fine-tuning Large Language Models (LLMs) with Reinforcement Learning with Verifiable Reward (RLVR) is the frequent degradation of multi-attempt performance.<n>This is often accompanied by catastrophic forgetting, where models lose previously acquired skills.<n>We argue that standard RLVR objectives lack a crucial mechanism for knowledge retention.
  arXiv  Detail & Related papers  (2025-09-09T06:34:32Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.