A Token-Level Analysis of Sampled-Token Reverse-KL On-Policy Distillation

Bing Shao; Jiazheng Zhang; Long Ma; Yujiong Shen; Senjie Jin; Xin Guo; Yuming Yang; Mingxu Chai

cs.LGAug 26, 2026

A Token-Level Analysis of Sampled-Token Reverse-KL On-Policy Distillation

Authors: Bing Shao, Jiazheng Zhang, Long Ma, Yujiong Shen, Senjie Jin, Xin Guo, Yuming Yang, Mingxu Chai,

Organizations: College of Computer Science and Artificial Intelligence, Fudan University

Abstract

On-policy distillation (OPD) supervises a student on its own trajectories with token-level signals from a frozen teacher, yet how a sampled loss allocates updates across tokens remains poorly understood. We analyze the gradient of the per-token K2 estimator of reverse KL with respect to the student logits. The $\ell_1$ norm of this gradient factorizes into the absolute teacher--student log-probability gap and a student-side softmax factor that grows as the sampled token becomes less likely under the student. In our math-distillation runs, these per-token norms are highly non-uniform: low-student-probability tokens account for a disproportionate share of their sum and are also enriched in large teacher--student gaps. As a lightweight intervention suggested by this analysis, we study Surprise-aware Reweighting (SuRe), a detached, bounded weighting rule that further amplifies this existing allocation. Across two Qwen3 student scales, SuRe improves several math metrics over vanilla OPD and shows no clear degradation on the selected out-of-domain benchmarks. Our primary contribution is therefore a gradient-level characterization of reverse-KL OPD trained with the K2 estimator, with SuRe as one empirical instantiation.

Figures & tables

Figure 1: From token-level gradient analysis to SuRe. (a) The same student-sampled token and prefix are scored by the student and frozen teacher. (b) Exact K2 student-logit gradient norm for a fixed sample; the conditional-mean diagnostic is observational. (c) SuRe applies detached, bounded weights that preserve per-token gradient direction. Curves show analytic weight functions.

Panel of Figure 2 — Figure 2: Signed post-training checkpoint shift. Curves are grouped by the checkpoint that generated each rollout; every realized token and prefix is rescored under both checkpoints. Base rollouts lean negative, whereas OPD rollouts exhibit a positive tail. Panels show a log-count overlay (a) and a signed tornado view (b); annotations in (b) count tokens with $|\Delta_{\mathrm{OPD-Base}}|\geq 2$ .

Figure 3: Why signed residuals rather than entropy/JSD. Panels (a,b) compare the teacher–OPD residual gap with JSD and OPD entropy; most tokens sit in a low-gap, low-uncertainty bulk. Panel (c) ranks tokens by each statistic and measures the share of the sum of gradient norms covered by each ranking. JSD and entropy use a top-50 approximation, computed by restricting each position’s full-distribution diagnostic to its 50 highest-probability candidate tokens.

Figure 4: Mean logit-gradient norms are larger at low-student-probability positions. Panel (a) shows the mean per-token logit-gradient

\ell_{1}

norm within student-probability bins, with shading for the 25th–75th percentiles. Panel (b) shows student probability distributions within absolute-residual deciles; the connected points indicate mean student probability.

Qwen3-1.7B-Base
	AIME24		AIME25		AMC23		MATH-500
Methods	avg@ $8$	pass@ $8$	avg@ $8$	pass@ $8$	avg@ $8$	pass@ $8$	avg@ $4$	pass@ $4$
- Base	4.17	16.67	4.17	16.67	30.00	65.00	56.35	76.80
- KD	7.08	20.00	3.33	13.33	30.94	67.50	56.45	75.60
- SeqKD	5.83	20.00	4.17	20.00	32.19	72.50	56.25	74.80
- Vanilla OPD	9.17	16.67	5.83	16.67	39.38	67.50	66.55	80.80

