PaperlayerView exact revision on arXiv

Source-linked AI summary

Dr. OPD: Learning What to Follow for Optimal On-Policy Distillation of Large Language Models

Zhenyu Wang, Tianze Wang, Linjun Zhang, Yifan Hu

arXiv:2609.38025v1cs.LGcs.AIcs.CLmath.OC

TL;DR

Vanilla OPD weights teacher signals equally despite their differing usefulness to student performance. Dr. OPD selects weighted supervision through bilevel optimization and iterative updates. It outperforms evaluated baselines across settings, including a 9.7-point math gain for Qwen3-1.7B over vanilla OPD.

  • Problem

    Vanilla OPD treats all teacher signals equally, although different token-level signals can have different effects on student performance.

    Source: §1 · p. 1

  • Method

    Dr. OPD formulates weighted distillation as bilevel optimization, selecting token weights by the resulting student’s reward and alternating closed-form weight updates with student updates.

    Source: §1 · pp. 2–3, §3.1 · p. 4

  • Results

    Dr. OPD outperforms evaluated baselines across settings; Qwen3-1.7B improves over vanilla OPD by +9.7 points on math and scores 35.0 vs. the Qwen3-4B teacher’s 34.1.

    Source: §1 · p. 3, §7 · p. 10

  • Takeaways & Limitations

    The comparison with OPD+GRPO and GRPD suggests that how outcome rewards are incorporated into distillation matters, not simply adding rewards.

    Source: §5.2 · p. 8

Abstract

from arXiv · show

On-policy distillation (OPD) trains a student on its own generated responses using dense, token-level supervision from a stronger teacher. Vanilla OPD treats all teacher signals equally, assuming that the teacher's supervision is equally important for every token. However, teacher signals at different tokens may have very different effects on the student's performance: some correct important reasoning errors, while others have little effect on the final answer. Motivated by this observation, we introduce Dr. OPD (OPD Done Right), which defines the optimal weighted OPD to maximize the student's performance. We formulate Dr. OPD as a bilevel optimization problem in which the student learns from weighted teacher supervision, while the weights are selected to maximize the expected reward of the resulting student. To solve Dr. OPD, we develop an efficient iterative solver that updates the token weights and student policy alternatively. At each round, it updates weights in closed form and then takes one gradient step on the resulting weighted OPD objective. Under regularity conditions, we show that this weighted update achieves a higher expected reward than a vanilla OPD update. Empirically, across strong-to-weak and same-size distillation on math and code, Dr. OPD consistently outperforms all evaluated baselines. In particular, in the strong-to-weak distillation setting, Dr. OPD improves average math performance by $9.7$ points over vanilla OPD, and enables the smaller student to surpass its larger teacher.

1 INTRODUCTION

Vanilla OPD treats every teacher signal equally, although signals can differ in how much they help the student. Dr. OPD learns token weights to improve student performance and reports gains across distillation settings.

1 INTRODUCTION

  • Vanilla OPD gives equal importance to all teacher signals, though a crucial reasoning correction and a wording preference can have different effects on the final answer.The paper motivates adapting each signal’s weight according to its usefulness for the student.Source: §1 · pp. 1–2
  • Dr. OPD assigns positive, token-specific weights to teacher signals; vanilla OPD is the special case with every weight equal to 1.A weight above 1 places more trust in the teacher at that token, while a weight between 0 and 1 places less.Source: §1 · p. 2
  • Dr. OPD formulates weighted distillation as bilevel optimization: the student trains under a given weight function, while weights are selected by the resulting student’s performance.The performance measure is an outcome reward based on generated responses.Source: §1 · p. 2
  • Dr. OPD improves Qwen3-1.7B over vanilla OPD by +9.7 points on math and surpasses the Qwen3-4B teacher, 35.0 vs. 34.1.Across the evaluated strong-to-weak and same-size settings, it outperforms the baselines.Source: §1 · p. 3, Figure 2 · p. 3
  • Its iterative solver alternates a closed-form weight update with a gradient step on the weighted OPD objective.It uses a single Jacobian–vector product to compute tokenwise gradient inner products efficiently.Source: §1 · pp. 2–3

