Reward-Aligned Distillation

Yilan Chen; Zian Su; Manish Prajapat; Dingmin Wang; Rein Houthooft; Rui Song; Vijay Lingam

cs.LGOct 10, 2026

Reward-Aligned Distillation

Authors: Yilan Chen, Zian Su, Manish Prajapat, Dingmin Wang, Rein Houthooft, Rui Song, Vijay Lingam

Organizations: UC San Diego · Purdue University · AWS AI Labs · Amazon

Abstract

Reinforcement learning with verifiable rewards improves the reasoning capabilities of large language models, but remains computationally expensive and sample inefficient. On-policy distillation (OPD) provides dense token-level supervision from a stronger teacher, yet standard OPD distills every token regardless of whether the teacher's preference actually improves task reward. Existing token-selection methods rely on proxies such as uncertainty or teacher--student divergence, which need not identify reward-improving supervision. We introduce Reward-Aligned Distillation (RAD), a bilevel framework that weights distillation signals based on their predicted effect on the student's post-update reward. RAD differentiates through a temporary optimizer step to score each token's supervision, then updates the token weights by projected ascent from the initial token weights. The resulting weights can strengthen, suppress, or reverse individual distillation-gradient contributions. At a matched training-response budget, RAD improves mean accuracy on six competition-math benchmarks over the strongest baselines by 6.5 percentage points for Qwen3-1.7B (non-thinking) and 2.5 points for Qwen3-4B-Instruct-2507. RAD-Simple, a simplified variant, achieves comparable math accuracy at the same training-response budget. On LiveCodeBench-v6, RAD reaches 54.69% and 58.67% pass@1 with Qwen3-1.7B (non-thinking) and OLMo3-7B-Instruct, respectively, exceeding the strongest baselines by 8.7 and 6.4 points. These results support reward-guided weighting of teacher supervision across model families and tasks.

Figures & tables

Figure 1: Reward-aligned distillation (RAD) consistently outperforms OPD and other baselines on mathematical reasoning and code generation tasks. The teachers are Qwen3-30B-A3B-Instruct-2507 for the Qwen student and OLMo3.1-32B-Instruct for the OLMo student. (a) Rollout-matched Qwen3-1.7B (non-thinking) results on both math and coding tasks ( Tables 1 and 8 ). (b) LiveCodeBench-v6 pass@1 for Qwen3-1.7B (non-thinking) and OLMo3-7B-Instruct, comparing RAD at step 100 with baselines at step 200. The grey bars show the strongest baselines, RLAD and GRPO, respectively; each RAD bar gives its gain over that baseline (pp). Dashed lines mark the corresponding teachers.

Panel of Figure 2 — Figure 2: One iteration of reward-aligned distillation (RAD), illustrated with SGD. (1) Starting from student $\bm{\theta}$ , assign the initial token weight $c$ to every token. (2) Take one temporary step on the weighted distillation objective to obtain $\bm{\theta}^{+}$ , as in Equation 3 . (3) Sample separate meta rollouts from $\bm{\theta}^{+}$ to estimate the post-update reward gradient. (4) Differentiate this reward through the temporary step to score each token using Equation 4 . Positive scores increase token weights and negative scores decrease them; projected ascent clips the updated weights to $[c-1,c+1]$ , as in Equation 5 . (5) Discard the temporary update and take one weighted distillation step from the original $\bm{\theta}$ . Reset the weights to $c$ for the next batch.

Student: Qwen3-1.7B (non-thinking)
Method	AIME24	AIME25	AIME26	HMMT25 Feb	HMMT25 Nov	Olympiad Bench	Mean	Time (h)
Teacher	$74.79_{\pm 3.72}$	$61.46_{\pm 3.53}$	$74.79_{\pm 4.56}$	$43.75_{\pm 5.38}$	$59.17_{\pm 5.59}$	$77.15_{\pm 0.52}$	$65.18_{\pm 1.66}$	—
Base	$12.08_{\pm 3.89}$	$11.04_{\pm 4.20}$	$8.13_{\pm 2.63}$	$6.04_{\pm 3.38}$	$3.75_{\pm 2.00}$	$39.30_{\pm 0.82}$	$13.39_{\pm 1.49}$	—
GRPO	$29.58_{\pm 5.64}$	$19.58_{\pm 3.89}$	$18.96_{\pm 5.98}$	$15.63_{\pm 2.82}$	$10.21_{\pm 4.93}$	$53.07_{\pm 0.93}$	$24.51_{\pm 1.46}$	5.12
OPD	$40.00_{\pm 5.65}$	$27.29_{\pm 4.44}$	$26.46_{\pm 5.83}$	$17.71_{\pm 3.67}$	$16.67_{\pm 6.24}$	$58.19_{\pm 1.17}$	$31.05_{\pm 2.01}$	14.28
TIP	$39.17_{\pm 4.93}$	$28.54_{\pm 4.40}$	$31.04_{\pm 5.98}$	$17.08_{\pm 3.09}$	$18.33_{\pm 4.08}$	$57.93_{\pm 1.18}$	$32.02_{\pm 1.70}$	34.09

