PaperlayerView exact revision on arXiv

Source-linked AI summary

Distillation as Probability Transport: Routed On-Policy Distillation

Tianle Xia, Lingxiang Hu, Yiding Sun, Linfang Shang, Ming Xu, Lan Xu, Ning Zheng, Wei Xu, Jie Jiang

arXiv:2609.08337v1cs.LGcs.CL

TL;DR

Sampled OPD provides scalar token credit but does not specify where probability should be redistributed. RouteOPD models this missing step as teacher-guided transport between student-excess sources and teacher-deficit destinations, using jointly realizable pairwise targets and demand-adaptive budgets. Across four settings and four reasoning benchmarks, it consistently outperforms sampled reverse-KL OPD and improves routing fidelity while reducing leakage.

  • Problem

    Sampled OPD reduces teacher feedback to scalar credit on individual student-generated tokens, leaving the corresponding probability redistribution unspecified.

  • Method

    RouteOPD decomposes teacher–student disagreement into student-excess sources and teacher-deficit destinations, couples explicit routes, and optimizes pairwise log-odds from a bounded teacher potential with demand-adaptive budgets.

  • Results

    RouteOPD improves the four-benchmark average over sampled reverse-KL OPD by 2.19–3.61 points across four teacher–student settings, with a mean gain of 2.70 points.

  • Takeaways & Limitations

    Explicitly modeling probability transport improves routing fidelity and reduces background leakage while consistently improving reasoning accuracy across the evaluated settings and benchmarks.

Abstract

from arXiv · show

On-policy distillation (OPD) transfers teacher knowledge on student-generated trajectories, but efficient sampled objectives reduce the teacher distribution to scalar credit on individual tokens. Such credit indicates whether a token should gain or lose probability, yet leaves the corresponding redistribution unspecified. We recast OPD as teacher-guided probability transport and propose RouteOPD (Routed On-Policy Distillation), which decomposes local teacher--student disagreement into student-excess sources and teacher-deficit destinations and couples them into explicit transport pairs. RouteOPD optimizes pairwise log-odds toward jointly realizable targets obtained from a bounded teacher potential, while adapting the transport budget to the concentration of teacher demand. This formulation directs updates toward teacher-preferred destinations and controls their magnitude within a single transport operator. Experiments across four teacher--student settings and four mathematical-reasoning benchmarks demonstrate that RouteOPD consistently outperforms sampled reverse-KL OPD, with improvements accompanied by higher routing fidelity and lower background leakage. These results demonstrate the effectiveness of explicitly modeling probability transport in on-policy distillation.

1 Introduction

OPD provides dense teacher feedback on student-generated states, but sampled scalar credit leaves probability redistribution to student softmax geometry. RouteOPD makes that redistribution explicit through teacher-guided excess-to-deficit transport and improves reasoning performance with higher routing fidelity and lower leakage.

  • Motivation: Sampled OPD identifies whether a visited token should gain or lose probability, but leaves the redistribution destination unspecified.Suppressing one token releases mass that must be assigned elsewhere in a normalized categorical policy.
  • Approach: RouteOPD decomposes teacher–student disagreement into student-excess sources and teacher-deficit destinations, then couples them into explicit routes.The method constructs source and destination marginals on a shared teacher–student top-k support.
  • Approach: Pairwise log-odds updates localize direct logit changes to the specified source and destination rather than unrelated background logits.The shared log-softmax background cancels in the pairwise gradient.
  • Approach: RouteOPD derives cycle-consistent targets from a bounded teacher potential and adapts transport magnitude to teacher-deficit concentration.Concentrated demand permits a wider correction budget, while diffuse demand yields a more conservative update.
  • Evidence: Across four teacher–student settings and four reasoning benchmarks, RouteOPD improves the four-benchmark average over sampled reverse-KL OPD by 2.19–3.61 points, with a mean gain of 2.70 points.It also exceeds full-vocabulary reverse KL by 1.24 points on average and fixed-mid routing by 1.20 points.
  • Evidence: Matched interventions raise routing fidelity from 35.4% to 90.4% and reduce full-vocabulary background leakage from 68.3% to 13.8%.Teacher-deficit destinations reach 64.43 Avg@16 versus 62.72 for random-background and 63.72 for rank/frequency-matched destinations.
  • Evidence: Sparse adaptive routing reaches 65.53 Avg@16 at 2.5% overhead, close to 65.61 at 12.5% overhead for full-vocabulary routing.The default top-32 union retains over 97% of teacher and student mass with 96.38% sign agreement.

