PaperlayerView exact revision on arXiv

Source-linked AI summary

Reward-Aligned Reweighting for On-Policy Distillation

Haofeng Xu, Junwei Su, Lansong Diao, Wenchao Zhou, Chuan Wu

arXiv:2609.35517v1cs.LG

TL;DR

Standard OPD treats local teacher preference as a proxy for correction utility, although a decision's value depends on the student's continuation. R2-OPD uses outcome agreement and disagreement magnitude to reallocate teacher supervision. Across seven mathematical reasoning benchmarks, it achieves the highest average accuracy among compared methods in both distillation settings.

  • Problem

    Standard OPD applies token-level corrections uniformly even though local teacher preference may not match student continuation value, and trajectory outcomes do not identify individual-decision utility.

    Source: §1 · pp. 1–2

  • Method

    R2-OPD continuously reweights teacher corrections using outcome agreement and disagreement magnitude, giving reward-aligned corrections greater relative influence while retaining dense feedback.

    Source: §1 · p. 2

  • Results

    Across seven mathematical reasoning benchmarks, R2-OPD has the highest average accuracy in both settings, gaining +3.5 and +2.4 over standard OPD for 1.7B and 4B students.

    Source: §1 · p. 2

  • Takeaways & Limitations

    Under sufficient reward–utility agreement conditions, supervision reallocation improves first-order task progress over uniform OPD.

    Source: §1 · p. 2

    The stated analysis identities assume a differentiable student, positive teacher and student probabilities, and a bounded generation horizon.

    Source: §2.1 · p. 2

Abstract

from arXiv · show

On-policy distillation (OPD) trains a student language model with dense feedback from a stronger teacher on student-generated trajectories. Yet standard OPD weights token-level distillation terms uniformly, implicitly treating local teacher preference as a proxy for correction utility. A decision's task value, however, depends on how the student completes the subsequent reasoning. This mismatch can cause imitation to suppress viable student strategies or reinforce paths the student cannot reliably execute. Verified trajectory outcomes provide complementary evidence about continuation quality, but do not directly identify the utility of individual decisions. We introduce Reward-Aligned Reweighting for On-Policy Distillation (R$^{2}$-OPD), which uses outcome agreement and the magnitude of teacher--student disagreement to continuously reallocate teacher supervision. It gives reward-aligned corrections greater relative influence while retaining dense feedback, moving beyond uniform imitation and hard filtering. Our analysis formalizes the mismatch between local teacher preference and student continuation value and establishes sufficient conditions for reallocation to improve first-order task progress over uniform OPD. Across seven mathematical reasoning benchmarks, R$^{2}$-OPD achieves the highest average accuracy among the compared training methods in both cross-size and same-size distillation. It outperforms standard OPD on all seven benchmarks, with average gains of 3.5 and 2.4 percentage points for 1.7B and 4B students, respectively. An extension to code generation yields an average gain of 1.6 percentage points over standard OPD. These results highlight outcome-guided supervision allocation as an effective way to translate dense teacher feedback into stronger student performance across model scales and task domains.

1 INTRODUCTION

The introduction argues that local teacher preference does not necessarily reflect the value of a decision under the student's own continuation policy. R2-OPD uses verified outcomes and disagreement magnitude to reallocate teacher supervision, with reported gains across reasoning and code tasks.

1 INTRODUCTION

  • Uniform token-level correction coefficients encode local teacher preference, but task success depends on student continuation value, creating risks of over-correcting successful paths and over-trusting failed ones.The paper formalizes that moving a next-token distribution toward the teacher can reduce reverse KL while lowering expected reward when the student continuation policy is fixed.Source: §1 · pp. 1–2
  • Verified trajectory outcomes provide complementary evidence about continuation quality but do not identify individual-decision utility, so hard filtering can discard feedback of uncertain value.The introduction motivates using outcomes to guide correction strength rather than treating outcome agreement as a definitive token-value label.Source: §1 · p. 2
  • R2-OPD continuously combines outcome agreement with teacher–student disagreement magnitude, giving reward-aligned corrections greater relative influence while retaining dense feedback and preserving correction signs and per-response absolute coefficient mass.The analysis establishes sufficient reward–utility agreement conditions for reallocation to improve first-order task progress over uniform OPD.Source: §1 · p. 2
  • Across seven mathematical reasoning benchmarks, R2-OPD achieves the highest average accuracy in both settings, gaining +3.5 and +2.4 over standard OPD for 1.7B and 4B students, respectively.The code-generation extension gains +1.6 over standard OPD.Source: §1 · p. 2

2 PRELIMINARIES

