DiffGate: Difficulty-Gated Teacher Guidance for On-Policy Distillation

Karn Tiwari; Varnith Chordia; Prathosh A P

cs.LGOct 3, 2026

DiffGate: Difficulty-Gated Teacher Guidance for On-Policy Distillation

Authors: Karn Tiwari, Varnith Chordia, Prathosh A P

Organizations: Indian Institute of Science, Bengaluru, India · Snap Inc.

Abstract

On-policy distillation (OPD) has emerged as a widely used paradigm for post-training large language models, reducing the train--test mismatch of conventional distillation by supervising the student on its own generated trajectories. However, existing OPD objectives remain largely token-local and outcome-agnostic, optimizing teacher--student agreement at each prefix despite reasoning quality being determined at the trajectory level. Reinforcement learning with verifiable rewards (RLVR), particularly Group Relative Policy Optimization (GRPO), provides complementary outcome-level supervision but suffers from sparse rewards and coarse credit assignment. We show that OPD and RLVR exhibit complementary blind spots: teacher signals provide dense local guidance but are weakly aligned with rollout correctness, whereas group-relative rewards capture task success but provide coarse token-level credit and vanish on all-failure groups. We introduce DiffGate, an outcome-gated objective that combines GRPO with selective, bounded teacher guidance. Teacher supervision is applied only to failed trajectories, scaled by group difficulty, and smoothly bounded to prevent extreme teacher--student discrepancies from dominating optimization. The verifier therefore determines which trajectories receive teacher guidance, while the teacher provides dense token-level update directions within those trajectories. Across Qwen3-0.6B and Qwen3-1.7B students, DiffGate improves code avg@8 over matched GRPO by $+1.7$ and $+1.8$ points and pass@8 by $+1.6$ and $+5.7$ points, respectively. On mathematics, avg@8 remains within $0.5$ points of GRPO while pass@8 improves by $+1.1$ and $+3.9$ points. Overall, DiffGate improves pass@8 across all four model--domain settings, demonstrating improved solution coverage under our evaluation protocol.

Figures & tables

Figure 1: Teacher disagreement alone does not reliably identify incorrect rollouts. For a distillation-only Qwen3-1.7B student evaluated with a frozen Qwen3-8B teacher: (a) average token-level teacher–student disagreement distinguishes correct from incorrect rollouts only at chance level on mathematics (AUROC $0.50$ and $0.49$ ) and weakly on code (AUROC $0.60$ ); (b) aggregate disagreement becomes more predictive mainly when it also reflects response length; and (c) within an incorrect rollout, large disagreement can occur on ordinary tokens rather than on the token associated with the decisive error. These observations motivate DiffGate to use the verifier to determine which trajectories receive teacher guidance, while the teacher provides the local token-level update direction. Further analysis is provided in Appendix Section B .

Figure 2: Overview of DiffGate. The student generates on-policy trajectories that receive sequence-level correctness feedback from a verifier and token-level guidance from a frozen teacher. DiffGate bounds the sampled-token teacher–student log-probability gap, applies teacher guidance only to failed trajectories, and scales its strength according to the group solve rate. The resulting token-level advantage combines local teacher guidance with the global GRPO outcome signal to update the student policy. See Algorithm 1 for details.

Student	Method	Minerva		Olympiad		AIME24		AIME25		HMMT25		Average		Science (OOD)
Student	Method	avg@8	pass@8	avg@8	pass@8	avg@8	pass@8	avg@8	pass@8	avg@8	pass@8	avg@8	pass@8	avg@8	pass@8
0.6B	Base	10.7	27.7	15.1	37.9	2.5	6.6	0.8	6.8	0.4	3.3	5.9	16.4	21.7	48.8
	OPD (forward KL)	14.1	35.7	23.5	49.8	3.8	15.6	3.8	16.4	0.8	7.6	9.2	25.0	29.4	59.1
	OPD (reverse KL)	18.2	36.7	29.7	54.7	2.5	15.4	7.9	19.9	1.7	9.1	12.0	27.2	30.1	58.2
	EOPD	17.1	35.2	29.9	53.9	6.2	17.8	5.4	15.1	2.1	6.3	12.2	25.7	30.2	58.1
	Uni-OPD	18.9	36.9	30.5	55.8	6.2	15.9	8.8	24.0	0.8	5.0	13.0	27.5	30.6	57.7