2 Related Work

Related work studies why OPD benefits from student-visited states, how teacher feedback and supervision are allocated, and how OPD updates evolve. RouteOPD instead targets the per-update softmax redistribution that determines where probability moves.

  • On-policy distillation: OPD addresses the mismatch between fixed offline prefixes and student-generated histories by querying teachers on student-visited states.Prior theory and empirical work examines when this state-distribution change matters.
  • OPD objectives: Prior OPD methods reformulate targets, extrapolate dense feedback, balance exploration, or calibrate teacher reliability.These approaches modify the feedback applied at a visited state rather than explicitly specifying redistribution destinations.
  • Supervision allocation: Other methods allocate supervision across problems, trajectories, token positions, or reasoning spans.Examples include frontier weighting, rollout allocation, informative-token selection, and critical-span routing.
  • Learning dynamics: Analyses connect OPD outcomes to thinking-pattern compatibility, progressive alignment, update-direction stabilization, and spectrally concentrated changes.These studies characterize when OPD succeeds and how its parameter changes accumulate.
  • Positioning: RouteOPD operates one level earlier by exposing per-update softmax redistribution and replacing its implicit destination rule with explicit excess-to-deficit routes.This distinguishes its focus from prior analyses of state coverage, feedback design, allocation, and parameter-space dynamics.

3 Method: RouteOPD

RouteOPD converts teacher–student disagreement on student-visited states into sparse, explicit source–destination updates. It uses pairwise log-odds, a shared bounded potential, and demand-adaptive budgets to make routes local, jointly realizable, and efficient.

  • Scalar feedback: Sampled reverse-KL suppresses or encourages the sampled token while student geometry determines where released probability goes.The teacher supplies the correction sign, not the recipient.
  • Explicit routes: RouteOPD forms excess sources and deficit destinations from separately renormalized teacher and old-student distributions on a sparse shared support.Their common mass is the available transport mass, and states below the threshold receive zero route loss.
  • Explicit routes: Independent product coupling preserves both source and destination marginals while allowing sampled routes without constructing a k^2 table.The method samples independent source–destination pairs from the normalized marginals.
  • Locality: For each route, pairwise log-odds changes directly raise the destination logit and lower the source logit, giving zero direct gradient to unrelated logits.Softmax normalization can still induce probability changes beyond these direct logit updates.
  • Target construction: A bounded teacher token potential produces jointly realizable pair targets instead of independently clipped edge demands.The shared potential makes all pair targets cycle-consistent and bounded.
  • Adaptive budget: The transport budget is conditioned on teacher-deficit concentration, allowing wider corrections for concentrated demand and conservative updates for diffuse demand.The concentration statistic is computed over the nonempty deficit set.
  • Implementation: After top-k extraction, routing costs O(k + m) per valid token and reuses teacher logits without an extra forward pass.The algorithm samples m route pairs and differentiates only current-student terms.

4 Experiments