Table 1: In-domain performance (%). For AIME24, AIME25, and AMC23 we report avg@

8

and pass@

8

; for MATH-500 we report avg@

4

and pass@

4

. Cells with deeper background color correspond to better performance within each model group.

Panel of Figure 5 — Figure 5: Training dynamics on Qwen3-1.7B-Base. Vanilla OPD (blue dashed) vs. SuRe at $\alpha{=}1.0$ (red). SuRe improves mean rollout score and increases early gradient norm while tracking OPD’s actor entropy.

Figure 6: Out-of-domain performance (pass@

1

, %). Bars compare Base, Vanilla OPD, and SuRe on CRUX, IFEval, and MMLU-Pro for both student scales. On the evaluated OOD tasks, OPD and SuRe are broadly comparable, but neither uniformly improves over the base; both are slightly below the base on 1.7B MMLU-Pro.

Figure 7: Ablations on AMC23 (Qwen3-1.7B-Base). (a)

\alpha

sweep for SuRe against vanilla OPD. (b) Orientation controls comparing SuRe, High-reweight, and Random-reweight. SuRe exceeds Random-reweight at

k{=}2,4,8

, while Random-reweight scores higher at

k{=}1

. SuRe exceeds High-reweight at

k{=}1,2,4

and ties it at

k{=}8

Appendix figures & tables9 assets

Supplementary material from the paper’s appendix.

Appendix

Hyper-parameter	Qwen3 (8B $\to$ 1.7B)	Qwen3 (8B $\to$ 4B)
Student model	Qwen3-1.7B-Base	Qwen3-4B-Base
Teacher model	Qwen3-8B	Qwen3-8B
Optimizer	AdamW	AdamW
Learning rate	$1\times 10^{-6}$	$1\times 10^{-6}$
LR warmup steps	10	10
Training epochs	2	2

Appendix

Table 3: Training hyperparameters across student scales for Vanilla OPD and SuRe. The only difference between the two methods is the per-token weight

w_{t}

in ( 16 ); the entries below are shared by the corresponding model-scale runs except for the memory-related knobs required by the larger student. Rows prefixed by “Framework” denote batching limits inherited from the training implementation. The extra reference-policy KL coefficient is separate from the K2 distillation objective.

Qwen3-1.7B-Base
	CRUX			IFEval			MMLU-Pro
Method	@1	@5	@10	@1	@5	@10	@1	@5	@10
Base	4.75	21.75	33.12	22.92	47.50	57.86	27.07	62.94	78.17
Vanilla OPD	35.38	63.00	70.62	27.73	44.18	53.97	26.44	57.32	69.91
SuRe ( $\alpha{=}0.2$ )	34.62	63.25	71.12	26.80	44.92	53.05	26.52	57.37	69.98
SuRe ( $\alpha{=}0.5$ )	34.75	60.88	68.50	27.17	44.92	52.68	26.86	57.31	69.56

Appendix

Table 4: Full OOD pass@

k

results (%). CRUX, IFEval, and MMLU-Pro are the OOD benchmarks reported in the main text. IFEval uses prompt-level strict accuracy and MMLU-Pro uses exact match.

Panel of Figure 8 — Figure 8: Training dynamics on Qwen3-4B-Base. Vanilla OPD vs. SuRe at $\alpha{=}1.0$ . SuRe maintains comparable actor entropy while changing the score and gradient-norm dynamics, providing an additional cross-scale view of the optimization behavior.

Item	Content
Task	MATH-500 example asking for primes $p$ such that $8x\equiv 1\pmod{p}$ has no solution.
Config	Qwen3-1.7B OPD student at step $222$ ; Qwen3-1.7B-Base as base; Qwen3-8B as teacher; temperature $1.0$ , top- $p=1.0$ , max length $8192$ .
Prompt	“Determine the sum of all such $p$ . Please reason step by step, and put your final answer within \boxed{} .”
Response sketch	The student argues that the congruence is solvable iff $\gcd(8,p)=1$ ; only the prime $p=2$ divides $8$ ; therefore the answer is $\boxed{2}$ .
Global pattern	$376$ response tokens; mean JS $(\pi_{\mathrm{OPD}},\pi_{\mathrm{base}})=0.052$ ; only $16.8\%$ of tokens have JS $>0.10$ .

