Transfer-Stratified On-Policy Distillation for RL-Improved Reasoning Teachers

1 Introduction

• On MATH, student gains can fall short of a GRPO teacher’s improvement, motivating the paper’s question of which teacher improvements transfer.Prior work cited in the paper reports stronger students from GRPO-trained teachers.Source: §1 · p. 1
• Teacher improvements transfer unevenly: problems where the teacher succeeds and the student does not create knowledge gaps, while partial student success creates precision gaps.When the teacher does not outperform the student, it offers no better direction, and matching it may be harmful.Source: §1 · p. 1, Figure 1 · p. 2
• TS-OPD stratifies problems by joint teacher–student success, routes knowledge and precision gaps to matched objectives with token gating, and skips problems without a teacher advantage.Its problem-level decomposition distinguishes which teacher improvements may transfer.Source: §1 · p. 2
• TS-OPD achieves the highest macro-average correctness and smallest benchmark-averaged transfer gap at all three student scales, while distilling only on selected problems.The comparison uses a GRPO-improved teacher.Source: §1 · p. 2

RL-improved teachers

• Sparse-to-Dense uses a GRPO-improved teacher with forward-KL warmup and reverse-KL OPD uniformly, but compares final students rather than inherited teacher gains or their location.Direct-OPD instead transfers the RL-induced policy shift as a dense reward.Source: §2 · p. 2

On-policy objectives

• GKD and MiniLLM use on-policy objectives, while AKL, ToDi, EOPD, and TrOPD adapt forward and reverse KL without choosing direction from per-problem teacher–student performance.The passage describes adaptive weights, probability ratios, teacher entropy, and token-level trust regions among these approaches.Source: §2 · p. 3

Selective distillation

• Selective methods use student-only or teacher-only signals for problem or token selection, whereas TS-OPD uses joint success to choose KL direction and adds response-local token gating.The passage contrasts TS-OPD with PACED, TREK, TGOPD, OPDVR, and TIP.Source: §2 · p. 3

3.1 Evaluation Metrics and the GRPO Transfer Gap

• The transfer gap is teacher gain minus student gain in percentage points, and does not isolate loss of GRPO-acquired knowledge because student improvement may include knowledge already present before GRPO.avg@K measures average single-response correctness; pass@K measures solution coverage within the sampling budget.Source: §3.1 · p. 3

3.2 Two Learning Gaps on Training Problems

• A marks zero observed student success despite teacher success, C marks lower but nonzero student success, P marks student advantage, and N marks equal observed rates.A and C denote knowledge and precision gaps; P and N show no observed teacher advantage and are not prioritized for distillation.Source: §3.2 · p. 4

3.3 Training Strata and Test Gain Shortfalls

• The diagnosis uses eight sampled responses from each model on each of 7,496 MATH training problems and examines test shortfalls on 120 AIME and HMMT problems.The test set uses 16 responses per model per problem.Source: §3.3 · p. 4, Table 1 · p. 4
• 18.08 and 12.50 pp are the test gain shortfalls for the 0.6B and 1.7B students, respectively; their macro-average shortfalls reverse to 5.49 and 6.14 pp.The reported training-set teacher-advantage fractions are 77.16% and 43.52%; the authors caution that strata describe supervision needs, not established predictors of test shortfalls.Source: Table 1 · p. 4, §3.3 · p. 4

3.4 Implications for OPD Supervision

• For A problems, correct student rollouts may be rare; for C problems, uniform supervision may miss decisive errors and spend budget on P and N, motivating targeted guidance that preserves existing correct behavior.The proposed analysis treats these as mechanism hypotheses requiring fixed-budget stratification controls and token-level analysis.Source: §3.4 · pp. 4–5

4 Transfer-Stratified On-Policy Distillation

• TS-OPD filters to A and C problems, routes A to forward KL and C to reverse KL, gates token positions, and applies an entropy brake on C.Strata are computed once from eight samples per model and remain fixed during training; routing-only removes token gating.Source: §4 · p. 5, §4.1 · p. 5