Table 1: Math accuracy (%): RAD at step 100 versus baselines and RAD-Simple at step 200. Values are mean

\pm

standard deviation over 16 sampled responses; Mean equally weights six benchmarks. Bold / underline mark the best/second-best trained means, with ties sharing a rank. Parentheses give gains of the higher RAD or RAD-Simple score in each column over the best baselines. Time denotes total training time in hours.

Figure 3: Math evaluation dynamics for Qwen3-1.7B (non-thinking) in the left panel and Qwen3-4B-Instruct-2507 in the right panel, averaged equally over AIME24/25 and HMMT2025-Feb/Nov. Distillation baselines tend to plateau or decline early, while GRPO improves slowly. RAD reaches a higher peak within 100 updates than these baselines achieve over 200 updates.

Figure 4: LiveCodeBench-v6 pass@1 (%) by difficulty, comparing RAD at step 100 with baselines at step 200. RAD outperforms the baselines across all difficulty levels for both student models. Numbers above RAD’s bars give its improvement over the best baseline in that tier (percentage points). Error bars show

\pm

1 standard deviation over four sampled responses; black lines mark available teacher means. Table 8 reports the full means and standard deviations.

Variant	Change from RAD	Acc. (%)	$\Delta$ (pp)
RAD	AdamW lookahead; separate meta rollouts	32.71	—
Negative score	Reverse the score: $I_{t}\mapsto-I_{t}$	0.36	$(-32.34)$
Shuffled weights	Permute weights within each response	6.35	$(-26.35)$
Positive only	Restrict weights to $[0,1]$	25.00	$(-7.71)$
Negative-bound warmup	Cosine lower bound: $0\to-1$ over 50 steps	31.20	$(-1.51)$
Pre-update gradient	Reward gradient before the temporary update	34.32	$(+1.61)$

Table 2: Qwen3-1.7B (non-thinking) math ablations at step 100 (

c=0

): mean accuracy over the four benchmarks (AIME24/25 and HMMT2025-Feb/Nov). Reversing scores or shuffling weights sharply reduces accuracy. Variants use separate meta rollouts unless rollout reuse is specified.

\Delta

(pp) denotes difference from RAD. Bold / underline : best/second-best displayed means.

Table 7

Figure 5: OPD 2 evaluation dynamics for Qwen3-4B-Instruct-2507, using Table 4 ’s benchmarks. Horizontal lines mark the initial student and teacher. OPD 2 + RAD approaches teacher accuracy at step 100.

Appendix figures & tables12 assets

Supplementary material from the paper’s appendix.

Appendix

Domain	Use	Dataset	# Problems	Responses / prompt	Metric
Math	Train	DeepMath-103K, level $\geq 6$	57,046	8	Binary correctness reward
Math	Eval	AIME24	30	16	pass@1 accuracy (Math-Verify)
Math	Eval	AIME25	30	16	pass@1 accuracy (Math-Verify)
Math	Eval	AIME26	30	16	pass@1 accuracy (Math-Verify)
Math	Eval	HMMT2025-Feb	30	16	pass@1 accuracy (Math-Verify)
Math	Eval	HMMT2025-Nov	30	16	pass@1 accuracy (Math-Verify)

Appendix

Table 5: Dataset statistics and evaluation metrics. Counts denote dataset examples or tasks; responses per prompt are sampled whenever a prompt is selected. RAD additionally samples an independent outer batch of 64 meta prompts with eight responses each (512 responses). Coding training and validation counts refer to the Eurus split before model-specific prompt-length filtering.

Figure 6: Initial token weight and negative-weight warmup comparison for Qwen3-1.7B (non-thinking). Each curve shows one run with seed 42 over the shared step-10–100 window. Accuracy equally averages AIME24/25 and HMMT2025-Feb/Nov, with 16 responses per problem. Fixed

c=0

without warmup uses the same run as Table 1 , restricted to these four benchmarks. All runs use AdamW lookahead, separate meta rollouts, weight learning rate

