PaperlayerView exact revision on arXiv

Source-linked AI summary

SAKI: Maximal-Coupling-Routed Teacher Supervision for On-Policy Distillation

Miteto Wei, Xiaohan Wang, Zehao Chen, Jiajun Chai, Sichao Liu, Li Wang, Haoyuan Xu, Zhaoyu Hu, Wei Lin, Guojun Yin

arXiv:2609.36601v1cs.AI

TL;DR

Weak students can visit teacher-misaligned prefixes during on-policy distillation, making the representativeness of supervision an important concern. SAKI combines a KL-constrained teacher-guided rollout with maximal coupling to route token-level supervision, and improves both reported metrics over the matched teacher-guided baseline at both student scales.

  • Problem

    Weak students may make compounding early mistakes that lead to prefixes the teacher would rarely generate, where conditional supervision may be less representative.

    Source: §1 · pp. 1–2

  • Method

    SAKI uses maximal coupling to route sampled-token reverse-KL supervision to accepted positions and teacher Top-1 supervision to correction positions.

    Source: §1 · p. 2

  • Results

    SAKI achieves the best macro Mean@8 and Pass@8 at both student scales, improving over TRB in both metrics for the 1.7B and 0.6B students.

    Source: §4.2 · p. 7, §5 · p. 9

  • Takeaways & Limitations

    Placement controls and fixed-prefix analysis support correction-based routing as a conflict-adaptive supervision signal.

    Source: §4.2 · p. 8, §5 · p. 9

    Online teacher verification remains an intrinsic cost, so throughput stays below student-only generation and depends on response length, correction rate, and committed tokens per verification wave.

    Source: §K.5 · p. 24

Abstract

from arXiv · show

On-policy distillation (OPD) reduces train-test state mismatch by training a student on its own generated trajectories, but weak students may visit teacher-misaligned prefixes where supervision is less representative. We introduce SAKI (Supervision Allocation with KL-constrained Interpolation), which combines a KL-constrained teacher-guided rollout with maximal coupling and reuses realized accept/correction events to route token-level supervision. Accepted positions retain sampled-token reverse-KL supervision, while correction positions receive direct supervision on the teacher's highest-probability token. Under maximal coupling, the correction probability is exactly TV(p_t, q_t), so the same trust-region radius controls rollout deviation and upper-bounds intervention and specialized-supervision frequency. We further implement an engine-resident speculative verifier that preserves the exact-q trajectory distribution and coupling semantics while improving matched-workload rollout throughput by 4.22x. Across seven mathematical reasoning benchmarks, SAKI improves the matched teacher-guided baseline in Mean@8 and Pass@8 for both 1.7B and 0.6B students. Placement controls and fixed-prefix analysis further support correction-triggered routing as a conflict-adaptive supervision signal.

1 INTRODUCTION

On-policy distillation trains on student-generated trajectories, but weak students can reach prefixes where teacher supervision is less representative. SAKI guides rollouts within a student-centered trust region and uses coupling events to route specialized supervision.

1 INTRODUCTION

  • Weak students can make compounding early mistakes that move rollouts toward prefixes the teacher would rarely generate, where conditional supervision may be less representative.This is a concern for reasoning distillation because early choices alter later trajectories and queried states.Source: §1 · pp. 1–2
  • TRB constructs a teacher-guided behavior distribution within a student-centered KL trust region, balancing movement toward the teacher against deviation from student behavior.This addresses where supervision is queried, while leaving open whether accepted and corrected positions should receive the same objective.Source: §1 · p. 2
  • SAKI retains sampled-token reverse-KL supervision at accepted positions and gives correction positions direct supervision on the teacher’s highest-probability token.The correction event is the routing signal, rather than an external token-selection heuristic.Source: §1 · pp. 2–3
  • Under maximal coupling, correction probability equals TV(p_t, q_t), the minimum possible intervention probability, and is bounded by ϵ/2 under the trust region.Thus the trust-region radius limits both rollout deviation and the frequency of correction-triggered supervision.Source: §1 · pp. 2–3
  • An engine-resident speculative verifier preserves exact-q trajectory distribution and coupling semantics while providing a 4.22× matched-workload speedup over the external-loop implementation.It commits proposals only through the first rejection, applies an exact residual correction there, and discards the invalid suffix.Source: §1 · pp. 2–3