The preliminaries define the prompt, response, and prefix notation, then describe sampled-token on-policy distillation and trajectory-level verifier outcomes. They distinguish local teacher preference signals from response-level success or failure signals.

2.1 NOTATION AND DISTILLATION OBJECTIVE

  • For prompt q and response o, s_t denotes the prefix (q, o_<t); the setup uses a trainable student, frozen teacher, fixed student rollout snapshot, and assumes differentiability, positive probabilities, and a bounded generation horizon.The response probability is expressed as the product of token probabilities across its prefixes.Source: §2.1 · p. 2

2.2 ON-POLICY DISTILLATION

  • Sampled-token OPD queries teacher log-probabilities at student-sampled tokens and uses a detached log-probability surrogate whose expected gradient matches the fixed-rollout reverse-KL objective at rollout parameters.The sampled advantage is an unbiased single-sample estimate of local reverse KL; standard OPD does not reweight its coefficients by outcomes.Source: §2.2 · p. 3

2.3 REINFORCEMENT LEARNING WITH VERIFIABLE REWARDS

  • With verifiable rewards, each sampled response receives a shared binary outcome, while the token-specific distillation signal A_t measures teacher preference for the sampled token.The response outcome indicates success or failure but does not identify any individual decision's utility.Source: §2.3 · p. 3

3 TEACHER PREFERENCE VERSUS STUDENT CONTINUATION VALUE

This section separates local agreement with the teacher from the value of an action under the student's continuation policy. It shows how outcome feedback can complement local preference when allocating supervision strength.

3 TEACHER PREFERENCE VERSUS STUDENT CONTINUATION VALUE

  • When only the next-action distribution is interpolated toward the teacher and the student's continuation policy is held fixed, a negative covariance between teacher preference and continuation value makes small positive interpolation reduce both reverse KL and expected reward.The analysis considers a finite action set with positive teacher and student probabilities.Source: §3 · p. 4
  • Local reverse KL decreases to first order whenever the distillation advantage is nonconstant, but the reward change depends on its covariance with student continuation value.Centered teacher preference determines the direction of first-order probability change; over-correction and over-trust shift probability away from above-average-value actions and toward below-average-value actions, though other actions can offset these contributions.Source: §3 · p. 4
  • Trajectory outcomes provide evidence about continuation value under student-generated continuations, complementing local teacher preference to guide the relative strength of teacher corrections.A disagreement between outcome and token-level preference does not by itself establish that an individual correction is harmful.Source: §3 · p. 4

4 REWARD-ALIGNED REWEIGHTING

R2-OPD continuously reallocates teacher correction strength using outcome agreement and disagreement magnitude, while preserving per-response absolute coefficient mass. Its task-progress benefit is conditional on reward-aligned gates covarying positively with task-projection utility.

4 REWARD-ALIGNED REWEIGHTING

  • R2-OPD adjusts token-level correction strength using outcome alignment and disagreement magnitude, then normalizes each response to preserve its original absolute coefficient mass.The gate adjusts relative correction weights; a common response-level factor restores the coefficient budget.Source: §4.1 · pp. 4–5

4.1 REWARD-ALIGNED DISTILLATION ADVANTAGES

  • Reward-aligned corrections have a positive product of signed outcome and distillation advantage; a sigmoid gate gives aligned corrections greater weight while retaining positive weight for finite sharpness.Disagreement is scaled relative to its typical within-response magnitude, and β controls gate sharpness.Source: §4.1 · pp. 4–5

4.2 TRAINING OBJECTIVE

  • Advantages and reweighting factors are detached, so gradients flow only through token log-probabilities and reward adds no separate policy-gradient term.The implementation applies the reweighted contributions through the OPD surrogate.Source: §4.2 · p. 5

4.3 STRUCTURAL PROPERTIES

  • The coefficient decomposition preserves the teacher correction's sign while outcome-dependent modulation attenuates over-correction on successful responses and over-trust on failed responses.For finite β, the reweighting retains nonzero corrections and their signs; increasing β shifts coefficient mass toward reward-aligned feedback when both groups carry mass, without implying monotonic changes to individual coefficients.Source: §4.3 · pp. 5–6

4.4 WHEN REALLOCATION IMPROVES TASK PROGRESS

  • Reweighting improves response-level first-order task projection over standard OPD if and only if gate and task-projection utility have positive covariance under the original coefficient-mass distribution.Reducing reward-inconsistent mass alone does not guarantee task progress; the gate and utility are analysis quantities, not training inputs.Source: §4.4 · p. 6

5 EXPERIMENTS