2 WEIGHTED ON-POLICY DISTILLATION

Weighted OPD generalizes vanilla OPD by weighting teacher-token matching according to positive, token-dependent weights. Its weighted divergence remains a valid distillation objective, and reduces to the ordinary KL divergence when all weights are 1.

2 WEIGHTED ON-POLICY DISTILLATION

  • Positive token weights define a weighted KL divergence that generalizes ordinary KL and controls the importance of matching the teacher’s token probabilities.An added term ensures nonnegativity under token-dependent weighting.Source: §2 · p. 3
  • The weighted KL divergence is zero exactly when student and teacher probabilities match for every token; setting all weights to 1 recovers ordinary KL.The added term is zero in that all-ones case.Source: §2 · p. 4
  • For any positive weight function, minimizing weighted OPD still pushes the student toward the teacher, and the objective is valid.The theorem states equality only when the student and teacher probabilities match at every token for almost every context under the reference distribution.Source: §2 · p. 4

3 OPD DONE RIGHT: FIND THE BEST WEIGHT FOR OPD

Dr. OPD chooses token-level teacher weights to maximize the reward of the resulting student, rather than treating every teacher signal equally. Its iterative solver alternates student updates with reward-informed weight updates.

3.1 FORMULATION OF DR. OPD

  • Dr. OPD is a bilevel optimization: weights maximize the resulting student’s reward, while the student minimizes weighted OPD under those weights.Weights are constrained pointwise to prespecified bounds containing 1.Source: §3.1 · p. 4

3.2 ITERATIVE SOLVER FOR DR. OPD

  • The solver alternates one weighted-OPD gradient step for the student with a weight update based on the updated student’s reward.Restricting weight updates to sampled response tokens makes each step finite-dimensional.Source: §3.2 · p. 4
  • Token credit measures whether increasing trust in a teacher signal would locally improve student reward, guiding whether its weight should rise or fall.Positive credit suggests a weight above 1; negative credit suggests a weight below 1.Source: §3.2 · p. 5
  • A reward lower-bound surrogate yields a closed-form weight update, avoiding the need to parameterize or train a separate weight model.The weighting-strength parameter controls how strongly token weights may be adjusted.Source: §3.2 · p. 6
  • Under the stated conditions, a nonuniform weight update achieves higher reward than vanilla OPD.The guarantee applies when the updated weights are not all 1.Source: §3.2 · p. 6

4 ALGORITHM

The practical algorithm estimates token credits from sampled responses and outcome rewards, updates token weights, then takes a weighted-OPD gradient step. A Jacobian–vector product avoids explicitly forming a parameter gradient for every token.

4 ALGORITHM

  • Each round samples multiple responses per prompt, obtains teacher signals and outcome rewards, estimates token credits, updates weights, and takes one weighted-OPD gradient step.For verifiable math and code, rewards can come from answer checking or test execution; general tasks can use learned reward models.Source: §4 · p. 6
  • For verifiable math and code tasks, the reward-gradient estimate uses group-relative advantages over prompt-response pairs.For a general reward, the gradient can be estimated by weighting each response’s log-probability gradient by its reward.Source: §4 · p. 6
  • A single Jacobian–vector product computes the token-credit directional derivatives without explicitly forming and storing a parameter gradient for every response token.Explicitly computing all token gradients is prohibitively expensive for large models.Source: §4 · pp. 6–7
  • The algorithm normalizes estimated credits by their RMS scale to stabilize weight updates, then uses the resulting weights in the student’s gradient step.The additional computations add little overhead compared with vanilla OPD.Source: §4 · p. 7

5 EXPERIMENTAL RESULTS

