PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Stanley Wei, Jiin Woo, Kaan Ozkara, Youngsuk Park

arXiv:2610.21234v1cs.LG

TL;DR

OPD commonly uses one teacher per rollout, leaving token-level teacher selection during student rollouts unexplored. This paper develops a theoretical framework for teacher routing and evaluates routing heuristics across domains, reporting improvements over vanilla OPD and teacher experts on some tasks. Its analysis predicts that different routing schemes may be beneficial in different scenarios.

  • Problem

    OPD typically uses a single teacher per prompt rollout or training run, while token-level teacher selection during rollouts remains an underused on-policy flexibility.

    Source: §1 · p. 2

  • Method

    The paper develops a log-linear theoretical framework that unifies confidence-based routing with a proposed representation-based routing family and analyzes their optimization and diversity implications.

    Source: §1 · p. 2, §6 · p. 11

  • Results

    Experiments across code, Triton kernel generation, math, and visual reasoning report pass@k improvements over vanilla OPD, while different routing schemes are predicted to be beneficial in different scenarios.

    Source: §1 · p. 2

  • Takeaways & Limitations

    Token-level teacher routing is an additional OPD design axis, with the framework offering predictions about when routing schemes may better suit student ability and output diversity.

    Source: §6 · p. 11

    The theoretical analysis fixes the prefix state and considers hard routing only, not the additive mixture models used in empirical experiments.

    Source: §4.2 · p. 5

Abstract

from arXiv · show

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.

1 Introduction

The paper frames token-level teacher selection as an on-policy flexibility and develops a framework for choosing compatible teachers. It considers teacher-student mismatch, teacher competence, and the diversity of student responses.

1 Introduction

  • The introduction contrasts OPD's dense token-level teacher feedback with reinforcement learning's single reward scalar for guiding gradient updates.It presents OPD as a proposed alternative to RL in post-training.Source: §1 · p. 1
  • The proposed framework treats distributional mismatch, teacher competence, and induced diversity over good student responses as compatibility criteria for OPD.The paper states that too much student-teacher mismatch makes the optimization landscape unfavorable.Source: §1 · p. 1
  • Token-level routing lets a student use different teachers at different rollout prefixes, such as a math teacher for reasoning before a kernel specialist handles implementation.This extends prompt-level routing, where prompts are assigned to domain experts.Source: §1 · p. 2
  • The paper develops a theoretical framework that separates teacher competence from the teacher-student gap and interprets confidence- and representation-based routing mathematically.It argues that either routing family may be more useful in different scenarios, and that representation-based routing may yield more diverse student generations.Source: §1 · p. 2
  • Experiments across code, Triton kernel generation, math, and visual reasoning report pass@k improvements over vanilla OPD, indicating potential for high-quality and diverse student responses.The authors also describe principled predictions of when routing families may be effective in practical setups.Source: §1 · p. 2

2 Preliminaries

The preliminaries formulate vanilla OPD around student-generated responses and explain how its practical gradient approximation yields a sampled-token advantage. The routing method later changes the teacher term at each sampled prefix.

2 Preliminaries

  • Vanilla OPD considers a prompt drawn from a prompt distribution and a response sampled from the student policy, with each prefix state predicting the next token.The student policy is parameterized by θ.Source: §2.1 · p. 2
  • In practice, OPD approximates the loss gradient, an approximation equivalent to moving the gradient inside the original expectation.The displayed gradient expressions are not needed to state this relationship.Source: §2.1 · p. 2
  • Token-level routing substitutes a selected teacher or soft mixture into the teacher term at each sampled prefix, holding the selection and mixture weights fixed during the update.This choice determines the sampled-token advantage.Source: §2.1 · p. 2
  • The approximation uses the sampled token's log-probability difference between teacher and student as the token-level advantage, termed sampled-token advantage OPD.The paper distinguishes this from a Top-K KL loss variant.Source: §2.1 · p. 3

3 Token-Level Teacher Selection

The paper routes teacher supervision at individual student-rollout prefixes, rather than fixing one teacher throughout training. It describes direct, confidence-based, and representation-based routing schemes.

