An RL View of OPD: Least Square Policy Distillation for Sample-Efficient LLM Reasoning

AI summary

Overview

• Research area: Large language model post-training, specifically on-policy distillation (OPD) and its connection to KL-regularized reinforcement learning, with a focus on sample-efficient reasoning improvement.
• Technical level: Advanced. The paper combines an exact mathematical reformulation of OPD, a Lagrangian-derived least-squares objective, and a regret/sample-complexity analysis, alongside empirical evaluations on mathematical reasoning benchmarks.
• One-sentence scope: The paper reinterprets on-policy distillation as token-level KL-regularized policy optimization, derives a least-squares distillation objective (LSPD) that supports optimistic exploration and off-policy data reuse, proves a logarithmic regret bound for an idealized version, and shows consistent gains on six math benchmarks across three Qwen3 teacher–student pairs.

What This Paper Is About

On-policy distillation trains a student model on its own generated trajectories using token-level teacher supervision, and is typically implemented with policy-gradient frameworks such as PPO or GRPO. Those implementations are on-policy: they need fresh, expensive rollouts throughout training and make limited use of previously collected data.

The paper asks whether viewing OPD through the lens of KL-regularized reinforcement learning can yield a distillation algorithm that explores better, reuses old rollout data, and preserves student policy diversity, while still improving reasoning accuracy.

Key Contributions

1. An RL-inspired policy distillation framework (LSPD). The paper shows that the reverse-KL OPD objective is an exact reformulation of token-level KL-regularized policy optimization with a teacher-induced log-ratio reward. Motivated by optimistic KL-regularized value-based RL, it derives a least-squares matching objective, adds a Huber-type robust penalty and an explicit entropy bonus, and shows that the resulting loss is naturally off-policy. A theoretical analysis of an idealized optimistic formulation gives a regret bound of order Õ(log K), where K is the number of update rounds.

2. Effective data reuse for sample-efficient distillation. LSPD-RB, the replay-buffer variant, reaches saturated performance in approximately 10 rollout batches, compared with more than 40 for LSPD with one update per batch, according to the paper's off-policy experiments.

3. Improved reasoning performance, diversity, and efficiency. Across six benchmarks and three teacher–student settings, LSPD improves Avg@16 and Pass@16 by +1.59 and +1.87 points on average over baselines. LSPD-RB further improves Pass@16 with only 10 training steps, while higher measured entropy and Pass@k curves indicate better policy diversity and solution coverage.

Main Findings

• Average benchmark performance. Averaged over all 18 model–benchmark combinations, LSPD achieves Avg@16 of 31.60 and Pass@16 of 53.30, outperforming the strongest baseline, EOPD, by +0.91 and +0.85 points, and standard OPD by +1.99 and +2.27 points. LSPD-RB achieves a comparable Avg@16 of 31.51 and a higher Pass@16 of 54.86, exceeding EOPD by +0.82 Avg@16 and +2.42 Pass@16, and OPD by +1.89 and +3.84 points.

• Ranking across settings. LSPD attains the best or tied-best Avg@16 on 11 of 18 settings and Pass@16 on 9 of 18 settings, and is top-2 on 17 of 18 and 16 of 18 settings respectively. Considering LSPD and LSPD-RB together, at least one achieves the best or tied-best result on 17 of 18 settings for both metrics, so the better-performing variant is top-2 across all 36 metric–setting combinations.

• Rollout efficiency. LSPD-RB achieves its results in only 10 training steps, whereas the baselines require more than 40 steps to converge stably. The abstract states the fully off-policy variant matches vanilla OPD using only the first 25% of rollout batches.

• Policy entropy during training. Tracking student policy entropy over 200 training steps for three teacher–student pairs, both LSPD and EOPD maintain entropy at approximately 0.5, consistently higher than standard OPD. LSPD raises entropy directly through the regularization term in its objective rather than EOPD's entropy-gated forward KL formulation.

