Adaptive Pullback: Reweighting Supervision at Penalty Jumps in On-Policy Distillation

Guangyu Liu; Siyuan Gan; Shilong Zhou; Linjian Meng; Boyan Wang; Jing Huo

cs.LGOct 20, 2026

Adaptive Pullback: Reweighting Supervision at Penalty Jumps in On-Policy Distillation

Authors: Guangyu Liu, Siyuan Gan, Shilong Zhou, Linjian Meng, Boyan Wang, Jing Huo

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

Abstract

On-policy distillation (OPD) supervises student-generated responses using sampled-token student-teacher log-ratios. A positive log-ratio penalizes a token that the student assigns higher probability to than the teacher, but its current value does not distinguish sustained high penalties from abrupt increases along a response. We examine these local changes by comparing consecutive blocks of tokens within the same response. Our diagnostic identifies penalty jumps, abrupt increases in average penalty that persist on held-out tokens. Across these jumps, penalties rise while teacher confidence decreases. We call this pattern penalty-confidence tension. Boundaries into high-penalty blocks do not exhibit comparable confidence changes, showing that local changes provide information beyond the current penalty value. Motivated by this observation, we propose Adaptive Pullback OPD (AP-OPD). At each detected jump, AP-OPD uses the preceding block as a local reference and uniformly scales all supervision coefficients in the following block so that its weighted average penalty matches the preceding block's original mean. This reference determines the scaling factor without an additional strength constant. AP-OPD improves performance on both mathematical reasoning and code generation.

Figures & tables

Figure 1: Penalty–confidence tension at penalty jumps: penalties rise while teacher confidence decreases. Adaptive Pullback uniformly scales post-jump supervision to match the preceding block’s original mean penalty. Numbers are illustrative.

D_{i}

: block mean penalty;

H

: teacher entropy;

P_{\max}

: teacher top-1 probability;

w

: supervision weight.

Figure 2: Held-out rise in token penalty

d_{t}

around penalty jumps, abrupt increases in block-averaged penalty detected with the

z=2

gate, and three control boundaries. The rise is relative to the mean penalty of the block before each boundary. Boundaries are selected on one half of each block’s tokens and the curves are measured on the other half. The curves use 8-token bins and average both detection directions; bands for Jump and High-magnitude boundary are 95% bootstrap intervals clustered by prompt. The boundary lies at offset 0, and

B=64

Figure 3: Teacher prediction changes across penalty jumps, abrupt increases in block-averaged penalty detected with the

z=2

gate, and control boundaries (

B=64

). (a) Teacher entropy and (b) teacher top-1 probability, averaged over all tokens in the block before and the block after each boundary, on 153 independent matched groups. Their changes indicate decreasing teacher confidence at jumps, accompanying the penalty increase in the observed penalty–confidence tension. Appendix B.4 also reports all 353 groups.

\Delta

is after minus before, and the vertical axis in (b) is truncated. (c) Drop in the student’s and the teacher’s log-probability of the sampled tokens (before minus after), measured on held-out tokens from 324 prompts; (a,b) use all block tokens. Error bars are 95% bootstrap intervals.

Method	MATH500	AMC23	Minerva	Olympiad	AIME24	AIME25	Macro
Teacher (Qwen3-8B-RL-Math)	86.80	74.22	35.75	50.19	32.08	21.67	50.12
Student (Qwen3-1.7B-Base)	42.15	24.06	11.67	16.63	3.54	1.67	16.62
Vanilla OPD	72.90	43.75	20.96	33.63	12.40	7.29	31.82
IW-OPD	71.60	43.28	22.15	32.31	11.35	7.19	31.31
FiRe-OPD	72.80	45.39	21.74	32.83	11.88	8.44	32.18
AP-OPD (top- $\alpha$ )	72.97	44.77	22.06	34.26	13.02	8.65	32.62

Table 1: Main results at step 200, reported as Avg@

k

in percent under a shared evaluation protocol, averaged over two training runs. The teacher and the student before distillation are reference evaluations and are not ranked. Best in bold, second underlined; Macro weights the six benchmarks equally.

Method	HE+	MBPP+	LCB	Macro
Teacher	83.69	72.95	32.29	62.97
Student	61.59	56.81	14.57	44.32
Vanilla OPD	70.43	63.89	23.57	52.63
IW-OPD	70.88	65.28	25.00	53.72
AP-OPD	73.17	65.15	25.71	54.68

Table 2: Code generation uses Avg@4 (%). Qwen3-4B-Instruct-2507 (Teacher) and undistilled Qwen3-1.7B (Student) are unranked references. HE+ is HumanEval+; LCB is LiveCodeBench v6. Macro is the unweighted mean; best trained scores are bold.

Variant	Offline			Trained
	Cov.	Unsel.	$q$	Macro	$\Delta$
