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
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 μB , achieving superior alignment with the teacher distribution across rollout steps.
Unbiased estimator
B
γ
Role
HETS
0
0
structural baseline
ResOPD
μB
0
our estimator
SC-TA
μB
1
boundary estimator
Joint oracle
μtail
γ⋆
unavailable optimum
Table 1: Coordinates in the two-control family.
Dataset
Models
Covariance Trace ↓
NSR ↓
Sampled
vOPDT-Top-k
HETS
ResOPD
Sampled
vOPDT-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 L2 norm of the full-rKL gradient. Sparse estimators ( vOPDT-Top-k , HETS, ResOPD) share an identical teacher budget of k=4 , with vOPDT-Top-k adopting a detached conditional-KL baseline over this support. Best results per row are in bold , and second-best are underlined .
Method
Benchmark Accuracy (%)
Aggregate (%)
DeepMath held-out (avg@2)
AIME26 (avg@16)
HMMT26 (avg@16)
OlymMATH (avg@2)
Avg. Score
Avg. AUC
▼ Original Models
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
–
▼ 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
vOPDT-Top-k
HETS
ResOPD ( ours )
Δ (%) ↓
Support-mass gap ( ×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 M2 ( ×10−3 )
–
10.55
11.73
10.67
−9.0%
Trajectory residual M2 ( ×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. Δ 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.
Method
Science Benchmarks (%)
Code Benchmarks (%)
Aggregate (%)
GPQA-D (avg@10)
SuperGPQA (avg@3)
LCB-v5 (avg@3)
LCB-v6 (avg@3)
Avg. Score
Avg. AUC
▼ Original Models
Student
45.25
32.46
21.76
24.38
30.96
–
Teacher
56.06
48.87
45.51
41.71
48.04
–
▼ 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
Δ
1
HETS
60.62
37.69
34.75
44.35
+0.11
ResOPD
56.88
39.02
37.50
44.47
4
HETS
59.17
36.93
34.50
43.53
+1.01
ResOPD
60.00
37.88
35.75
44.54
16
HETS
57.92
39.39
35.75
44.35
+0.70
ResOPD
59.38
39.02
36.75
45.05
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). Δ 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τp for HETS and ResOPD, characterizing the scaling coefficient of the head-versus-tail partition. (b) Tail log-mass mismatch reduction ( HETS−ResOPD , measured via Δ∣log(τp/τ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 (%)
HETS
ResOPD
Δ
HETS
ResOPD
Δ
1×10−6
45.94
47.10
+1.16
46.12
47.57
+1.45
2×10−6
45.94
46.61
+0.67
45.69
46.17
+0.48
5×10−6
42.54
46.35
+3.81
43.88
47.77
+3.89
Table 8: Learning-rate sensitivity of HETS vs. ResOPD. Δ denotes the improvement of ResOPD over HETS.
Dataset
VSC/VResOPD
Δ
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
Teacher model
Qwen3.5-27B
Data
Training dataset
1,024 prompts sampled from DeepMath-103K
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
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 (%) ↑
NSR ↓
HETS
ResOPD
HETS
ResOPD
GSM8K 0.8B/4B
1
16
75.0%
74.8%
−0.5
20.5
5.446
4.310
4
40
93.1%
92.9%
27.3
52.5
3.940
2.573
16
136
97.7%
97.9%
62.6
75.3
2.027
1.337
64
520
99.3%
99.7%
85.1
89.5
0.807
0.570
HMMT26 4B/27B
1
16
89.1%
89.3%
−1.0
32.0
6.656
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×10−6 solid, 2×10−6 dashed, 5×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 ( M2 ). While HETS exhibits large gradient-norm excursions and growing mass/residual discrepancies as the learning rate scales up to 5×10−6 , ResOPD maintains stable dynamics and stays close to its low-rate trajectories throughout optimization.
Language model distillation transfers a larger model’s capabilities to smaller or specialized models using outputs, probabilities, or internal representations.
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.
Yongliang Miao, Shuang Liu, Yanguang Liu +2
The Chinese University of Hong Kong, Shenzhen · Carnegie Mellon University · New Jersey Institute of Technology +1
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.
Yuqian Fu, Haohuan Huang, Kaiwen Jiang +4
State Key Laboratory of Multimodal Artificial Intelligence Systems, CASIA · School of Artificial Intelligence, UCAS · Fudan University +1
On-policy distillation (OPD) transfers teacher knowledge to a student through feedback on student-generated rollouts, but generating complete rollouts and computing per-token teacher feedback are computationally expensive. Existing truncation methods save computation by stopping rollouts early. However, they do not compensate for truncated token gradients, which biases the gradient estimates and changes Vanilla OPD's original training objective. Furthermore, their truncation decisions rely on assumptions about supervision quality or criteria that may not generalize across settings. We introduce Stochastic Truncated On-Policy Distillation (STOPD), which accelerates OPD by randomly truncating student rollouts and correcting the resulting sampling bias with blockwise inverse-probability weighting (IPW). At each training step, we sample a stopping position for each rollout before generation and let the student generate until that position or natural termination. We obtain teacher feedback on every generated token and use blockwise IPW to compensate for the lower probability of retaining later tokens while controlling the weights. We prove that this weighting preserves Vanilla OPD's gradient sum in expectation and derive one-step update-error bounds for SGD, Adam, and Muon. We evaluate STOPD across multiple teacher-student model pairs, task domains, and optimizers. STOPD reduces generated tokens by 25-48% while maintaining accuracy comparable to Vanilla OPD, achieving 1.3-2.2x training speedups.