Experiments compare Dr. OPD with OPD baselines across strong-to-weak and same-size Qwen3 settings on math and code. Dr. OPD leads the reported average results, and a 1.7B student surpasses its larger teacher in average math performance.

5.1 EXPERIMENTAL SETUP

  • The experiments use four Qwen3 student–teacher configurations: larger Instruct teachers for Base or Instruct students, and RL-trained teachers with same-size Base or Instruct students.The strong-to-weak teacher is Qwen3-4B; same-size teachers are RL-trained Qwen3-4B models, and all teachers are frozen.Source: §5.1 · p. 7, §5 · p. 7
  • Training uses DeepMath-103K for math and Eurus-RL-Code for code, with 50 post-training steps; evaluation reports best checkpoints using avg@16 for math and avg@4 for code.Baselines are OPD, ExOPD, GRPD, and OPD+GRPO; math benchmarks are AIME24, AIME25, AMC, Minerva Math, and OlympiadBench, while code benchmarks are HumanEval+, MBPP+, and LiveCodeBench.Source: §5.1 · pp. 7–8

5.2 DISTILLATION RESULTS

  • 1.8 math points and 6.4 code points: Dr. OPD improves over vanilla OPD for the 1.7B-Base student, with fewer successful rollouts making reward-based token credits less informative.Despite its lower starting accuracy, the 1.7B-Base student achieves the best average performance among the compared methods.Source: §5.2 · p. 8, Table 1 · p. 8
  • 35.0 vs. 34.1 average math performance: Dr. OPD lets the 1.7B student surpass its 4B teacher and leads average math and code results for both students.The 35.0 and 34.1 values compare the 1.7B student with the larger teacher.Source: Table 1 · p. 8, §5.2 · p. 8
  • 2.3 and 3.0 points over vanilla OPD: Dr. OPD leads both same-size setups and surpasses the RL-trained teacher on all five Instruct benchmarks and four of five Base benchmarks.For the Base student, average performance rises from 34.1 to 36.0.Source: Table 2 · p. 9, §5.2 · p. 9
  • Every tested positive λ improves over vanilla OPD on all five benchmarks and average performance; λ = 0.4 performs best overall.λ = 0 reduces to vanilla OPD; Table 3 varies λ for the Qwen3-1.7B student distilled from Qwen3-4B.Source: Table 3 · p. 9, §5.2 · p. 9

6 RELATED WORKS

Prior work studies on-policy distillation objectives, token-level trust and weighting, and learner-performance-based reweighting. Dr. OPD applies reward sensitivity to token-level teacher supervision, distinguishing its continuous weighting objective from sparse masking and token-logit alignment approaches.

On-policy distillation and trust weighting

  • OPD research spans divergence objectives, reverse-KL distillation, reasoning post-training, multi-teacher transfer, reward extrapolation, and teacher–student compatibility.Related weighting approaches use agreement, distribution discrepancy, compatibility, batch budgets, or outcome feedback to select or gate supervision.Source: §6 · pp. 9–10

Related approaches to selective supervision

  • SparseOPD studies how little supervision suffices, while Dr. OPD assigns continuous weights by each signal’s effect on student performance; Unmasking OPD instead aligns signals with an estimated success-improving gradient in token-logit space.Dr. OPD uses parameter-space credits and a clipped weighting rule derived from a reward lower bound relative to vanilla OPD.Source: §6 · p. 10

Bilevel optimization, reweighting, and influence

  • Dr. OPD evaluates reward sensitivity after a one-step OPD update, linking token-level supervision to bilevel optimization, data reweighting, and influence methods.Its finite weight updates come from a lower bound on reward improvement.Source: §6 · p. 10

7 CONCLUSION

Dr. OPD selects weighted distillation objectives to improve student task performance and solves this bilevel problem with alternating weight and student updates. Across evaluated setups, it outperforms existing baselines and can enable students to surpass their teachers.

