Compatibility Matters: A Framework for Token-Level Teacher Selection in On-Policy Distillation

Stanley Wei; Jiin Woo; Kaan Ozkara; Youngsuk Park

cs.LGOct 15, 2026

Compatibility Matters: A Framework for Token-Level Teacher Selection in On-Policy Distillation

Authors: Stanley Wei, Jiin Woo, Kaan Ozkara, Youngsuk Park

Organizations: Princeton University · Amazon Annapurna Labs

Abstract

On-policy distillation (OPD) has recently been introduced as a promising paradigm for continual learning, where a student policy can be supervised at the token-level. Typically, OPD is performed using a single teacher per prompt rollout (or even per training run). However, this ignores a whole other axis of flexibility that is available due to the on-policy nature - token-level teacher selection during rollouts. Building on this insight, we theoretically explore what makes a teacher and a student compatible for OPD, and leverage our findings to derive several empirical heuristics for teacher selection. In particular, our framework captures several desiderata for OPD: compatible teacher-student reasoning patterns, teacher competence, and induced student diversity. Theoretically, we categorize several types of teacher routing schemes, and we prove that under different scenarios, different routing schemes are likely to be beneficial for optimization and student diversity. We back up our findings with extensive experimental results in domains such as math, coding, kernel generation, and visual reasoning, where we find upwards of 10 percentage points of improvement on pass@ $k$ when compared to vanilla OPD and the teacher experts on certain tasks, indicating both capability improvement and diversity over high-quality responses.

Figures & tables

Figure 1: Triton-only entropy routing (

\tau=0.1

; iteration 300). Here, we plot the score gap of the kernel teacher minus the thinking teacher. The full trace is included in Section H.3 , where we can observe that the blue (positive) regions on this plot correspond roughly to parts of the student rollout where Triton code is being generated, and the red parts correspond to general reasoning or thinking behavior to plan the kernel code.

	Joint fast 0		Math-portion	Kernel-portion
Method	pass@1	pass@16	pass@1	pass@1
Initial student model
Qwen3-4B	1.79	15.04	15.08	12.04
Teachers
Math RL expert	2.87	19.92	23.85	11.46
Dr.Kernel-8B	1.42	14.06	6.74	17.57

Table 1: AIME 2025

\times

KernelBench L2 (%). Joint fast 0 requires correctness of both components in one response (higher is better); we also report the math-portion and kernel-portion breakdown for completeness. Note that for the sampled advantage Hellinger routing, the scores on each individual section recover roughly the teacher performance, and together yield the highest joint score.

Panel of Figure 2 — Figure 2: Abridged rollout on algorithmic coding task by a student trained with feature routing, at step 300. Here, the text is colored by the code-minus-math routing score. Note that the coloring demonstrates that the reasoning part of the rollout leans more towards the math teacher, and the Python implementation is markedly closer to the code teacher. Omitted text is marked by $[\ldots]$ , and the full per-token score gap and complete rollout are in Figure 3 and Section H.2 .

Initial student model
	Fast 0
Method	pass@1	pass@16
Qwen3-4B	9.63	53.00
Teachers
Dr.Kernel-8B	12.63	78.00
Without token-level routing

Table 2: Triton-only evaluation on 100 KernelBench Level 2 tasks (%). Fast 0 measures kernel correctness.

Initial student model
Method	HumanEval+		MBPP+
	pass@1	pass@64	pass@1	pass@64
Qwen3-4B	76.28	92.42	58.98	76.78
Teachers
Code RL expert	81.15	93.59	58.38	77.29
Math RL expert	78.24	93.08	60.45	77.23

Table 3: Algorithmic coding correctness (%). Although our feature-matching selection does not achieve the highest pass@1, it shows improvements for pass@64, indicating the benefits of diversity of this router.

Appendix figures & tables2 assets

Supplementary material from the paper’s appendix.

Appendix

Initial student model
Method	WeMath	MathVista-mini	CharXiv RQ
Qwen3-VL-2B-Instruct	60.97	62.70	29.10
Teachers
Qwen3-VL-4B-Instruct	76.52	74.20	39.30
Without token-level routing
OPD (visual teacher)	66.31	66.40	32.00

Appendix

Table 4: Visual reasoning question-level pass@1 (%).

Figure 3: Code-only final-layer feature routing (

\tau=0.1

; step 300). We plot the code-minus-math score gap here against response token index. Here, the dashed colored segments show the mean gaps during explanation and implementation.

On-Policy Distillation

On-policy distillation trains a model on its own outputs using teacher guidance, aligning supervision with the states it encounters at deployment.

See all 507 papers

Spend Teacher Tokens Where They Matter: Success-Referenced On-Policy Distillation
Oct 2, 2026Xiang Chen, Futao Su, Kong Wang +2Language Model Distillation On-Policy Distillation
Tongji University · Xiaohongshu · Shanghai Jiao Tong University +1
Not All Disagreement Is Learnable: Token Teachability in On-Policy Distillation
May 26, 2026Yuanyi Wang, Su Lu, Yanggan Gu +6Selective Knowledge Distillation Language Model Distillation
The Hong Kong Polytechnic University, PolyU · Hong Kong Polytechnic University Daya Bay Technology and Innovation Research Institute
TIP: Token Importance in On-Policy Distillation
Apr 15, 2026Yuanda Xu, Hejian Sang, Zhengze Zhou +3Knowledge Distillation On-Policy Distillation

Oct 2, 2026J/K move · Enter open · S save

Spend Teacher Tokens Where They Matter: Success-Referenced On-Policy Distillation

Xiang Chen, Futao Su, Kong Wang, Jiayi Chen+1

Tongji University · Xiaohongshu · Shanghai Jiao Tong University+1

Language Model Distillation On-Policy Distillation Knowledge Distillation Selective Knowledge Distillation

Open paper AskPDF arXiv

On-policy distillation (OPD) combines student-generated rollouts with dense token-level supervision from a teacher, but providing such supervision for every rollout requires substantial teacher computation. We introduce Success-Referenced On-Policy Distillation (SR-OPD), which reduces this cost by selecting which prompts and rollouts receive teacher supervision. When the student produces both successful and failed rollouts for the same prompt, a successful rollout can serve as a natural reference for selecting failed rollouts. SR-OPD therefore focuses on such prompts and prioritizes failed rollouts whose hidden-state trajectories show sustained divergence from a successful reference, while accounting for estimated teacher-input cost. Across three teacher-student pairs and six mathematical reasoning benchmarks, SR-OPD uses only 3.46-5.02% of the teacher-input tokens required by Vanilla OPD in the one-pass setting while maintaining comparable reasoning performance. Under a controlled setting matched to 5% of Vanilla OPD's teacher-input budget, further experiments support both key design choices: focusing supervision on prompts with both successful and failed rollouts, and using successful rollouts to guide failure selection. These results indicate that a student's own successful behavior can serve as a useful reference for allocating teacher supervision under a fixed teacher-input budget.

Abstract

On-Policy Distillation

Explore similar work

Spend Teacher Tokens Where They Matter: Success-Referenced On-Policy Distillation

Not All Disagreement Is Learnable: Token Teachability in On-Policy Distillation

TIP: Token Importance in On-Policy Distillation