Across four teacher–student settings and four mathematical-reasoning benchmarks, RouteOPD consistently improves performance over sampled-RKL while explicit routing controls identify teacher-deficit destinations, adaptive budgeting, and sparse support as important design choices. The gains coincide with stronger routing fidelity, lower background leakage, practical efficiency, and more repaired reasoning errors.

  • Overall performance: RouteOPD improves every teacher–student setting, raising the four-benchmark average by 2.19–3.61 points over sampled-RKL, with a 2.70-point mean gain.Matched bootstrap intervals remain uniformly positive.
  • Routing controls: Teacher-deficit destinations reach 64.43 Avg@16, exceeding rank/frequency-matched destinations at 63.72 and random-background destinations at 62.72 under matched controls.The same controls also show 64.43 for the full pairwise update, compared with 63.55 for source-only and 63.84 for destination-only updates.
  • Routing controls: Adaptive budgeting reaches 65.53 Avg@16 versus 64.43 for fixed-mid routing, with its advantage increasing from −0.10 to +2.70 points as teacher-deficit concentration rises.Reversing the concentration rule reaches 63.88, supporting concentration-aware allocation rather than uniformly larger updates.
  • Mechanism and training evidence: Adaptive RouteOPD reaches 90.4% routing fidelity and 13.8% full-vocabulary background leakage, compared with 35.4% and 68.3% for sampled-RKL.Across adaptive checkpoints, the routing signature strengthens while AIME24/25 Avg@16 rises from 27.19 to 43.65; the fixed bank reports 92.7% median fidelity versus 33.1% for sampled-RKL.
  • Mechanism and training evidence: The routing signature emerges throughout training: accuracy rises as routing fidelity and destination specificity improve and then saturate, rather than appearing at one isolated checkpoint.Figure 5 compares accuracy, fidelity, and destination specificity during training.
  • Efficiency: At k = 32, sparse routing retains 97.79% of student mass and 97.13% of teacher mass, reaches 65.53 Avg@16 at 2.5% overhead, and approaches full-vocabulary routing’s 65.61 at 12.5%.The k = 32 setting is identified as the practical Pareto knee.
  • Behavioral outcomes: On AIME24, correct-response share rises to 52.5% under RouteOPD from 48.1% under sampled-RKL and 50.2% under fixed-mid routing, with repaired failures primarily being reasoning errors.Among 42 sampled-RKL failures repaired, 34 are reasoning errors; format validity remains 95.21% and clipping falls from 4.47% to 2.34%.

5 Conclusion

The paper formalizes routed distillation with probability transport whose targets are bounded, cycle-consistent, and jointly realizable. Its theory explains sampled OPD’s implicit redistribution and establishes locality and unbiased estimation properties for routed updates.

  • Implicit redistribution: These source and non-source components sum to zero, reflecting probability conservation but not an explicit destination choice.
  • Implicit redistribution: Sampled OPD suppresses a negatively credited source while increasing every non-source logit according to the student’s current probabilities.The source component decreases in proportion to 1 − πold(at | st), while each non-source component increases in proportion to πold(v | st).
  • Coupling properties: The product coupling preserves exact excess and deficit marginals, maximizes entropy, and supports unbiased Monte Carlo estimation without materializing the full source–destination table.
  • Direct-logit locality: Routed pairwise log-odds updates have no direct gradient on unrelated logits, although normalization and shared parameters can still alter unrelated probabilities.
  • Jointly realizable targets: A shared bounded teacher potential produces valid edge targets that are jointly integrable across all routed pairs.Every target is a destination-minus-source potential difference, so one token potential realizes all pairwise targets simultaneously.
  • Jointly realizable targets: Independent edge clipping can violate four-cycle consistency and therefore fail to correspond to any shared set of token scores.The supplied example gives D11 + D22 = 2B but D12 + D21 = 1.8B.

A.5 Implementation and Numerical Conventions