(a) Comparison with magnitude-cap
Vanilla OPD	–	–	–	31.55	–
Magnitude-cap	7.17	0.00	13.75	31.40	$-0.15$
AP-OPD (top- $\alpha$ )	7.19	0.00	14.63	32.50	$+0.95$
AP-OPD ( $z$ -gate)	6.49	28.71	11.60	32.50	$+0.95$

Table 3: Ablations, same number of training steps and evaluation protocol as Table 1 , first run only. Macro: macro Avg@

k

;

\Delta

: difference from vanilla OPD. Offline statistics on the analysis rollouts are reported in percent: Cov., response-mean share of tokens in selected blocks; Unsel., share of responses with no selected boundary;

q

, response-mean fraction of penalty mass removed.

Appendix figures & tables13 assets

Supplementary material from the paper’s appendix.

Appendix

Setting	Value
Student initialization	Qwen3-1.7B-Base
Frozen teacher	Qwen3-8B-RL-Math; see Table 5
Training prompts	DAPO-Math-17K; 17,398 prompts
Prompt / response limit	2,048 / 8,192 tokens
Steps	200, from the base checkpoint
Batch	128 prompts, one rollout each; one update per step

Appendix

Table 4: Training configuration shared by our mathematical-reasoning runs.

Setting	Value
Initialization	Qwen3-8B-Base
Training prompts	DAPO-Math-17K
Policy optimization	GSPO
Steps	200
Prompt / response limit	2,048 / 4,096 tokens
Batch	512 prompts, 16 responses each; mini-batch of 32 prompts

Appendix

Table 5: Training configuration of the frozen mathematical-reasoning teacher.

Benchmark	Problems	$k$	Metric
MATH500	500	4	Avg@4
AMC23	40	16	Avg@16
Minerva	272	4	Avg@4
OlympiadBench	675	4	Avg@4
AIME24	30	16	Avg@16
AIME25	30	16	Avg@16

Appendix

Table 6: Evaluation benchmarks and samples per problem.

Setting	Value
Student initialization	Qwen3-1.7B
Frozen teacher	Qwen3-4B-Instruct-2507
Training prompts	Eurus-2-RL-Data code subset; 12,800 prompts
Prompt / response limit	2,048 / 8,192 tokens
Steps	100, one epoch
Batch	128 prompts, one rollout each; one update per step

Appendix

Table 7: Training configuration of our code-generation runs.

Rule	Selected/resp.	Unselected (%)	Token cov. (%)	$q$ (%)
Top- $\alpha$ ( $5\%$ )	1.85	0.00	7.19	14.63
$z=1.5$	3.93	8.98	13.37	20.04
$z=2$ (AP-OPD)	1.91	28.71	6.49	11.60
$z=3$	0.36	75.78	1.10	2.56
All $\Delta_{i}>0$	–	–	–	36.11
Magnitude-cap	1.85	0.00	7.17	13.75

Appendix

Table 8: Selection rules on the analysis rollouts at

B=64

: selected boundaries per response, share of responses without a selection, share of tokens in selected blocks, and removed penalty fraction

q

. Token coverage and

q

are averaged equally over responses;

q

uses responses with nonzero penalty mass. Among the

z

-gate thresholds, only

z=2

is used for training; the other

z

values are reported for reference. Magnitude-cap selections are not penalty jumps.

Detect	$-$ High-mag. bound.	$-$ Rand. pos.	$-$ Rand. bound.
odd	0.019 [0.006, 0.031]	0.041 [0.025, 0.058]	0.135 [0.105, 0.168]
even	0.036 [0.016, 0.057]	0.050 [0.034, 0.066]	0.145 [0.116, 0.177]

Appendix

Table 9: Held-out rise

U_{0:64}

z=2

jumps minus that of each matched control, for each detection direction on the analysis set at

B=64

, with 95% intervals.

Rollouts	$U_{0:64}$	$U_{64:128}$	Ratio
Analysis set	0.190 [0.172, 0.209]	0.096 [0.081, 0.111]	50.3%
Vanilla, 300 steps	0.188 [0.183, 0.194]	0.101 [0.096, 0.106]	53.6%
Paired rule, 300 steps	0.190 [0.184, 0.196]	0.103 [0.098, 0.108]	54.5%

Appendix

Table 10: Held-out rise of the first and second blocks after

z=2

jumps and their ratio, on the analysis set and on two further training runs. Both rises use the pre-jump block as the reference and are measured with the split-half protocol, averaging both detection directions. A positive second-block rise relative to this reference is compatible with returning to the level of position-matched random boundaries in Figure 2 .

