PaperlayerView exact revision on arXiv

Source-linked AI summary

Self-Distilled Policy Gradient

Yifeng Liu, Shiyuan Zhang, Yifan Zhang, Quanquan Gu

arXiv:2606.04036v1cs.LG

TL;DR

Sparse sequence-level verifier rewards make token-level credit assignment difficult, while privileged self-distillation can provide dense supervision but may be noisy or biased by unavailable information. SDPG combines exact full-vocabulary on-policy reverse-KL distillation with verifier-based policy optimization and reference-policy KL regularization, using gating and scheduling controls. The framework reports better performance and stability than RLVR and self-distillation baselines on LLM reasoning tasks.

  • Problem

    RLVR provides sparse sequence-level rewards, and self-distillation can suffer from noisy strong teacher signals or privileged information unavailable at inference.

  • Method

    SDPG combines binary verifier rewards, exact full-vocabulary privileged OPD, group-relative policy optimization, reference-policy KL regularization, positive-advantage gating, and a warmup-decay distillation schedule.

  • Results

    SDPG improves performance and stability over baseline algorithms, with SDPG variants outperforming GRPO and RLSD across benchmarks and reaching high-reward plateaus earlier than GRPO.

  • Takeaways & Limitations

    Combining verifier rewards with OPD-equivalent dense token signals improves credit assignment while retaining RLVR’s exploration and selection benefits.

  • Takeaways & Limitations

    Privileged OPD can remain biased by information unavailable at inference, motivating decay of the distillation coefficient near training’s end.

Abstract

from arXiv · show

On-policy self-distillation, where a language model conditions on privileged context to supervise its own generations, is a promising source of dense supervision for sparse-reward reinforcement learning. Actually, it can be instantiated as an auxiliary full-vocabulary student-to-teacher reverse Kullback-Leibler divergence loss. We therefore propose SDPG, a self-distilled policy-gradient framework that combines group-relative verifier advantages with normalized standard deviation, exact full-vocabulary on-policy self-distillation, as well as reference-policy KL regularization. Empirically, SDPG improves stability and performance over RLVR and self-distillation baselines. The code is available at https://github.com/lauyikfung/SDPG.

1 Introduction

RLVR provides effective outcome supervision for reasoning models but remains sparse at the token level and unstable under early negative advantages. SDPG combines verifier rewards with privileged self-distillation, reference-policy regularization, and stabilizers to provide denser supervision and improve training.

  • RLVR uses rule-based outcome rewards but provides sparse sequence-level supervision across tokens and can be unstable under early negative advantages.
  • On-policy distillation supplies dense token-level guidance, while traditional approaches often require a larger, stronger teacher with substantial memory cost.
  • Self-distillation uses one model as both deployable student and privileged teacher, with the teacher additionally conditioned on demonstrations, answers, or reasoning paths.
  • SDPG integrates exact full-vocabulary privileged OPD into KL-regularized policy optimization, combining sparse verifier outcomes with dense context-conditioned teacher supervision.
  • The framework applies positive-advantage gating and a warmup-decay distillation schedule, and reports improvements over GRPO and self-distillation baselines.

2 Background

RLVR commonly uses GRPO to optimize verifier rewards, but sequence-level advantages make token credit assignment sparse. On-policy self-distillation addresses distribution mismatch with privileged conditioning, while SDPG uses full-vocabulary reverse-KL distillation and UKL reference regularization.

  • RLVR and GRPO: A rule-based verifier assigns scalar outcome rewards, commonly 1 for correct and 0 for incorrect generated reasoning sequences.
  • RLVR and GRPO: GRPO samples output groups from a frozen rollout policy and normalizes group rewards to compute sequence-level advantages without a separate value network.
  • RLVR and GRPO: Applying one scalar advantage to every token makes credit assignment sparse, while many early negative advantages can degrade foundational language capabilities under PPO clipping.
  • On-policy self-distillation: On-policy self-distillation evaluates losses on student-sampled prefixes and uses the same model as deployable student and privileged teacher under additional context.
  • On-policy self-distillation: The method uses student-to-teacher full-vocabulary reverse KL, whose detached-teacher gradient has an equivalent local policy-gradient form on fixed sampled prefixes.
  • KL regularization: Unnormalized KL adds a mass-correction term for measures that may not be normalized and is used for policy regularization against a fixed reference model.

3 Self-Distilled Policy Gradient