2 RELATED WORK

Prior work addresses on-policy state mismatch through student-generated trajectories and teacher-guided rollouts, while speculative verification can accelerate generation. This paper uses speculative verification to realize a separate teacher-guided distribution and reuse coupling events for supervision routing.

2.1 ON-POLICY DISTILLATION AND TEACHER-GUIDED ROLLOUTS

  • On-policy distillation trains on student-induced states; TRB adds a student-centered KL-constrained teacher-guided behavior distribution while retaining reverse-KL supervision at visited prefixes.The paper uses the rollout’s accept/correction information to route specialized supervision.Source: §2.1 · p. 3

2.2 SPECULATIVE DECODING AND DISTILLATION

  • Speculative verification accelerates autoregressive generation by verifying draft tokens in parallel, and Draft-OPD also uses verification outcomes as training structure [Lei et al., 2026].Here, speculative execution realizes a separate teacher-guided intermediate distribution, with maximal-coupling decisions reused to route token-level supervision.Source: §2.2 · p. 3

3 METHOD

SAKI combines a student-centered teacher-guided rollout with maximal coupling, then routes token-level supervision according to realized accept/correction events. Its exact-q speculative verifier preserves the sequential sampling semantics while batching verification work.

3 METHOD

  • At each prefix, maximal coupling realizes the trust-region behavior distribution and exposes an accept/correction event that determines the rollout token and local objective.The residual correction token controls the next prefix, while teacher-mode supervision controls the parameter update.Source: §3 · p. 3, Figure 1 · p. 4

3.1 PROBLEM SETUP

  • The setup distinguishes the frozen rollout student distribution and fixed teacher distribution from the trainable student; retained proposals use the sampled-token reverse-KL signal.The reverse-KL signal acts on the student–guided overlap mass min{p, q}.Source: §3.1 · p. 4

3.2 TEACHER-GUIDED TRUST-REGION ROLLOUT

  • TRB geometrically interpolates between student and teacher distributions under a student-centered KL constraint; ϵ = 0 recovers q_t = p_t, while a sufficiently large region permits q_t = T_t.Source: §3.2 · p. 4

3.3 MAXIMAL COUPLING AND CORRECTION EVENTS

  • Maximal coupling realizes q with the minimum possible intervention probability, and its correction event is proposal-dependent rather than a scalar conflict score.The correction event has explicit distributional meaning, rather than acting as an arbitrary gate.Source: §3.3 · p. 5

3.4 COUPLING-ROUTED SUPERVISION

  • SAKI keeps sampled-token reverse-KL updates on accepted positions and switches corrections to teacher-mode supervision, with ϵ annealed to zero in the main schedule.Teacher-mode supervision outperforms stochastic teacher sampling by +0.62 Mean@8 and +1.56 Pass@8 on the 1.7B student.Source: §3.4 · p. 6

3.5 EFFICIENT ENGINE-RESIDENT EXACT-q ROLLOUT

  • The engine verifies draft blocks in batch, commits the longest accepted prefix, applies an exact residual correction at the first rejection, and discards later speculative states.This preserves the same autoregressive law and coupling events as sequential exact-q sampling.Source: §3.5 · pp. 6–7

4 EXPERIMENTS

SAKI improves macro Mean@8 and Pass@8 over TRB at both student scales, while controlled analyses support correction-triggered routing and an engine-resident exact-q verifier. The experiments use matched training and evaluation protocols across seven mathematical reasoning benchmarks.

4.1 EXPERIMENTAL SETUP

  • The matched protocol uses Qwen3-0.6B-Base and Qwen3-1.7B-Base students, a shared teacher and training setup, and student-only evaluation across seven benchmarks.Each student scale uses the same checkpoints, prompts, 200-step budget, rollout batch, and eight responses per prompt; evaluation samples eight responses per problem.Source: §4.1 · p. 7

