PaperlayerView exact revision on arXiv

Source-linked AI summary

$β$-OPSD: Deriving with Policy Optimization, Training with Self-Distillation

Jiawei Xu, Minghui Liu, Juzheng Zhang, Tom Goldstein, Furong Huang

arXiv:2607.28582v1cs.LG

TL;DR

OPSD can be brittle because direct teacher matching leaves its regularization and update strength implicit. The paper introduces β-OPSD, derives interpolated policy targets, and trains them through efficient logit distillation with return-to-go credit assignment. Across mathematical reasoning benchmarks, β-OPSD consistently improves stability and reasoning performance over vanilla OPSD and other reported baselines.

  • Problem

    OPSD can be brittle and often requires substantial engineering, while direct teacher matching provides no explicit control over proximity to a reference policy.

  • Method

    β-OPSD treats β as a KL-regularization parameter, derives a geometric reference-to-teacher target, and approximates it with scheduled token-level logit interpolation and return-to-go credit assignment.

  • Results

    β-OPSD consistently improves optimization stability and reasoning performance over vanilla OPSD; on Qwen3-1.7B, avg@12 improves by up to 9.16 percentage points across three mathematical benchmarks.

  • Takeaways & Limitations

    Smooth target paths and trajectory-level credit assignment provide effective principles for on-policy distillation while retaining its practical efficiency.

  • Takeaways & Limitations

    The exact sequence-level optimum contains an intractable normalizer, so training uses a tractable local approximation at each decoding prefix.

Abstract

from arXiv · show

On-policy self-distillation (OPSD) is a promising approach to improve reasoning language models, but it remains brittle in practice: making it work reliably often requires substantial engineering effort. We identify a structural source of this difficulty: vanilla OPSD is precisely the $β=1$ member of a broader policy-optimization family, where $β$ weights the KL penalty anchoring the student to a reference policy. This equivalence turns $β$ from an implicit value fixed at one into a controllable regularization parameter, yielding a more general formulation that trades off proximity to a reference policy against privileged teacher guidance. We introduce $β$-OPSD and derive its optimal policy as a geometric interpolation between the reference policy and the privileged teacher. Directly optimizing this objective with reinforcement learning, however, would be costly and high-variance. Rather than optimize the RL objective directly, we turn its closed-form solution into a distillation target. Each value of $β$ selects a target along the reference-to-teacher path, which we implement efficiently by mixing their token-level logits. In this way, inexpensive distillation approximates the solution of expensive policy optimization. Return-to-go credit assignment further aligns token updates with the sequence-level objective while retaining the simplicity of OPSD. Experiments on mathematical reasoning benchmarks show that $β$-OPSD consistently outperforms vanilla OPSD, improving optimization stability and downstream reasoning performance. Our results provide a principled route from self-distillation to policy optimization and back without sacrificing the efficiency that makes OPSD practical.

1. Introduction

OPSD improves reasoning by distilling privileged teacher guidance on student-generated trajectories, but its effectiveness can be brittle. β-OPSD makes the teacher–reference trade-off explicit, derives interpolated targets, and adds sequence-level credit assignment while retaining efficient distillation.

  • Motivation: OPSD trains on student-sampled trajectories while a privileged teacher supplies supervision unavailable at deployment.Unlike offline distillation, the teacher responds to the student’s actual intermediate states.
  • Motivation: Direct teacher matching fixes the teacher as the sole target and provides no explicit control over movement away from the current or initial policy.
  • Core idea: Vanilla OPSD is precisely the β = 1 member of a KL-regularized policy-optimization family, making reference regularization explicit and controllable.The family trades off privileged teacher guidance against proximity to a reference policy.
  • Core idea: β-OPSD schedules targets smoothly from the reference policy toward the privileged teacher instead of moving abruptly toward the teacher.Larger β values produce more conservative updates that remain closer to the reference policy.
  • Training: Return-to-go credit assignment propagates future student–target mismatch back to earlier tokens, better aligning token updates with the sequence-level objective.This retains token-level distillation while accounting for how early tokens affect later prefixes.
  • Results: 9.16 percentage points: β-OPSD improves avg@12 over vanilla OPSD across AIME 2024, AIME 2025, and HMMT 2025 on Qwen3-1.7B.

