Technical Summary: Preference-Based Self-Distillation (PBSD)

Problem Statement

On-Policy Distillation (OPD) offers an efficient alternative to Reinforcement Learning (RL) for post-training large language models (LLMs) by providing dense, token-level training signals. However, standard OPD relies on a stronger, external teacher model, creating bottlenecks in computational efficiency and scalability due to the need for separate architectures and shared vocabularies.

To address this, on-policy self-distillation has emerged, where the same model serves as both teacher and student, with the teacher instantiated by conditioning the model on additional privileged context (e.g., expert demonstrations or reasoning traces). Despite its potential, existing self-distillation methods suffer from two critical limitations:

1. Training Instability and Performance Degradation: Current approaches rely on direct Kullback-Leibler (KL) divergence alignment to the context-augmented teacher. Recent analyses suggest this can actively degrade reasoning performance over time by suppressing epistemic verbalization (uncertainty, self-checking), leading to overly confident, shorter reasoning traces.
2. Lack of Explorative Diversity: Treating a context-augmented version of the same model as the teacher assumes the teacher is a strong oracle. However, since the teacher and student share identical parameters, the supervision often lacks the diversity and exploratory value that a truly stronger external model would provide.

The work argues that simply copying the teacher distribution via KL alignment is suboptimal, as it ignores the latent reward structure of the responses.

Methodology: Preference-Based Self-Distillation (PBSD)

The authors propose Preference-Based Self-Distillation (PBSD), a framework that reframes on-policy self-distillation through a reward-regularized perspective, rather than using fixed teacher KL alignment.

1. Reward-Regularized Objective

Instead of minimizing only the reverse KL divergence between the student $\pi_\theta$ and the teacher $\pi_{\text{teach}}$, PBSD optimizes a joint objective that balances KL alignment with reward maximization:
$$\max_{\pi_\theta} \mathbb{E}_{x \sim \mathcal{D}} \mathbb{E}_{y \sim \pi_\theta(\cdot|x)} [r(x, y)] - \beta D_{\text{KL}}(\pi_\theta(\cdot|x) \parallel \pi_{\text{teach}}(\cdot|x))$$
where $r(x, y)$ is a latent target reward and $\beta$ controls the regularization strength.

Theoretical Insight: The authors derive that the analytical optimum of this objective is not the teacher distribution itself, but a reward-tilted teacher distribution:
$$\pi^*(y|x) = \frac{\pi_{\text{teach}}(y|x) \exp(r(x, y)/\beta)}{Z(x)}$$
This target policy is provably superior to the original teacher under the reward-regularized objective, provided the reward varies over the responses supported by the teacher. The teacher serves as a support distribution (inductive bias), while the reward reweights probability mass toward higher-quality behaviors.

2. Optimization via Preference Learning

Since the latent reward $r(x, y)$ is unobserved, PBSD optimizes the objective using preference-based learning (inspired by DPO).

• Teacher Construction: The teacher $\pi_{\text{teach}}(y|x)$ is defined as the same base model conditioned on privileged context $c$ (e.g., $\pi(y|x, c)$).
• Preference Pairs: For each input $x$, a positive sample $y^+$ is generated from the teacher (better conditioned) and a negative sample $y^-$ from the current student (on-policy).
• Loss Function: The method maximizes the probability that the teacher response is preferred over the student response using a Bradley-Terry model:
  $$\max_{\theta} \mathbb{E} \left[ \log \sigma \left( \beta \left( \log \frac{\pi_\theta(y^+|x)}{\pi_{\text{teach}}(y^+|x)} - \log \frac{\pi_\theta(y^-|x)}{\pi_{\text{teach}}(y^-|x)} \right) \right) \right]$$
  This formulation effectively optimizes the "preference gap" between teacher and student samples while maintaining on-policy sampling, avoiding the need for an explicit reward model.

3. Statistical Analysis

The work provides a statistical analysis of the induced preference learning problem. It finds that contextual self-distillation is theoretically preferable to learning from an external teacher when:

• The positive samples generated by the teacher provide diverse score-gap directions (covering informative directions beyond the student's current distributions).
• The preference margins between teacher and student remain moderate.
  The analysis shows that local sample complexity is determined by the smallest eigenvalue of the empirical information matrix. Contextual self-distillation naturally maintains moderate margins (avoiding saturation) and offers richer directions than pure self-copying, leading to tighter estimation bounds.

Main Contributions

1. Reward-Regularized Objective: Introduction of a self-distillation objective that balances proximity to the teacher with reward maximization, and proof that the optimal target is a reward-reweighted teacher distribution rather than the teacher itself.
2. PBSD Framework: Proposal of a novel, preference-based on-policy self-distillation framework that optimizes preference gaps between teacher and student samples without requiring an external teacher or reward model.
3. Theoretical Justification: Provision of a statistical analysis establishing conditions under which contextual self-distillation outperforms distillation from external teachers, specifically regarding the trade-off between score-gap diversity and margin moderation.
4. Empirical Validation: Demonstration of consistent improvements across multiple model sizes (1.7B, 4B, 8B) on benchmarks for mathematical reasoning and tool use.

Experimental Results

The authors evaluated PBSD on mathematical reasoning (OpenThoughts dataset, evaluated on AIME 2024/2025, HMMT 2025) and tool use (ToolAlpaca-based setup) using Qwen3 models.

• Performance: PBSD consistently achieved the strongest average performance across all model sizes and tasks.
  • On Qwen3-8B, PBSD achieved a mathematical average of 65.2, outperforming the base model (+3.7) and the strongest OPSD baseline (+1.0).
  • On Qwen3-4B and 1.7B, PBSD also achieved the highest averages, outperforming baselines such as GRPO, DAPO, and OPSD.
  • In tool use, PBSD achieved the best Top-1 accuracy across all sizes (e.g., 72.0% for 8B).
• Stability: Unlike OPSD, which often peaks early and then declines, PBSD showed stable training progress throughout the entire 500-step budget.
• Token Efficiency: PBSD retained the token efficiency of self-distillation (one rollout per prompt) while achieving performance comparable to or better than RL-based methods (such as GRPO), which require massive token budgets.
• Ablation Studies:
  • Fixed Teacher: It was found that keeping the teacher model fixed at initialization is sufficient and effective; updating the teacher every 5 steps offered no clear advantage.
  • Exploration: PBSD preserved exploratory reasoning behavior (moderate completion lengths) better than direct KL alignment methods and avoided collapsing into overly short, overly confident traces.

Significance and Claims

The work claims that PBSD closes a fundamental gap in post-training paradigms by combining the token efficiency of self-distillation with the stability and reward awareness of preference optimization.

• Beyond KL Alignment: The work challenges the assumption that self-distillation must be a direct distribution alignment problem. It posits that the teacher should define the support of useful behaviors, while optimization should selectively reinforce higher-quality responses through reward reweighting.
• Scalability: By eliminating dependence on external teacher models or reward models, PBSD offers a scalable path to democratizing powerful open-source agents, particularly in scenarios involving multi-step reasoning and tool use.
• Robustness: The framework provides a more robust target distribution that preserves exploratory reasoning and allows for more stable optimization compared to previous self-distillation baselines.

The authors conclude that reward-aware self-distillation represents a promising paradigm for the stable and scalable post-training of reasoning-oriented LLMs, effectively bridging the gap between dense supervision and reward optimization without external dependencies.