Preference-Based Self-Distillation: Beyond KL Matching via Reward Regularization

• Standard on-policy distillation uses a stronger external teacher for dense token-level signals, while self-distillation reuses the base model under added prompt context.Table 1 characterizes PBSD as combining on-policy optimization, dense supervision, token efficiency, and reward awareness without an external teacher or scorer.Source: §1 · pp. 1–2, Table 1 · p. 2
• KL-based self-distillation can destabilize training, degrade reasoning, and suppress uncertainty, self-checking, and error correction in reasoning traces.Prompt-only teacher supervision may also lack the diversity and exploratory value of a stronger external model.Source: §1 · p. 2
• PBSD combines teacher–student preference gaps with on-policy student sampling to learn a reward-reweighted target rather than directly matching the teacher.Its reward-regularized objective balances closeness to the teacher with reward maximization.Source: §1 · p. 3
• The analytic optimum reweights the teacher distribution by latent reward and improves over the original teacher under the same objective.The paper also analyzes when contextual self-distillation yields more favorable comparison geometry than external-teacher distillation.Source: §1 · p. 3
• PBSD achieves the strongest overall performance across mathematical reasoning and tool-use benchmarks at multiple model scales, with improved stability over prior self-distillation baselines.The method also retains self-distillation's token-efficiency benefits.Source: §1 · p. 3

• Standard on-policy distillation trains on student-sampled trajectories with dense teacher signals, but depends on an external teacher and can face vocabulary, output-space, and distributional-support mismatch.Source: §2 · p. 3
• Self-distillation reuses the base model as teacher and student, usually with added context, but existing methods largely retain KL-based teacher matching.PBSD instead targets a reward-reweighted teacher distribution.Source: §2 · p. 3
• Unlike RLHF and DPO-style preference optimization, PBSD constructs teacher–student preference pairs specifically for contextual self-distillation rather than using externally annotated pairs or general online updates.Source: §2 · p. 3

3.1 A Reward-Regularized Objective for On-Policy Distillation

• The objective combines KL matching to the teacher with latent-reward maximization, favoring higher-reward responses within teacher-supported behavior.The teacher is a reference policy, not a target to copy uniformly.Source: §3.1 · p. 4
• The objective’s analytic optimum reweights the teacher distribution by exp(r(x, y)/β), and is never worse than the teacher under the same objective.The inequality is strict when reward varies across teacher-supported responses on a positive-measure set of inputs.Source: §3.1 · p. 4

3.2 Solving the Objective via Preference-Based Optimization

• PBSD estimates the reward-aware target through teacher–student preference pairs, with a teacher response preferred over an on-policy response from the current student.The preference margin uses student-to-teacher log-probability ratios for the positive and negative responses.Source: §3.2 · p. 5
• The teacher is a context-augmented policy that may share the student’s base model; the distinction comes from conditioning information, not necessarily model capacity.The additional context is used to construct the teacher signal, while the student remains conditioned on the prompt alone.Source: §3.2 · p. 5
• Each online update maximizes sampled pairwise likelihood, treating the teacher response as positive and the current student response as negative.Gradients flow through the optimized student policy, increasing relative likelihood of teacher-generated responses over the student’s sampled responses.Source: §3.2 · p. 5

4.1 Pairwise MLE and Informative Comparisons

• PBSD’s finite-sample local sample complexity is governed by the smallest eigenvalue of the empirical information matrix.The theorem assumes a locally well-specified pairwise logistic model, a locally stable Gauss–Newton Hessian, bounded score-gap features, and positive minimum eigenvalue.Source: §4.1 · p. 7
• Under a frozen local comparison model, bounded score gaps, and positive full population information, the local estimation rate is O(n^-1/2) with high probability for sufficiently large n.The empirical information is bounded away from zero under those conditions.Source: §4.1 · p. 7
• More diverse teacher positives can add score-gap directions, while moderate teacher–student margins preserve logistic curvature and strengthen the empirical Hessian.The curvature is maximized at zero margin and decays as the absolute margin grows.Source: §4.1 · p. 7
• Contextual self-distillation may offer more favorable comparison geometry because teacher and student share a base model, whereas external teachers can create larger shifts and saturated margins.Richer score-gap directions and moderate margins jointly increase the empirical information matrix’s minimum eigenvalue and tighten the estimation bound.Source: §4.1 · p. 7

5.1 Experiment Setup

• Experiments test three Qwen3 scales on math and tool use, using a shared 500-step budget and Avg@12 for math versus top-1 accuracy for tool use.Math uses up to 30K OpenThoughts examples; the ToolAlpaca-based setup uses 4,046 training and 94 test examples.Source: §5.1 · p. 8

5.2 Main Results