2. Method

The method embeds OPSD in a KL-regularized policy-optimization family, whose optimal policy interpolates geometrically between a reference policy and privileged teacher. It then approximates that sequence-level solution with scheduled local logit interpolation and return-to-go-weighted distillation.

  • Policy-optimization formulation: β-OPSD uses a teacher-to-reference log-ratio reward together with a KL penalty weighted by β to balance teacher preference and reference proximity.β > 0 controls how strongly the learned policy remains close to the reference.
  • Policy-optimization formulation: At β = 1, reference-policy terms cancel and the objective becomes the standard sequence-level reverse KL from the student to the privileged teacher.
  • Optimal targets: The closed-form optimal policy is a geometric interpolation between the reference policy and privileged teacher, with β selecting a point along that path.For β ≥ 1, β = 1 gives the teacher endpoint and β →∞ approaches the reference endpoint.
  • Optimal targets: A scheduled coefficient creates a smooth curriculum whose early targets are closer to the reference and later targets incorporate stronger teacher guidance.This avoids abrupt projection onto the privileged teacher.
  • Distillation target: Because the sequence-level normalizer is intractable, β-OPSD approximates the optimum at each decoding prefix by interpolating reference and teacher logits.The local target is computed with a weighted sum of logits followed by a softmax.
  • Distillation target: Rather than directly optimize costly, high-variance RL, training matches the RL-derived interpolant target with an OPSD-shaped distillation objective.When βk = 1, the interpolant reduces to the privileged teacher and standard OPSD is recovered.
  • Credit assignment: β-OPSD differs from vanilla OPSD through two changes: scheduled logit-interpolant targets and return-to-go weights for sequence-level credit assignment.
  • Credit assignment: Return-to-go weights replace myopic local token weights so earlier actions receive credit for future target mismatch.With γ = 1, the estimator recovers the exact sequence-level gradient; practice uses γ < 1 to control long-generation return magnitude.

3. Related Work

Related work spans smoother on-policy supervision, KL-regularized policy optimization, curriculum-like interpolation, and future-aware credit assignment. β-OPSD combines these ingredients into a single distillation procedure with a principled target path from student to teacher.

  • Related work synthesis: β-OPSD combines smoother supervision, reference-regularized optimization, and return-based credit assignment into one distillation procedure.
  • On-policy distillation and self-distillation: On-policy distillation trains students on their own trajectories with dense token-level teacher feedback, avoiding offline exposure-bias mismatch but requiring reliable supervision on unfamiliar prefixes.
  • RLHF and reference-regularized policy optimization: KL-regularized reinforcement learning and RLHF optimize reward while constraining the policy toward a reference model, with closed-form solutions connecting rewards, references, and optimal distributions.
  • Interpolation, curricula, and teacher-guided learning: Curriculum and imitation-learning methods smooth optimization by gradually changing supervision, while related student–teacher mixtures address distillation’s distributional gap.
  • Interpolation, curricula, and teacher-guided learning: β-OPSD’s logit-interpolant target forms a distributional curriculum whose early targets stay near the student and later targets move toward the privileged teacher.
  • Policy gradients and credit assignment: Policy-gradient work motivates return-to-go credit assignment, which replaces myopic local token-KL updates with signals reflecting future distributional mismatch.

4. Experiments