The implementation validates response states and tokenizer compatibility, constructs routes from detached top-k statistics, and limits incremental routing overhead to O(k + m) per valid token. Experiments use matched training and compute conditions across teacher–student settings.

  • Validity and routing: A valid response state is generated and unpadded, token-aligned with finite teacher and student scores, and associated with a nonempty token support.
  • Validity and routing: Route construction uses detached top-k scores on the teacher–student token union, while learner updates gather differentiable probabilities only for sampled pair IDs.
  • Validity and routing: O(k + m) per valid token is the incremental route overhead: O(k) for detached routing statistics and O(m) for pair sampling.
  • Numerical conventions: Default routing uses k = 32 and m = 2, with Bmin = log 1.2, Bmax = log 1.5, and κ = 1.0.
  • Compatibility: Teacher–student pairs are included only when vocabulary IDs, token maps, special tokens, serialization, and tokenized prefixes match exactly.
  • Experimental controls: Training methods share prompts, rollout seeds, response masks, token budgets, optimizer steps, and checkpoint schedules; routed methods reuse teacher scores without extra queries.

B.3 Evaluation and Uncertainty

Evaluation covers four mathematical-reasoning benchmarks with common decoding and answer-processing protocols. Uncertainty is quantified through paired problem-level bootstrap intervals, while checkpoint and contamination controls support reproducible comparisons.

  • Compatibility: Token-level comparisons require identical ordered vocabulary IDs and compatible serialized prefixes.
  • Aggregation: Each problem receives equal weight, with 16 decoded responses averaged before aggregation and exact-answer correctness computed after deterministic normalization.
  • Uncertainty: Paired 95% confidence intervals bootstrap problems while retaining all decoded samples for each selected problem and matching identifiers across methods.
  • Controls: Checkpoint evaluation uses a common protocol, while exact hashes and n-gram screening check training–evaluation overlap.

C.1 Per-Setting and Per-Benchmark Results

Across settings, adaptive RouteOPD improves benchmark accuracy over sampled reverse-KL, full-vocabulary reverse KL, and fixed-mid routing. Its gains coincide with improved routing fidelity and reduced background leakage, while controls test destination identity and pairwise coupling.

  • Per-setting results: All twelve lower confidence bounds are positive, with adaptive RouteOPD gaining +2.19 to +3.61 Avg@16 points over sampled-RKL across settings.
  • Per-setting results: +1.20 macro improvement over fixed-mid routing isolates the contribution of conditioning update magnitude on teacher demand.
  • Checkpoint trajectories: Adaptive RouteOPD raises AIME24/25 Avg@16 from 27.19 at initialization to 43.65 at the final checkpoint.
  • Checkpoint trajectories: Routing fidelity rises from 78.2% to 90.4% as full-vocabulary leakage falls from 25.6% to 13.8%.
  • Routing diagnostics: Destination controls compare teacher-deficit, uniform-deficit, random-background, student-proportional, and rank- or frequency-matched permutations while holding source and target magnitude fixed.
  • Routing diagnostics: Pairwise, source-only, and destination-only updates are compared with and without mismatch-mass weighting to test whether coupling both sides matters beyond token selection or loss scale.

D.4 Adaptive Transport Budget

Adaptive budgeting assigns larger transport budgets when teacher-deficit demand is concentrated, improving accuracy beyond fixed and reversed controls. Sparse routing preserves nearly full support mass while offering the best measured accuracy–cost trade-off.

  • Destination choice: 64.43 Avg@16 for teacher-deficit destinations exceeds 63.72 for rank/frequency-matched destinations and 62.72 for random backgrounds.The destination block holds source identity, target construction, mismatch weighting, and fixed-mid budget constant.
  • Budget calibration: 65.53 Avg@16 is achieved by adaptive concentration budgeting, versus 64.12–64.43 for fixed budgets and 63.88 for reversed concentration.The ordering rejects a simple larger-update explanation: adaptive budgeting performs best when teacher demand is concentrated.
  • Joint realization: Shared potentials eliminate the 27.6% four-cycle violations produced by independent edge clipping, making pair targets jointly realizable.The shared construction also improves target-realization diagnostics under fixed and adaptive budgets.
  • Budget calibration: Adaptive routing is nearly tied with fixed-mid routing in the lowest concentration bin, but its advantage grows to +2.70 points in the highest bin.The adaptive gains across concentration bins are −0.10, +0.40, +1.30, and +2.70 points.

E.3 Routing Fidelity, Leakage, and Target Realization

