Technical Summary: $\beta$-OPSD

1. Problem Statement

On-policy self-distillation (OPSD) has emerged as a promising technique for improving the reasoning capabilities of large language models (LLMs). In standard OPSD, a student model learns from its own sampled trajectories while receiving supervision from a "privileged teacher" (the same model architecture with access to ground-truth solutions or verified traces) available only during training.

However, the authors identify that vanilla OPSD is often brittle in practice, requiring substantial engineering effort to function reliably. The core issue is that standard OPSD treats the teacher as the unique optimization target, implicitly fixing the regularization strength that anchors the student to its initial or reference policy. This rigid formulation lacks a mechanism to control how aggressively the student moves away from its reference policy toward the teacher, often leading to instability or suboptimal convergence.

2. Methodology

The paper proposes $\beta$-OPSD, a framework that reframes on-policy self-distillation as a broader family of KL-regularized policy optimization problems. The methodology proceeds in three main stages: deriving the theoretical target, converting it to a practical distillation objective, and refining the credit assignment.

2.1. Theoretical Derivation: $\beta$-OPSD as Policy Optimization

The authors demonstrate that vanilla OPSD is precisely the $\beta = 1$ member of a KL-regularized policy optimization family. They define the objective function $J_\beta(\theta)$ as:
$$J_\beta(\theta) = \mathbb{E}_{y \sim \pi_\theta} \left[ \log \frac{p_T(y|x,c)}{\pi_{\text{ref}}(y|x)} \right] - \beta D_{KL}(\pi_\theta(\cdot|x) \parallel \pi_{\text{ref}}(\cdot|x))$$
Where:

• $\pi_\theta$: The student policy.
• $p_T$: The privileged teacher policy (conditioned on problem $x$ and privileged info $c$).
• $\pi_{\text{ref}}$: A reference policy (e.g., the initial student).
• $\beta$: A controllable regularization parameter.

Key Insight: When $\beta=1$, the reference terms cancel out, recovering the standard reverse KL minimization toward the teacher. For $\beta > 1$, the objective becomes more conservative, penalizing deviations from the reference policy more heavily.

2.2. Optimal Policy and Geometric Interpolation

The authors derive the closed-form optimal policy $\pi^\star_\beta$ for this objective:
$$\pi^\star_\beta(y|x,c) = \frac{\pi_{\text{ref}}(y|x)^{1-1/\beta} p_T(y|x,c)^{1/\beta}}{Z_\beta(x,c)}$$
This solution represents a geometric interpolation between the reference policy and the privileged teacher. The parameter $\beta$ (specifically its inverse $1/\beta$) acts as a weight determining the position along this path.

• $\beta = 1 \implies$ Target is the teacher.
• $\beta \to \infty \implies$ Target approaches the reference policy.

2.3. Practical Training: Scheduled Logit Interpolation

Directly optimizing the RL objective is computationally expensive due to the intractable normalizer $Z_\beta$ and high-variance gradients. Instead, $\beta$-OPSD converts the theoretical solution into an efficient distillation target.

• Local Approximation: Since the sequence-level normalizer is intractable, the method approximates the optimal policy at each decoding prefix $h_t$ by linearly interpolating the logits of the reference and teacher:
  $$\tilde{p}_{\beta_k}(\cdot | h_t, c) = \text{softmax}\left( \left(1 - \frac{1}{\beta_k}\right) z_{\text{ref}} + \frac{1}{\beta_k} z_T \right)$$
• Curriculum Scheduling: Rather than fixing $\beta$, the method employs a schedule where the teacher weight $w_k = 1/\beta_k$ increases linearly from $w_{\text{start}}$ to $w_{\text{end}}$ over training steps. This creates a smooth curriculum, starting with targets closer to the student and gradually incorporating stronger teacher guidance.

2.4. Return-to-Go Credit Assignment

Standard token-level OPSD updates are "myopic," treating each token's loss independently. This ignores how an early token choice affects future prefixes and the subsequent mismatch with the target.
To address this, $\beta$-OPSD introduces return-to-go (RTG) credit assignment. The loss for a token $t$ is weighted by the cumulative future mismatch (discounted return) along the trajectory:
$$G_{t, \gamma} = \sum_{s=t}^{T} \gamma^{s-t} \rho_s$$
where $\rho_s$ is the log-ratio mismatch at step $s$. This ensures that early tokens receive credit (or blame) for the distributional mismatch they induce later in the sequence, aligning the token-level updates with the sequence-level objective.

3. Key Contributions

1. Theoretical Unification: The paper establishes that vanilla OPSD is the $\beta=1$ case of a broader KL-regularized policy optimization family, making the implicit regularization choice explicit and controllable.
2. $\beta$-OPSD Framework: It introduces a generalized objective where the optimal policy is a geometric interpolation between the reference and teacher, realized practically via scheduled logit interpolation.
3. Efficient Training Algorithm: The authors propose a method to approximate the expensive RL-derived solution using inexpensive token-level distillation, avoiding the need for complex RL machinery like advantage estimation.
4. Sequence-Level Alignment: The integration of return-to-go credit assignment corrects the myopic nature of standard token-KL updates, providing an unbiased estimator for the sequence-level objective.

4. Experimental Results

The method was evaluated on mathematical reasoning benchmarks (AIME 2024, AIME 2025, HMMT 2025) using Qwen3 models ranging from 1.7B to 8B parameters.

• Performance Gains: $\beta$-OPSD consistently outperformed vanilla OPSD, SFT, and GRPO baselines.
  • On Qwen3-1.7B, $\beta$-OPSD improved the average performance (avg@12) by 5.74 percentage points over vanilla OPSD, with specific gains of 9.16 points on AIME 2024.
  • On larger models (4B and 8B), it maintained competitive or superior performance, improving the overall average by 1.76 and 1.66 points respectively.
• Ablation Studies:
  • Target Distribution: Replacing the fixed teacher target with the scheduled logit-interpolant target yielded significant gains (e.g., +6.03 points on AIME 2024), confirming the benefit of a smooth transition.
  • Credit Assignment: Using return-to-go weights instead of local token gradients improved performance (e.g., +5.55 points on AIME 2025), validating the importance of future-aware updates.
  • Reference Construction: The best performance was achieved using a dynamic student (current model) and a fixed teacher (privileged teacher), suggesting that tracking the evolving student distribution while maintaining a stable supervisory anchor is optimal.

5. Significance and Claims

The paper claims to provide a "principled route from self-distillation to policy optimization and back." Its primary significance lies in:

• Stability: By treating $\beta$ as a controllable parameter, the method addresses the brittleness of vanilla OPSD, allowing for smoother optimization that trades off proximity to the reference policy against teacher guidance.
• Efficiency: It achieves the benefits of policy optimization (sequence-level objectives, KL regularization) without the computational cost of full RL, retaining the efficiency of standard distillation.
• Generalizability: The approach offers a mechanism to reduce the distributional gap between student and teacher, which is particularly beneficial for smaller models where this gap is larger.

The authors conclude that smooth target paths and trajectory-level credit assignment are effective principles for on-policy distillation, offering a robust alternative to direct teacher imitation.