• Across three Qwen3 scales, PBSD has the strongest average results in math and the strongest tool-use result at every scale.The reported averages and tool-use results span Qwen3-8B, Qwen3-4B, and Qwen3-1.7B.Source: §5.2 · p. 8, Table 2 · p. 9

5.3 Sample Efficiency and Training Dynamics

• PBSD stays stable and reaches the strongest final AIME25 performance, while OPSD peaks early and then declines.Figure 3A plots Avg@12 against training steps.Source: Figure 3 · p. 9, §5.3 · p. 9
• PBSD achieves a favorable AIME25 accuracy–token trade-off, using far fewer tokens than GRPO while remaining competitive with OPSD.Figure 3B plots accuracy against cumulative generated tokens on a symlog x-axis.Source: Figure 3 · p. 9, §5.3 · p. 10
• On 30 selected AIME25 problems, OPSD and PBSD recover substantial base-model majority-vote errors; refreshing the teacher every five steps shows no clear benefit.The fixed teacher at initialization is already effective in this setting.Source: §5.3 · p. 10

6 Limitations

• When privileged context is weak, noisy, or only marginally informative, the induced teacher can be weak and the benefit of self-distillation is reduced.The reported experiments focus on reasoning-centric domains and structured tool use; broader open-ended behaviors warrant further study.Source: §6 · p. 10

7 Conclusion

• PBSD learns a reward-reweighted target policy instead of uniformly imitating the teacher, while combining self-distillation’s token efficiency with improved training stability.The paper also provides a local statistical analysis of when contextual self-distillation yields more favorable comparison geometry than external-teacher distillation.Source: §7 · p. 10

A Additional Discussion of Limitations

• Weak or marginal privileged context can yield a weak teacher and reduce self-distillation benefits; broader open-ended behaviors remain for further study.PBSD targets the setting where the teacher is a better-conditioned version of the same base model, rather than a strong external oracle.Source: A Additional Discussion of Limitations · p. 14

B Appendix Overview

• Appendix D derives the reward-tilted teacher optimum and its objective-level comparison, Appendix E details the statistical analysis, and Appendix F collects experimental protocols and additional analyses.Appendix F includes broader-domain transfer, context sensitivity, checkpoint behavior, and teacher-size comparisons.Source: B Appendix Overview · p. 15

C Appendix for Related Work

• Prior self-distillation methods largely use KL-based matching, while reward-aware reweighting of the teacher distribution remains an open question that PBSD addresses.The related-work discussion distinguishes PBSD’s contextual-teacher focus from DPO as a broader preference-optimization alternative to RLHF.Source: C Appendix for Related Work · p. 16

D.1 Details for the Reward-Regularized Distillation Objective

• For β > 0, strict concavity makes the reward-regularized objective’s optimizer unique and recoverable from first-order conditions.The derivation applies the probability-normalization constraint using a Lagrange multiplier.Source: §D.1 · p. 17
• The optimum exponentially reweights the teacher distribution by response reward, so the teacher serves as a reference measure rather than the final target.The derivation normalizes the reward-weighted teacher probabilities to form the target distribution.Source: §D.1 · p. 17
• The reward-tilted policy is never worse than the teacher under the same objective, and is strictly better when reward varies across teacher-supported responses.Source: §D.1 · p. 18

E Appendix for Theoretical Analysis

• The theoretical analysis asks which teacher-generated positive and on-policy student negative comparison properties strengthen PBSD’s policy-estimation guarantee.It focuses on teacher-supported directional diversity and moderation of preference margins.Source: E Appendix for Theoretical Analysis · p. 19

E.1 Details for Online PBSD Objective

• The online-objective appendix derives the sample-level gradient and local Hessian form used in the main-text PBSD analysis.It identifies a Gauss–Newton component and averages pairwise contributions to obtain the Hessian expression.Source: §E.1 · p. 19

E.2 Details for Pairwise MLE and Informative Comparisons

• The proof uses local regularity conditions, including bounded score-gap features and stable curvature near the target parameter.These conditions support the local inverse-Hessian argument.Source: §E.2 · pp. 19–20
• Violating bounded-feature or curvature assumptions can increase score noise or parameter perturbations, so the result is a local guarantee rather than a global robustness claim.Source: §E.2 · p. 20
• Positive population information requires non-saturated teacher–student margins and score-gap coverage across locally identifiable directions.The condition rules out comparisons that are uninformative either through saturated margins or insufficient directional coverage.Source: §E.2 · p. 21
• The strengthened result freezes the comparison distribution at the target parameter, rather than analyzing the full adaptive online trajectory.The frozen pairs draw positives from the teacher and negatives from the student policy at the target parameter.Source: §E.2 · pp. 20–21
• For fixed d, μ, and failure probabilities, the local bound is O(n^-1/2) under the frozen-model population-information condition.The bound vanishes as the number of comparison pairs n grows.Source: §E.2 · pp. 21, 23

