PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Shangzhe Li, Yuxiao Yang, Tianrun Yu, Kaixiang Zhao, Xiaoyun Wang, Taylor W. Killian, Weitong Zhang

arXiv:2609.35505v1cs.LG

TL;DR

OPD’s RL formulation leaves exploration and data efficiency underexplored, while common policy-based implementations require fresh rollouts and offer limited historical data reuse. The paper introduces LSPD, using optimistic KL-regularized value-based learning principles for distillation. Across mathematical reasoning benchmarks and teacher–student settings, it reports improved reasoning performance, stronger Pass@k scaling, and greater rollout efficiency, with its replay-buffer variant comparable to vanilla OPD using only the first 25% of rollout batches.

  • Problem

    OPD’s algorithmic implications for exploration and data efficiency remain underexplored, and common policy-based implementations require fresh rollouts while offering limited historical data reuse.

    Source: §1 · p. 1

  • Method

    LSPD uses optimistic KL-regularized value-based learning principles, robust quadratic matching of student and teacher log-probabilities, and entropy regularization to support off-policy trajectory reuse.

    Source: §1 · pp. 1–2, §4.2 · p. 4, §4.3 · p. 5

  • Results

    Across six benchmarks and three teacher–student settings, LSPD improves Avg@16 by +1.59 points on average over baselines.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The results show how value-based RL principles can improve policy distillation by preserving policy diversity and enhancing sample efficiency through off-policy learning.

    Source: §7 · p. 10

    The OPD-to-KL-regularized-RL formulation assumes teacher and reference policies assign positive probability to tokens in the student’s support at each prefix.

    Source: §4.1 · p. 3

Abstract

from arXiv · show

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.

1 Introduction

The paper addresses OPD’s limited exploration and historical data reuse by introducing an RL-inspired distillation framework. It reports improved reasoning, rollout efficiency, and a logarithmic regret guarantee for an idealized formulation.

1 Introduction

  • The paper identifies exploration and data efficiency as underexplored implications of OPD’s RL formulation, noting that policy-based implementations rely on fresh rollouts and offer limited historical data reuse.Source: §1 · p. 1
  • The idealized optimistic formulation has a sharp O(log K) regret bound, where K counts update rounds.Source: §1 · p. 2
  • LSPD draws on optimistic KL-regularized value-based RL, matching student and teacher log-probabilities with a robust quadratic loss and explicit entropy regularization.Its objective supports multiple updates per rollout and reuse of trajectories from earlier student policies.Source: §1 · pp. 1–2
  • Approximately 10 rollout batches suffice for LSPD-RB to reach saturated performance, compared with more than 40 for LSPD with one update per batch.The comparison is from the paper’s off-policy experiments.Source: §1 · p. 2
  • +1.59 points in Avg@16 and +1.87 points in Pass@16 on average over baselines across six benchmarks and three teacher–student settings.The reported gains are averaged across the settings described.Source: §1 · p. 2

2 Related Works

Related work spans language-model distillation and reinforcement learning for LLM post-training. The paper situates its contribution at their intersection: quadratic matching of student and teacher log-probabilities, motivated by optimistic KL-regularized RL.

2 Related Works

  • The paper connects reverse-KL distillation to optimistic KL-regularized RL and introduces a maximum-entropy least-square objective for off-policy data reuse.Source: §2 · p. 2
  • Prior distillation work ranges from teacher-generated-data forward KL to student-trajectory and reverse-KL methods, with later work exploring divergence choices, data efficiency, adaptive reuse, contrastive objectives, and other extensions.Examples include DistiLLM [Ko et al., 2024] and DistiLLM-2 [Ko et al., 2025].Source: §2 · p. 2
  • RL post-training includes preference-based alignment and verifiable-reward approaches, while IPO uses a squared loss on policy-to-reference log-likelihood ratios rather than matching student and teacher log-probabilities.Source: §2 · p. 2

3 Preliminaries

The preliminaries model LLM post-training at the token level as a contextual bandit. They relate KL-regularized policy optimization to maximum-entropy RL and describe OPD as minimizing token-level reverse KL on student-generated prefixes.

3 Preliminaries

  • LLM post-training is formulated as a contextual bandit in which a query and generated-token prefix form the context, the next token is an action, and a reward function is maximized.Source: §3 · p. 3
  • KL-regularized RL optimizes reward while regularizing policy divergence from a fixed reference; with a uniform reference, the divergence term becomes negative entropy.Source: §3 · p. 3
  • OPD minimizes token-level reverse KL between teacher and student policies over prefixes sampled from the student’s own rollouts.The paper notes that common implementations use on-policy optimization such as PPO.Source: §3 · p. 3