• Pass@k scaling. In a Qwen3-4B-Base student distilled from a Qwen3-8B non-thinking teacher, LSPD achieves the best or tied-best Pass@32 and Pass@64 on all three tested benchmarks, reaching 85.54, 43.33, and 46.67 Pass@32 and 89.16, 46.67, and 50.00 Pass@64 on AMC23, AIME24, and AIME25. Averaged across the three benchmarks, LSPD reaches Pass@32/64 of 58.51/61.94 versus 57.00/60.43 for EOPD, a gain of +1.51 and +1.52 points.

• Off-policy update count. Increasing the number of optimization steps per rollout batch from N=1 to N=4, 16, and 64 improves sample efficiency, though gains shrink beyond N=4. With N ≥ 4, performance typically saturates after roughly 30 training steps, whereas N=1 requires more than 50 steps. LSPD-RB saturates within approximately 10 training steps; for LSPD-RB, each newly collected rollout batch is followed by 256 optimization steps sampled from the buffer.

• Entropy regularization ablation. Removing entropy regularization changes Avg@16 only modestly (down 0.17 points for LSPD and 0.75 for LSPD-RB) but degrades Pass@16 substantially (down 2.26 and 1.64 points). Averaged across both methods and all benchmarks, the degradation is 0.46 points in Avg@16 versus 1.95 points in Pass@16, and Pass@16 decreases in all 12 method–benchmark comparisons without entropy regularization.

• Theory results. Under an assumption that the optimal log-ratio reward lies in the reward class, the idealized algorithm satisfies Regret(K) = O(d_E log(N_R K)), where d_E and N_R are complexity measures on the order of Õ(d) for a d-dimensional linear reward. A remark states this implies Õ(ε⁻¹) sample complexity for a mixture policy whose reverse-KL divergence to the optimal policy is at most ε, compared with the Õ(ε⁻¹) sample complexity Foster et al. (2024) establish for MLE with a deterministic expert under a different error criterion. The provided paper text is truncated mid-sentence in this remark, so the full comparison is not available.

Methodology in Plain English

The authors begin by rewriting the standard OPD loss. Instead of treating the teacher's token-level log-probability as a fixed advantage signal inside a PPO-style update, they define a reward as the teacher's log-probability minus a reference policy's log-probability at the same prefix. With that substitution, minimizing the reverse KL to the teacher becomes exactly the same form as maximizing a reward minus a KL penalty to a reference policy — the classical KL-regularized RL objective, with the regularization weight set to one. This is presented as an exact reformulation, not an extra penalty term.

Having reframed OPD as an RL problem, the authors import a technique from value-based RL: optimism under uncertainty. After collecting K rounds of rollouts, they form a least-squares confidence set over reward functions that are statistically consistent with all accumulated prefix–token pairs, then pick the pointwise most favorable reward inside that set and update the token distribution in closed form. A Lagrangian relaxation of this step turns the pointwise optimization into a quadratic matching problem between the student's and the teacher's log-probabilities.

That relaxation is what makes the method practical. The practical loss keeps the quadratic log-probability matching but replaces the raw squared error with a Huber-type robust penalty for stability, and adds an explicit entropy bonus to counter the tendency of distillation to collapse onto a narrow set of high-probability responses. Because the objective matches log-probabilities rather than requiring gradients of the current policy's samples, it can be optimized on trajectories produced by earlier student policies.

Two variants result. LSPD clears the buffer after each rollout batch and performs one or more on-policy or semi-on-policy updates. LSPD-RB keeps a replay buffer of all previously collected query–response pairs and samples minibatches from it, making training fully off-policy.

Experiments use three Qwen3 teacher–student pairs — 8B to 4B-Base, 4B to 1.7B-Base, and 1.7B to 0.6B-Base — with teacher thinking mode disabled. Training uses DAPO-Math-17K, and evaluation covers MATH-500, Minerva, Olympiad-Bench, AMC23, AIME24, and AIME25, reporting averages over 16 generations against KD, OPD, and EOPD baselines.

Why This Matters