Table 1: Mathematical reasoning results. We report avg@8 and pass@8, with GRPO serving as the outcome-only baseline. Bold and underlined values indicate the best and second-best student results, respectively. The science (OOD) columns report average performance on GPQA, MMLU-STEM, and OlympiadBench-Physics (Table 6 ), evaluated using the mathematics-trained students without additional science training.

Student	Method	MBPP+		HumanEval+		LCB		Average
Student	Method	avg@8	pass@8	avg@8	pass@8	avg@8	pass@8	avg@8	pass@8
0.6B	Base	30.1	53.4	22.6	45.5	6.1	14.1	19.6	37.7
	OPD (forward KL)	37.8	63.7	31.4	60.8	7.7	18.7	25.6	47.7
	OPD (reverse KL)	46.0	65.8	33.4	56.5	10.4	20.8	29.9	47.7
	EOPD	46.0	65.3	38.5	58.8	10.3	21.5	31.6	48.6
	Uni-OPD	45.2	65.5	36.9	58.1	10.8	20.5	31.0	48.0

Table 2: Code generation results. We report avg@8 and pass@8 on MBPP+, HumanEval+, and LiveCodeBench. Bold and underlined values indicate the best and second-best student results, respectively.

Figure 3: Why DiffGate helps. Maths, 1.7B student; all methods use the same data, teacher, and step budget. (a) Student gradient norm over training: the bar denotes the median and the whisker spans minimum to maximum, shown on a log scale. (b) Policy entropy, shown as a five-step rolling mean; the dotted line marks $0.1$ . (c) The two components of DiffGate’s update on failed rollouts: the applied teacher term per token (dark blue) and the magnitude of the GRPO advantage per token (light blue). The GRPO advantage is zero for an all-failure group and, with $n=16$ , is approximately $0.25$ or larger for groups with $1$ – $8$ successes and approximately $1.10$ or larger for groups with $9$ or more successes, ignoring the small $\varepsilon$ term; values are weighted by the measured frequency of each group type. The grey region indicates this approximate lower-bound contribution. (d) Applied teacher term per response token under difficulty-gated and additive weighting on the same rollouts.

Figure 4: Performance over training. avg@8 (top) and pass@8 (bottom) across training checkpoints for the 1.7B student on the code benchmarks in Table 2 . Step $0$ denotes the base model (dotted line), enabling comparison of how each method improves accuracy and solution coverage.

Appendix figures & tables20 assets

Supplementary material from the paper’s appendix.

Appendix

Notation	Description
General setup
$\mathcal{D}$	Distribution over training prompts.
$x$	Prompt sampled from $\mathcal{D}$ .
$\pi_{\theta}$	Student policy parameterized by $\theta$ .
$\pi_{\theta_{\mathrm{old}}}$	Frozen behavior policy used to generate trajectories for the current policy update.
$\pi_{\mathrm{T}}$	Fixed teacher policy.

Appendix

Table 3: Summary of notation. Symbols used for the policy setup, outcome supervision, teacher guidance, and DiffGate optimization.

Group	Setting	Value
Models	Student	Qwen3-0.6B / Qwen3-1.7B (also Qwen3-4B, App. E.2 )
	Teacher	Qwen3-8B, frozen (also Qwen3-14B, App. E.2 )
	Mode	Non-thinking; shared chat template
Optimization	Optimizer	AdamW
	Learning rate	$3\times 10^{-6}$ , cosine schedule, no warmup
	Prompts per step	$128$ ; PPO mini-batch $32$

Appendix

Table 4: Training and evaluation configuration. Shared across methods unless otherwise specified. Training uses

128

prompts per step. The mathematics and code training sets contain

21{,}928

and

4{,}082

prompts, respectively; the reported budgets correspond to approximately

1.00

and

1.94

dataset passes. GRPO-based methods differ only in their teacher-guidance configuration.

