Organizations: State Key Laboratory of Novel Software Technology, Nanjing University, Nanjing, China · Shanghai Artificial Intelligence Laboratory, Shanghai, China
On-policy distillation (OPD) has recently emerged as a popular post-training paradigm for large language models (LLMs), providing an efficient way to transfer the knowledge and capabilities of teacher models into student models. However, teacher guidance on student-generated prefixes is not always reliable. Training should optimize the model to generate responses that are more likely to be correct, or equivalently, to get higher outcome rewards. But during OPD, the teacher model may provide guidance that discourages the student from moving toward correct trajectories or moves the student toward incorrect ones, which is misaligned with outcome reward. Such misaligned guidance is unreliable, as it would mislead the optimization process and ultimately degrade model performance. To mitigate misaligned teacher guidance, we propose Reward-Aligned On-Policy Distillation (RA-OPD). The key insight is to keep only trajectories whose induced updates move the student toward correct trajectories or discourage the student from moving toward incorrect ones. Specifically, for each sampled trajectory, RA-OPD checks whether its trajectory-level distillation return is consistent with its outcome reward and then filters out the misaligned trajectories. RA-OPD selects more reliable trajectories to improve student model performance without requiring additional computational cost. We evaluate RA-OPD on math and code benchmarks using models from the Qwen3 family and the DeepSeek-R1 family. Across seven math benchmarks and three code benchmarks, RA-OPD significantly outperforms standard OPD and other tested OPD variants.
Figures & tables
Figure 1: Mean avg@ k across seven math reasoning benchmarks for Qwen3-8B-Base distilled from Qwen3-8B-Base-GRPO-Math. RA-OPD achieves the strongest average performance among the tested OPD methods.
Figure 2: RA-OPD detects unreliable teacher guidance by comparing each trajectory’s trajectory-level distillation return with its outcome reward and filters out reward-misaligned trajectories before policy optimization.
avg@ k
Student
Method
AIME24
AIME25
AIME26
AMC
MATH
Minerva
Olympiad
Avg.
Qwen3-4B-Base
Base Model
8.02
5.31
6.15
30.80
54.07
22.79
25.22
21.77
OPD
21.04
19.17
14.27
57.83
84.67
34.33
53.44
40.68
ExOPD
25.73
20.42
18.65
58.81
83.65
36.63
52.63
42.36
Uni-OPD
21.67
18.12
18.85
58.81
83.88
37.04
53.52
41.70
RA-OPD
27.08
23.75
21.77
65.36
87.98
39.71
55.52
45.88
Table 1: Math reasoning results (%) for the Qwen3 family. We report avg@ k and pass@ k , with k=32 for AIME and AMC and k=8 for all other benchmarks. Within each student block, the best and second-best results are shown in bold and underlined, respectively.
avg@ k
Student
Method
AIME24
AIME25
AIME26
AMC
MATH
Minerva
Olympiad
Avg.
DeepSeek-R1- Distill-Qwen-7B
Base Model
54.58
40.00
47.50
82.83
94.27
37.50
67.15
60.55
OPD
57.19
46.15
55.31
86.60
95.15
38.51
72.09
64.43
ExOPD
59.79
48.75
57.29
86.94
95.43
38.79
71.85
65.55
Uni-OPD
60.83
51.77
60.21
86.07
95.55
38.79
72.65
66.55
RA-OPD
65.62
55.31
64.17
88.70
96.53
41.13
73.93
69.34
Table 2: Math reasoning results (%) for the DeepSeek model series, with Skywork-OR1-Math-7B as the teacher and DeepSeek-R1-Distill-Qwen-7B as the student. Results follow the avg@ k and pass@ k evaluation protocol of Table 1 ; the best and second-best results are shown in bold and underlined, respectively.
Appendix figures & tables5 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 3: Mean avg@ k across the seven math benchmarks for the ablations in Table 4 . RA-C and RA-I mask only correct-negative and incorrect-positive conflicts, respectively, whereas RA-Inv retains only these two conflict types.
Method
Rule
avg@ k
pass@ k
OPD
All
44.34
57.66
RA-C
C-negative
45.31
60.63
RA-I
I-positive
45.83
59.75
RA-Inv
Conflicts
43.41
56.80
RA-OPD
Both
49.43
61.56
Appendix
Table 4: Reward-alignment ablations on Qwen3-8B-Base, averaged across the seven math benchmarks (%). RA-C and RA-I mask only correct-negative ( R=1,G<0 ) and incorrect-positive ( R=0,G>0 ) conflicts, respectively, whereas RA-Inv retains only these two conflict types.
Figure 4: Actor entropy for RA-OPD and standard OPD during training. The left and right panels show Qwen3-8B-Base and DeepSeek-R1-Distill-Qwen-7B, respectively.
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 (R2-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, R2-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.
Haofeng Xu, Junwei Su, Lansong Diao +2
The University of Hong Kong · Alibaba Group · University of Science and Technology of China
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.
On-policy distillation (OPD) has become a central post-training tool for large language models (LLMs), providing dense per-token teacher supervision along the student's own rollouts. In this work, we identify a common structural cause underlying OPD, which we call prefix failure. Under prefix failure, dense per-token supervision induces a bimodal teacher mixture and fragmented gradients that token-level loss truncation or reweighting fail to address. We thus propose Trajectory-Refined Distillation (TRD), a trajectory-level correction method that revises the student's rollout under the teacher guidance while within on-policy support. By correcting problematic prefixes before distillation, TRD mitigates prefix failure at its source. Moreover, TRD improves exploration by exposing the student to alternative valid derivations under teacher guidance, even when the original rollouts are already correct. TRD can also be applied to on-policy self-distillation (OPSD), a parameter-sharing variant that uses the student model conditioned on privileged information as the teacher. TRD improves average math accuracy while remaining competitive on code. Under OPSD, it solves 9/23 AMOBench questions left unsolved by the Qwen3-8B base model, nearly doubling the strongest dense-KL baseline. Code is available at https://github.com/louieworth/trd.
Li Jiang, Haoran Xu, Yichuan Ding +1
McGill University · Mila Quebec AI Institute · UT Austin