4.2 MAIN RESULTS

  • Ours achieves the best macro Mean@8 and Pass@8 at both student scales, beating TRB by 1.1/2.9 points for 1.7B and 1.2/2.0 points for 0.6B.Mean@8/Pass@8 improve from 27.9/44.6 to 29.0/47.5 for 1.7B, and from 17.2/33.6 to 18.4/35.6 for 0.6B.Source: §4.2 · p. 7
  • For the 1.7B student under the matched per-trajectory teacher-mode budget, Ours reaches 29.00/47.50 Mean@8/Pass@8, exceeding TV-Weighted-TM (28.24/46.77) by 0.76/0.73 points, respectively.Source: §4.2 · p. 8

4.3 MECHANISTIC ANALYSIS: PERSISTENT ALIGNMENT AND CONFLICT-ADAPTIVE TEACHER SUPPORT

  • At step 200, correction-triggered supervision retains 4.26-point C1 and 2.94-point C16 gains over TRB after 149 RKL-only steps.Figure 2(a) marks the correction-active phase and step 51; the gains were 4.63 points in C1 and 4.54 in C16 at step 50.Source: §4.3 · p. 8, Figure 2 · p. 9
  • Relative to Random-TM, the C16 advantage grows from +0.25 points in the lowest-conflict quartile to +1.32 in the highest, with a Q4–Q1 contrast of +1.07 points.The Q4–Q1 contrast for teacher-supported non-argmax tokens is +1.28 points; the corresponding C1 interaction is not significant.Source: §4.3 · p. 9, Figure 2 · p. 9

4.4 EXACT-q ROLLOUT EFFICIENCY

  • The engine-resident backend reaches 3,276 committed tokens/s versus 776 for the external loop, a 4.22× speedup under the matched workload.The workload uses 64 prompts, eight responses per prompt, maximum length 7,168, and K = 8.Source: §4.4 · p. 9, Table 3 · p. 9

5 CONCLUSION

The conclusion presents SAKI as a maximal-coupling method for routing teacher-mode supervision and reports gains at both student scales. It also describes the supplementary implementation materials and the authors’ review of AI-assisted work.

5 CONCLUSION

  • SAKI combines a teacher-guided trust-region policy with maximal coupling to route teacher-mode supervision at correction events, and achieves the best macro Mean@8 and Pass@8 at both student scales.Placement controls and fixed-prefix analysis support correction-based, conflict-adaptive routing.Source: §5 · p. 9
  • The authors report that they independently verified mathematical arguments and reviewed and tested AI-assisted code and text edits.They take responsibility for the final content and artifacts.Source: AI USE STATEMENT · p. 10
  • Supplementary code provides a reference implementation of SAKI’s core components, with appendices covering configuration, derivations, and systems validation.Appendix A reports training and evaluation configuration; Appendices G, I, and K cover coupling, exact execution, and systems implementation and validation.Source: REPRODUCIBILITY STATEMENT · p. 10

A TEACHER CONSTRUCTION, MODEL COMPATIBILITY, AND TRAINING DETAILS

The setup trains Qwen3 Base students against a Qwen3-4B-Base-GRPO teacher on DAPO-Math-17K. It uses Base-to-Base pairs to align token-level behavior and holds the student training protocol constant across scales.

A TEACHER CONSTRUCTION, MODEL COMPATIBILITY, AND TRAINING DETAILS

  • The teacher is initialized from Qwen3-4B-Base and trained with GRPO on processed DAPO-Math-17K using the reported recipe and a step-by-step answer format.The teacher uses a released rule-based mathematical reward, not a learned reward model; format reward and reference-model KL regularization are disabled.Source: A TEACHER CONSTRUCTION, MODEL… · p. 12
  • Base-to-Base teacher–student pairs align vocabulary, termination behavior, and response masking for exact token-level rollout and maximal coupling.The released Base and post-trained checkpoints use different default termination-token conventions.Source: A TEACHER CONSTRUCTION, MODEL… · p. 12
  • The 0.6B and 1.7B students share a training protocol and 17,917 DAPO-Math prompts over 200 optimization steps, with parameter count the primary difference.Source: A TEACHER CONSTRUCTION, MODEL… · p. 12

B DERIVATION OF THE TRUST-REGION GEOMETRIC BRIDGE

This section reproduces the geometric bridge underlying TRB and gives its boundary behavior as the trust-region constraint varies.

B DERIVATION OF THE TRUST-REGION GEOMETRIC BRIDGE

C SINGLE-SAMPLE ESTIMATION OF REVERSE KL