The paper's main impact is conceptual: it shows that the increasingly popular on-policy distillation recipe is not just "RL-like" but exactly a KL-regularized RL problem, which opens the door to the algorithmic toolkit of that literature — optimism, replay, entropy regularization, and regret analysis. That matters for researchers who currently implement OPD with PPO or GRPO and inherit their on-policy sample costs.

Potential real-world applications (the paper does not report deployment studies; these follow from the method's properties):

• Cheaper distillation pipelines for reasoning models. Cut the number of expensive rollout batches needed to reach a target accuracy, reducing GPU time for producing smaller reasoning-capable models.
• On-device or edge model compression. The 1.7B-to-0.6B setting shows gains at small scale, which is the regime relevant to shipping models to constrained hardware.
• Diverse-answer systems. For applications such as verifier-guided search, best-of-n sampling, or tool use where multiple candidate solutions are generated and filtered, the Pass@k improvements matter more than single-sample accuracy.
• Continual updating of deployed models. A replay-buffer objective that reuses historical trajectories fits scenarios where new data arrives slowly and each rollout is costly.

Industry relevance: the method is evaluated on Qwen3 models from 8B down to 0.6B parameters, uses a publicly available training set, and the authors provide code at https://github.com/UNCSciML/LSPD. Reduced rollout count translates directly into lower training cost, which is the dominant expense in LLM post-training.

Future Directions

• Closing the gap between theory and practice. The regret guarantee applies to an idealized optimistic algorithm, while the practical method uses a Huber-robust quadratic matching loss plus an entropy bonus. Whether the practical loss retains the Õ(log K) regret rate is not established.
• Testing beyond mathematical reasoning. All reported experiments use DAPO-Math-17K and six math benchmarks. The paper reports out-of-domain evaluation in its appendix, but the main text does not describe those results, so generality to non-math reasoning is unresolved here.
• Understanding the hyperparameter landscape. The confidence radius, the regularization parameter, the Lagrange multiplier, and the entropy weight all enter the derivation, and the paper reports an ablation only for the entropy term.
• Comparing against the MLE sample-complexity benchmark. The theoretical discussion contrasts the reverse-KL guarantee with the Õ(ε⁻¹) rate for maximum likelihood estimation under a deterministic expert. The provided text is truncated mid-comparison, so the full implications of that contrast remain to be worked out.

Target Audience

Researchers and graduate students working on LLM post-training, reinforcement learning for language models, or knowledge distillation, particularly those with an interest in the theory of KL-regularized RL and in sample efficiency. Practitioners who train reasoning models with PPO- or GRPO-style on-policy distillation will find the off-policy replay variant the most directly applicable part. Readers without a background in policy optimization, KL regularization, or regret analysis will find the theoretical sections difficult, though the experimental tables and the plain description of the loss are accessible.

Authors’ abstract

We study on-policy distillation (OPD) through the lens of reinforcement learning, establishing a connection between the reverse-KL objective in OPD and KL-regularized policy optimization. Building on this connection, we introduce Least-Square Policy Distillation (LSPD), an RL-inspired framework that brings optimistic exploration and off-policy data reuse from value-based RL into policy distillation. LSPD preserves policy diversity through exploration while improving rollout efficiency by repeatedly learning from previously collected trajectories. Our theoretical analysis connects LSPD to optimistic value-based learning and shows that its idealized formulation achieves a sharp $\tilde{\mathcal O}(\log K)$ regret bound under online exploration. Empirically, LSPD consistently outperforms existing distillation baselines across six mathematical reasoning benchmarks and diverse teacher-student settings, with average gains of +1.59 points in Avg@16. Remarkably, through Pass@k evaluations up to k=64, we found that LSPD better preserves policy diversity by achieving stronger performance as k grows. Its fully off-policy variant achieves comparable performance to vanilla OPD using only the first 25% of rollout batches. Together, these results provide an RL perspective on OPD that offers both a principled interpretation and a practical route toward more effective and rollout-efficient language model distillation.