PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Minjae Oh, Sangjun Song, Gyubin Choi, Yunho Choi, Yohan Jo

arXiv:2605.07865v1cs.LGcs.AIcs.CL

TL;DR

On-Policy Distillation is unstable because its lightweight single-sample estimator has high gradient variance. vOPD uses a closed-form reverse-KL control-variate baseline, consistently improving reasoning performance over base OPD while matching full-vocabulary stabilization at lower cost.

  • Problem

    OPD offers dense token-level supervision but remains unstable because its single-sample Monte Carlo estimator has high variance and its stabilization recipe is underdeveloped.

  • Method

    vOPD treats OPD as policy-gradient reinforcement learning and subtracts the per-token negative reverse KL as a detached, closed-form value baseline.

  • Results

    Across mathematics and science reasoning benchmarks, vOPD consistently improves over base OPD, with average accuracy gains up to +3% and matching full-vocabulary OPD while reducing wall-clock time up to 57.7%.

  • Takeaways & Limitations

    vOPD provides a principled, efficient stabilization of On-Policy Distillation while retaining the single-token backward pass and reducing estimator variance.

  • Takeaways & Limitations

    Experiments are limited to models up to 7B parameters, and reported wall-clock results depend on the implementation.

Abstract

from arXiv · show

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.

1 Introduction

On-Policy Distillation offers dense token-level supervision and can match RLVR accuracy, but its single-sample estimator remains vulnerable to gradient variance. vOPD addresses this instability with a control-variate baseline while preserving unbiased gradients and improves accuracy across reasoning benchmarks.

  • Motivation: RLVR assigns credit across thousands of intermediate reasoning tokens using one sparse scalar reward, requiring large rollouts and prolonged training.This sparse supervision arises because the final answer receives the only reward signal.
  • Motivation: OPD instead minimizes reverse KL with dense token-level signals, enabling faster training while empirically matching RLVR accuracy.Because OPD is on-policy and reward-driven, it can use a single-sample Monte Carlo estimator within standard RL pipelines.
  • Method: vOPD subtracts a control-variate baseline from the policy-gradient estimator to reduce variance without biasing the gradient in expectation.The approach builds on baseline subtraction used in actor-critic methods, including PPO, GRPO, and RLOO.
  • Results: vOPD delivers up to +3% absolute average accuracy over base OPD across six mathematics and science reasoning benchmarks.The evaluation covers four Qwen3 and Olmo-3 models and benchmarks including MATH500, Minerva Math, AMC23, AIME24/25, SciKnowEval, and GPQA-Diamond.
  • Results: vOPD improves accuracy by up to +6.2% on MATH500 compared with base OPD.The paper reports these results alongside consistent improvements over baseline methods.

2 Preliminaries

OPD minimizes the student-to-teacher reverse KL using student-generated samples, but its efficient single-sample estimator has high variance. Full-vocabulary and top-k variants reduce variance at the cost of computation or bias, motivating baseline subtraction as an unbiased variance-reduction tool.

  • OPD: OPD minimizes the reverse KL divergence between the student πθ and teacher πT using responses sampled from the student during generation.On-policy sampling provides an unbiased KL estimator and mitigates exposure bias between training and test-time generation.
  • OPD: The base OPD objective is commonly optimized with a single-sample Monte Carlo estimate and policy-gradient RL using the detached reward rt(ct, yt).The reward is log πT(yt | ct) − log πθ(yt | ct), treated as a fixed scalar without gradient flow through it.
  • OPD: Base OPD is computationally efficient because its backward pass touches only the sampled token, but its estimator has high variance and causes training instability.This limitation motivates variance-stabilizing alternatives discussed in the preliminaries.
  • Variance-reduction variants: OPDfull-V computes the exact per-token expectation over the vocabulary, achieving zero variance for a context but requiring a full-vocabulary backward pass at every token.For Qwen3, the vocabulary size is approximately |V| ≈150k.
  • Variance-reduction variants: OPDtop-k restricts KL computation to the student’s k most likely tokens, reducing computation relative to full-vocabulary OPD but introducing gradient bias and only marginal reported gains.Its backward pass flows through k tokens per position rather than one sampled token, so it is heavier than base OPD.
  • Policy-gradient baseline: Policy-gradient baselines subtract bt(ct) from the reward to form an advantage, preserving unbiasedness when independent of yt while reducing variance through the value function.The canonical value function is Vπθ(ct) = Eyt∼πθ(·|ct)[rt(ct, yt)].