Appendix

Table 5: A compact view of the case. The example is a correct OPD rollout. We keep only the information needed to interpret the subsequent token table, rather than listing every logged field.

Where	Sampled token	OPD vs. base	Interpretation
After “analyze the given congruence”	newline after “:”	rank $5\to 1$ , $\Delta_{\mathrm{OPD-Base},t}=+2.95$	The final OPD checkpoint ranks the transition into displayed-equation format more highly than Base.
At the key concept “invertible modulo $p$ ”	emphasis marker before “invertible”	rank $6\to 1$ , $\Delta_{\mathrm{OPD-Base},t}=+6.82$	The largest endpoint difference occurs where the solution introduces the modular-inverse idea.
At the restatement “which primes $p$ …”	the transition token “the”	rank $5\to 1$ , $\Delta_{\mathrm{OPD-Base},t}=+2.90$	Some large endpoint differences occur on discourse or fluency tokens, not only mathematical symbols.
At an intermediate mention of the answer	token “2”	rank $2\to 2$ , $\Delta_{\mathrm{OPD-Base},t}=-2.70$	The final OPD checkpoint assigns this token lower probability than Base; the final answer remains correct, but the local confidence profile differs.

Appendix

Table 6: Four readable token-level changes from the case. The point is not that every large change is a mathematical token, but that the endpoint difference is localized: a few reasoning, formatting, and transition positions differ strongly instead of the whole response moving uniformly.

Hyper-parameter	KD	SeqKD	GRPO
Data source	DeepMath hard	Teacher rollouts	DeepMath hard
Teacher model	Qwen3-8B	Qwen3-8B	–
Supervision	Token-level $\pi_{T}$	Sequence sampled from $\pi_{T}$	Verifiable answer reward
Loss	Forward KL on teacher tokens	SFT cross-entropy on teacher seq.	GRPO clipped policy gradient
Rollout sampling	off-policy (teacher)	off-policy (teacher)	on-policy (student)
Generation temperature	–	1.0	1.0

Appendix

Table 7: Baseline-specific hyperparameters for KD, SeqKD, and GRPO on Qwen3-1.7B-Base. Entries shared with OPD/SuRe (see Table 3 ) are omitted.

	AIME24		AIME25		AMC23		MATH-500
Method	avg@ $8$	pass@ $8$	avg@ $8$	pass@ $8$	avg@ $8$	pass@ $8$	avg@ $4$	pass@ $4$
GRPO	8.33	23.33	4.17	20.00	41.56	70.00	65.80	79.00
SuRe ( $\alpha{=}1.0$ )	9.58	23.33	7.08	16.67	43.12	75.00	67.65	80.80

Appendix

Table 8: Comparison with on-policy RL on Qwen3-1.7B-Base (%). GRPO is a reward-based RL baseline trained with verifiable answer matching; SuRe is purely a distillation objective with no reward signal. AIME24, AIME25, and AMC23 compute avg@

8

/pass@

8

from only the first

8

32

rollouts per problem, in generation order; MATH-500 computes avg@

4

/pass@

4

from all

4

rollouts. Bold marks the better entry within each column.

	AIME24		AIME25		AMC23		MATH-500
Training objective	avg@ $8$	pass@ $8$	avg@ $8$	pass@ $8$	avg@ $8$	pass@ $8$	avg@ $4$	pass@ $4$
Vanilla OPD	9.17	16.67	5.83	16.67	39.38	67.50	66.55	80.80
Mean-normalized SuRe	10.00	20.00	4.58	10.00	41.88	72.50	69.20	82.20
Exact-shuffled	10.00	23.33	5.83	23.33	38.44	70.00	67.95	81.00
Exact rank-reversed	–	–	–	–	–	–	67.00	81.40
Mean-normalized uplift-only	9.58	26.67	4.17	13.33	39.38	72.50	68.40	81.40