Metric	Control	All 353 groups	Independent 153
$H$ (nats)	random	0.300 [0.275, 0.342]	0.311 [0.263, 0.365]
$H$ (nats)	high-magnitude	0.172 [0.122, 0.265]	0.261 [0.212, 0.317]
$P_{\max}$	random	$-$ 0.099 [ $-$ 0.112, $-$ 0.091]	$-$ 0.102 [ $-$ 0.116, $-$ 0.088]
$P_{\max}$	high-magnitude	$-$ 0.061 [ $-$ 0.085, $-$ 0.050]	$-$ 0.083 [ $-$ 0.098, $-$ 0.069]

Appendix

Table 11: DID of the teacher’s entropy and top-1 probability, computed as the post-minus-pre change at jumps minus that at control boundaries, with 95% intervals.

Mean of the two runs
Model	MATH500	AMC23	Minerva	Olympiad	AIME24	AIME25	Macro
Vanilla OPD	72.90	43.75	20.96	33.63	12.40	7.29	31.82
IW-OPD	71.60	43.28	22.15	32.31	11.35	7.19	31.31
FiRe-OPD	72.80	45.39	21.74	32.83	11.88	8.44	32.18
AP-OPD (top- $\alpha$ )	72.97	44.77	22.06	34.26	13.02	8.65	32.62
AP-OPD ( $z$ -gate, $B=64$ )	73.08	46.72	21.42	32.94	12.19	8.02	32.39

Appendix

Table 12: Per-benchmark Avg@

k

in percent at step 200 under the same evaluation protocol. Mean: arithmetic mean of the two runs, calculated from unrounded scores. Macro: equal-weight mean over the six benchmarks; the additional ablations use run 1 only. Best in bold, second underlined within each block; ties at the reported precision share the same marking.

Contrast	MATH500	AMC23	Minerva	Olympiad	AIME24	AIME25	Macro
$z$ -gate $-$ Vanilla	$-0.75$	$+2.81$	$+1.10$	$+0.04$	$+0.63$	$+1.88$	$+0.95$
	[ $-$ 2.40, 0.80]	[ $-$ 0.47, 6.25]	[ $-$ 1.19, 3.31]	[ $-$ 1.44, 1.48]	[ $-$ 1.04, 2.29]	[ $-$ 0.42, 5.00]	[0.02, 1.90]
Top- $\alpha$ $-$ Vanilla	$-0.45$	$+2.19$	$+1.19$	$+1.11$	$+1.04$	$+0.63$	$+0.95$
	[ $-$ 2.05, 1.10]	[ $-$ 0.94, 5.47]	[ $-$ 0.92, 3.31]	[ $-$ 0.33, 2.52]	[ $-$ 1.67, 4.17]	[ $-$ 1.25, 2.71]	[0.01, 1.90]
Magnitude-cap $-$ Vanilla	$-0.95$	$+0.78$	$+1.65$	$-0.74$	$-2.50$	$+0.83$	$-0.15$
	[ $-$ 2.60, 0.65]	[ $-$ 1.56, 3.28]	[ $-$ 0.46, 3.86]	[ $-$ 2.22, 0.70]	[ $-$ 5.00, $-$ 0.21]	[ $-$ 1.46, 3.33]	[ $-$ 0.99, 0.70]

Appendix

Table 13: Paired differences in percentage points among our run-1 models, with 95% paired bootstrap intervals.

Model	Clipped steps	Resp. length	Student entropy	Weighted	Unweighted	Mean
	(1–50 / 51–200)	(first / last 50)	(first / last 50)	bound. (%)	resp. (%)	weight
Vanilla OPD	43 (36 / 7)	1464 / 2437	0.764 / 0.168	–	–	1
AP-OPD (top- $\alpha$ )	30 (30 / 0)	1441 / 2432	0.757 / 0.174	6.52	1.57	0.952
AP-OPD ( $z$ -gate)	33 (33 / 0)	1485 / 2408	0.816 / 0.173	6.79	24.22	0.945
Magnitude-cap	32 (31 / 1)	1440 / 2434	0.726 / 0.175	6.52	1.33	0.961

Appendix

Table 14: Run-1 training dynamics over 200 steps, reported as means of logged values. Clipped steps have pre-clipping gradient norm above 1.

Rule	MATH500	AMC23	Minerva	Olympiad	AIME24	AIME25	Macro	Clipped	Mean $w$
Vanilla OPD	69.90	43.12	20.31	30.67	11.46	6.25	30.29	–	1
Paired ( $+$ pre, $-$ post)	70.00	41.88	20.77	30.33	11.67	7.71	30.39	32 / 100	1.006
Pre-jump $\times 4$	69.15	41.88	20.50	31.26	12.29	6.25	30.22	88 / 100	1.200
Post-jump $\times 0.3$	70.25	44.38	21.23	31.11	12.29	6.88	31.02	30 / 100	0.955

Appendix

Table 15: Alternative reweighting rules at step 100, reported as Avg@