Experiments compare R2-OPD across student sizes and task domains, with controlled ablations. It leads the mathematical comparisons and also outperforms standard OPD in code generation.

5 EXPERIMENTS

  • The study evaluates reward-aligned supervision allocation across student sizes and task domains and uses controlled ablations to examine its role.Source: §5 · p. 6

5.1 EXPERIMENTAL SETTINGS

  • The experiments assess student sizes and task domains using a shared frozen teacher and DAPO-Math-17k, with β fixed at 0.001 for mathematical distillation.Evaluation reports avg@16 on AIME24/25 and avg@4 on AMC23/24, MATH500, MinervaMath, and OlympiadBench; the main 1.7B cross-size comparison uses three seeds.Source: §5.1 · pp. 6–7, Figure 1 · p. 7

5.2 MAIN RESULTS

  • 46.2 ± 0.4% for the 1.7B student and 50.8% for the 4B student accompany the highest average accuracy in both settings and wins over standard OPD on all seven benchmarks.The average gains over standard OPD are +3.5 and +2.4, respectively.Source: Table 1 · p. 7, §5.2 · p. 8

5.3 EXTENSION TO CODE GENERATION

  • R2-OPD reaches 61.2% average avg@4 in code generation, beating vanilla OPD by +1.6 and OPDVR by +1.0.It outperforms both methods on HumanEval+ and LiveCodeBench.Source: Table 2 · p. 8, §5.3 · p. 8

5.4 ABLATION STUDIES

  • Ablations favor the full outcome- and magnitude-dependent configuration: reversed alignment drops average accuracy by 3.8, while removing alignment, magnitude, or mass normalization drops it by 2.4, 1.6, and 0.9.Permuting outcomes drops accuracy by 1.7; the tested gate-sharpness values from 10^-4 to 10^-1 retain strong performance, and adding GRPO yields an average gain of +0.8.Source: Table 3 · p. 8, §5.4 · p. 8

5.5 TRAINING DYNAMICS

  • R2-OPD maintains a higher smoothed training-reward trajectory over most of training while retaining higher student policy entropy in the middle and later stages.Raw reward observations fluctuate and overlap, so the trend does not imply an advantage at every update.Source: Figure 2 · p. 9, §5.5 · p. 9

6 RELATED WORK

Related work spans on-policy distillation, selective token supervision, and outcome-guided distillation with RLVR. R2-OPD is positioned as continuously reweighting distillation coefficients using outcome alignment and disagreement magnitude.

On-policy distillation of language models

  • On-policy distillation work addresses reverse-KL transfer, student-generated contexts, divergence choices, rollout reuse, and optimization; sampled-token OPD treats teacher–student log-probability gaps as per-token rewards.ExOPD amplifies teacher–reference log-ratio rewards relative to KL regularization.Source: §6 · p. 9

Selective token supervision in OPD

  • Selective-supervision methods choose or reweight tokens using disagreement, uncertainty, support compatibility, advantage regions, or outcome-informed directions and confidence gating.Source: §6 · p. 9

Outcome-guided distillation and RLVR

  • Outcome-guided methods combine or sequence objectives, route samples, or use verifier feedback; SG-OPD routes by sign consistency and OPDVR masks outcome-inconsistent corrections, while R2-OPD continuously reweights coefficients.Source: §6 · p. 9

7 CONCLUSION

The conclusion frames effective on-policy distillation as depending on how teacher feedback is allocated, not only where it is collected.

7 CONCLUSION

  • R2-OPD addresses the mismatch between local teacher preference and student continuation value by reweighting corrections with outcome alignment and disagreement magnitude while preserving dense feedback, correction signs, and per-response absolute coefficient mass.The analysis gives sufficient reward–utility agreement conditions for improved first-order progress over uniform OPD; experiments report gains over standard OPD in mathematical reasoning and code generation.Source: §7 · p. 10

A IMPLEMENTATION DETAILS

The appendix documents training configuration, policy-loss implementation, evaluation reporting, and ablation settings.

A.1 TRAINING CONFIGURATION

  • The training configuration shares optimization and rollout settings between mathematical reasoning and code generation, uses the same prompt budget for controlled comparisons, and runs on one eight-GPU H800 machine.The listed software environment includes Ubuntu 22.04.5, Python 3.11.15, PyTorch 2.11.0, and CUDA 12.9.Source: §A.1 · p. 15, Table 4 · p. 15