Appendix

Table 9: Matched weighting controls on Qwen3-1.7B-Base (%). AIME24, AIME25, and AMC23 report avg@

8

/pass@

8

; MATH-500 reports avg@

4

/pass@

4

. Rank-reversed was evaluated only on MATH-500.

Seed	Vanilla OPD	Original SuRe
42	66.55	67.65
43	66.40	67.50

Appendix

Table 10: Second-seed check on MATH-500 avg@

4

(%). Original SuRe uses the submitted, unnormalized weighting rule.

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

Revisiting On-Policy Distillation: Empirical Failure Modes and Simple Fixes
Mar 26, 2026Yuqian Fu, Haohuan Huang, Kaiwen Jiang +4Language Model Distillation On-Policy Distillation
State Key Laboratory of Multimodal Artificial Intelligence Systems, CASIA · School of Artificial Intelligence, UCAS · Fudan University +1
Unbiased Top- $k$ Estimation for On-Policy Distillation
Sep 28, 2026Linjian Meng, Siyuan Gan, YuHan Li +5Language Model Distillation On-Policy Distillation
Shanghai Artificial Intelligence Laboratory, Shanghai, China · State Key Laboratory of Novel Software Technology, Nanjing University, Nanjing, China
Understanding Off- vs On-Policy Distillation: A Tale of Distinct Training Objectives
Sep 29, 2026Qiwei Di, Xuheng Li, Kaixuan Ji +3Multi-Teacher Knowledge Distillation On-Policy Distillation
Department of Computer Science, University of California, Los Angeles, CA 90095, USA

Mar 26, 2026J/K move · Enter open · S save

Revisiting On-Policy Distillation: Empirical Failure Modes and Simple Fixes

Yuqian Fu, Haohuan Huang, Kaiwen Jiang, Jiacai Liu+3

State Key Laboratory of Multimodal Artificial Intelligence Systems, CASIA · School of Artificial Intelligence, UCAS · Fudan University+1

Language Model Distillation On-Policy Distillation

Open paper AskPDF arXiv

On-policy distillation (OPD) is increasingly used in LLM post-training because it can leverage a teacher model to provide dense supervision on student rollouts. The standard implementation, however, usually reduces distribution matching to a sampled-token log-ratio, which can make the learning signal fragile on long rollouts whose prefixes drift away from the teacher's typical support. We revisit this formulation from both theoretical and implementation perspectives. Theoretically, token-level OPD is biased relative to sequence-level reverse-KL minimization, but admits a substantially tighter worst-case variance bound; a controlled synthetic study further shows that stronger future-reward coupling increases gradient variance and destabilizes training. Empirically, we identify three failure modes of sampled-token OPD: imbalanced token-level supervision, unreliable teacher guidance on student-generated prefixes, and tokenizer or special-token mismatch. These findings motivate teacher top-K local support matching, a truncated reverse-KL objective that compares teacher and student distributions over a teacher-supported token set at each prefix, together with top-p rollout sampling and special-token masking. Across single-task reasoning and multi-task benchmarks spanning agentic and reasoning settings, this objective improves optimization stability and yields a +19.8% performance gain over standard sampled-token OPD baselines, providing a practical recipe for more stable on-policy distillation.

Abstract

Language Model Distillation

Explore similar work

Revisiting On-Policy Distillation: Empirical Failure Modes and Simple Fixes

Unbiased Top-kkk Estimation for On-Policy Distillation

Understanding Off- vs On-Policy Distillation: A Tale of Distinct Training Objectives

Unbiased Top- $k$ Estimation for On-Policy Distillation