ResOPD: Tail Residualization for Sparse On-Policy Distillation

Penghui Yang; Long Xing; Xuanlang Dai; Ziyu Liu; Kai Chen; Yuhang Zang

cs.LGOct 4, 2026

ResOPD: Tail Residualization for Sparse On-Policy Distillation

Authors: Penghui Yang, Long Xing, Xuanlang Dai, Ziyu Liu, Kai Chen, Yuhang Zang

Organizations: Shanghai Jiao Tong University · Shanghai Artificial Intelligence Laboratory · University of Science and Technology of China · Fudan University

Abstract

On-policy distillation (OPD) trains a student model on self-generated trajectories, but transmitting dense teacher distributions across long reasoning traces creates prohibitive communication and memory bottlenecks. Practical systems therefore rely on sparse teacher interfaces, typically transmitting either the sampled-token score or a small Top- $k$ distribution. However, this sparse setting faces a fundamental dilemma: sampled-token estimators are unbiased but suffer from severe gradient variance, whereas directly optimizing Top- $k$ objectives introduces systematic bias. To improve this trade-off, we propose ResOPD (On-Policy Distillation with Tail Residualization), which provides unbiased full-vocabulary reverse KL gradient estimation under on-policy sampling, with substantial variance reduction under the same sparse payload in the evaluated settings. ResOPD aggregates the unobserved vocabulary into an observable coarse tail event, computes its exact aggregate gradient, and samples only the fine-grained within-tail residual, which requires no additional teacher queries or forward passes. Extensive experiments demonstrate that ResOPD substantially reduces gradient variance, stabilizes online training dynamics, and improves downstream performance across the evaluated settings. These results establish ResOPD as an efficient, plug-and-play variance reduction primitive for sparse on-policy distillation.

Figures & tables

Figure 1: (a) Sparse OPD: the teacher model only sends the Top- $k$ head probabilities and the sampled token score, leaving the tail distribution unobserved. (b) Bias vs. variance trade-off: while deterministic Top- $k$ surrogates suffer from systematic gradient bias, ResOPD provides an unbiased gradient estimator and substantially reduces gradient variance. (c) Downstream evaluation: the radar chart shows benchmark scores across mathematics, coding, and science; while the relative ranking varies by task, ResOPD achieves the best overall performance.

Figure 2: Overview of ResOPD. (a) When the student sampled-token falls in the head, ResOPD uses only the exact coarse gradient; upon a tail hit, it adds a residual correction. (b) ResOPD suppresses between-event variance compared to HETS by centering tail samples around $\bm{\mu}_{B}$ , achieving superior alignment with the teacher distribution across rollout steps.

Unbiased estimator	$\bm{B}$	$\gamma$	Role
HETS	$\bm{0}$	$0$	structural baseline
ResOPD	$\bm{\mu}_{B}$	$0$	our estimator
SC-TA	$\bm{\mu}_{B}$	$1$	boundary estimator
Joint oracle	$\bm{\mu}_{\mathrm{tail}}$	$\gamma^{\star}$	unavailable optimum

Table 1: Coordinates in the two-control family.

Dataset	Models	Covariance Trace $\downarrow$				NSR $\downarrow$
Dataset	Models	Sampled	$\text{vOPD}_{\text{T-Top-}k}$	HETS	ResOPD	Sampled	$\text{vOPD}_{\text{T-Top-}k}$	HETS	ResOPD
GSM8K	0.8B / 4B	2.146	1.730	1.559	1.018	5.422	4.370	3.940	2.573
GSM8K	9B / 27B	0.215	0.192	0.107	0.068	6.202	5.560	3.101	1.955
GSM8K	4B / 27B	0.333	0.297	0.183	0.108	6.058	5.400	3.335	1.964
HMMT26	4B / 27B	0.208	0.189	0.092	0.052	6.589	5.977	2.914	1.630
DeepMath	4B / 27B	0.183	0.169	0.087	0.047	7.352	6.785	3.500	1.907

Table 2: Exact frozen-prefix covariance trace and NSR. Covariance traces of student logit-gradients are computed via full action enumeration and averaged across rollouts; NSR normalizes each trace by the squared

L_{2}

norm of the full-rKL gradient. Sparse estimators (

\text{vOPD}_{\text{T-Top-}k}

, HETS, ResOPD) share an identical teacher budget of

k=4

, with

\text{vOPD}_{\text{T-Top-}k}

adopting a detached conditional-KL baseline over this support. Best results per row are in bold , and second-best are underlined .