SDPG combines verifier-based outcome optimization, exact full-vocabulary on-policy self-distillation, and reference-policy KL regularization. It gates and schedules distillation to align privileged supervision with verifier-endorsed rollouts while controlling noise and late-training over-constraint.

  • 3.1 KL-Regularized Policy Optimization with Outcome and OPD: SDPG combines binary-verifier outcome optimization, gated full-vocabulary OPD on sampled prefixes, and KL anchoring to a fixed reference policy.The training loop samples responses from the unprivileged rollout policy, computes group-relative verifier advantages, and minimizes the combined objective.
  • 3.1 KL-Regularized Policy Optimization with Outcome and OPD: The objective adds the reward-based policy-gradient loss, a β(k)-weighted gated OPD loss, and reference-policy KL regularization.When β = 0, the method reduces to the corresponding RPG-style objective; when α = 0, it removes the reference-policy anchor.
  • 3.3 Full-Vocabulary Distillation on Sampled Prefixes: SDPG uses exact full-vocabulary student-to-teacher KL on each sampled prefix rather than a sampled-token approximation.The privileged teacher is conditioned on additional context, while the sampled prefixes come from the unprivileged rollout policy.
  • 3.3 Full-Vocabulary Distillation on Sampled Prefixes: The reverse-KL OPD gradient is locally identical to a detached-sampling policy-gradient update with a centered teacher/student log-ratio advantage.This is a gradient identity; SDPG still minimizes the explicit full-vocabulary KL objective for more accurate gradient estimation.
  • 3.5.2 β Scheduler: Positive advantage gating restricts OPD to verifier-endorsed rollouts, while β warmup and end-of-training decay limit noisy or over-constraining privileged supervision.The gate vanishes when all rewards in a group are identical, and decay phases out distillation after useful information has been internalized.

4 Experiments

The experiments evaluate SDPG and baseline algorithms on Qwen3-4B mathematical reasoning tasks using standardized training and benchmark settings. SDPG variants achieve stronger performance and more stable training dynamics than the baselines.

  • Experiment Settings: Experiments train Qwen3-4B on DAPO-Math-17k and evaluate AIME2024, AIME2025, and AMC23.The setup uses 13.9k English training samples and Gemini 2.5 Pro to generate privileged information.
  • Experiment Settings: Training runs for 400 steps with 8 responses per prompt, batch size 128, and learning rate 1 × 10−6.Optimization uses AdamW, gradient clipping, bfloat16 mixed precision, vLLM rollouts, and 8 NVIDIA H100 GPUs.
  • Training Dynamics: Figure 3 tracks benchmark accuracy, group-relative reward, actor entropy, and average response length across training.Its top row covers AIME24, AIME25, and AMC23; its bottom row covers the three training dynamics.
  • Experiment Results: Both SDPG-URKL and SDPG-UFKL reach higher final accuracies than baselines, while SDPG-UFKL avoids RLSD’s entropy collapse.The comparison spans the three benchmark panels and the actor-entropy trajectory.
  • Experiment Results: SDPG opens its accuracy gap over GRPO within the first 50 steps and reaches the high-reward plateau several hundred steps earlier.SDPG-UFKL also maintains higher actor entropy than RLSD, whose entropy collapses toward zero by step 250.

5 Related Work

Related work addresses sparse reinforcement-learning supervision through verifiable rewards, process-level signals, on-policy distillation, and self-distillation. These approaches trade off token-level density, annotation or memory costs, and teacher–student compatibility.

  • Reinforcement Learning with Verifiable Rewards: RLVR rewards correct final answers with an automatic verifier, avoiding human preference labels but potentially giving all tokens the same sparse outcome signal.Process reward models and step-level estimators provide intermediate signals, but human step annotations can be expensive.
  • On-Policy Distillation: On-policy distillation trains a student on its own trajectories while a teacher supplies dense token-level targets through KL-based objectives.Traditional OPD commonly uses a larger external teacher.
  • On-Policy Distillation: External teachers impose substantial memory usage and may provide mismatched guidance because of model heterogeneity.This motivates alternatives that use the same model under privileged knowledge.
  • Self-Distillation: Self-distillation evaluates the student policy with privileged knowledge such as solution paths or environmental feedback, then transfers that information back into the student.OPSD uses full KL divergence, OPCD decouples the on-policy strategy, and TRRD incorporates the teacher policy into an importance ratio.

6 Conclusion

The paper presents SDPG as a combination of verifier-based RLVR and exact full-vocabulary on-policy self-distillation. Its policy-gradient interpretation links reverse-KL distillation to a centered log-ratio token advantage while retaining the explicit KL objective.

  • Contribution: SDPG combines verifier-based RLVR with exact full-vocabulary on-policy distillation.The framework also uses binary verifier rewards alongside the distillation signal.
  • Interpretation: Reverse-KL OPD has a local policy-gradient interpretation with a centered log-ratio token advantage at a fixed sampled prefix.SDPG still implements the explicit full-vocabulary KL rather than replacing it with sampled-token optimization.
  • Implications: The combined objective improves credit assignment while retaining RLVR’s exploration and selection benefits.On language-model reasoning tasks, the proposed algorithms achieve better performance and stability than baseline algorithms.

A.1 Proof of Proposition 3.1

