Self-Distilled Policy Gradient

Technical Summary: Self-Distilled Policy Gradient (SDPG)

Problem Statement
Reinforcement Learning with Verifiable Rewards (RLVR) has become a standard for post-training Large Language Models (LLMs) on reasoning tasks like mathematics and code generation. Algorithms such as Group Relative Policy Optimization (GRPO) optimize against rule-based outcome rewards (e.g., correct/incorrect answers) without requiring human preference annotations. However, RLVR faces two primary limitations:

1. Sparse Rewards: The reward signal is applied at the sequence level, making it difficult to assign credit to specific reasoning steps (tokens), leading to inefficient exploration.
2. Training Instability: Early training stages often suffer from negative advantages, causing instability and potential degradation of foundational language capabilities.

While On-Policy Distillation (OPD) has been proposed to provide dense, token-level supervision, traditional approaches rely on external, larger teacher models, which incur high memory costs and introduce distribution mismatches. Recent "self-distillation" methods, where the model acts as both student and teacher using privileged context (e.g., ground-truth solutions), address the memory issue but risk reinforcing locally plausible tokens on globally incorrect trajectories if not properly constrained. Furthermore, pure self-distillation can lead to mode collapse and reduced exploration.

Methodology: SDPG Framework
The paper proposes Self-Distilled Policy Gradient (SDPG), a framework that integrates exact full-vocabulary on-policy self-distillation with KL-regularized policy optimization. The core objective combines three components:

1. Outcome-Based Policy Gradient ($L_{out}$):
   SDPG retains the sparse binary reward signal from a verifier. It utilizes a group-relative advantage ($A_{out}$) normalized within a group of rollouts to stabilize learning, similar to GRPO but formulated as a direct REINFORCE-style surrogate loss without PPO-style importance ratio clipping.

2. Exact Full-Vocabulary On-Policy Distillation ($L_{OPD}$):
   Instead of sampling tokens to approximate the teacher signal, SDPG computes the exact reverse Kullback-Leibler (KL) divergence between the student distribution ($p_t$) and a privileged teacher distribution ($q_t$) conditioned on additional context $c$ (e.g., solutions or reasoning paths).

   • Gradient Identity: The authors derive a key theoretical result (Proposition 3.1): the gradient of the reverse-KL full-vocabulary loss with respect to the student policy is locally identical to a detached-sampling policy gradient update. The "advantage" in this update is a centered log-ratio of the teacher and student probabilities.
   • Gating Mechanism: To prevent the privileged teacher from reinforcing incorrect trajectories, the distillation loss is gated by the outcome advantage. Distillation is only applied when the rollout is verified as correct ($A_{out} > 0$).

3. Reference Policy Regularization ($L_K$):
   To prevent the policy from drifting too far from its initial capabilities or collapsing, SDPG anchors the updated policy to a fixed reference policy ($\pi_{ref}$). The paper employs Unnormalized KL (UKL) divergence (forward or reverse) for this regularization, which offers more elegant and symmetric gradients compared to standard normalized KL, particularly when dealing with unnormalized measures.

Stabilization Strategies
To manage the noise of the privileged signal and the trade-off between exploration and exploitation, SDPG introduces two stabilizers:

• Positive Advantage Gating: Ensures distillation only occurs on trajectories endorsed by the verifier.
• Warmup-Decay Schedule ($\beta(k)$): The distillation coefficient $\beta$ is gradually increased (warmup) to allow the model to find correct trajectories before relying on distillation, and then decayed (decay) toward the end of training. This prevents the privileged signal from over-constraining the model once it has internalized the necessary knowledge, encouraging final-stage exploration.

Key Contributions

1. Theoretical Derivation: The paper derives an exact local policy-gradient form for reverse-KL full-vocabulary on-policy distillation, showing it is equivalent to a detached-sampling update with a centered log-ratio advantage.
2. Novel Objective: It proposes a unified objective combining binary outcome rewards, gated full-vocabulary privileged distillation, and unnormalized KL regularization against a fixed reference.
3. Stabilization Techniques: The introduction of positive-advantage gating and a warmup-decay schedule for the distillation coefficient to balance dense supervision with exploration.

Experimental Results
The authors evaluated SDPG on mathematical reasoning benchmarks (AIME2024, AIME2025, AMC23) using Qwen3-4B.

• Performance: Both SDPG variants (using Unnormalized Forward KL and Unnormalized Reverse KL) outperformed baseline GRPO and the self-distillation baseline RLSD across all benchmarks.
• Stability: SDPG variants reached high-reward plateaus significantly faster than GRPO. Notably, SDPG maintained higher actor entropy throughout training compared to RLSD, which exhibited entropy collapse (a sign of mode collapse) around step 250.
• Efficiency: SDPG methods produced response lengths that were sufficient for multi-step reasoning but more concise than the verbose outputs of GRPO.

Significance
The paper claims that SDPG successfully bridges the gap between sparse-reward RLVR and dense-supervision self-distillation. By treating self-distillation as a policy gradient update and rigorously controlling the privileged signal through gating and scheduling, SDPG achieves superior stability and performance. It demonstrates that dense, token-level supervision derived from the model's own privileged context can effectively guide reasoning models without the memory overhead of external teachers or the instability of pure self-distillation.