Experiments evaluate β-OPSD on Qwen3 models and mathematical reasoning benchmarks, then isolate the effects of interpolant targets, return-to-go credit assignment, references, and schedules. β-OPSD generally improves over Vanilla OPSD, with especially large gains for Qwen3-1.7B.

  • Experimental setup: β-OPSD evaluates Qwen3-1.7B, Qwen3-4B, and Qwen3-8B on AIME 2024, AIME 2025, and HMMT 2025 using avg@12.The student generates solutions from problems, while the privileged teacher additionally conditions on ground-truth solutions during training.
  • Main results: β-OPSD improves overall average performance over Vanilla OPSD across all three Qwen3 scales.Evaluations use the 100-step checkpoint from the default 200-step training schedule.
  • Main results: 9.16, 5.27, and 2.78 points are the Qwen3-1.7B gains over Vanilla OPSD on AIME 2024, AIME 2025, and HMMT 2025, respectively.The corresponding average-performance improvement is 5.74 points.
  • Main results: 1.76 and 1.66 points are the overall-average gains over Vanilla OPSD for Qwen3-4B and Qwen3-8B, respectively.β-OPSD slightly underperforms Vanilla OPSD on some individual benchmarks but achieves the strongest overall benchmark-suite performance.
  • Ablations: The scheduled β-OPSD target improves avg@12 over direct teacher distillation by 6.03, 5.30, and 1.67 points on AIME 2024, AIME 2025, and HMMT 2025.Both variants use the same return-to-go estimator, isolating the target-distribution effect.
  • Ablations: Return-to-go credit assignment improves avg@12 over the local token-gradient objective by 1.12, 5.55, and 3.61 points on the three benchmarks.The variants share the same target and sampling distribution, isolating the credit-assignment benefit.
  • Ablations: The dynamic-student and fixed-teacher construction performs best across all three benchmarks when teacher weight is fixed at 0.5.This setup tracks the evolving student while retaining a stable privileged teacher anchor.
  • Ablations: The 0.5 → 0.8 schedule gives the strongest overall result among the schedules considered.Other schedules can perform better on individual benchmarks, indicating that teacher-guidance strength and timing affect performance.

5. Conclusion

β-OPSD frames on-policy distillation as trajectory-level policy optimization, with vanilla OPSD as the β = 1 case and geometric interpolants linking a reference student to a privileged teacher. Return-to-go credit assignment and smooth targets support consistent improvements over vanilla OPSD, SFT, and GRPO on competition-level mathematical reasoning benchmarks.

  • β-OPSD interprets on-policy distillation as trajectory-level policy optimization rather than merely local token imitation.
  • Vanilla OPSD is the β = 1 teacher endpoint of a reference-regularized reinforcement-learning family.
  • The family motivates tractable token-level logit interpolants along a geometric path from a reference student to the privileged teacher.
  • Return-to-go credit assignment accounts for future distributional mismatch and provides an unbiased gradient estimator for the sequence-level objective.
  • β-OPSD consistently improves over vanilla OPSD, SFT, and GRPO on competition-level mathematical reasoning benchmarks.

B.1. Proof of Proposition 2.1

The proof rewrites the β-OPSD objective as a teacher KL term plus reference-policy regularization, then shows that β = 1 removes the reference term and recovers standard OPSD.

  • Expanding the objective yields a teacher-matching term and a reference-policy KL penalty weighted by β − 1.
  • At β = 1, the reference-regularization term cancels, leaving the standard sequence-level OPSD objective.

B.2. Proof of Proposition 2.2

The proof expands the β-OPSD objective and identifies its maximizer as the normalized distribution combining teacher and reference log probabilities with β-dependent weights.

  • Expanding the β-OPSD objective expresses it as an expectation of teacher, student, and reference log probabilities.
  • The candidate optimum combines the reference log probability weighted by β − 1 with the teacher log probability, minus β times a normalizer.
  • Substitution separates a π-independent term from a nonnegative KL term, which is minimized exactly at the proposed policy.

B.3. Proof of the Return-to-Go Gradient Estimator

The proof establishes that return-to-go credit assignment produces an unbiased estimator of the sequence-level reverse-KL gradient when the interpolant target is fixed and γ = 1.

  • The estimator proof treats the interpolant target as fixed during the current gradient update.
  • The token-level discrepancy is defined as the student log probability minus the interpolant target’s log probability.
  • Autoregressive factorization, score-function differentiation, and zero expected score remove past terms through causality.
  • The resulting return-to-go estimator is unbiased for the sequence-level gradient.
  • For γ = 1, the practical surrogate has the desired gradient up to the normalization factor 1/T; γ < 1 gives a discounted approximation with reduced long-horizon magnitude and variance.

C.1. Training Details