4.1 Strata and KL Direction

• For A problems, forward KL can raise teacher-preferred tokens that the student rarely selects, whereas reverse KL barely moves such tokens; C problems instead need rejected choices corrected.This fixed-prefix gradient argument motivates the direction choice but does not describe the gated update or guarantee a correct response.Source: §4.1 · pp. 5–6

4.2 A Layer: Gated Forward KL

• A-layer gating averages forward KL only at positions with high teacher–student disagreement and teacher mass above 0.9 on a small set of candidate tokens.The gate keeps the top 20% of positions by disagreement; the teacher candidate count is set to 6.Source: §4.2 · p. 6

4.3 C Layer: Gated Reverse KL and Entropy Brake

• C-layer reverse KL penalizes sampled tokens the teacher rejects at selected positions, while an entropy brake counters reverse KL's concentration that can remove correct solutions and lower pass@K.The C gate selects the top 20% of positions with positive discrepancy; unselected positions are not updated.Source: §4.3 · p. 6

5 Experiments

• Experiments span two teacher lineages, three student sizes, eight benchmarks, and direct GRPO, uniform OPD, routing-only, and TS-OPD objectives.Runs use one epoch over 7,496 MATH training problems; per-problem analyses use seven benchmarks because MATH-test has K=1.Source: §5 · p. 7

5.1 Teacher GRPO Is Necessary but Not Sufficient

• With the GRPO-8B teacher, TS-OPD has the best avg@K and smallest transfer gap at every size, exceeding Direct GRPO by 2.84, 0.97, and 2.71 points.pass@K is also highest for TS-OPD or TS-OPD (routing only) at every size, while uniform OPD does not reliably help.Source: §5.1 · p. 7, Table 2 · p. 8

5.2 Only Part of the Teacher’s GRPO Gain Reaches the Student

• On the 567 problems improved by teacher GRPO, TS-OPD improves 50.1%, 62.1%, and 68.4% across student sizes, versus 37.1%, 59.5%, and 58.1% for same-problem Direct GRPO.TS-OPD newly solves the largest share at every size, yet about a third of teacher-improved problems still do not improve even at 4B.Source: §5.2 · p. 7, Figure 3 · p. 9
• Forward KL newly solves more knowledge-gap problems, while reverse KL makes fewer precision-gap problems worse; routing combines these strengths and token gates mainly improve precision-gap reliability.Routing newly solves 47.0%, 56.2%, and 71.7% of knowledge-gap problems; gates improve reliability over routing alone by 5.4, 2.0, and 1.0 points.Source: §5.2 · p. 8

5.3 Filtering Helps Only with Routing and Gating

• On student-advantage problems, every method makes about half or more worse, while over 89% of the larger no-gap stratum remains unchanged.Skipping changes the share made worse in that no-gap stratum by at most 1.3 points.Source: §5.3 · pp. 8–9
• Filtering alone loses 1.3 to 5.6 avg@K points, whereas routing reduces the transfer gap by 2.09, 2.69, and 5.78 points and gating lowers it further.Filtering lowers knowledge-gap acquisition, especially when A∪C runs take only 27 of 118 steps; routing and gating recover gains through acquisition and precision-gap updates.Source: Table 3 · p. 8, §5.3 · p. 9

• Teacher improvement and student transfer are distinct, and useful dense supervision requires separating acquisition from consolidation needs.Source: §6 · p. 9
• TS-OPD combines problem-level routing, response-local token gates, and an entropy brake; its gains depend on the full design.Source: §6 · p. 10
• Across two teacher lineages and three student sizes, the routed objective is most useful with a GRPO-improved teacher, motivating evaluation of teachable strata alongside endpoint accuracy.Source: §6 · p. 10

Scope

• Experiments cover one domain, one model family, one teacher size, and one teacher-improvement algorithm, with students up to 4B.Transfer beyond mathematical reasoning, to students near the teacher's size, or to teachers improved by other RL methods is untested.Source: Limitations · p. 10