3 Token-Level Teacher Selection

  • Token-level selection chooses which teacher supervises each student prefix, offering combinatorially more rollout-level choices than fixing one teacher throughout training.The paper formalizes a claim that some routing scheme cannot do worse than vanilla OPD under its stated reward-suboptimality criterion.Source: §3 · p. 3
  • The routing framework scores teachers at each prefix and uses a temperature to define routing weights; it covers hard selection and soft mixtures.The paper considers C = 2 teachers in both its theory and experiments.Source: §3 · p. 3
  • Direct section routing assigns a domain teacher to the corresponding rollout segment, such as a thinking teacher for reasoning and a kernel specialist for implementation.The example concerns writing kernel code.Source: §3 · p. 3
  • Confidence-based routing scores a teacher by its negative entropy at the prefix, using teacher confidence alone rather than the student policy.The paper describes this confidence as a proxy for successful continuation rewards, more accurate when the teacher is correctly confident.Source: §3 · p. 3
  • Representation-based routing scores teacher-student compatibility through policy-distribution or feature similarity at the prefix.The paper uses negative Hellinger distance or negative squared feature distance; feature matching is especially natural for shared architectures, while Hellinger distance is used otherwise.Source: §3 · p. 3

4 Theoretical Analysis of Teacher Routing Schemes

The analysis frames teacher selection around teacher competence, the student–teacher imitation gap, and diversity, then distinguishes when representation alignment or confidence is useful. Its theoretical results connect these routing choices to optimization and retention of successful response modes under stated assumptions.

4.1 Decoupling Teacher Competence and the Student-Teacher Gap

  • The framework separates teacher competence from the student–teacher gap and treats token-level KL as a tractable surrogate for the imitation gap OPD minimizes.The performance-difference analysis motivates examining competence and optimization toward a teacher separately.Source: §4.1 · pp. 4–5

4.2 Unifying Representation- and Entropy-Based Routing

  • Teacher–student feature alignment and teacher-head magnitude jointly shape accessible gradient signal, complementing representation-based and confidence-based routing.With equal head magnitudes, greater alignment roughly maximizes the initial gradient norm; with similarly scaled alignments, greater teacher norm does so and roughly corresponds to greater confidence.Source: §4.2 · pp. 5–6

4.3 Representation-Based Teacher Routing

  • With equal teacher-head norms, lower feature discrepancy yields a larger gradient norm and, under suitable assumptions, faster learning-time guarantees.The stated separation requires a sufficiently small learning rate.Source: §4.3 · p. 6

4.4 Confidence-Based Routing

  • When teachers share features, lower entropy corresponds to a larger gradient signal and, under suitable assumptions, a learning-time advantage.The proof sketch connects greater head magnitude to larger logit gaps and lower entropy.Source: §4.4 · pp. 6–7

4.5 Overlap with Good Solutions and Coverage

  • Under uniformly distributed high-quality modes, a more aligned teacher gives a stronger guarantee on how many modes the optimized student retains.The coverage discussion adapts an undercoverage perspective for successful modes, while experiments report higher pass@k even when pass@1 is not necessarily higher.Source: §4.5 · p. 7
  • If successful modes are uniform and both teachers assign each mode probability c/M, the theorem lower-bounds distinct successful modes observed in k samples using a teacher-dependent error term.That error increases with feature discrepancy; the analysis compares lower bounds rather than actual mode counts.Source: §4.5 · p. 7

5 Experimental Results

Experiments examine token-level teacher routing across math–Triton generation, Triton generation with thinking mode, and algorithmic coding; visual-reasoning results are deferred to the appendix. Across these settings, the strongest routing scheme varies by task, with results linked to teacher–student feature alignment.

5.1 A Synthetic Math-Triton Task

  • 29.30% pass@16: Hellinger routing achieved the best joint fast0 score on the synthetic math–Triton task, nearly 10 percentage points above vanilla OPD on either teacher, entropy routing, and section routing.The joint score requires both the math answer and kernel to be correct in one response; Table 1 also reports the component scores.Source: §5.1 · p. 8, Table 1 · p. 9

5.2 Triton-only Task with Thinking Mode

  • Entropy-based routing performed best on fast0 in the Triton-thinking setup, outperforming Hellinger-based routing there.The evaluation uses 100 KernelBench Level 2 problems; fast0 measures kernel correctness.Source: §5.2 · p. 10, Table 2 · p. 10

5.3 Algorithmic Coding

  • Feature-matching routing improved pass@64 over the baselines on algorithmic coding, although it did not necessarily achieve the highest pass@1.The authors relate this pattern to preserving coverage of successful responses.Source: Table 3 · p. 10, §5.3 · p. 11