7 CONCLUSION

  • Dr. OPD formulates optimal weighted OPD as a bilevel problem: the student trains under a weight function, and the weights are selected according to the resulting student’s reward.The lower-level problem trains the student; the upper-level problem selects weights.Source: §7 · p. 10
  • An iterative solver alternates closed-form weight updates with a gradient update on the resulting weighted OPD objective.Source: §7 · p. 10
  • Across all evaluated experimental setups, Dr. OPD consistently outperforms existing baselines and can enable students to surpass their teachers.Source: §7 · p. 10

A.1 SHARED TRAINING CONFIGURATION

The experiments use shared training infrastructure and frozen teachers across methods, while training data and prompts vary by task and student. A common stopping rule is applied during training and evaluation.

A.1 SHARED TRAINING CONFIGURATION

  • All methods use the same trainer, rollout engine, outcome verifier, and frozen teacher within each student–teacher setting, isolating the training-objective comparison.Runs use one node with 8×H100 GPUs; learning rates, batch sizes, rollout settings, and sequence limits are shared.Source: §A.1 · p. 14

A.2 TRAINING DATA

  • Math training uses DeepMath-103K with difficulty bands varying by student; code uses Eurus-RL-Code, except the Base student uses easier APPS and TACO subsets.The easier subsets increase groups with both correct and incorrect responses, which provide nonzero group-relative advantages.Source: §A.2 · p. 14

A.3 PROMPT TEMPLATES AND STOPPING RULES

  • Most students use their own chat template, while Qwen3-1.7B-Base uses a fixed four-shot math prompt; all methods share the same stopping rule in training and evaluation.Responses are truncated at the first completed boxed answer or next-problem delimiter, otherwise at end-of-sequence or 12,288 tokens.Source: §A.3 · pp. 14–15

B IMPLEMENTATION OF DR. OPD

The implementation constructs token credits from group-relative reward signals and uses them to weight the OPD actor update. It smooths and optimizer-scales the reward direction, then computes per-token derivatives with a forward-mode JVP.

B.2 OPTIMIZER-SCALED REWARD DIRECTION

  • The implementation smooths the raw GRPO reward-gradient estimate and scales it coordinate-wise using the actor optimizer’s transformation.The reward direction uses reward momentum and the actor’s stored second moment; weights default to one if required state is unavailable.Source: B IMPLEMENTATION OF DR. OPD · p. 15, §B.2 · p. 16

B.1 GROUP-RELATIVE REWARD SIGNAL

  • Group-relative reward advantages are zero when all responses in a group have identical rewards, and the reward direction normalizes each response contribution by its valid-token length.The implementation uses K = 8 responses per prompt.Source: §B.1 · pp. 15–16

B.3 TOKEN CREDITS, RMS SCALE, AND WEIGHT CLIPPING

  • Forward-mode JVP computes token-level directional derivatives without materializing a parameter gradient for each token.The student replica is synchronized with the actor, dropout is disabled, and derivatives use automatic differentiation.Source: §B.3 · p. 16
  • The credit scale is a token-weighted RMS across valid response tokens for one prompt, without mean subtraction; a prompt with all-zero credits assigns its valid tokens weight one.Prompt and padding tokens are excluded, while valid response tokens with zero credit are included.Source: §B.3 · p. 16

B.4 ACTOR UPDATE

  • The actor update passes detached weighted teacher signals as token advantages to a PPO-style objective, with weights and teacher signals held fixed during the update.At the reference policy, the ascent gradient is the token average of the weighted teacher corrections.Source: §B.4 · p. 17

C PROOFS AND THEORETICAL DETAILS

The theoretical analysis states assumptions for weighted distillation and the reward-improvement guarantee. It fixes a reference distribution for each optimization round and imposes integrability and local smoothness conditions.

C PROOFS AND THEORETICAL DETAILS

  • The proofs establish weighted-distillation properties and a reward-improvement guarantee after specifying the reference distribution for one optimization round.Source: C PROOFS AND THEORETICAL DETAILS · p. 17