\gamma=10^{4}N_{1}/N

, and no further weight normalization. The initialization variants clip to

[c-1,c+1]

; cosine initialization reaches

c=0

at step 100. Negative-weight warmup keeps

c=0

and the upper bound at 1 while lowering the lower bound from 0 to

-1

by step 50 using a cosine schedule. Dots are observed evaluations; labels show step-100 accuracy.

Student: Qwen3-1.7B (non-thinking)
Method	Easy (322)	Medium (383)	Hard (350)	Overall (1,055)
Teacher	$96.27_{\pm 0.79}$	$69.71_{\pm 1.21}$	$26.07_{\pm 0.24}$	$63.34_{\pm 0.38}$
Base (student)	$60.02_{\pm 1.17}$	$16.45_{\pm 0.26}$	$1.79_{\pm 0.71}$	$24.88_{\pm 0.35}$
GRPO	$66.07_{\pm 2.64}$	$23.11_{\pm 1.13}$	$4.07_{\pm 0.47}$	$29.91_{\pm 1.14}$
OPD	$81.37_{\pm 1.54}$	$41.38_{\pm 1.40}$	$10.79_{\pm 1.02}$	$43.44_{\pm 0.70}$
TIP	$82.69_{\pm 1.08}$	$40.34_{\pm 2.16}$	$9.71_{\pm 0.90}$	$43.10_{\pm 0.52}$

Appendix

Table 8: Full LiveCodeBench-v6 pass@1 (%; mean

\pm

standard deviation over four sample positions), comparing RAD at step 100 with baselines at step 200. Bold/underline: best/second-best trained means; parentheses: gains of RAD in each column over the best evaluated baseline (pp, before rounding).

Figure 7: Coding-validation dynamics for Qwen3-1.7B (non-thinking). Each point reports pass@1 over all 1,024 validation problems for one training run per method. Curves stop at the last available evaluation. We compare RAD at step 100 with baselines at step 200; the RAD curve is dotted beyond step 100 and extends to step 200.

	Step 100					Step 200
Method	APPS	CodeContests	Codeforces	TACO	Overall	Overall
GRPO	17.61	17.51	15.45	7.85	13.67	18.16
OPD	28.87	25.99	21.95	11.52	20.51	21.97
TIP	27.46	23.08	26.83	12.57	20.21	22.27
RLAD	37.32	26.79	31.71	15.18	24.51	23.34
ExOPD	19.01	20.16	16.26	10.21	15.82	19.73

Appendix

Table 9: Qwen3-1.7B (non-thinking) coding-validation pass@1 (%). Step-100 scores are shown by dataset; available step-200 scores are shown overall. Overall weights all 1,024 problems equally. Bold/underline: best/second-best reported value per column, with ties sharing a rank; dashes: unavailable. Parentheses give RAD’s gain over the best baseline at the same step (pp, before rounding). RAD at step 100 matches baselines at step 200 in training responses.

Method	Step 100	Step 200
GRPO	34.28	35.84
OPD	33.20	33.01
TIP	27.44	25.68
RLAD	34.28	33.98
RAD (ours)	40.23	—
$\Delta$ vs. best baseline	$(+5.96)$	$(+4.39)$

Appendix

Table 10: OLMo3-7B-Instruct coding-validation pass@1 (%) over 1,024 Eurus problems, with one response per problem and one training run per method. Bold/underline: best/second-best reported value per column, with ties sharing a rank; dashes: unavailable. Parentheses give RAD’s step-100 gain over the best available baseline at the indicated step (percentage points, before rounding).

Figure 8: Coding-validation dynamics for OLMo3-7B-Instruct. Each point reports pass@1 over all 1,024 Eurus validation problems. Curves stop at the last available evaluation. We compare RAD at step 100 with all four baselines at step 200.

Figure 9: SWE-bench Verified at a 65,536-token context and 75 tool turns. Solid bars show pass@1 with

\pm 1

standard deviation across the three rollouts; hatched bars show pass@3. Horizontal lines mark the pass@1 of the base student and the teacher.

Method	Steps	Student GPUs	Qwen3-1.7B (non-thinking), math	OLMo3-7B-Instruct Coding
OPD	200	6	14.28	136.51
TIP	200	6	34.09	166.79
RLAD	200	6	14.11	123.81
ExOPD	200	6	51.02	—
OPD 2	200	4	33.99	—
RAD	100	6	16.66	101.61

Appendix

Table 11: Total training time (hours) for RAD at step 100 versus baselines at step 200, matching generated training responses. RAD-Simple is included only for math, at step 200. Checkpoint saves are included; periodic validation, startup, and restart downtime are excluded.