4 Methodology and Analysis

The methodology recasts OPD as KL-regularized policy optimization, then develops an optimistic, data-reusing distillation objective and a robust entropy-regularized practical variant.

4.1 On-Policy Distillation as KL-Regularized Reinforcement Learning

  • OPD is an exact token-level counterpart of KL-regularized policy optimization with η = 1 and a teacher-induced reward.The reward measures the teacher’s next-token log-likelihood relative to a fixed reference; the formulation assumes teacher and reference assign positive probability to tokens in the student’s support.Source: §4.1 · pp. 3–4

4.2 Least Square Policy Distillation

  • Unlike PPO’s updates driven by current reward estimates, the proposed optimistic update explicitly accounts for uncertainty to encourage exploration toward potentially better tokens.The method is motivated by optimistic reward estimation in KL-regularized reinforcement learning.Source: §4.2 · p. 4
  • LSPD estimates reward from accumulated prefix–token pairs and chooses optimistic token distributions using favorable reward estimates consistent with the data.The formulation builds a least-squares confidence set over reward functions and supports updates from historical rollouts.Source: §4.2 · pp. 4–5

4.3 Practical Implementation

  • LSPD’s practical objective combines robust student–teacher log-probability matching with an entropy bonus, replacing the pointwise optimistic term.The implementation uses a Huber-type penalty to reduce the influence of excessively large log-probability differences.Source: §4.3 · p. 5
  • LSPD supports multiple updates per rollout batch, while LSPD-RB additionally reuses responses generated by earlier policies through a replay buffer.The paper distinguishes the on-policy or semi-on-policy variant from the fully off-policy replay-buffer variant.Source: §4.3 · p. 5, §4.2 · p. 5

5 Experiments

Experiments compare LSPD with distillation baselines across mathematical reasoning tasks, model settings, and evaluations of accuracy, diversity, and rollout efficiency. LSPD performs strongly across benchmarks, while replay-buffer training and entropy regularization relate to sample efficiency and solution coverage.

5.1 Experimental Setup

  • The experiments use three Qwen3 teacher–student settings and six mathematical reasoning benchmarks, comparing KD, OPD, and EOPD with LSPD variants.Training uses DAPO-Math-17K; the six evaluations include MATH-500, Minerva, Olympiad-Bench, AMC23, AIME24, and AIME25.Source: §5.1 · pp. 5–6

5.2 Main Results

  • Across 18 model–benchmark combinations, LSPD averages 31.60 Avg@16 and 53.30 Pass@16, beating EOPD by +0.91 and +0.85 points, respectively.Table 1 also reports LSPD-RB at 31.51 Avg@16 and 54.86 Pass@16, using 10 training steps versus at least 40 for other methods.Source: §5.2 · p. 6, Table 1 · p. 7
  • LSPD maintains student entropy comparable to EOPD and higher than OPD across three teacher–student pairs, consistent with greater policy diversity.LSPD uses direct entropy regularization with robust log-probability matching, while EOPD uses an entropy-gated forward-KL term.Source: §5.2 · p. 6, Figure 1 · p. 7
  • For a Qwen3-4B-Base student distilled from a Qwen3-8B non-thinking teacher, LSPD achieves best or tied-best Pass@32/64 on all three benchmarks: 85.54/89.16 AMC23, 43.33/46.67 AIME24, and 46.67/50.00 AIME25.Source: §5.2 · p. 6

5.3 Off-policy and LSPD-RB Replay Buffer Results

  • LSPD-RB reaches saturated performance in about 10 rollout batches; semi-on-policy LSPD needs more, with N = 1 taking over 50 steps on AIME24 and AIME25.Increasing updates per batch improves sample efficiency, but gains become substantially smaller beyond N = 4.Source: Figure 2 · p. 8, §5.3 · p. 8
  • Removing entropy regularization reduces Pass@16 by 2.26 points for LSPD and 1.64 for LSPD-RB, while Avg@16 falls by 0.17 and 0.75 points, respectively.Across both methods and all benchmarks, Pass@16 decreases in all 12 comparisons; Avg@16 changes are comparatively small and mixed.Source: §5.3 · pp. 8–9, Table 3 · p. 9

6 Theoretical analysis

The theoretical analysis bounds reverse-KL regret for an optimistic formulation under stated assumptions. Its fast-rate guarantee depends on bounded rewards and student–teacher coverage.

Assumption 6.1

  • Assumption 6.1 requires a bounded reward class, with B < ∞; choosing the initial student as reference assumes student–teacher coverage.The text relates coverage dependence to reverse-KL minimization, including OPD and LSPD.Source: §6 · p. 9

