Technical Summary: vOPD (On-Policy Distillation with Control Variate Baseline)

1. Problem Statement

On-Policy Distillation (OPD) has emerged as a dominant post-training paradigm for Large Language Models (LLMs), particularly in reasoning domains. Unlike Reinforcement Learning with Verifiable Rewards (RLVR), which relies on sparse terminal rewards, OPD minimizes the reverse Kullback-Leibler (KL) divergence between a student and a teacher using dense, token-level signals. This allows for faster training and effective handling of long Chain-of-Thought (CoT) tasks.

However, OPD faces a critical stability issue in practice. The standard implementation uses a single-sample Monte Carlo estimator to approximate the gradient. This approach suffers from high gradient variance, leading to unstable training dynamics. Existing stabilization methods attempt to mitigate this but introduce significant trade-offs:

• Full-vocabulary OPD ($OPD_{full-V}$): Computes the KL divergence over the entire vocabulary. While this eliminates variance, it incurs prohibitive computational overhead (requiring a backward pass over the full vocabulary, e.g., ~150k tokens) at every step.
• Top-k OPD ($OPD_{top-k}$): Restricts the KL calculation to the top-$k$ most likely tokens. While lighter, this introduces bias into the gradient by truncating the support, and empirical results show it yields only marginal performance gains over the unstable base OPD.

The core challenge is to stabilize OPD training without sacrificing the computational efficiency of the single-sample estimator or introducing bias into the objective.

2. Methodology: vOPD

The authors propose vOPD (On-Policy Distillation with a control variate baseline), which reframes OPD as a policy-gradient Reinforcement Learning (RL) problem and applies a standard variance-reduction technique: subtracting a control variate baseline.

2.1 Theoretical Foundation

In policy-gradient RL, subtracting a baseline $b_t(c_t)$ from the reward $r_t$ yields an advantage function $a_t = r_t - b_t$. Crucially, if the baseline is independent of the sampled action, the gradient remains unbiased, but its variance is reduced. The canonical choice for a baseline is the value function $V^{\pi}(c_t)$.

The key insight of vOPD is that for OPD, the value function admits a closed-form solution that requires no additional critic model or inference:
$$V^{\pi_\theta}(c_t) = \mathbb{E}_{y_t \sim \pi_\theta} [r_t(c_t, y_t)] = -D_{KL}(\pi_\theta(\cdot|c_t) \parallel \pi_T(\cdot|c_t))$$
Here, the reward is defined as $r_t = \log \pi_T - \log \pi_\theta$. The expected reward under the student distribution is simply the negative reverse KL divergence between the student and teacher at that context.

2.2 Algorithmic Implementation

The vOPD gradient estimator is formulated as:
$$\nabla_\theta J_{vOPD} = \mathbb{E} \left[ \sum_{t} \left( r_t(c_t, y_t) + D_{KL}(\pi_\theta \parallel \pi_T) \right) \nabla_\theta \log \pi_\theta(y_t|c_t) \right]$$

• Unbiasedness: The baseline (the KL term) is computed based on the distributions $\pi_\theta$ and $\pi_T$ but is detached from the gradient flow regarding the sampled token $y_t$. Thus, the gradient remains unbiased.
• Efficiency: The KL baseline is computed during the forward pass (which is already required for the reward calculation). The backward pass still only touches the single sampled token, preserving the computational efficiency of base OPD.
• Variance Reduction Mechanism: The baseline acts as a regularizer. When the student and teacher distributions strongly disagree (high KL), the reward $r_t$ is large and negative. The baseline $D_{KL}$ is large and positive, effectively canceling out the negative reward and dampening the resulting heavy-tailed gradients that cause instability.

2.3 Top-k Approximation ($vOPD_{top-k}$)

To further reduce the forward-pass cost of computing the full-vocabulary KL, the authors propose approximating the baseline using only the student's top-$k$ tokens.

• Unlike $OPD_{top-k}$, which uses the top-$k$ support to replace the loss objective (introducing bias), $vOPD_{top-k}$ uses the top-$k$ approximation only for the baseline.
• Since the baseline is detached from the sampled token, the gradient remains unbiased regardless of $k$.
• Empirically, even small $k$ (e.g., $k=5$) provides sufficient variance reduction to stabilize training.

3. Key Contributions

1. Closed-Form Value Baseline: The paper demonstrates that the OPD value function is exactly the negative reverse KL divergence, making it available directly from the forward pass without extra models.
2. Unbiased Stabilization: vOPD stabilizes training by reducing gradient variance while maintaining the unbiased nature of the single-sample Monte Carlo estimator, unlike top-k variants that bias the objective.
3. Computational Efficiency: vOPD retains the single-token backward pass of base OPD. The $vOPD_{top-k}$ variant further reduces forward-pass costs, offering a superior accuracy-to-compute tradeoff compared to full-vocabulary methods.
4. Theoretical Analysis: The authors provide a derivation showing that variance reduction is proportional to the square of the KL divergence, confirming that the method specifically targets the high-mismatch tokens responsible for instability.

4. Experimental Results

The authors evaluated vOPD on four models (Qwen3-1.7B, 4B, and Olmo-3-7B) across six reasoning benchmarks (MATH500, Minerva Math, AMC23, AIME24/25, SciKnowEval, GPQA-Diamond).

• Performance Gains: vOPD consistently outperforms base OPD.
  • On MATH500, vOPD achieved absolute gains of up to +6.2% over base OPD.
  • Across all benchmarks, vOPD variants showed an average improvement of +3% over base OPD.
  • vOPD performance matched or slightly exceeded $OPD_{full-V}$ (the most expensive method) while being significantly faster.
• Comparison with Stabilization Variants:
  • $OPD_{top-k}$ yielded only marginal gains (e.g., +0.4% average) over base OPD, confirming the bias issue.
  • vOPD variants significantly outperformed $OPD_{top-k}$ and matched $OPD_{full-V}$.
• Efficiency:
  • vOPD reduced wall-clock time by up to 57.7% compared to $OPD_{full-V}$.
  • $vOPD_{top-k}$ was the most efficient, offering the best accuracy-to-compute tradeoff.
• Stability Metrics:
  • Gradient Norms: vOPD maintained gradient norms 1–2 orders of magnitude lower than base OPD, confirming the suppression of noisy, destabilizing gradients.
  • Hyperparameter Sensitivity: Performance was stable across different values of $k$ (5 to 100), with even coarse approximations ($k=5$) providing sufficient stabilization.

5. Significance and Claims

The paper claims that vOPD bridges the gap between Reinforcement Learning and Knowledge Distillation by providing a principled, efficient approach to stable On-Policy Distillation.

• Mechanism: It identifies high-gradient variance from single-sample estimators as the root cause of OPD instability and solves it using a control variate baseline derived from the problem's own structure (the KL divergence).
• Practical Impact: It offers a drop-in replacement for base OPD that requires no additional models (critics) or expensive backward passes, making it suitable for industrial-scale post-training where compute efficiency is paramount.
• Generalization: The method is shown to be robust across different model scales, architectures (Qwen, Olmo), and reasoning domains (mathematics and science).

The authors conclude that vOPD acts as a principled regularizer on destabilizing reward tokens, transforming OPD from an unstable heuristic into a reliable training paradigm without compromising the efficiency that made it attractive in the first place.