F Appendix for Experiment

• The experiment appendix follows the OPSD protocol where applicable and details datasets, training, evaluation, case studies, and ablations.Source: F Appendix for Experiment · p. 25

F.1 Datasets

• Experiments use Qwen3 instruct models at 1.7B, 4B, and 8B scales across mathematical reasoning and tool use.Source: §F.1 · p. 25
• Math training samples up to 30K OpenThoughts problem–solution pairs and evaluates on AIME 2024, AIME 2025, and HMMT 2025 using Avg@12.Source: §F.1 · p. 25
• The ToolAlpaca-based tool-use study follows Shenfeld et al. [21], pairs user queries with tool/API information, and evaluates correct tool-call generation on a split with 4,046 training and 94 test examples.Source: §F.1 · p. 25

F.2 Data Visualization

• The data examples illustrate the key–value format consumed by training, including tool requests and mathematical problems with reasoning traces and answers.Source: §F.2 · pp. 25–26

F.3 Detailed Training Configuration

• In PBSD, the on-policy student rollout is the negative sample and the better-conditioned context-augmented teacher supplies the positive sample.Source: §F.3 · p. 27
• Methods share aligned training settings where applicable, while PBSD is matched to OPSD in optimization and rollout configuration to focus comparison on the objective.PBSD replaces token-level KL matching with pairwise preference-based self-distillation and keeps its teacher fixed at the initial checkpoint.Source: §F.3 · pp. 27–28

F.4 Detailed Evaluation Configuration

• Math evaluation samples 12 responses per prompt and reports Avg@12, while tool-use evaluation reports top-1 accuracy on a 94-example test split.Source: §F.4 · p. 28

F.5 Case-by-Case Study

• The AIME25 case study compares majority-vote correctness for the base student, expert-demonstration teacher, OPSD, and PBSD across 30 selected problems.The table marks each problem-method result as correct or incorrect.Source: Table 3 · p. 28, §F.5 · p. 28

F.6 Capability Gain from Expert Demonstrations

• The hard AIME24 study records problem-level accuracy and completion lengths for the base student and demonstration-conditioned teacher, plus completion lengths for OPSD and PBSD.The comparison is designed to assess whether PBSD preserves longer completions while improving accuracy.Source: Table 4 · p. 29, §F.6 · p. 29

F.7 Teacher Update Frequency Study

• The teacher-update ablation compares fixed initialization against refreshing the teacher every five steps on AIME24, AIME25, HMMT25, and their average.Source: §F.7 · p. 29, Table 5 · p. 29

F.8 Sensitivity to Privileged Context

• With sparse or imperfect teacher context, PBSD remains close to the default setting on Qwen3-4B mathematical reasoning.The compared contexts are ground-truth reasoning plus answer, answer only, and noisy reasoning plus ground-truth answer.Source: §F.8 · p. 30, Table 6 · p. 30

F.9 Broader-Domain Transfer Beyond Math and Tool Use

• PBSD achieves the strongest performance on HumanEval+, LiveCodeBench, and IFEval, while remaining close to the best method on MBPP+.Source: §F.9 · p. 30, Table 7 · p. 30

F.10 Longer-Horizon Training and Checkpoint Selection

• Peak and last-checkpoint AIME25 results are close, and training beyond 500 steps does not materially change the conclusions.Table 8 reports the latest, best-within-500, and later-checkpoint results.Source: §F.10 · p. 30, Table 8 · p. 30

F.11 OOD Capability Preservation and KL-to-Base Diagnostics

• PBSD preserves baseline-level performance on broader capabilities while improving the target math domain, with OOD token-level KL of 0.08 indicating proximity to the base model.The evaluation compares the math-trained Qwen3-4B checkpoint with its corresponding base model.Source: §F.11 · pp. 30–31, Table 9 · p. 31

F.12 β Sensitivity

• PBSD remains stable across a reasonably wide range of β, with smaller values slightly better in this setting and β = 0.1 a competitive default.The sensitivity analysis is on Qwen3-4B mathematical reasoning.Source: §F.12 · p. 31, Table 10 · p. 31

F.13 Comparison to Larger Teachers and Teacher–Student KL

• The larger external teacher improves AIME25 slightly, while the contextual self-teacher does better on AIME24 and HMMT25; their teacher–student KL values are 0.188 and 0.258, respectively.The KL comparison is for a Qwen3-8B teacher with privileged context c versus a larger teacher without privileged context.Source: §F.13 · p. 31, Table 11 · p. 31