p=\frac{1}{n}\sum_{i=1}^{n}R_{i},\qquad A_{i}^{\mathrm{GRPO}}=\frac{R_{i}-\bar{R}}{\sigma_{R}+\varepsilon},\qquad\bar{R}=\frac{1}{n}\sum_{j=1}^{n}R_{j}

Appendix

Algorithm 1 DiffGate: Difficulty-Gated On-Policy Distillation

Arm	Rollouts / prompt	Math 0.6B	Math 1.7B	Code 0.6B	Code 1.7B
		171 steps		62 steps
OPD (forward KL)	1	4	4.7	1.5	5
OPD (reverse KL)	1	4.1	4.5	1.4	1.5
EOPD	1	4.2	4.7	4	4
Uni-OPD	8	15	17	14	13
GRPO	16	36	39	16	19

Appendix

Table 5: Training cost. Wall-clock hours on

4\times

A100-80GB for the runs behind main Tables 1 and 2 .

Figure 5: Hyperparameter sweep, 1.7B student. Teacher strength $\lambda$ (left) and difficulty sharpness $\gamma$ (right), for mathematics (top) and code (bottom); avg@8 on the left axis, pass@8 on the right. Every point is DiffGate at a different setting, so this figure varies one constant rather than comparing arms. Hollow markers are the reported configuration $\lambda=\gamma=1$ .

Method	GPQA		MMLU-STEM		OlympiadBench-Physics		Average
	avg@8	pass@8	avg@8	pass@8	avg@8	pass@8	avg@8	pass@8
Qwen3-0.6B mathematics student
Base	13.0	50.8	48.5	86.2	3.6	9.4	21.7	48.8
OPD (forward KL)	26.1	77.3	58.5	88.5	3.6	11.4	29.4	59.1
OPD (reverse KL)	26.5	74.2	59.0	87.8	4.8	12.6	30.1	58.2
EOPD	26.1	74.3	59.7	88.1	4.7	12.0	30.2	58.1

Appendix

Table 6: Out-of-distribution transfer to Science. avg@8 / pass@8 of mathematics-trained students on GPQA, MMLU-STEM and OlympiadBench-Physics, with no science training; the final columns average the three and are the science columns of Table 1 . Bold / underline : best / second-best per student.

Method	Minerva		Olympiad		AIME24		AIME25		HMMT25		Average
Method	avg@8	pass@8	avg@8	pass@8	avg@8	pass@8	avg@8	pass@8	avg@8	pass@8	avg@8	pass@8
Base (Qwen3-4B)	38.8	52.2	51.2	71.7	20.4	43.3	17.5	35.6	7.5	24.6	27.1	45.5
OPD (forward KL)	38.9	56.3	48.5	70.4	20.8	36.5	15.0	32.4	7.9	19.1	26.2	43.0
OPD (reverse KL)	39.0	54.5	49.4	70.3	20.8	40.3	17.5	39.2	8.8	20.8	27.1	45.0
DiffGate	41.8	54.3	53.3	72.1	22.1	47.1	21.2	40.6	11.7	21.7	30.0	47.2
Teacher 14B	43.6	57.4	55.4	74.8	30.0	49.2	20.0	34.5	9.2	22.1	31.6	47.6

Appendix

Table 7: Mathematical reasoning with a Qwen3-4B student and Qwen3-14B teacher. Avg@8 and pass@8 across five mathematics benchmarks. Bold and underline denote the best and second-best student results, respectively.

Method	MBPP+		HumanEval+		LCB		Average
Method	avg@8	pass@8	avg@8	pass@8	avg@8	pass@8	avg@8	pass@8
Base (Qwen3-4B)	62.4	78.5	73.9	89.8	26.0	41.1	54.1	69.8
OPD (forward KL)	64.0	78.7	67.8	86.2	22.8	38.1	51.6	67.7
OPD (reverse KL)	66.9	77.6	71.8	85.7	25.6	39.7	54.8	67.7
DiffGate	67.1	78.7	77.1	89.5	27.6	43.3	57.3	70.5
Teacher 8B	70.8	81.6	80.0	90.2	28.1	43.1	59.6	71.7