k

in percent from one training run with 8K evaluation responses. Clipped: steps whose pre-clipping gradient norm exceeds 1; Mean

w

: mean token weight.

Rollouts	Bound	Removed	Jumps kept	Rise kept	Overlap	Expected
Analysis set	matched	11.6	99.0	89.8 [86.1, 93.8]	14.5 [13.1, 15.9]	12.5
Analysis set	default	27.0	96.3	71.1 [67.0, 75.8]	32.3 [30.9, 33.6]	28.6
Vanilla, 300 steps	matched	13.2	98.9	86.1 [84.8, 87.4]	13.3 [13.1, 13.6]	13.1
Vanilla, 300 steps	default	30.1	95.9	68.0 [66.7, 69.4]	31.9 [31.6, 32.2]	30.2
Paired rule, 300 steps	matched	13.4	98.7	84.2 [82.9, 85.5]	13.7 [13.5, 14.0]	13.1
Paired rule, 300 steps	default	30.3	95.5	66.4 [65.0, 67.8]	32.5 [32.1, 32.8]	30.3

Appendix

Table 16: Token-level clipping on the analysis set and the two 300-step runs at

B=64

and

z=2

. Values are percentages; intervals are 95%. Matched:

\kappa=3.131

, which removes as much penalty mass as AP-OPD on the analysis set; default: REOPOLD’s

\kappa=1.720

. Removed: penalty mass removed by clipping, averaged over responses. Jumps kept: original jumps detected again after clipping. Rise kept: held-out rise

U_{0:64}

after clipping relative to no clipping. Overlap: share of AP-OPD’s removal that lies above the bound. Expected: share of all penalty mass above the bound, the overlap expected if AP-OPD’s removal were unrelated to token size. Overlap and Expected pool mass over tokens.

Language Model Distillation

Language model distillation transfers a larger model’s capabilities to smaller or specialized models using outputs, probabilities, or internal representations.

See all 508 papers

Your Teacher Can't Help You Here: Combating Supervision Fidelity Decay in On-Policy Distillation
May 29, 2026Yanjiang Liu, Jie Lou, Xinyan Guan +7CoT Distillation Language Model Distillation
University of Chinese Academy of Sciences · 2Chinese Information Processing Laboratory Institute of Software, Chinese Academy of Sciences 2University of Chinese Academy of Sciences, Beijing, China · 3Xiaohongshu
Reward-Aligned Reweighting for On-Policy Distillation
Sep 28, 2026Haofeng Xu, Junwei Su, Lansong Diao +2RL for Knowledge Distillation Language Model Distillation
The University of Hong Kong · Alibaba Group · University of Science and Technology of China
Reward-Gated On-Policy Distillation
Jul 4, 2026Mohammad Sadegh Akhondzadeh, Vijay Lingam, Atula Tejaswi +3Language Model Distillation On-Policy Distillation
University of Cologne · AWS AI · UT Austin +1

May 29, 2026J/K move · Enter open · S save

Your Teacher Can't Help You Here: Combating Supervision Fidelity Decay in On-Policy Distillation

Yanjiang Liu, Jie Lou, Xinyan Guan, Yuqiu Ji+6

University of Chinese Academy of Sciences · 2Chinese Information Processing Laboratory Institute of Software, Chinese Academy of Sciences 2University of Chinese Academy of Sciences, Beijing, China · 3Xiaohongshu

CoT Distillation Language Model Distillation On-Policy Distillation Teacher-Student Learning RL Fine-Tuning

Open paper AskPDF arXiv

On-policy distillation transfers reasoning capabilities by training a student model on its own generated trajectories using token-level feedback from a teacher. However, we identify a critical bottleneck, Supervision Fidelity Decay (SFD): as student-generated prefixes lengthen, the teacher's next-token distribution becomes less confident and less discriminative. Consequently, the teacher-dependent corrective signal in reverse-KL distillation weakens, causing student drift to compound across long reasoning chains. To mitigate SFD, we introduce Lookahead Group Reward (\ours{}). Building on the insight that next-step teacher confidence reflects the discriminative strength of future reverse-KL supervision, \ours{} evaluates the student's top-K candidate tokens by the teacher confidence they induce at the subsequent step and assigns a group-normalized reward. To maintain computational efficiency, we further design an entropy-triggered tree-attention mechanism. Across six math and code benchmarks, \ours{} improves mean@8 by 2.57 points over OPD for a 7B student, with gains increasing in longer-generation and reaching +4.92 points on AIME-26 at 39k tokens.

Abstract

Language Model Distillation

Explore similar work

Your Teacher Can't Help You Here: Combating Supervision Fidelity Decay in On-Policy Distillation

Reward-Aligned Reweighting for On-Policy Distillation

Reward-Gated On-Policy Distillation