$\blacktriangledown$ Original Models
Method	Benchmark Accuracy (%)				Aggregate (%)
Method	DeepMath held-out (avg@2)	AIME26 (avg@16)	HMMT26 (avg@16)	OlymMATH (avg@2)	Avg. Score	Avg. AUC
Qwen3.5-4B Student	48.41	56.25	35.23	31.50	42.85	–
Qwen3.5-27B Teacher	89.29	80.62	62.88	63.25	74.01	–
$\blacktriangledown$ Alternative Sparse OPD (Biased)
Teacher Top- $k$ fKL	50.20	54.58	35.04	34.50	43.58	42.80

Table 3: Qwen3.5-27B-to-4B OPD results. DeepMath held-out and OlymMATH are evaluated using avg@2, while AIME26 and HMMT26 adopt avg@16 due to their smaller problem sets.

Diagnostic Metric	Top- $k$ rKL	$\text{vOPD}_{\text{T-Top-}k}$	HETS	ResOPD ( ours )	$\Delta$ (%) $\downarrow$
Support-mass gap ( $\times 10^{-3}$ )	11.17	3.53	5.17	3.50	$-32.3\%$
$\|$ Tail log-mass mismatch $\|$	0.715	0.858	0.723	0.322	$-55.5\%$
Token residual $M_{2}$ ( $\times 10^{-3}$ )	–	10.55	11.73	10.67	$-9.0\%$
Trajectory residual $M_{2}$ ( $\times 10^{-5}$ )	–	2.76	3.50	2.65	$-24.3\%$

Table 4: Aggregated online diagnostics across training (lower is better). Values represent the average over all 192 optimization steps.

\Delta

denotes the relative percentage reduction of ResOPD compared to HETS.

Figure 3: Training-time diagnostics (64 steps/epoch; lower is better). Panels (a) and (b) illustrate the student–teacher mismatch in Top-

k

support mass and tail log-mass across vOPD, HETS, ResOPD, and direct Teacher Top-

k

rKL. Panels (c) and (d) compare vOPD, HETS and ResOPD via the squared residual between the binary Top-

k

indicator and its student-predicted probability, evaluated at the token level and trajectory level, respectively. Detailed formulations are provided in Appendix D.1 .

Figure 4: Teacher negative log-likelihood (NLL) across training in the Math OPD setup. Evaluations are computed over the four core benchmark families with identical weights to the main suite: AIME26, HMMT26, DeepMath held-out, and OlymMATH. Lower values indicate higher fidelity to the Teacher distribution.

$\blacktriangledown$ Original Models
Method	Science Benchmarks (%)		Code Benchmarks (%)		Aggregate (%)
Method	GPQA-D (avg@10)	SuperGPQA (avg@3)	LCB-v5 (avg@3)	LCB-v6 (avg@3)	Avg. Score	Avg. AUC
Student	45.25	32.46	21.76	24.38	30.96	–
Teacher	56.06	48.87	45.51	41.71	48.04	–
$\blacktriangledown$ Unbiased rKL Estimators
Sampled-token OPD	50.30	36.87	43.31	39.43	42.48	41.21

Table 5: Distillation results from Qwen3-30B-A3B-Instruct to Qwen3-4B. Evaluations use avg@10 for GPQA-D and avg@3 for SuperGPQA, LCB-v5, and LCB-v6. AUC integrates performance across the fixed training trajectory (Base/10/25/50/100).

$k$	Method	AIME26	HMMT26	OlymMATH	Average	$\Delta$
1	HETS	60.62	37.69	34.75	44.35	$+0.11$
1	ResOPD	56.88	39.02	37.50	44.47	$+0.11$
4	HETS	59.17	36.93	34.50	43.53	$+1.01$
4	ResOPD	60.00	37.88	35.75	44.54	$+1.01$
16	HETS	57.92	39.39	35.75	44.35	$+0.70$
16	ResOPD	59.38	39.02	36.75	45.05	$+0.70$

Table 6: Performance under different Top-

k

support sizes (matched one-epoch budget). Evaluation setups follow Tab. 3 (avg@16 for AIME26/HMMT26 and avg@2 for OlymMATH).

\Delta

denotes the average improvement of ResOPD over HETS.

Figure 5: Online diagnostic metrics across teacher Top-

k

support sizes. (a) Between-event variance factor

P\tau_{p}