5.4 Interpreting Effective Routing Schemes

  • Hellinger- or feature-based routing performed better on code-only and synthetic Triton-math tasks, while entropy-based routing performed better on Triton-thinking.The authors associate feature alignment with settings where teacher and student heads have similar magnitudes, and entropy with settings where teacher features are similarly misaligned with student features.Source: §5.4 · p. 11

6 Conclusion

The conclusion frames token-level teacher routing as an additional axis for on-policy distillation and identifies open questions about its downstream benefits.

6 Conclusion

  • Token-level teacher routing adds a new axis for OPD, with open questions about using diversity to warm-start RL and making teacher selection more reward-aware.The conclusion presents both as directions for future work.Source: §6 · p. 11

A Related Work

The related work situates OPD among RL, distillation, and teacher-selection research. The appendix proofs analyze how compatibility, teacher competence, and feature alignment relate to OPD optimization and diversity.

A Related Work

  • Prior OPD work includes dense teacher supervision on student rollouts, self-distillation, multi-domain experts, and analyses of compatibility and prefix drift.Existing teacher-selection approaches include training-long teacher choice and confidence-based token routing, while the paper extends arbitration to token-level schemes for different settings.Source: A Related Work · p. 16

B Proofs of the Feature-Matching Teacher Selection

  • The proof roadmap connects teacher competence and compatibility through the performance difference lemma, then develops gradient, optimization-time, and coverage results for feature matching.It also states that token-level teacher selection can realize potential unavailable from using only a single teacher.Source: B Proofs of the Feature-Matching… · p. 17

B.1 Why Minimize Token-Level KL? A Performance-Difference View

  • The appendix derives a bound on the imitation gap using token-level KL terms, the quantity OPD directly optimizes.This surrogate supports analyzing teacher choice through student–teacher compatibility.Source: §B.1 · pp. 17–18

B.2 Proofs for Gradients at Initialization

  • At initialization, greater teacher–student feature alignment corresponds to a larger gradient norm in the stated nonnegative-inner-product regime.The proof notes that, under suitable assumptions, this initialization relationship yields a separation in optimization rate.Source: §B.2 · p. 20

B.3 Proofs for Learning-Time Separation

  • The learning-time analysis bounds the updates needed to reach a target loss using the initial gradient, subject to smoothness and learning-rate conditions.The stated bounds require a Lipschitz gradient and 0 < ηΛ < 1.Source: §B.3 · pp. 21–22

B.4 Proofs for Diversity and Undercoverage

The diversity analysis links representation alignment to bounds on how well a student covers high-quality response modes. Helper lemmas establish technical constants used in the guarantees.

B.4 Proofs for Diversity and Undercoverage

  • At a KL minimum, the diversity theorem bounds student probabilities for each good mode under assumptions of uniformly distributed teacher modes and controlled error.The analysis defines diversity as the expected number of distinct successful modes observed in k independent samples; the error bound decreases with alignment, strengthening the diversity lower bound.Source: §B.4 · pp. 24–25

B.5 Helper Lemmas

  • The helper lemmas establish finite smoothness and probability-control constants used in the theoretical guarantees.The appendix states that these problem-relevant parameters are finite; it gives separate lemmas for reverse-KL smoothness and explicit probability-control constants.Source: §B.5 · pp. 25–26

C Proofs of Confidence-Based Routing Selection

The confidence-based routing analysis compares teacher entropy and initial student gradients in a shared-feature setting. Under its assumptions, lower teacher entropy corresponds to a larger teacher head magnitude and gradient norm.

C Proofs of Confidence-Based Routing Selection

  • Under the stated shared-feature, equal-alignment, and common-direction assumptions, H1 < H2 if and only if W1 > W2, and the gradient comparison is positive if and only if a1 > a2.The section motivates entropy-based routing as a way to analyze learnability, with reward handled separately through the performance difference lemma.Source: C Proofs of Confidence-Based Routing… · p. 27

D Unifying Compatibility- and Confidence-Based Teacher Selection

Teacher selection balances optimization signal, teacher competence, and student–teacher compatibility. Routing can use confidence or representation-based scores computed from the student’s current prefix.