RouteOPD improves routing fidelity and reduces background leakage relative to sampled-RKL, while shared-potential targets remain coherent through optimization. The routing signature persists beyond ideal logits but is attenuated by parameter sharing and optimizer state.

  • Distributional evidence: 92.7% median routing fidelity and 11.2% median leakage are achieved by adaptive RouteOPD, compared with 33.1% and 70.1% for sampled-RKL.The matched permutation has 21.7% median fidelity, while fixed teacher-deficit routing lies between the alternatives.
  • Target realization: 100% target completion with near-zero leakage occurs in exact pair-isolated and potential-wide probes, confirming the algebraic construction.Scale matching reduces completion to 92.6%, showing that realized optimization steps can attenuate the ideal target geometry.
  • Target realization: Fresh AdamW reaches 88.1% completion, saved optimizer state reaches 85.4%, and SGD reaches 90.2%, showing attenuation during parameter-space optimization.The routing signature survives shared-parameter and optimizer effects, but not without degradation.
  • Mechanism stratification: Higher disagreement supplies more source mass, while concentrated teacher demand yields clearer destinations with higher fidelity and lower leakage.Mismatch and teacher entropy provide complementary stratifications of the transport signal.

F Training Dynamics and Stability

Adaptive RouteOPD’s training trajectory reflects sustained improvements in accuracy and routing diagnostics rather than unstable optimization. Its default sparse configuration also provides a practical efficiency advantage over full-vocabulary routing.

  • Training trajectories: Figure 8 tracks AIME24/25 Avg@16, mismatch mass, routing fidelity, and full-vocabulary background leakage across sampled-RKL, fixed-mid routing, and adaptive RouteOPD.These trajectories test whether the endpoint comparison develops steadily during training.
  • Stability: Gradient norms converge similarly across methods, while adaptive RouteOPD ends with shorter responses, a lower clip ratio, and 93.4% routed-update eligibility.Eligibility rises from 84.1% at initialization, and policy entropy moves toward the teacher endpoint of 1.310.
  • Efficiency: The k = 32, m = 2 default adds 2.5% end-to-end time and 2.0 GB peak memory while preserving the same teacher-call count.Full-vocabulary routing costs 12.5% end-to-end and 9.4 GB additional memory for only 0.08 more Avg@16 points.
  • Paired outcomes: Against sampled-RKL, RouteOPD repairs 42 failures and loses 21 previously correct responses in matched AIME24 comparisons.The paired analysis separates these transitions by identical benchmark problem and decoding-sample identifiers.

H.2 Error Taxonomy and Annotation

The error taxonomy uses deterministic checks followed by blinded semantic judging, with precedence rules enforcing one label per incorrect response. RouteOPD’s main observed change is fewer reasoning errors, supported by paired transitions and stable formatting validity.

  • Annotation procedure: Across 1,920 responses, two independent blinded judging passes achieve 92.4% raw agreement and Cohen’s κ = 0.864, with 146 cases adjudicated.Deterministic correctness and format checks precede semantic judging.
  • Assignment rules: The first decisive error determines the mutually exclusive semantic label, while clipped but parseable responses remain separate from format errors.This precedence rule prevents later surface properties from overriding earlier mathematical failures.
  • Error distributions: RouteOPD has 130 reasoning errors versus 151 for sampled-RKL and 142 for fixed-mid routing, while arithmetic, premature-stop, and format counts remain similar.The aggregate distribution therefore aligns with the paired transition analysis rather than broad relabeling.
  • Error transitions: Reasoning repairs account for 34 of RouteOPD’s 42 sampled-RKL wrong-to-correct transitions, while 15 of 21 regressions become reasoning errors.The net transition pattern is concentrated in mathematical reasoning rather than formatting.
  • Generation diagnostics: RouteOPD combines stable format validity with a lower clipping rate than the OPD baselines.These diagnostics support the observed reasoning-error transitions.
Loading 2609.08337v1…