3 Control Variate Baseline for OPD

vOPD stabilizes on-policy distillation by subtracting a closed-form reverse-KL value baseline from OPD’s single-sample policy-gradient estimator. The baseline preserves unbiasedness, reduces variance without extra backward passes, and can be approximated on top-k support to lower compute.

  • vOPD baseline: vOPD uses OPD’s RL interpretation to subtract a control variate baseline, producing an unbiased, lower-variance estimator without additional backward passes.The method addresses OPD’s high variance while retaining computational efficiency.
  • Closed-form value function: The OPD value function is exactly the negative per-step reverse KL, computed from already-available student and teacher distributions without a learned value network or extra forward pass.The full-vocabulary baseline is evaluated during the forward pass and detached from gradients.
  • Unbiased gradient: vOPDfull-V has the same expected gradient as OPD while backpropagating only through the single sampled token, because the baseline KL does not propagate gradients through the vocabulary.This preserves the lightweight backward computation of base OPD.
  • Variance reduction: Variance reduction is largest when DKL(πθ∥πT)^2 is large at high-mismatch tokens, damping the heavy-tailed negative gradients that destabilize OPD.The token-level reward damping directly reduces gradient variance.
  • Top-k approximation: vOPDtop-k approximates the baseline KL on the student’s top-k support, reducing the O(|V|) exact-KL cost while keeping the baseline detached and the gradient unbiased.Unlike OPDtop-k, which changes the optimization target and biases the gradient, vOPDtop-k uses the approximation only as a baseline subtracted from the reward.

4 Experiments

Experiments evaluate vOPD across mathematical and scientific reasoning, model scales, size-mismatched pairs, and model families. vOPD consistently improves over base OPD while retaining stable, efficient training through its control-variate baseline.

  • Experimental setup: Experiments compare vOPDtop-k and vOPDfull-V with base OPD, OPDfull-V, and OPDtop-k across model scale, size mismatch, and model family.The primary setting distills Qwen3-1.7B into Qwen3-1.7B-Base, with additional Qwen3-4B, Qwen3-0.6B-Base, and Olmo-3-7B configurations.
  • Mathematical reasoning: +6.2% absolute gains on MATH500 are achieved by vOPDtop-k and vOPDfull-V in the Qwen3-1.7B-Base setting, with average gains above +3% over base OPD.At 4B, both vOPD variants gain around +4% on MATH500 and around +2.5% on average; in the Olmo-3-7B family, vOPDtop-k reaches an average of 33.1%.
  • Mathematical reasoning: 21.1% average performance is achieved by vOPDfull-V in the Qwen3-0.6B-Base experiment, while vOPDtop-k reaches 20.0%, both on par with OPDfull-V.These gains remain consistent across model scales, size-mismatched teacher-student pairs, and model families.
  • Stability and efficiency: 1–2 orders of magnitude lower gradient norms are maintained by vOPD than by base OPD, while training remains stable and reaches higher accuracy.The wall-clock comparison reports vOPDtop-k and vOPDfull-V at modest overhead, whereas OPDfull-V is most expensive because of its full-vocabulary backward pass; accuracy is stable for k ∈{5, 20, 50, 100} and the full-vocabulary baseline.
  • Scientific reasoning: Around +4% gains over base OPD are obtained by OPDfull-V and both vOPD variants on SciKnowEval chemistry and GPQA-Diamond, while OPDtop-k shows little gain.The scientific-reasoning experiment trains Qwen3-1.7B into Qwen3-1.7B-Base using SciKnowEval chemistry and evaluates on its test set and GPQA-Diamond.

5 Related Work