D Unifying Compatibility- and Confidence-Based Teacher Selection

  • Closer student–teacher feature alignment raises ρ_i toward 1 and increases gradient norm, holding the teacher head norm fixed.The analysis also notes that lower entropy increases the teacher head norm and, holding teacher features fixed, increases gradient norm.Source: D Unifying Compatibility- and… · p. 28
  • Gradient comparisons capture imitation pressure, while teacher competence remains in a separate reward-gap term; their trade-off is open.The discussion says it remains an open question how to quantify the trade-off between teacher competence and optimization pressures.Source: D Unifying Compatibility- and… · p. 29
  • Confidence routing can select a confidently wrong teacher, and low-entropy imitation may improve pass@1 while worsening pass@k.The passage gives bad post-training data as one possible source of confident errors and reports observing this pass@1/pass@k pattern in experiments.Source: D Unifying Compatibility- and… · p. 29
  • A confidently correct teacher may still be difficult to imitate when its thinking patterns mismatch the student; the ideal case combines correctness and compatibility.In the theory, this ideal corresponds to ψ = ϕ with a very large teacher head.Source: D Unifying Compatibility- and… · p. 29
  • Representation-based routing can force imitation of an undesirable reward distribution, leaving competence-versus-signal calibration an open problem.The stated failure case assumes a teacher post-trained toward that reward distribution.Source: D Unifying Compatibility- and… · p. 29

E.2 Routing Scores and Teacher Targets

  • Routing scores use the same prefix as the supervised next token, without lookahead; weights are detached and held fixed during the update.The resulting target is either the selected teacher under hard routing or an arithmetic mixture under soft routing.Source: §E.2 · p. 30
  • Direct section routing assigns teachers by prompt tags or thinking-block boundaries, with the closing tag routed according to the preceding state.This rule is used for the synthetic math–Triton and Triton-only settings.Source: §E.2 · p. 31
  • Confidence routing scores a teacher by negative entropy of its top-8 renormalized distribution, independent of the student.The passage follows H-OPD.Source: §E.2 · p. 31

E.1 Implementation of OPD Objectives and Gradients

The implementation uses a semi-gradient approximation to OPD and evaluates sampled-token advantage and top-k KL objectives. The latter removes token-level sampling noise on its restricted support but leaves probability mass outside that support without direct signal.

E.1 Implementation of OPD Objectives and Gradients

  • The reverse-KL gradient approximation omits future time-step effects that a prefix may influence.The implementation differentiates per-token KL terms while holding sampled prefixes and teacher targets fixed.Source: §E.1 · pp. 29–30
  • Sampled-token advantage treats the advantage and student log probability as constants, differentiating only the trailing student log probability.With routed teachers, the loss matches Equation (1) up to a 1/|y| factor; prompt and padding tokens are excluded from averaging.Source: §E.1 · pp. 29–30
  • Top-k KL mixes teachers’ individually renormalized top-K distributions over their union support and compares them with the student restricted to that support.The experiments use K = 8; the reverse KL is averaged over rollout length.Source: §E.1 · p. 30
  • Top-k KL avoids token-level sampling noise through a full expectation on restricted support, but gives no direct signal to student mass outside it.The experiments use sampled-token advantage and top-k KL; detailed comparison of the implementations is left to future work.Source: §E.1 · pp. 29–30

F Training Details and Hyperparameters

Training covers algorithmic coding, synthetic math–Triton, visual reasoning, and Triton-only settings under a shared optimization recipe. The prompts, teachers, and training conditions differ by task.

F Training Details and Hyperparameters

  • Algorithmic coding uses 25,276 Eurus-2-RL prompts, with 1,024 held out for validation, and a math expert trained on DeepMath for 500 iterations.The code prompts are the same prompts used to train the code RL expert.Source: F Training Details and Hyperparameters · p. 32
  • Synthetic math–Triton training pairs KernelBook programs with distinct shuffled DeepMath problems, yielding 8,077 training pairs and 509 validation pairs after filtering.The source programs deduplicate to 8,884; composite prompts longer than 4,096 tokens are dropped.Source: F Training Details and Hyperparameters · p. 32
  • Visual reasoning uses 54,840 image problems; the vision-language teacher receives images, while the text-only teacher receives appended image descriptions.The student is Qwen3-VL-2B-Instruct, with Qwen3-VL-4B-Instruct and Qwen3-4B-Instruct-2507 as teachers.Source: F Training Details and Hyperparameters · p. 32
  • All settings train for 300 updates (400 for visual reasoning), using 128 prompts per batch, one rollout per prompt, and AdamW at 10^-6.The shared recipe uses temperature 1, top-p 1, no top-k restriction, no warmup, and gradient clipping at 1; thinking mode is enabled only for Triton-only.Source: F Training Details and Hyperparameters · p. 32