The proof establishes a local gradient identity for reverse-KL full-vocabulary OPD and derives unbiased surrogate losses for KL regularization under rollout sampling. It emphasizes that the explicit KL objective and its policy-gradient interpretation are distinct.

  • Reverse-KL OPD: For a fixed prefix, the proof treats student and teacher branches, sampled-prefix distributions, and policy-gradient coefficients as detached quantities.The local identities are evaluated at the current iterate under these detached-surrogate assumptions.
  • Reverse-KL OPD: Reverse-KL full-vocabulary OPD yields a centered log-ratio advantage in the student-side gradient.The centering follows from adding a state-dependent baseline without changing the gradient.
  • Reverse-KL OPD: The policy-gradient form interprets the explicit full-vocabulary KL objective rather than replacing it with a sampled-token objective.Sampled-token Monte Carlo estimators can nevertheless approximate the full-vocabulary expectation using detached student samples.
  • Normalized KL Terms: The appendix derives corresponding differentiable surrogate terms for both forward-KL and reverse-KL regularization.These derivations use expectations over the current policy together with teacher and reference-policy ratios.
  • Normalized KL Terms: Rollout-based KL regularization requires derived surrogate losses because the gradient of an expectation over the policy does not equal the expectation of the gradient.The proof states that directly using the original KL loss forms in this setting can produce biased estimates.

A.3 Analysis in one-step off-policy settings

The practical implementation accounts for rollout-policy drift and decomposes the distillation gradient into direct path-wise and score-function components. Because the score-function contribution is negligible relative to the outcome advantage and adds variance, the implementation uses the direct OPSD-style approximation.

  • One-step off-policy rollouts require importance ratios because the rollout model can differ from the current policy after within-step updates.The discrepancy introduces stale importance weights and changes the exact on-policy loss.
  • The final near on-policy objective combines clipped reward advantages with the corresponding KL surrogate loss.The clipping hyperparameters can match GRPO and RLSD or use asymmetric thresholds as in DAPO.
  • The full OPSD-style distillation objective produces both direct path-wise and score-function gradient paths.The direct path-wise term is the implementation used in modern RL frameworks.
  • The implementation omits the score-function term because β · DKL is negligible relative to the sequence-level outcome advantage and its variance outweighs its theoretical benefit.The retained distillation term therefore acts almost as a local shaping constraint.

B Reinforcement Learning with Self-Distillation

RLSD uses privileged self-distillation to redistribute token-level credit within verifier-grounded policy updates, whereas SDPG optimizes a separate full-vocabulary reverse-KL objective. This distinction preserves the teacher's distribution-matching role in SDPG rather than limiting it to advantage reweighting.

  • RLSD uses privileged information gain to reweight the GRPO token advantage according to the sign of the sequence-level verifier advantage.Teacher-favored tokens receive more positive credit on successful trajectories, while the ratio is inverted for negative-advantage trajectories.
  • RLSD clips and interpolates the evidence weight with the original GRPO advantage using λrlsd and ϵw.λrlsd = 0 recovers uniform GRPO advantage, while λrlsd = 1 gives fully reweighted RLSD advantage.
  • RLSD's surrogate replaces the uniform sequence-level GRPO advantage with a token-dependent advantage.This makes privileged information a credit-redistribution mechanism rather than an auxiliary KL objective.
  • Unlike SDPG, RLSD does not optimize a separate full-vocabulary distribution-matching OPD loss.The privileged model changes only the magnitude of the verifier-grounded policy-gradient update, not its sign.

C Analysis on On-policy Context Distillation

SDPG contrasts with pure on-policy context distillation by combining full-vocabulary reverse-KL self-distillation with verifier rewards and reference-policy regularization. The ablation evidence assigns dense supervision to OPD and update stability to the policy KL anchor.

  • SDPG applies reverse-KL full-vocabulary OPD between the deployable student distribution and detached privileged teacher distribution.The implementation uses the explicit objective DKL(pt∥¯qt), while its centered log-ratio form describes the local policy-gradient signal.
  • Pure self-distillation lacks a binary outcome verifier and applies teacher signals even on trajectories the verifier would reject.SDPG addresses these differences with binary outcome rewards, positive-advantage gating, and reference-policy KL regularization.
  • Removing β eliminates the OPD term, recovering RPG with binary outcome reward but without self-distillation.Removing α instead leaves binary reward and full-vocabulary OPD while removing policy KL regularization.
  • Removing β loses the early-training accuracy advantage on AIME24 and AIME25, while removing α causes response lengths around 2,000 tokens and rising entropy.The ablations support complementary roles for dense OPD supervision and policy-KL stabilization.

D.2 Robustness Across Model Scales: Qwen3-1.7B

Experiments on Qwen3-1.7B reproduce SDPG's advantage at a smaller scale. SDPG variants lead most performance columns and maintain more stable training behavior than GRPO, RLSD, and pure self-distillation.

  • On Qwen3-1.7B, SDPG-UFKL achieves the highest score in five of six Last/Best columns, while SDPG-URKL leads AIME25.Together, the SDPG variants occupy the top two positions in every column except AIME24 Best.
  • OPCD training becomes unstable after step 250, with AIME24 accuracy dropping from 0.13 to 0.02 and response length collapsing below 300 tokens.Reward also turns sharply negative during this degradation.
  • Figure 6 compares AIME24, AIME25, AMC23, reward, entropy, and response length across baseline algorithms and SDPG variants.The benchmark panels use pass@1 mean@32, while the lower panels show training dynamics.
  • SDPG-URKL and SDPG-UFKL outperform GRPO and RLSD across all three benchmarks.The accuracy gap opens within the first 50 steps and persists throughout training.
Loading 2606.04036v1…