The section connects sampled reverse-KL updates to accepted proposals, then compares placement controls under the same exact-q rollout. Across the seven-benchmark suite, the reported results favor correction-triggered routing over the matched placement controls.

C SINGLE-SAMPLE ESTIMATION OF REVERSE KL

  • Under exact student sampling, the sampled log-ratio estimates reverse KL and supplies an unbiased score-function gradient coefficient.On teacher-guided rollouts, accepted-proposal RKL updates are restricted to the student–guided overlap, while diverging states receive the direct teacher-mode target.Source: C SINGLE-SAMPLE ESTIMATION OF REVERSE KL · p. 14

D ADDITIONAL PLACEMENT-CONTROL RESULTS

  • Table 6 presents matched supervision-placement controls for the 1.7B student on the same sevenbenchmark suite: Random-TM controls supervision budget, while TV-Weighted-TM favors positions with larger local student–guided disagreement.Source: Table 6 · p. 15

C SINGLE-SAMPLE ESTIMATION OF REVERSE KL

  • Mean@8 gains occur on six of seven benchmarks for the 1.7B student and all seven for the 0.6B student; macro Mean@8 and Pass@8 improve at both scales.Figure 3 also decomposes the 1.7B gain into count-matched random placement and the additional gain from placing updates at maximal-coupling corrections.Source: Figure 3 · p. 15

D ADDITIONAL PLACEMENT-CONTROL RESULTS

  • TV-Weighted-TM samples valid positions without replacement using local TV as its weight, while the correction mask sets only the per-trajectory supervision budget.The selected positions receive teacher-mode loss; remaining valid positions retain RKL, and the exact-q rollout is unchanged.Source: D ADDITIONAL PLACEMENT-CONTROL RESULTS · p. 15

E ADDITIONAL FIXED-PREFIX RESULTS

The fixed-prefix analysis compares correction-triggered teacher-mode supervision with matched TRB and random-placement runs on frozen positions. It reports that the reduction in teacher-to-student KL persists after correction-triggered updates stop.

E ADDITIONAL FIXED-PREFIX RESULTS

  • The fixed probe contains 2,048 positions—1,024 corrections and 1,024 matched accepted positions—across 213 prompts, with prefixes and initial statistics frozen.The three matched 1.7B runs use the same initialization, teacher, exact-q rollout, K = 8, annealing schedule, horizon, and seed.Source: E ADDITIONAL FIXED-PREFIX RESULTS · pp. 15–16
  • At step 51, the trust-region radius reaches zero, so qt = pt and correction-triggered teacher-mode updates cease.Source: E ADDITIONAL FIXED-PREFIX RESULTS · p. 16
  • The reduction in D_KL(T∥S) persists after correction-triggered supervision is disabled; the accompanying entropy reduction indicates greater concentration around teacher-supported mass.Table 7 reports the complete fixed-prefix statistics underlying Figure 2(a).Source: E ADDITIONAL FIXED-PREFIX RESULTS · p. 16, Table 7 · p. 16

F CORRECTION-POSITION TEACHER SUPERVISION

At correction positions, the method uses the teacher’s Top-1 token as a hard-label target, while teacher-sampled supervision is a stochastic alternative. In the reported 1.7B setting, teacher-mode supervision performs better than that alternative.

F CORRECTION-POSITION TEACHER SUPERVISION

G PROOF OF PROPOSITION 1

The proof establishes that maximal coupling produces the exact q marginal, and the following sections analyze its intervention behavior, rollout endpoints, and engine-resident implementation. The implementation preserves the target trajectory distribution while moving verification into the inference runtime.

G PROOF OF PROPOSITION 1

  • Maximal coupling produces the exact q marginal for the generated token.The proof sums the accepted and correction cases to obtain the unconditional output distribution.Source: G PROOF OF PROPOSITION 1 · pp. 17–18

G.1 MINIMAL-INTERVENTION PROPERTY

  • Pinsker’s inequality and the trust-region constraint bound the intervention rate, while maximal coupling attains the minimum intervention probability.The proof identifies equality under the maximal-coupling construction.Source: §G.1 · p. 18