Theorem 6.2

  • With high probability, Algorithm 2 has regret O(dE log(NRK)) under Assumptions C.1 and 6.1.The theorem states that dE and NR are each of order Õ(d) when the reward is a d-dimensional linear function.Source: §6 · p. 9

Reverse-KL OPD and forward-KL SFT

  • The comparison gives MLE an Õ(ε^-1) sample complexity for an ε-small imitation gap with a deterministic expert, while noting that the error criteria differ.The passage contrasts this with reverse-KL fast rates under a coverage condition and forward-KL SFT’s mass-covering tendency.Source: §6 · p. 9

7 Conclusion

The conclusion presents LSPD as an RL-inspired approach connecting reverse-KL distillation to KL-regularized policy optimization. Its experiments and theoretical analysis support off-policy reuse, stronger reasoning results, and rollout efficiency.

7 Conclusion

  • LSPD connects reverse-KL distillation to KL-regularized policy optimization through a teacher-induced reward.The connection motivates an optimistic least-squares formulation that enables trajectory reuse with a purely off-policy replay buffer.Source: §7 · pp. 9–10
  • Across six mathematical reasoning benchmarks and multiple teacher–student settings, experiments show improved reasoning performance, stronger Pass@k scaling, and greater rollout efficiency.LSPD-RB achieves performance comparable to vanilla OPD using only the first 25% of rollout batches.Source: §7 · p. 10
  • The conclusion links the theoretical guarantee and empirical findings to value-based RL principles for preserving policy diversity and improving sample efficiency through off-policy learning.Source: §7 · p. 10

A Hyperparameters and Implementation Details

The implementation uses a Huber-type penalty to stabilize LSPD training, compares methods under fixed experimental settings, and applies a step-by-step reasoning prompt with boxed final answers.

A.1 Details on Huber-Type Penalty in Practical Implementation

  • The Huber-type penalty preserves quadratic matching for moderate log-probability discrepancies and switches to linear growth for large discrepancies to reduce outlier influence and improve stability.The threshold c separates the quadratic and linear regimes.Source: §A.1 · p. 13

A.2 Training Details

  • Experiments reproduce OPD, EOPD, and KD under the same setup, holding batch size, rollout and evaluation configurations, and optimizer settings fixed.Runs use four NVIDIA RTX PRO 6000 GPUs with 96GB VRAM each.Source: §A.2 · p. 13

A.3 Prompt Template

  • Training and evaluation prompts ask the model to reason step by step and place its final answer within a box.The prompt template follows TTRL.Source: §A.3 · p. 13

B Additional Experimental Results

Additional experiments assess LSPD beyond mathematical reasoning and examine its hyperparameter settings and Huber penalty. LSPD and LSPD-RB improve average out-of-domain benchmark scores over EOPD, while the Huber penalty also improves Avg@16 results.

Out-of-Domain Results

  • LSPD-RB averages 56.84 across four out-of-domain benchmarks, +1.84 points over EOPD at 55.01; LSPD averages 56.02, +1.02 points over EOPD.LSPD-RB leads on GPQA, MMLU, and HumanEval, while LSPD leads on WinoGrande.Source: B Additional Experimental Results · p. 13

Training Hyperparameters

  • For off-policy experiments without a replay buffer, the number of optimizer updates per rollout iteration varies across N=1, 4, 16, and 64.Table 4 reports training hyperparameters for LSPD and LSPD-RB across experiments.Source: Table 4 · p. 14

Ablation on Huber-Type Robust Penalty

  • Across AMC23, AIME24, and AIME25, the Huber-type penalty improves Avg@16 by an average of +0.50 points.The comparison uses a Qwen3-1.7B-Base student, a Qwen3-4B non-thinking teacher, and LSPD-RB with or without the penalty.Source: B Additional Experimental Results · p. 14
  • Only 0.345% of evaluated tokens enter the Huber penalty's linear region; 99.655% remain in its quadratic region.The analysis uses about 335K tokens from 128 responses sampled from 32 DAPO-Math-17K prompts.Source: B Additional Experimental Results · p. 14

C Theoretical Analysis and Proofs

The appendix analyzes LSPD as token-level optimistic least-squares learning, reducing the sequential proof to contextual-bandit reasoning. It relates the regularized objective to reverse KL and establishes sharp one-step and confidence results under stated assumptions.