G Evaluations

Evaluations report pass@1 and task-specific higher-k metrics across synthetic math–Triton, Triton-only, algorithmic coding, and visual reasoning. Each task uses its stated benchmark set and sampling protocol.

G Evaluations

  • Pass@k uses an unbiased estimator from n independent responses per task, with shared seeds across models and the same sample pool for every k.The report uses pass@1 and pass@16 for math–Triton and Triton-only, pass@1 and pass@64 for coding, and pass@1 only for visual reasoning.Source: G Evaluations · p. 32
  • Synthetic math–Triton evaluates 3,000 AIME 2025–KernelBench pairs and requires both the math answer and kernel to be correct for joint fast0.The math and kernel portion columns report their respective correctness separately.Source: G Evaluations · p. 32
  • Triton-only evaluation uses 100 KernelBench Level 2 tasks, 16 responses per task, and a 32,768-token response cap.Correctness is fast0, and the reported metrics are pass@1 and pass@16.Source: G Evaluations · p. 33
  • Algorithmic coding evaluation uses 164 HumanEval+ and 378 MBPP+ problems, drawing 128 responses per problem.The reported metrics are pass@1 and pass@64.Source: G Evaluations · p. 33
  • Visual reasoning reports pass@1 on WeMath, MathVista-mini, and CharXiv RQ, using question-level accuracy or single sampled responses as specified.WeMath averages 16 sampled responses per question; MathVista-mini and CharXiv use one, with final answers graded by an LLM judge.Source: G Evaluations · p. 33

H Additional Experiments and Analyses

This section presents additional experiments and analyses, including Hellinger-routing results on visual reasoning benchmarks and example rollouts for algorithmic coding and Triton-thinking tasks. It also describes Triton focal-loss implementations that compute element-wise loss and return the sum, with one stated aim of reducing kernel launches and memory reads and writes.

H.1 Visual Reasoning

  • 68.69% on WeMath and 67.40% on MathVista-mini: Hellinger routing beats entropy routing by 0.67 and 0.70 points, respectively, and ties H-OPD at 33.20% on CharXiv RQ.It also outperforms single-teacher OPD on all three benchmarks.Source: §H.1 · p. 33

H.2 Code-only feature routing trace

  • The coding-task example rollout accompanies a plot of code-minus-math score gaps by response token, with mean gaps shown separately for explanation and implementation.The example is from the algorithmic coding task referenced in Figures 2 and 3; Figure 3 uses code-only final-layer feature routing at τ = 0.1 and step 300.Source: §H.2 · p. 33, Figure 3 · p. 34

H.3 Triton-only entropy routing trace

  • The Triton-thinking example rollout is the task for which teacher score gaps are plotted in Figure 1.Source: §H.3 · p. 35
  • The implementation is intended to reduce kernel launches and memory reads/writes.Source: §H.3 · p. 39
  • Although inputs and targets are 4D tensors, their dimensions are not immediately important: the kernel treats them as 1D, so the tensors can be passed as 1D.The launch uses BLOCK_SIZE = 1024.Source: §H.3 · pp. 39, 42, 44
  • The Triton kernel computes focal loss per element, stores the results, and the Python wrapper returns their sum.The computation clamps inputs, calculates cross-entropy and alpha balancing, then applies the (1 - pt)^gamma factor.Source: §H.3 · pp. 41–42
  • The kernel clamps inputs to [eps, 1 - eps] before taking logarithms, keeping the log arguments positive.The passage notes that an incorrectly handled log(1 - clamped) could produce NaNs.Source: §H.3 · pp. 42–43
  • The Triton kernel clamps inputs, computes cross-entropy and focal weighting, stores element-wise losses, and returns their sum while fusing operations to reduce memory traffic and kernel launches.The implementation uses a 1D grid with BLOCK_SIZE 1024 and requires CUDA tensors.Source: §H.3 · pp. 44–45

Related papers

Loading 2610.21234v1…