A.2 POLICY LOSS AND OPTIMIZATION PROTOCOL

  • The policy loss averages per-response token means, masks prompts and padding, and uses detached snapshot log probabilities with fixed teacher scores and reweighting coefficients during accumulation.Truncated responses retain standard OPD coefficients, while zero-mass responses retain zero coefficients; outcomes enter pure distillation through the gate without an added reward-derived advantage.Source: §A.2 · pp. 15–16

A.3 EVALUATION AND REPORTING

  • The main mathematical comparison reports three-seed mean and sample standard deviation for 1.7B cross-size distillation and single-run results for 4B same-size distillation.The seven-benchmark macro-average is computed within each run before cross-seed aggregation.Source: §A.3 · p. 16

B ABLATION EXPERIMENTS

B.1 ABLATION DEFINITION

The section defines ablation and outcome-branch controls for R2-OPD and reports their comparison alongside the gate-sharpness analysis. The controls test which parts of reweighting matter while retaining teacher supervision.

B.1 ABLATION DEFINITION

  • A reversed-alignment control flips the gate’s relative emphasis without changing correction signs or verifier labels.The outcome-permuted control instead shuffles outcome labels across complete responses while preserving advantages, scales, token masks, and per-response coefficient mass.Source: §B.1 · pp. 16–17

B.2 OUTCOME-BRANCH CONTROLS

  • Success-only reweights successful responses while retaining standard OPD on failures; Failure-only does the reverse, without discarding either branch’s supervision.Both controls use the full gate and per-response normalization on the selected branch, without rescaling for branch frequency.Source: §B.2 · pp. 17–18
  • Full R2-OPD achieves the highest seven-benchmark average of 46.7, exceeding Success-only and Failure-only by 1.2 and 2.2, respectively.The comparisons reflect branch-level effects under the observed outcome distribution, not causal token-level utility.Source: Table 5 · p. 18, §B.2 · p. 18

C HYPERPARAMETER SENSITIVITY

  • Figure 3 compares inconsistent coefficient-mass share and mathematical accuracy across gate sharpness β, marking β = 10−3 and standard OPD at β = 0.The fixed-rollout analysis holds advantages, outcomes, response-local scales, and token masks constant across β; accuracy uses one run per setting.Source: Figure 3 · p. 18

C.1 SENSITIVITY OF TASK PERFORMANCE

The section examines gate sharpness, training behavior, computational cost, and the scope of the gradient analysis. It reports robust accuracy across tested sharpness values while clarifying that coefficient alignment and rollout-point identities do not alone establish reward ascent or equal runtime.

C.1 SENSITIVITY OF TASK PERFORMANCE

  • 46.7% is the highest observed accuracy in the sweep at β = 10−3, rising from approximately 42.8% at β = 10−8.Sampled values from 10−4 to 10−1 achieve roughly 46.1–46.7%, and all tested settings exceed the displayed OPD reference; stronger gating offers no consistent improvement over the default.Source: §C.1 · pp. 18–19

C.2 REWARD-INCONSISTENT MASS ALLOCATION

  • DB(10−3) ≳ 0.50 lower-bounds mean response-normalized coefficient displacement across the full batch, including unchanged fallback responses.This is not a pooled mass ratio or parameter-gradient displacement, and it does not identify which token ranges account for accuracy gains.Source: §C.2 · p. 20

D.1 RESPONSE LENGTH

  • R2-OPD and standard OPD both reach approximately 5.5k response tokens by step 40 and remain closely matched thereafter.The similar length trajectories provide evidence against longer training rollouts as the primary explanation for performance gains.Source: §D.1 · p. 20, Figure 5 · p. 20

D.2 EVALUATION TRAJECTORIES ACROSS ALL BENCHMARKS

  • R2-OPD shows early and sustained gains on several benchmarks, with visible separation on AIME24, AIME25, AMC24, and OlympiadBench.Most methods improve rapidly within 40–60 training steps, then change more gradually.Source: §D.2 · p. 20, Figure 6 · p. 21

E COMPUTATIONAL COST

  • Matched rollout batches and update schedules require no extra model passes, trainable parameters, or inference-time components for reweighting.Outcome verification may add task-dependent cost, and unchanged model-pass counts alone do not establish wall-clock equivalence or negligible overhead.Source: E COMPUTATIONAL COST · pp. 20, 22

F.1 STATEWISE REVERSE-KL GRADIENT

  • A nonpositive mean coefficient does not imply uniform suppression because the update depends on its association with the action-dependent score function.The score function’s expectation is zero.Source: §F.1 · p. 22

F.2 FIRST-ORDER EQUIVALENCE OF THE SAMPLED SURROGATE

  • The sampled-surrogate identity establishes gradient equivalence at rollout parameters, not equality of loss values or differentiation through state visitation.The equivalence also requires the sampling distribution to agree with the rollout policy used to define the distillation advantage.Source: §F.2 · p. 23