C.1 NOTATION AND REGULARITY CONDITIONS

  • The analysis assumes bounded positive weights, holds the reference context distribution and weights fixed when differentiating the round’s distillation objective, and refreshes the reference between rounds.The feasible weights satisfy wmin ≤ w ≤ wmax, with 0 < wmin < 1 < wmax < ∞.Source: §C.1 · p. 17
  • The reference measure is an unnormalized token-occupancy measure that includes context visitation and sampled-token probability, and its total mass is the expected response length.Source: §C.1 · p. 17
  • Regularity conditions require finite expected response length, integrability and differentiability of the weighted objective, and a locally smooth reward with Lipschitz gradient.A sufficiently small student step keeps responses in the neighborhood where smoothness holds; the stated condition is ληLRM ≤ 1.Source: §C.1 · pp. 17–18
  • Unlike fixed-reference distillation loss, reward is evaluated under the student’s own response distribution, subject to differentiation and summability conditions.Source: §C.1 · p. 18

C.2.1 PROOF OF THEOREM 3

The supplement establishes properties of weighted divergence and develops local analyses for choosing and evaluating token weights. The reward and loss guarantees depend on stated conditions, including positive weights and, for one-step loss descent, uniform smoothness.

C.2.1 PROOF OF THEOREM 3

  • The corrected weighted divergence is nonnegative and is zero only when student and teacher distributions agree almost everywhere; with unit weights it reduces to vanilla OPD's reverse KL loss.For token-dependent weights, the linear correction prevents the weighted log-ratio sum from becoming negative; it cancels when weights are token-independent.Source: §C.2.1 · p. 18

C.2.2 DERIVATION OF THE SAMPLED-TOKEN GRADIENT

  • With weights and the reference sampling law fixed, the sampled-token gradient retains the full conditional expectation; at the reference parameters, it reduces to the weighted token gradients.The local token correction follows the teacher–student log-ratio unless the token's score gradient is zero; descent applies to the aggregate loss.Source: §C.2.2 · p. 19

C.2.3 DESCENT OF THE LOWER-LEVEL OBJECTIVE

  • A one-step response decreases the fixed-reference loss when the objective has a uniformly L_D-Lipschitz gradient and ηL_D ≤ 1.The loss-smoothness condition is required for this descent statement, not for the reward bounds.Source: §C.2.3 · p. 19

C.3 BILEVEL FORMULATION WITH MULTIPLE INNER SOLUTIONS

  • Exact teacher matching makes all feasible positive weightings share zero-loss minimizers; the local response analysis does not establish convergence to the global bilevel solution.When exact matching is infeasible, weighting can change the best attainable compromise and the learning trajectory before convergence.Source: §C.3 · p. 19

C.4 CREDIT AS SENSITIVITY TO THE WEIGHT FUNCTION

  • Credit measures sensitivity per unit reference measure; an atom's coordinate derivative additionally includes its mass, and zero-measure weight changes leave the objective unchanged.Source: §C.4 · p. 20

C.5 PROOF OF PROPOSITION 4

  • The reward-change bound is exact at the vanilla weights, and its first variation there agrees with the credit expression.Source: §C.5 · p. 20

C.6 CLIPPED WEIGHTS AND PROOF OF THEOREM 5

  • The surrogate's unique feasible maximizer is the clipped weight 1 + λc; the reward gain is strictly positive exactly when credit is nonzero on a set of positive measure.If credit vanishes almost everywhere, the update recovers vanilla OPD.Source: §C.6 · p. 21

D SUPPLEMENTARY EXPERIMENTAL RESULTS

The supplementary experiments show training-reward dynamics across students and tasks, including same-size distillation and changes in weighting strength. Dr. OPD maintains a training-reward advantage in the same-size setting, while larger weighting strengths can substantially alter response length.

D SUPPLEMENTARY EXPERIMENTAL RESULTS

Related papers

Loading 2609.38025v1…