C Theoretical Analysis and Proofs

  • The proof adapts contextual-bandit analysis to conditional next-token distributions, summing per-prefix bounds along student-generated responses using the KL chain rule.It requires no transition model, action-value function, or Bellman recursion; the reduction relies on the special reward r*=log(π*/πref).Source: C Theoretical Analysis and Proofs · p. 14

C.1 Setup and Reverse-KL Identity

  • The analysis assumes a finite reward class and bounded response length, with the policy fixed during each rollout and its token observations available for the next update.Teacher feedback is modeled as conditionally unbiased with sub-Gaussian noise under Assumption C.1.Source: C Theoretical Analysis and Proofs · pp. 14–15, §C.1 · p. 15
  • The regularized objective is a reverse-KL objective to the target policy, and the autoregressive KL chain rule gives J(π*)=0.The regret in Eq. 6.1 is the cumulative objective gap, not a token-averaged surrogate.Source: §C.1 · pp. 16–17

C.2 Objective Decomposition and the Sharp One-Step Bound

  • Under optimism, Lemma C.9 gives a sharp one-step bound, which is summed across tokens and averaged over prefixes generated by the student policy.The proof first establishes the bound at each fixed prefix.Source: §C.2 · p. 19

C.3 Least-Squares Confidence and Uniform Optimism

  • Under Assumptions C.1 and 6.1, finite-class least-squares confidence contains the target reward, and uniform optimism holds with probability at least 1−δ/2.The cumulative eluder complexity makes response-length dependence explicit without a transition-estimation or Bellman-error term.Source: §C.1 · p. 17, §C.3 · pp. 19–20

C.4 Proof of Theorem 6.2

The proof establishes a logarithmic regret order by bounding token-level log-ratio reward errors and summing fixed-prefix bandit inequalities. It then interprets the covering-number dependence and contrasts the guarantee with NAILGUN under different feedback assumptions and error criteria.

C.4 Proof of Theorem 6.2

  • The proof controls cumulative reverse-KL regret using immediate token log-ratio rewards and fixed-prefix bandit inequalities, without propagating errors through an MDP.The argument handles observed-to-predictable regret with a concentration step while retaining the logarithmic regret order.Source: §C.4 · p. 21

C.5 Covering-Number Interpretation

  • The log-policy and induced reward classes have identical uniform covering numbers because the fixed reference term cancels in pairwise differences.The cancellation also makes their token-level uncertainty values and cumulative eluder complexities coincide.Source: §C.5 · p. 22
  • A finite uniform ε-net preserves the logarithmic regret form, with covering-number dependence evaluated on prefix–token pairs at the stated resolution.The extension assumes the empirical minimum is attained and makes additional dependence on K and T through the covering resolution and dE explicit.Source: §C.5 · p. 22

C.6 Comparison with On-Policy Distillation under Noisy Expert Feedback

  • The guarantee relies on conditionally unbiased feedback and does not establish that vanilla reverse-KL distillation fails under the paper’s assumptions.NAILGUN addresses corruption, whereas the paper’s analysis gives a finite-sample guarantee for optimistic least-squares learning (Sriraman et al., 2026).Source: §C.6 · p. 23
  • Unlike ordinary on-policy sampling, Algorithm 2 uses optimistic token selection and least-squares log-ratio reward estimation to obtain the stated fast-learning guarantee.The comparison is with NAILGUN, which updates a weighted aggregate of deterministic policies using noisy expert labels (Sriraman et al., 2026, Algorithm 3).Source: §C.6 · p. 23
  • Theorem 6.2 bounds high-probability cumulative reverse-KL regret, while NAILGUN bounds expected final-policy Hellinger error; their guarantees are complementary, not a strict sample-complexity improvement.Both have a 1/K output-error factor up to complexity and logarithmic terms, and the analysis accommodates stochastic targets under bounded log-ratio rewards.Source: §C.6 · p. 23

D Generated Samples

This section presents representative LSPD and LSPD-RB samples from AMC23, AIME24, and AIME25 in a Qwen3-8B non-thinking teacher to Qwen3-4B-Base student setup.

D Generated Samples

  • Representative generated samples use AMC23, AIME24, and AIME25 problems with a Qwen3-8B non-thinking teacher and Qwen3-4B-Base student.The section presents samples from both LSPD and LSPD-RB.Source: D Generated Samples · p. 23
  • A question asks for xy when real numbers x and y are both greater than 1 and satisfy logarithmic conditions; the excerpt then begins a solution to find xy.Source: D Generated Samples · p. 27
  • A question asks for integer bases b > 9 such that 17b divides 97b; the cited solution checks b = 21 and b = 49 and says both satisfy the condition.Source: D Generated Samples · pp. 28–29

Related papers

Loading 2609.35505v1…