Appendix

Table 8: Code generation with a Qwen3-4B student and Qwen3-8B teacher. Avg@8 and pass@8 on MBPP+, HumanEval+, and LiveCodeBench. Bold and underline denote the best and second-best student results, respectively.

$n=16$ rollouts per prompt
Method	Mathematics				Code
	0.6B		1.7B		0.6B		1.7B
	avg@8	pass@8	avg@8	pass@8	avg@8	pass@8	avg@8	pass@8
GRPO	14.2	28.4	21.4	35.2	31.3	48.2	43.4	55.7
DiffGate	13.7	29.5	21.1	39.1	33.0	49.8	45.2	61.4
$n=4$ rollouts per prompt

Appendix

Table 9: Effect of rollout group size. GRPO and DiffGate trained with

n{=}16

and

n{=}4

rollouts per prompt, holding data, teacher, optimizer, and budget fixed; averages over the benchmark sets of Tables 1 and 2 . The

n{=}16

rows are the runs of Tables 1 and 2 . Bold marks the better of the two arms within each group size.

Held-out accuracy
Method	Mathematics				Code
	0.6B		1.7B		0.6B		1.7B
	avg@8	pass@8	avg@8	pass@8	avg@8	pass@8	avg@8	pass@8
DiffGate, bounded $2\tanh(\Delta/2)$	13.7	29.5	21.1	39.1	33.0	49.8	45.2	61.4
DiffGate, unbounded (raw $\Delta$ )	12.7	28.2	19.9	38.0	31.9	49.0	43.7	60.3
bounded $-$ unbounded	$+1.0$	$+1.3$	$+1.2$	$+1.1$	$+1.1$	$+0.8$	$+1.5$	$+1.1$

Appendix

Table 10: Bounded versus unbounded teacher term. Top: held-out avg@8 / pass@8 of the two DiffGate variants. Bottom: median and maximum student gradient norm over training.

Figure 6: Student gradient norm over training. Bars: median. Whiskers: minimum to maximum step, with the maximum printed above. Log scale.

Objective	Mathematics				Code
	0.6B		1.7B		0.6B		1.7B
	avg@8	pass@8	avg@8	pass@8	avg@8	pass@8	avg@8	pass@8
GRPO + OPD (additive, $\gamma{=}0$ )	11.3	26.7	18.7	35.0	31.7	49.7	43.8	60.0
DiffGate ( $\gamma{=}1$ )	13.7	29.5	21.1	39.1	33.0	49.8	45.2	61.4

Appendix

Table 11: Difficulty gating against the additive baseline. Both arms are GRPO plus the same bounded teacher term on failed trajectories, matched on data, teacher,

n{=}16