for HETS and ResOPD, characterizing the scaling coefficient of the head-versus-tail partition. (b) Tail log-mass mismatch reduction (

\text{HETS}-\text{ResOPD}

, measured via

\Delta|\log(\tau_{p}/\tau_{q})|

Rollout	Method	DeepMath	AIME26	HMMT26	OlymMATH	Avg.	AUC
Default	Naive Top- $k$ rKL	53.17	56.04	35.42	35.00	44.91	45.45
	HETS	53.17	59.17	36.93	34.50	45.94	46.12
	ResOPD	54.76	60.00	37.88	35.75	47.10	47.57
Full	Naive Top- $k$ rKL	45.44	55.42	35.80	34.25	42.72	44.96
	HETS	49.80	57.50	37.88	35.50	45.17	45.68
	ResOPD	57.14	60.42	37.88	36.00	47.86	46.79

Table 7: First-epoch Math performance under default vs. full-support rollouts. Benchmark scores and Avg. denote step-64 accuracy (%); AUC (%) measures the normalized learning curve area over steps

\{0,6,32,64\}

. Default : top-

p=0.95

, top-

k=20

; Full : top-

p=1.0

, top-

k=-1

Learning Rate	Final Score (%)			Trajectory AUC (%)
Learning Rate	HETS	ResOPD	$\Delta$	HETS	ResOPD	$\Delta$
$1\times 10^{-6}$	45.94	47.10	$+1.16$	46.12	47.57	$+1.45$
$2\times 10^{-6}$	45.94	46.61	$+0.67$	45.69	46.17	$+0.48$
$5\times 10^{-6}$	42.54	46.35	+3.81	43.88	47.77	+3.89

Table 8: Learning-rate sensitivity of HETS vs. ResOPD.

\Delta

denotes the improvement of ResOPD over HETS.

Dataset	$\mathcal{V}_{\mathrm{SC}}/\mathcal{V}_{\mathrm{ResOPD}}$	$\Delta$	ResOPD-Lower Prefixes
DeepMath	0.512	$-0.126$	0 6.8%
MegaScience	0.532	$-0.207$	11.3%
SynLogic	0.558	$-0.161$	18.0%

Table 9: SC-TA vs. ResOPD covariance trace ordering across natural prompt distributions.

Method	Base (Step 0)	Step 6	Step 32	Step 64	Step 128	Step 192
HETS	42.85	48.86	45.04	45.94	46.60	45.72
ResOPD	42.85	48.23	47.95	47.10	47.71	50.12
SC-TA	42.85	47.66	46.69	47.26	48.27	46.33

Table 10: Matched evaluation curves across training updates (%). Evaluated on Qwen3.5-27B-to-4B distillation across AIME26/HMMT26 (avg@16) and held-out DeepMath/OlymMATH (avg@2).

Appendix figures & tables7 assets

Supplementary material from the paper’s appendix.

Appendix

Category	Hyperparameter	Configuration / Value
Models	Student model	Qwen3.5-4B
Models	Teacher model	Qwen3.5-27B
Data	Training dataset	1,024 prompts sampled from DeepMath-103K
Data	Max prompt length	2,048
Sampling	Enable thinking	True
	Max response length	32,768 (training and evaluation)

Appendix

Table 11: Primary Math OPD experimental configurations and hyperparameters.

Category	Hyperparameter	Configuration / Value
Models	Student model	Qwen3-4B
Models	Teacher model	Qwen3-30B-A3B-Instruct-2507
Data	Training dataset	30,000 deduplicated prompts;
		Math/Code/Science: 10,000 each
	Max prompt length	1,024
Sampling	Enable thinking	False

Appendix

Table 12: Cross-domain Qwen3 OPD experimental configurations and hyperparameters.

Figure 6: Step-wise learning curves corresponding to Tab. 3 on DeepMath held-out (left) and external competitions (right) across 192 optimizer steps.

Figure 7: Evaluation trajectories during cross-domain transfer distillation from Qwen3-30B-A3B-Instruct to Qwen3-4B across optimizer steps (0, 10, 25, 50, and 100).

Dataset Student/Teacher	$k$	Payload (B/token)	Student mass	Head Hit	Var. Removed (%) $\uparrow$		NSR $\downarrow$
Dataset Student/Teacher	$k$	Payload (B/token)	Student mass	Head Hit	HETS	ResOPD	HETS	ResOPD
GSM8K 0.8B/4B	1	16	$75.0\%$	$74.8\%$	$-0.5$	$\mathbf{20.5}$	5.446	$\mathbf{4.310}$
	4	40	$93.1\%$	$92.9\%$	27.3	$\mathbf{52.5}$	3.940	$\mathbf{2.573}$
	16	136	$97.7\%$	$97.9\%$	62.6	$\mathbf{75.3}$	2.027	$\mathbf{1.337}$
	64	520	$99.3\%$	$99.7\%$	85.1	$\mathbf{89.5}$	0.807	$\mathbf{0.570}$
HMMT26 4B/27B	1	16	$89.1\%$	$89.3\%$	$-1.0$	$\mathbf{32.0}$	6.656	$\mathbf{4.479}$

Appendix

Table 13: Support-size (

k

) sweep under fixed Teacher Top-

k

payloads. “Var. Removed” measures the empirical covariance-trace reduction relative to sampled-token OPD. Payload accounts for

k

Top-

k

entries plus the sampled token (INT32 ID and FP32 log-probability each). Student mass and Head Hit denote the predicted and empirical probabilities of landing in the Top-

k

support, respectively. Best variance reduction and NSR per setting are in bold .

Figure 8: Comprehensive frozen-prefix support-size ( $k$ ) diagnostics across benchmarks. ResOPD consistently reduces the dimensionless NSR across all evaluated support sizes on GSM8K and HMMT (log scale). Axis annotations detail the transmitted Teacher support size $k$ , the corresponding communication payload (bytes/token), and the empirical head-hit probability. Error bars denote 95% trajectory-level bootstrap confidence intervals.

Figure 9: Training dynamics and learning-rate sensitivity of HETS ( orange ) vs. ResOPD ( blue ) on DeepMath across three learning rates (

1\times 10^{-6}

solid,

2\times 10^{-6}

dashed,

5\times 10^{-6}

dotted). (A) Pre-clipping global gradient norm (inset displays full-scale unclipped peaks over early steps). (B) Support-mass gap (

Q-P

) on the transmitted Teacher Top-4 tokens. (C, D) Token-level and trajectory-level second residual moments (

M_{2}

). While HETS exhibits large gradient-norm excursions and growing mass/residual discrepancies as the learning rate scales up to

5\times 10^{-6}

, ResOPD maintains stable dynamics and stays close to its low-rate trajectories throughout optimization.

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

Look Before You Select: Rethinking Vocabulary Sparsification in On-Policy Distillation
Sep 28, 2026Yongliang Miao, Shuang Liu, Yanguang Liu +2Language Model Distillation On-Policy Distillation
The Chinese University of Hong Kong, Shenzhen · Carnegie Mellon University · New Jersey Institute of Technology +1
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
Does On-Policy Distillation Need to Complete Every Rollout?
Oct 15, 2026Zhengxin Zhang, Maosen Tang, Ning Wang +6On-Policy Distillation Language Model Distillation
Cornell University · Georgia Tech · Alibaba Group

Sep 28, 2026J/K move · Enter open · S save

Look Before You Select: Rethinking Vocabulary Sparsification in On-Policy Distillation

Yongliang Miao, Shuang Liu, Yanguang Liu, Yandong Bai+1

The Chinese University of Hong Kong, Shenzhen · Carnegie Mellon University · New Jersey Institute of Technology+1

Language Model Distillation On-Policy Distillation Memory-Efficient Training

Open paper AskPDF arXiv

On-policy distillation (OPD) uses teacher correction on student-generated responses. Full-vocabulary correction can provide important corrections even for tokens that the student assigns low probability, but backpropagating through all token logits becomes memory-intensive for long sequences. Existing memory-saving approaches estimate corrections from sampled tokens or restrict supervision to the student's TopK tokens, introducing sampling noise or changing the full-vocabulary correction. We introduce SparseOPD, which uses full-vocabulary teacher correction to determine which corrections matter before selecting the token logits to differentiate. SparseOPD first constructs the full-vocabulary correction without retaining its backward graph, then selects tokens by correction magnitude rather than student probability. Signed residual compensation preserves the total promoting and suppressing correction mass, while correction-aware budget allocation distributes the sparse support across positions. Finally, the update backpropagates only through the selected token logits. Across six task--scale settings spanning mathematics, chemistry QA, and multimodal reasoning, SparseOPD outperforms Sampled Token and TopK in task-average accuracy and matches or exceeds Full Vocabulary. Gradient cosine similarity reaches 99% on 4B mathematics, while 8K full-parameter profiling shows 70.5% lower backward memory.

Abstract

Language Model Distillation

Explore similar work

Look Before You Select: Rethinking Vocabulary Sparsification in On-Policy Distillation

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

Does On-Policy Distillation Need to Complete Every Rollout?