Vanilla OPSD and β-OPSD share a fixed training setup spanning hardware, optimization, parameter-efficient tuning, and on-policy generation.

  • Shared training configuration: Training uses four NVIDIA RTX A6000 GPUs with four-way data parallelism, colocated single-GPU vLLM generation, and up to 200 optimization steps.An eight-hour wall-time limit is also imposed.
  • Optimization and parameter-efficient tuning: Optimization uses a 5 × 10^-6 learning rate, 0.1 gradient-norm clipping, per-device batch size 1, and eight accumulation steps for an effective batch size of 32.The effective batch size is computed as 4 GPUs × 1 sample per GPU × 8 accumulation steps = 32.
  • Optimization and parameter-efficient tuning: Models use gradient checkpointing, bfloat16 precision, FlashAttention-2, and LoRA adapters with rank r = 64 and scaling αLoRA = 128.Adapters are applied to q_proj, k_proj, v_proj, o_proj, gate_proj, up_proj, and down_proj.
  • On-policy generation: On-policy trajectories allow 20,000-token sequences and 1,024-token completions, sampled with temperature 1.1, top-p = 0.95, and top-k = 20.These settings are shared by Vanilla OPSD and β-OPSD to isolate the distillation objective and Look-Ahead credit assignment.

C.2. Evaluation Details

Evaluation covers three Qwen3 model sizes and three mathematical reasoning benchmarks, using repeated stochastic candidate generation and pass@12 and avg@12 metrics.

  • Evaluation pipeline: Evaluation uses Qwen3-1.7B, Qwen3-4B, and Qwen3-8B checkpoints with corresponding LoRA adapters on AIME 2024, AIME 2025, and HMMT 2025.Adapters are loaded on top of the original Qwen3 base models.
  • Decoding configuration: For each problem, evaluation generates k = 12 independent candidate solutions using thinking mode, temperature 0.6, top-p = 0.95, and top-k = 50.The evaluation batch size is 32, with no separate completion-length cap in the evaluation script.
  • Scoring and metrics: pass@12 measures whether at least one of 12 solutions is correct, while avg@12 measures average correctness across all 12 solutions.Both metrics are reported in percentage points after multiplying raw fractions by 100.
  • Inference configuration: Inference uses tensor-parallel size 1, GPU-memory utilization 0.9, one evaluation repeat, and checkpoint sweeps at training steps 50, 75, 100, and 200.One Slurm evaluation job is submitted for each checkpoint–dataset pair.

D. Sampling from Logit Interpolants

Mixed student–teacher sampling constructs proposals by interpolating student and teacher distributions, then corrects the resulting off-policy mismatch with importance sampling. Results show comparable gains but added sensitivity and cost, so student on-policy sampling remains the default.

  • Proposal construction: The proposal distribution mixes the stop-gradient student and privileged teacher with teacher weight η: m̄θ,η = (1 − η)π̄θ + ηpT.η = 0 recovers student on-policy sampling, while η = 1 samples entirely from the privileged teacher.
  • Distribution correction: Because the objective is defined under the student trajectory distribution, mixed proposals introduce distribution mismatch corrected with per-decision importance sampling.All mixed-sampling variants use β-OPSD with discount factor γ = 0.99.
  • Surrogate objective: The displayed surrogate uses student log-probabilities weighted by a return-to-go term for the logit-interpolant target.The passage identifies the student mismatch with the logit-interpolant target and the corresponding optimized surrogate.
  • Experimental settings: Mixed-sampling experiments evaluate fixed η values of 0.2, 0.5, and 0.8, plus linear schedules that increase or decrease η during training.Here η denotes teacher weight, while 1 − η denotes student weight.
  • Results: At 200 steps, every evaluated schedule improves avg@12 over Vanilla OPSD on both AIME benchmarks, with best gains of 8.06, 5.55, and 3.06 points on AIME 2024, AIME 2025, and HMMT 2025.These gains are reported across all checkpoints for the mixed student–teacher sampling variants.
  • Practical trade-offs: No proposal schedule dominates across benchmarks, and mixed sampling adds dual-model decoding cost plus importance-sampling variance; therefore student on-policy sampling is the default.Dual-model sampling requires combining student and teacher distributions at every decoding step and has nearly double implementation complexity and generation overhead.
Loading 2607.28582v1…