On-policy distillation has become an important post-training approach for long-chain-of-thought reasoning, while control-variate baselines connect vOPD to established policy-gradient reinforcement-learning methods. In LLM reinforcement learning, learned critics and relative baselines are common alternatives for advantage estimation.

  • On-Policy Distillation: OPD became an important LLM post-training component for long Chain-of-Thought reasoning, with GKD and MiniLLM establishing it as an alternative to standard distillation.Later work popularized token-level Monte Carlo OPD, studied practical recipes, and integrated OPD into large-scale post-training systems.
  • Policy-Gradient Baselines: Control-variate baselines are central to on-policy policy-gradient reinforcement learning, underpinning actor-critic methods, advantage estimation, A3C, and PPO with GAE.This places vOPD’s baseline within a well-established reinforcement-learning variance-reduction framework.
  • LLM Reinforcement Learning: In LLM reinforcement learning, early RLHF systems used PPO with learned value models, whereas reasoning-oriented RLVR methods often use simpler relative baselines.GRPO and RLOO construct relative baselines from rewards across multiple sampled responses, while SPO studies alternative single-stream estimators.

6 Conclusion, Limitations, and Future Work

The paper introduces vOPD as an unbiased control-variate formulation that reduces single-sample OPD variance while retaining a single-token backward pass. Its evaluation is limited in scale and implementation-specific wall-clock results, motivating larger-scale validation and efficiency, teacher-access, and self-distillation extensions.

  • Conclusion: vOPD reduces the variance of OPD’s single-sample Monte Carlo estimator without changing the original objective.It uses the student–teacher negative reverse KL as a detached value baseline.
  • Conclusion: vOPD preserves an unbiased policy-gradient estimator while retaining base OPD’s single-token backward pass.The formulation uses a detached value baseline.
  • Limitations and Future Work: Experiments are limited to models up to 7B parameters, so validating vOPD at larger scale remains future work.The paper identifies larger-scale validation as a next step.
  • Limitations and Future Work: Wall-clock results depend on the implementation and are not definitive, while optimizing vOPDtop-k for greater speed than vOPDfull-V remains open.The proposed efficiency direction specifically targets vOPDtop-k relative to vOPDfull-V.
  • Limitations and Future Work: Future work could extend vOPD to self-distillation settings because the method requires access to a stronger teacher.The paper describes self-distillation as an interesting direction.

A Theoretical Derivations · A.1 Unbiasedness of Baseline Subtraction · A.2 Variance Reduction and the Optimal Baseline

The appendix establishes that subtracting an action-independent baseline preserves the policy gradient’s unbiasedness. It then derives baseline subtraction as a variance-reduction method and identifies the optimal scalar baseline, motivating value-function baselines in practice.

  • A Theoretical Derivations: The derivations define context c_t = (x, y_<t) and per-token reward r_t(c_t, y_t) = log π_T(y_t | c_t) − log π_θ(y_t | c_t).
  • A Theoretical Derivations: The score-function identity follows from the log trick, ∇_θ log π_θ(v | c_t) = ∇_θπ_θ(v | c_t)/π_θ(v | c_t).
  • A.1 Unbiasedness of Baseline Subtraction: Subtracting an action-independent baseline b_t(c_t) from the per-token reward leaves the policy gradient unbiased.
  • A.1 Unbiasedness of Baseline Subtraction: Linearity of expectation reduces the difference between baseline-subtracted and original gradients to a baseline-weighted expectation.
  • A.1 Unbiasedness of Baseline Subtraction: Because b_t(c_t) is independent of y_t, the score-function identity makes the remaining sum zero, proving the unbiasedness result.
  • A.2 Variance Reduction and the Optimal Baseline: For the per-step gradient estimator, the baseline is chosen to minimize tr(Var[g(b)]).
  • A.2 Variance Reduction and the Optimal Baseline: Since E[g(b)] is baseline-independent, minimizing variance is equivalent to minimizing the second moment E[∥g(b)∥2], yielding a convex quadratic in b.
  • A.2 Variance Reduction and the Optimal Baseline: The resulting optimal scalar baseline is a weighted value function; using the value function itself as the baseline motivates the practical choice in § 3.1.