Stage	Seconds	Share
Distillation rollout pipeline	211.1	35.2%
Student log-probabilities and advantages	11.9	2.0%
Lookahead stage	283.2	47.2%
Weight computation and student update	87.5	14.6%
Rollout weight synchronization	5.3	0.9%
Checkpoint saves	0.6	0.1%

Appendix

Table 12: RAD math training-time breakdown for Qwen3-1.7B (non-thinking), averaged over the same 100 updates as Table 11 . Entries give seconds per training step and each stage’s share of recorded training-step time. Checkpoint saves are amortized over all steps; validation is excluded and reported separately in the text.

Method	Step	Training (h)	Budget used	Accuracy (%)
Student: Qwen3-1.7B (non-thinking) ; budget 12.09 h
GRPO	290	11.98	99.1%	$26.98_{\pm 2.22}$
OPD	160	11.72	96.9%	$23.23_{\pm 2.20}$
TIP	70	12.19	100.9%	$24.27_{\pm 2.62}$
RLAD	170	12.08	99.9%	$25.10_{\pm 2.67}$
ExOPD	50	12.40	102.6%	$27.34_{\pm 2.70}$

Appendix

Table 13: Math accuracy for Qwen3-1.7B (non-thinking) at comparable training times, relative to RAD at step 90. RAD-Simple uses its latest evaluated checkpoint within this budget. Accuracy is the mean over AIME24/25 and HMMT2025-Feb/Nov, with standard deviation across 16 sample-position accuracies. Budget used is the fraction of RAD’s cumulative training time. Bold/underline mark the best/second-best means; parentheses give RAD’s gain over the strongest displayed baseline (pp).

Figure 10: Evaluation during extended training for Qwen3-1.7B (non-thinking). Left: math; right: coding. Math averages AIME24/25 and HMMT2025-Feb/Nov with 16 responses per problem. Coding reports pass@1 on all 1,024 Eurus validation problems. RAD is shown through step 90; GRPO is shown through step 450 in both math and coding. The vertical line marks step 90.

RL for Knowledge Distillation

See all 12 papers

Aug 28, 2026

When Teacher Guidance Misleads: Reward-Aligned On-Policy Distillation

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.

Siyuan Gan, Yuhan Li, Xiran Wang +5

State Key Laboratory of Novel Software Technology, Nanjing University, Nanjing, China · Shanghai Artificial Intelligence Laboratory, Shanghai, China

Language Model Distillation On-Policy Distillation Knowledge Distillation

Ask PDF

Sep 28, 2026

Reward-Aligned Reweighting for On-Policy Distillation

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.

Haofeng Xu, Junwei Su, Lansong Diao +2

The University of Hong Kong · Alibaba Group · University of Science and Technology of China

RL for Knowledge Distillation Language Model Distillation On-Policy Distillation Knowledge Distillation

Ask PDF

Aug 19, 2026

Beyond Imitation: Filtering On-Policy Distillation by Reasoning Progress

On-policy distillation (OPD) has emerged as an effective framework for post-training language models by pairing student-generated trajectories with dense token-level supervision from a teacher. However, OPD implicitly assumes that teacher-derived rewards are an appropriate proxy for reasoning progress, and therefore treats all teacher feedback equally during policy optimization. While in practice, this assumption does not always hold. We observe that teacher-derived rewards often conflict with genuine reasoning progress, as reasoning steps with clear reasoning advancement may still receive lower distillation rewards, simply due to deviation from teacher's outputs. To address this mismatch, we propose Reasoning-Progress-Aware Reward Filtering for On-Policy Distillation (R2-OPD), which constructs two within-trajectory rankings of reasoning spans, one from teacher-derived rewards and the other from independently estimated progress reward. Distillation rewards are selectively suppressed whenever the two rankings disagree, reducing supervision that conflicts with reasoning progress while preserving effective teacher guidance. Our approach shows consistent improvement over standard OPD especially regarding reasoning performances.

Chen Yang, Haiyuan Wan, Rengrong Xiong +2

The Hong Kong University of Science and Technology (Guangzhou) · Tsinghua University · Zhejiang University +1

Language Model Distillation On-Policy Distillation RL for Language Model Reasoning

Ask PDF

Abstract

RL for Knowledge Distillation

Explore similar work

When Teacher Guidance Misleads: Reward-Aligned On-Policy Distillation

Reward-Aligned Reweighting for On-Policy Distillation

Beyond Imitation: Filtering On-Policy Distillation by Reasoning Progress