F.4 TOKEN-LEVEL ALIGNMENT AND ITS SCOPE

  • Token-level outcome alignment alone does not imply reward ascent because shared-parameter aggregation introduces cross-token inner products.The task-progress criterion additionally depends on correction utility; the aligned/inconsistent coefficient partition describes dense coefficients, not the R2-OPD update rule.Source: §F.4 · p. 23, §F.3 · p. 23

G.2 PROOF OF PROPOSITION 1

The proofs distinguish local teacher preference from student continuation value, show how outcome evidence can guide reweighting without identifying single-decision utility, and illustrate the mismatch with a two-route counterexample.

G.2 PROOF OF PROPOSITION 1

  • Both first-order quantities decrease for sufficiently small positive interpolation, establishing the local result proved in Proposition 1.The proof uses a negative variance term and smoothness.Source: §G.2 · p. 25

G.3 CENTERED PREFERENCES AND THE TWO FAILURE MODES

  • Centered preferences determine normalized action-probability shifts; neither a raw preference sign nor an individual coefficient determines the aggregate parameter update.Actions above and below average student value shift in opposite directions, while the net effect depends on probability-weighted contributions and shared token parameters.Source: §G.3 · p. 25

G.4 WHAT TRAJECTORY OUTCOMES REVEAL

  • Trajectory outcomes connect to local value only in expectation, and a sampled outcome need not identify an individual correction’s utility.The outcome compares sampled reward with a fixed threshold rather than the student’s expected value, and includes continuation-sampling noise; reweighting therefore does not guarantee positive reward change.Source: §G.4 · p. 26

G.5 AN ANALYTIC TWO-ROUTE COUNTEREXAMPLE

  • The illustrative two-route construction shows that matching teacher route choices does not transfer the teacher’s ability to complete those routes.It demonstrates the possibility of preference–value mismatch, not its prevalence in language models or that R2-OPD resolves this decision process.Source: §G.5 · p. 26

H CORE PROOFS AND THEIR SCOPE

  • The coefficient-space setup fixes each complete response’s advantages, outcome, and positive scale as the reweighting parameter varies.The stated setup also assumes positive total absolute advantage mass and exact normalization.Source: H CORE PROOFS AND THEIR SCOPE · p. 26

H.1 PROOF OF PROPOSITION 2

  • The sigmoid-gate identity holds pointwise for finite sharpness, and the reweighted coefficient retains its sign for nonzero advantages.The sign result does not determine a token’s probability change under a shared-parameter batch update.Source: §H.1 · p. 27

H.2 PROOF OF PROPOSITION 3

The proofs establish per-response sign and mass preservation with monotone reallocation, then give conditions for positive relative progress while distinguishing those guarantees from reward ascent and optimizer behavior.

H.2 PROOF OF PROPOSITION 3

  • As sharpness increases, inconsistent coefficient mass decreases while aligned mass increases, with total absolute mass preserved.The proof states boundary cases where all nonzero advantages are aligned or inconsistent; monotonicity concerns group mass, not individual coefficients.Source: §H.2 · pp. 27–28

H.3 PROOF OF PROPOSITION 4

  • A positive per-response projection gain is sufficient when aligned tokens have greater mean utility and the weighted within-group covariance is nonnegative.These are additional utility conditions, not consequences of outcome alignment alone.Source: §H.3 · p. 28

H.4 RELATIVE PROGRESS AND REWARD ASCENT

  • A positive relative reward difference does not alone ensure reward ascent from the original policy; it must also overcome any negative baseline projection.These bounds concern specified deterministic directions, not arbitrary stochastic or adaptive-optimizer updates.Source: §H.4 · p. 29

H.5 GRADIENT COMPATIBILITY AND TASK PROGRESS

  • Positive gradient compatibility permits aligned-surrogate descent for a sufficiently small step, but does not establish reward ascent or fix update magnitude.Cross-token inner products determine whether gradient contributions reinforce or cancel.Source: §H.5 · p. 29

I IMPLEMENTATION CONVENTIONS

  • Implementation equivalence is scoped to the rollout-point policy-loss gradient; inverse-length weighting and adaptive-optimizer displacements do not inherit the stated reward guarantees.The analysis does not automatically extend its reverse-KL identity to inverse-length weighting, and finite-step guarantees for direct gradient updates do not cover arbitrary AdamW displacements.Source: I IMPLEMENTATION CONVENTIONS · p. 30
Loading 2609.35517v1…