A.3 Variance Reduction of vOPD

vOPDfull-V reduces per-step estimator variance by subtracting the value baseline b_t = −D_KL(π_θ(· | c_t) ∥ π_T(· | c_t)) while preserving the estimator’s expectation. The analysis motivates this reduction through weak correlation between sampled student log-probability gradients and the teacher-driven reward.

  • Variance reduction: vOPDfull-V’s per-step variance reduction follows by comparing second moments because both estimators have the same expectation.The derivation states that the variance difference equals the difference of second moments.
  • Control variate baseline: The vOPD value baseline is b_t = −D_KL(π_θ(· | c_t) ∥ π_T(· | c_t)).This baseline is the student-to-teacher per-token negative reverse KL.
  • Weak-correlation approximation: A weak-correlation approximation is plausible because sampled student tokens lie in high-probability regions, concentrating log π_θ(y_t | c_t) and ∥∇log π∥2.The reward r_t varies mainly with the teacher log-probability term, which lacks structural dependence on the student probability at y_t.

B Experiment Settings

The experiments used single-GPU systems pairing NVIDIA H200 NVL or A100 NVL GPUs with specified Intel Xeon CPUs. Training runs lasted 2–6 hours, depending on model and hardware.

  • Hardware: Experiments ran on a single NVIDIA H200 NVL or A100 NVL GPU.The H200 NVL or A100 NVL was used as the sole GPU for each experiment.
  • Hardware: The GPU systems used Intel Xeon Gold 6530 or Gold 6230 CPUs at 2.10GHz, respectively.The CPUs were paired with the H200 NVL and A100 NVL systems, respectively.
  • Training runtime: Training runs took between 2h for the 1.7B model on A100 and 6h for the 7B model on H200.Runtime varied with model size and GPU type.

B.1 Training Settings · B.2 Prompts

The paper uses a lightweight OPD training setup with 2048-token responses, temperature-1.0 student sampling, LoRA, and TRL’s MiniLLM2 implementation, while evaluation uses lower-temperature top-p sampling. Mathematical and scientific reasoning follow official prompt formats, requiring boxed answers and constrained JSON outputs, respectively.

  • B.1 Training Settings: Training uses a maximum response length of 2048 tokens, following evidence that 2–3K tokens suffice for OPD.The setting is motivated by recent OPD work on response-length requirements.
  • B.1 Training Settings: Student sampling during training uses temperature 1.0, following recent OPD practice.This temperature is specified for student sampling rather than evaluation.
  • B.1 Training Settings: The implementation adopts parameter-efficient LoRA training and builds on TRL’s MiniLLM2 implementation.Both choices are reported as implementation settings for training.
  • B.1 Training Settings: Evaluation uses sampling temperature 0.6, top-p 0.9, and vLLM for accelerated inference.These settings follow official Qwen3 and Olmo-3 guidelines.
  • B.2 Prompts: Mathematical and scientific reasoning use official prompts throughout training and evaluation.The paper states that both prompt formats follow official prompts.
  • B.2 Prompts: The math prompt asks the model to reason step by step and place its final answer within \boxed{}.The template begins with the problem, followed by the reasoning and answer-format instruction.
  • B.2 Prompts: The scientific reasoning prompt requires exactly one JSON object with the key "answer" and a single choice letter as its value.The template gives {"answer": "C"} as an example.

C Extended Experiment Results

The extended results report teacher-model accuracy, Qwen3-1.7B-to-Qwen3-0.6B-Base distillation results, and comprehensive hyperparameter-ablation benchmark results. Table 5 additionally presents mathematical reasoning benchmark comparisons with best and second-best performances distinguished.

  • The extended experiments include teacher-model accuracy results.
  • They report distillation results for Qwen3-1.7B into Qwen3-0.6B-Base from § 4.2.
  • The section provides comprehensive benchmark results for the hyperparameter ablations in Figure 2.
  • Table 5 presents additional mathematical reasoning benchmark results, marking the best performance in bold and the second-best performance with underlining.
Loading 2605.07865v1…