H TEACHER–STUDENT INTERPRETATION OF THE RESIDUAL DISTRIBUTION

  • Corrections replace student-excess probability mass with bridge-favored mass on tokens above a prefix-dependent teacher-to-student likelihood-ratio threshold.The positive residual is supported only on tokens whose likelihood ratio exceeds that threshold.Source: H TEACHER–STUDENT INTERPRETATION OF… · p. 19

I EXACTNESS OF ENGINE-RESIDENT BLOCK VERIFICATION

  • Engine-resident block verification preserves sequential maximal coupling’s autoregressive distribution by committing through the first rejection and discarding later speculative states.The correction is sampled from the exact residual distribution.Source: I EXACTNESS OF ENGINE-RESIDENT BLOCK… · pp. 19–20

J.1 ϵ = 0: RECOVERY OF STANDARD STUDENT ROLLOUT; J.3 TEACHER ENDPOINT

  • At ϵ = 0, the objective reduces to ordinary student-rollout RKL; as the trust-region radius grows, the bridge interpolates toward the teacher endpoint.The teacher-mode supervision term vanishes at ϵ = 0.Source: §J.1 · p. 20, §J.3 · p. 21

J.2 CORRECTION ACTIVITY UNDER ANNEALING

  • Correction activity decreases as ϵ is annealed, remains below the Pinsker bound, and vanishes at ϵ = 0 at step 51.Figure 4 compares observed correction probability with the upper bound induced by the current trust-region radius.Source: §J.2 · p. 20, Figure 4 · p. 21

K.1 WHY EXACT-q ROLLOUT IS EXPENSIVE; K.2 FROM TOKEN-WISE TO EXTERNAL BLOCK VERIFICATION

  • Naive per-token verification reached ∼5 tokens/s versus ∼3.02K for student-only generation; K = 8 blocks reduce service calls to approximately O(L/K) when largely accepted.Online teacher distributions are required during generation, and the earlier external coordinator still handled bridge construction, dense-logit transfer, and coupling.Source: §K.1 · p. 21, §K.2 · p. 22

K.3 ENGINE-RESIDENT q BACKEND; K.4 MODEL AND CACHE LIFECYCLE

  • The final backend places the recurrent block loop and teacher alongside the student draft in the engine; student synchronization is only approximately 1.1% of representative step time.The q-engine remains resident and cached states from previous student versions are invalidated after updates.Source: §K.3 · p. 22, §K.4 · p. 23

K.5 CORRECTNESS VALIDATION AND LIMITATIONS

The backend is validated across bridge, coupling, training-state, and routing behavior. The section also specifies how correction supervision is applied and notes that online teacher verification remains an intrinsic cost.

K.5 CORRECTNESS VALIDATION AND LIMITATIONS

  • Validation checks bridge normalization and endpoints, coupling samples and first-rejection commits, training-state synchronization, and the zero-radius student-only route.At exactly zero effective trust-region radius, the correction mask is identically zero and no teacher q-forward is performed.Source: §K.5 · p. 23
  • Online teacher verification remains costly: throughput stays below student-only generation and depends on response length, correction rate, and committed tokens per verification wave.The optimized system does not eliminate the intrinsic cost of online teacher verification.Source: §K.5 · p. 24

L ADDITIONAL OBJECTIVE AND IMPLEMENTATION DETAILS

  • Correction positions receive no RKL contribution; accepted-position RKL and correction-position teacher-mode terms have unit weight and share valid-token normalization.No additional correction-loss coefficient is used.Source: L ADDITIONAL OBJECTIVE AND… · p. 24
  • The sampled residual correction token determines subsequent prefixes, but the direct correction supervision target is the teacher’s Top-1 token and may differ.The residual token is the actual response token at that position, not the direct training target.Source: L ADDITIONAL OBJECTIVE AND… · p. 24
  • The rollout bridge and maximal coupling use the full vocabulary; the teacher Top-1 token is only a correction-supervision target.The Top-1 target does not truncate or otherwise modify the rollout distribution.Source: L ADDITIONAL OBJECTIVE AND… · p. 24
  • Reported benchmark evaluation generates from the current student alone, so performance measures the distilled student rather than the teacher-guided rollout engine.The validation procedure specifies student-only generation.Source: L ADDITIONAL OBJECTIVE AND… · p. 24
Loading 2609.36601v1…