rollouts, and step budget. They differ only in that term’s weight: constant (additive,

\gamma{=}0

) against the failed fraction of the group (DiffGate,

\gamma{=}1

Objective	Mathematics				Code
	0.6B		1.7B		0.6B		1.7B
	avg@8	pass@8	avg@8	pass@8	avg@8	pass@8	avg@8	pass@8
GRPO + bounded OPD (all rollouts)	11.9	27.4	20.2	36.8	31.4	49.5	43.2	61.3
GRPO + OPD (additive, all rollouts, raw $\Delta$ )	13.5	29.9	21.0	38.1	31.6	47.8	42.0	59.0
GRPO + bounded OPD (difficulty-weighted, all rollouts)	13.3	28.4	19.8	36.5	32.4	49.7	45.8	61.7
DiffGate	13.7	29.5	21.1	39.1	33.0	49.8	45.2	61.4

Appendix

Table 12: Selector and additive controls. All rows add the teacher term to GRPO; the first three apply it to every rollout. Bold and underlined : best and second-best per column.

Figure 7: Teacher-side training dynamics: GRPO, OPD and DiffGate. Columns are the four model–domain settings. GRPO and DiffGate use $n{=}16$ rollouts and the outcome reward, DiffGate adding the gated teacher term; the two OPD baselines use one rollout per prompt and no reward. All four share data, teacher, and step budget. Faint traces are per-step values, bold lines are a five-step rolling mean. Rows: mean $\Delta$ over all response tokens, and the teacher’s log-likelihood of the student’s own samples.

Group-level reward statistics
Quantity	Mathematics, 0.6B			Mathematics, 1.7B			Code, 0.6B			Code, 1.7B
	Start	GRPO	DiffGate	Start	GRPO	DiffGate	Start	GRPO	DiffGate	Start	GRPO	DiffGate
All-fail groups, $A_{\mathrm{GRPO}}\equiv 0$ (%)	32.4	7.5	10.8	35.6	2.6	2.6	35.4	22.3	19.6	13.7	9.2	5.7
All-pass groups, $A_{\mathrm{GRPO}}\equiv 0$ (%)	22.6	60.9	49.1	32.1	70.8	62.1	24.0	36.5	37.4	49.8	55.8	56.3
Prompts with no group signal (%)	55.0	68.5	59.8	67.7	73.5	64.8	59.4	58.8	57.0	63.5	65.0	62.1
Token-weighted mean $A_{\mathrm{GRPO}}$	-0.010	-0.051	-0.030	-0.010	-0.052	-0.037	+0.004	-0.065	-0.023	-0.014	-0.053	-0.023

Appendix

Table 13: Advantage and teacher-signal statistics over training. “Start” is the mean of the two arms at step 1; GRPO and DiffGate report averages over the final ten steps. Reward-uniform groups (all-fail or all-pass) have zero group-relative advantage. We measure teacher-signal statistics only on failed rollouts, and we average over the final block of all generated response tokens.

Figure 8: Reward-based routing of teacher guidance. the rows show: (1) the ratio of the all-token mean of $-\Delta$ to the failed-rollout mean of $|\Delta|$ , with the dotted line marking a ratio of $1$ ; (2) the fraction of the applied teacher contribution assigned to all-failure groups, together with the fraction of prompts belonging to such groups (dashed); (3) the effective gate factor, defined by the fraction of trajectories on which the gate fires multiplied by their mean difficulty weight; and (4) for DiffGate, the fraction of failed-rollout tokens on which the teacher signal opposes the outcome advantage.

Teacher disagreement and reward-signal availability
Quantity	Mathematics, 0.6B		Mathematics, 1.7B		Code, 0.6B		Code, 1.7B
	GRPO	DiffGate	GRPO	DiffGate	GRPO	DiffGate	GRPO	DiffGate
Correct/failed disagreement ratio	$0.99{\to}2.41$	$0.83{\to}1.34$	$0.85{\to}3.48$	$0.80{\to}1.57$	$1.10{\to}1.45$	$1.13{\to}1.29$	$1.68{\to}2.07$	$1.62{\to}1.72$
Teacher dose on GRPO-silent groups (%)	$55{\to}58$	$52{\to}52$	$78{\to}41$	$69{\to}26$	$69{\to}64$	$63{\to}58$	$58{\to}56$	$52{\to}39$
GRPO-silent groups, share of prompts (%)	$26{\to}8$	$28{\to}11$	$36{\to}3$	$31{\to}3$	$40{\to}22$	$32{\to}20$	$15{\to}9$	$13{\to}6$
Applied teacher guidance

Appendix

Table 14: Teacher-side diagnostics over training. Each entry reports start

\rightarrow

end, where start” averages the first five training steps and end” averages the final ten. For GRPO, teacher-derived quantities are computed only for diagnostic purposes and are never used for optimization; rows that require an applied teacher update are therefore left blank. The disagreement block reports the raw teacher–student log-probability difference

\Delta

before bounding, evaluated on teacher-scored tokens. The gradient row instead reports the median/maximum gradient norm over the full training run.

Figure 9: Training reward, response length and policy entropy. Rows: verifier pass rate of the $n{=}16$ rollouts on each step’s training prompts; mean response length in tokens; mean token entropy of the student. Faint traces are per-step values, bold lines a five-step rolling mean. DiffGate is a dense retrain; the remaining arms are the original runs. The reverse-KL arm’s solve rate is a single-sample diagnostic, not the quantity it optimises.

Base-student band	$n$	Teacher pass@16	Base	GRPO	DiffGate	$\Delta$	Wins (DiffGate:GRPO)
Code, 0.6B student (994 problems; teacher fails on 50% of the never-solved band, 31% overall)
Never solved (0/16)	616	50%	0.00	0.86	0.79	$-0.07$	16:21
1–4 / 16	140	96%	2.13	6.46	6.64	$+0.18$	22:17
5–11 / 16	132	98%	8.08	11.22	12.35	$+1.13$	33:17 †
12–16 / 16	106	98%	14.09	14.74	15.28	$+0.55$	9:3
Code, 1.7B student (994 problems; teacher fails on 58% of the never-solved band, 30% overall)

Appendix

Table 15: Performance across base-student difficulty bands. Problems are grouped by the number of successful responses produced by the base student out of

16

. “Teacher pass@16” denotes the fraction of problems in each band for which the Qwen3-8B teacher produces at least one correct response under the same sampling protocol. Base, GRPO, and DiffGate report the mean number of successful generations out of

16

, and

\Delta

denotes DiffGate minus GRPO. “Wins” counts problems favoring DiffGate versus GRPO for which a two-proportion

z

-test comparing their

16

sampled outcomes satisfies

|z|\geq 2

, with DiffGate listed first.

Figure 10: Performance over training. avg@8 (top) and pass@8 (bottom) on the benchmark sets of Tables 1 and 2 ; step $0$ is the base model (dotted line). GRPO and DiffGate were evaluated at steps $50$ and $62$ on code, and GRPO at steps $50$ , $100$ , $150$ and $171$ on mathematics; their other points are filled in as described in the text.

Explore similar work

Verify Before You Distill: Prompt-Level Teacher Gating for On-Policy Distillation
Sep 2, 2026Zhiwei Zhang, Zechen Sun, Fei Zhao +6On-Policy Distillation RL Post-Training
AllSpark Team
When Teacher Guidance Misleads: Reward-Aligned On-Policy Distillation
Aug 28, 2026Siyuan Gan, Yuhan Li, Xiran Wang +5Language Model Distillation On-Policy Distillation
State Key Laboratory of Novel Software Technology, Nanjing University, Nanjing, China · Shanghai Artificial Intelligence Laboratory, Shanghai, 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

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

Verify Before You Distill: Prompt-Level Teacher Gating for On-Policy Distillation

Zhiwei Zhang, Zechen Sun, Fei Zhao, Kang Peng+5

AllSpark Team

On-Policy Distillation RL Post-Training

Open paper AskPDF arXiv

On-policy distillation (OPD) accelerates post-training by providing dense token-level supervision from a frozen teacher on the student's own rollouts. Vanilla OPD applies this supervision uniformly across prompts, without checking whether the teacher is reliable for each prompt. Because reverse KL is mode-seeking, a confidently wrong teacher can induce a strong yet misleading update. Distributional proxies, such as entropy or teacher-student likelihood agreement, measure uncertainty or agreement but do not directly verify outcome correctness. We introduce Teacher-Gated On-Policy Distillation (TGOPD), built on the principle that teacher reliability should be verified at the prompt level before dense supervision is admitted. TGOPD estimates reliability from a small set of verifier-scored teacher probes and routes each prompt exclusively to dense OPD when the reliability check passes or to verifier-grounded GRPO otherwise. Across 4B and 35B students in mathematics, code, and instruction following, TGOPD outperforms Vanilla OPD in all six single-domain settings and achieves higher seven-benchmark averages at both scales under multi-domain training. By using otherwise-idle teacher capacity for reliability estimation, TGOPD also reduces teacher-side compute waste in asynchronous OPD, increasing teacher-node GPU utilization from 9.8% to 78.9% in the measured 4B single-domain run.

Abstract

Explore similar work

Verify Before You Distill: Prompt-Level Teacher Gating for On-Policy Distillation

When Teacher Guidance Misleads: Reward-Aligned On-Policy Distillation

Reward-Gated On-Policy Distillation