Statistical strength

• Most runs use a single training seed, so small differences may not be robust and per-problem outcome categories should be read as trends.Source: Limitations · p. 10

Strata and budget

• Strata are fixed from eight calibration samples per problem; observed student failure may not mean a problem is unsolvable, and labels are not updated as students improve.Source: Limitations · p. 10

Design choices and cost

• Hyperparameter ablations use only the 0.6B student with the GRPO-8B teacher, and screening plus teacher top-K scoring adds training costs.The transfer gap measures a gain shortfall, not how much GRPO-acquired knowledge is lost.Source: Limitations · p. 10

A.1 Calibrating the Entropy Brake

• The entropy-brake coefficient is calibrated before training from student responses to C problems and depends on the teacher–student pair and K_tok.The procedure uses up to 50 responses, their first 400 tokens, and positions in the top 20% of H_S−H_T to set β; Table 4 lists pair-specific values.Source: §A.1 · p. 14, Table 4 · p. 14

A.2 Reverse KL from Sampled Tokens

• At selected positions, the sampled token's r_t-weighted student log-probability gradient estimates the reverse-KL gradient, so minimizing Eq. 7 lowers reverse KL there.Source: §A.2 · p. 14

A.3 Training Algorithm

• Training samples calibration responses to form the problem sets, calibrates β, then draws batches from D+ and scores student-generated responses with the teacher.The algorithm applies stratum-specific gates, stops gradients through them, averages the losses over the batch, and makes one update.Source: §A.3 · p. 14

Models and data

• The three Qwen3 students are 0.6B, 1.7B, and 4B; teachers are base-8B and GRPO-8B, and training uses 7,496 MATH problems.TS-OPD trains on each pair's retained A∪C pool, while Direct GRPO, OPD-r, and OPD-f also have full-set runs; every run makes one pass and evaluates its final checkpoint.Source: B Experimental Setup and Compute · p. 15

Framework

• Training uses slime with Megatron-LM and SGLang, with thinking mode disabled for both training and evaluation.Source: B Experimental Setup and Compute · p. 15

Evaluation

• Evaluation uses MATH-test, MATH500, Minerva, OlympiadBench, AIME24/25, and HMMT-Feb25/Nov25, with sample counts varying by benchmark.Truncated or unparsable responses count as incorrect.Source: B Experimental Setup and Compute · p. 15

Compute

• All training and evaluation runs used 1,728 NVIDIA B300 GPU hours in total.Source: B Experimental Setup and Compute · p. 15

C Training Configurations

• Table 5 lists training hyperparameters, with shared settings applying to all runs.Source: Table 5 · p. 15

• Table 6 lists training arms, and Table 7 reports one-setting-at-a-time TS-OPD ablations for the 0.6B student with the GRPO-8B teacher.Table 2 also reports four OPD objectives with the base-8B teacher.Source: Table 7 · p. 16, Table 6 · p. 16
• Figure 4 categorizes outcomes across all 1,566 test problems after training, comparing students with S0 and showing the teacher from T0 to T1 in its first row.Source: Figure 4 · p. 17
• Tables 8–11 classify post-training outcomes relative to the raw student and examine teacher-improved problems, teacher-not-better problems, and knowledge-gap and precision-gap problems.Table 11 notes that knowledge-gap problems start at 0 and precision-gap problems start above 0.Source: Table 8 · p. 17, Table 9 · p. 17, Table 10 · p. 18
• Table 12 reports transfer gaps by benchmark for K≥8, with each student's final row giving the macro average used in Table 2.Source: Table 12 · p. 19
• Tables 13–15 report per-benchmark results for 0.6B, 1.7B, and 4B students, with macro averages weighting benchmarks equally.Direct GRPO uses no teacher.Source: Table 13 · p. 19, Table 14 · p. 20, Table 15 · p. 20
• Table 16 gives per-benchmark results for the raw students and both teachers.Source: Table 16 · p. 20