PaperlayerView exact revision on arXiv

Source-linked AI summary

Not Every Token Is Worth Distilling: Selective Supervision for Direct-OPD

Yibo Zhao, Zixuan Yang, Yunshi Lan, Xiang Li

arXiv:2609.29142v1cs.LGcs.AI

TL;DR

Direct-OPD transfers an RL-induced policy shift from a small model to a larger student, but its reward can remain fixed while the probability mass supporting the shift vanishes. S2D-OPD selects student-sampled states by teacher–reference JSD and masks low-divergence supervision. It improves held-out accuracy over dense Direct-OPD in seven of eight settings and matches it in the eighth.

  • Problem

    Direct-OPD’s log-ratio measures relative policy change but can stay fixed as both checkpoints’ probability mass on student candidates vanishes, leaving uncertainty about whether every state is worth distilling.

    Source: §1 · p. 2

  • Method

    S2D-OPD ranks student-sampled states by teacher–reference JSD and masks Direct-OPD supervision at low-divergence states, retaining the top 10% per response.

    Source: §1 · p. 2

  • Results

    Seven of eight teacher–student settings improved in held-out accuracy over dense Direct-OPD, with a mean gain of 0.95 points (95% CI 0.40–1.54).

    Source: §1 · p. 2

  • Takeaways & Limitations

    The exact construction shows that the Direct-OPD reward and update can remain fixed as candidate mass and teacher–reference divergences vanish, while small JSD bounds behavioral change.

    Source: §1 · p. 2

    JSD measures the magnitude, not correctness, of the teacher–reference shift, so retained supervision still depends on the quality of the teacher’s RL.

    Source: G LIMITATIONS · p. 18

Abstract

from arXiv · show

Direct On-Policy Distillation (Direct-OPD) transfers reinforcement-learning-induced policy improvements from a small model to a larger student by using the token-level log-ratio between post-RL and pre-RL checkpoints as dense supervision on the student's own rollouts. This transfer rewards the policy shift at every state, yet the log-ratio measures only relative change: it can stay fixed even as the probability mass that both checkpoints assign to the student's candidate tokens vanishes. Through an exact construction, we show that the Direct-OPD reward and its update can remain unchanged while the Jensen-Shannon divergence (JSD) and both KL directions between the checkpoints vanish with this mass, and we note that a small JSD bounds how much the teacher's behavior changed. Motivated by this analysis, we propose Selective Supervision for Direct-OPD (S$^2$D-OPD), which ranks student-sampled states by their teacher-reference JSD and masks Direct-OPD supervision at low-divergence states, retaining only the top 10% of states per response. Across two teacher pairs and four student models ranging from 1.7B to 8B parameters, S$^2$D-OPD improves held-out accuracy over dense Direct-OPD on AIME and HMMT benchmarks in seven of eight settings and matches it in the eighth, without extra forward passes. Our code is available at https://anonymous.4open.science/r/S2D-OPD-8868.

1 INTRODUCTION

Direct-OPD transfers a small model’s RL-induced policy shift to a larger student, but its token-level log-ratio can ignore how much probability supports that shift. The paper motivates divergence-guided selection and reports improved held-out accuracy over dense Direct-OPD in seven of eight settings.

1 INTRODUCTION

  • Large-scale RL remains costly as model size grows, while OPD can transfer capabilities efficiently but requires a teacher stronger than the student.Direct-OPD and Proxy-OPD offer a weak-to-strong route by transferring the RL-induced shift from a small model.Source: §1 · pp. 1–2
  • Seven of eight teacher–student settings improved in held-out accuracy over dense Direct-OPD, with a mean gain of 0.95 points (95% CI 0.40–1.54).S2D-OPD matched dense Direct-OPD in the remaining setting.Source: §1 · p. 2
  • S2D-OPD masks Direct-OPD supervision at low-divergence states, retaining the highest-JSD states within each response.It retains the top 10% of states per response and requires no extra forward passes.Source: §1 · p. 2
  • An exact construction shows the log-ratio reward and local gradient can remain fixed as candidate probability mass, teacher–reference JSD, and both KL directions vanish.The paper also shows that JSD bounds how much the teacher’s behavior can change at a state.Source: §1 · p. 2
  • Within settings, performance broadly rises with JSD percentile: the top bin beats a uniformly sampled 10% subset, while the lowest bin can degrade and eventually collapse the student.Across teacher pairs, the lower-overall-JSD pair benefits more from masking.Source: §1 · p. 2

2 RELATED WORK

Related work addresses RL cost and transfer through training algorithms, systems, and on-policy distillation. This paper focuses on transferring an RL-induced policy shift and selecting student-visited positions using teacher–reference divergence.

2 RELATED WORK

  • OPD trains on student-sampled prefixes using teacher next-token distributions, but primarily transfers the teacher’s policy rather than improvements already present in it.ExOPD extrapolates a teacher improvement, whereas Direct-OPD and Proxy-OPD transfer a reward-optimized checkpoint’s log-ratio to its pre-RL reference (Yang et al., 2026; Feng et al., 2026; Fu et al., 2026a).Source: §2 · p. 2
  • Existing on-policy distillation primarily learns from the teacher’s policy, limiting transfer to improvements already present there; this work ranks student-visited prefixes by teacher–reference JSD to select supervision positions.Token-level log-ratios capture relative changes but are insensitive to the absolute probability mass supporting those changes.Source: §2 · pp. 2–3
  • Selective-distillation criteria draw on student uncertainty, teacher confidence, or teacher–student disagreement, while this work scores the pre-RL-to-post-RL shift at student-visited prefixes.Unlike OPD2’s sign-agreement gate, the criterion uses teacher–reference JSD while retaining Direct-OPD updates at selected positions (Heo et al., 2026).Source: §2 · p. 3

3 PRELIMINARIES

The setup distinguishes a pre-RL reference, a post-RL teacher, and a larger student whose own sampled responses define the training states. Direct-OPD uses the teacher–reference policy shift as a reward and regularizes the student toward its initialization.

3 PRELIMINARIES

  • The reference is a pre-RL policy, the teacher is obtained from it through outcome-based RL, and the larger student samples responses that define states as prompt and preceding tokens.The experiments use publicly released teacher-side checkpoints, so the authors run no RL themselves.Source: §3 · p. 3
  • At each student-visited state, Direct-OPD uses the teacher–reference token log-ratio as a dense reward; it is positive when RL raised a token’s probability and negative when RL lowered it.The objective applies this reward on student-visited states with KL regularization to the student initialization.Source: §3 · p. 3
  • In practice, Direct-OPD evaluates rewards on the student’s top-K candidates and weights them by renormalized student probabilities.The local reward-gradient contribution is computed with the candidate set fixed; the student-anchor KL supplies a separate regularization gradient.Source: §3 · p. 3

4 METHOD

The method identifies a probability-mass mismatch in Direct-OPD and uses teacher–reference JSD to select the states receiving supervision. Selection is performed within each response, and the divergence score reuses probabilities already computed by Direct-OPD.

4 METHOD

  • Direct-OPD applies log-ratio supervision at every valid student-response position, while S2D-OPD selects positions using teacher–reference divergence.The section develops the mismatch and then specifies the divergence-guided selection procedure.Source: §4 · p. 3

4.1 THEORETICAL MOTIVATION: PROBABILITY-MASS MISMATCH

  • With the student and conditional candidate distributions fixed, an exact construction shows that teacher–reference JSD and both KL directions vanish with shared candidate mass, while candidate rewards and the Direct-OPD update stay fixed and nonzero.This isolates probability mass as a factor the Direct-OPD update ignores.Source: §4.1 · p. 4
  • At any state, small teacher–reference JSD bounds the teacher’s behavioral change, so low-divergence masking removes states where behavior changed little even though the Direct-OPD update can be large.The theoretical results do not establish that removing these states improves transfer; the paper tests that in JSD-percentile training analyses.Source: §4.1 · p. 4

4.2 DIVERGENCE-GUIDED STATE SELECTION

  • S2D-OPD scores student-sampled states by teacher–reference JSD and retains the highest-scoring states within each response.The score is evaluated on top-K student candidates plus a residual token for the remaining probability mass; this coarsening cannot increase JSD.Source: §4.2 · p. 4, Figure 1 · p. 5
  • A low score can reflect either little mass on the student’s candidates or substantial candidate mass distributed similarly by both checkpoints.In either case, a small score implies little teacher behavioral change on the candidates; differences among tail tokens outside the candidate set do not affect the score.Source: §4.2 · pp. 4–5
  • The method retains the top-scoring valid positions within each response, applies the mask to both Direct-OPD objective terms, and requires no extra forward passes.At least one state is retained per response; masked states receive neither reward nor student-anchor supervision.Source: §4.2 · p. 5

5 EXPERIMENTS

The experiments evaluate whether masking low-JSD states improves Direct-OPD transfer, why state selection matters, and whether gains persist across design choices. Across the main comparison, top-10% JSD selection improves held-out accuracy in seven of eight settings and ties in one.

5.2 MASKING LOW-DIVERGENCE STATES IMPROVES TRANSFER

  • 0.95 points: S2D-OPD raises mean held-out accuracy from 46.36% to 47.31% over dense Direct-OPD, improving seven of eight settings and tying one.The paired-problem bootstrap gives a 95% CI of 0.40–1.54 points and the exact one-sided sign-flip test gives p = 6.4 × 10^-4.Source: Table 1 · p. 6, §5.2 · p. 6

5.1 EXPERIMENTAL SETUP

  • Two teacher pairs and four students are evaluated with top-10% JSD retention per response, selecting checkpoints on AIME 2024/2025 and testing on held-out AIME 2026 and HMMT.Students train on Skywork-OR1-RL-Data; the held-out benchmarks postdate release of all student models and are absent from their training data.Source: §5.1 · p. 6

5.2 MASKING LOW-DIVERGENCE STATES IMPROVES TRANSFER

  • S2D-OPD matches or exceeds dense Direct-OPD during the first 60 steps for all four students, then maintains higher accuracy and steadier late training.Under the JustRL pair, dense Direct-OPD repeatedly falls from peaks for Qwen3-4B and Qwen3-8B, and falls below Qwen3-4B's initialization by step 300.Source: §5.2 · p. 6, Figure 2 · p. 7

5.3 LOW-DIVERGENCE SUPERVISION IS REDUNDANT AND HARMFUL

  • 52.9: the top JSD percentile bin's peak validation accuracy exceeds the random 10% control (51.2) and dense Direct-OPD (51.3); gains track log mean training JSD (R2 = 0.97).At an equal 10% budget, random selection is on par with dense supervision, while the top-bin result shows that which states are retained matters.Source: §5.3 · p. 7, Figure 3 · p. 7
  • 0.044 JSD falls below the 0.067 cutoff: S2D-OPD masks the largest-log-ratio state and retains one with larger redistribution (JSD 0.262).The masked state has a sampled-token probability drop from 0.12 to 2 × 10^-5, yet only 0.12 total-variation change in the teacher distribution.Source: Figure 5 · p. 8, §5.3 · p. 8
  • The three lowest JSD bins finish below the initial student, and the two lowest collapse; the top bin moves toward the teacher while preserving reference overlap.Top-K overlap suggests that lower bins move the student away from the reference without a commensurate gain in teacher alignment.Source: §5.3 · p. 8
  • 1.0–1.7 points: masking gains more held-out accuracy under lower-JSD QuestA than under JustRL (0.0–1.2), suggesting its benefit comes from filtering low-divergence supervision.Mean JSD is higher for JustRL than QuestA over both all states and retained states.Source: §5.3 · p. 8, Figure 4 · p. 8

5.4 SELECTIVE MASKING IS ROBUST ACROSS DESIGN CHOICES

  • 52.2–53.0 peak validation accuracy across 5%–20% retention beats dense Direct-OPD (51.3) and random 10% (51.2); forward KL tracks JSD, while batch ranking matches per-response selection.Reverse KL is weaker on Qwen3-4B. Batch-versus-response peak accuracy is 52.3 versus 52.9 under JustRL and 51.7 versus 51.7 under QuestA.Source: §5.4 · p. 9, Figure 6 · p. 9

6 CONCLUSION

S2D-OPD applies Direct-OPD supervision only to the highest-JSD states, improving or matching dense Direct-OPD while using a selected-position update and inherited adaptive KL control.

6 CONCLUSION

  • 10% retention improved held-out accuracy over dense Direct-OPD in seven of eight settings and tied in the eighth, without extra forward passes.The conclusion also reports that dense supervision was largely redundant and that low-divergence supervision could harm the student when trained on alone.Source: §6 · p. 9
  • The findings suggest weak-to-strong policy transfer depends not only on what is transferred but also on where it is applied.Adaptive retention rates and domains beyond mathematical reasoning are identified as natural next steps.Source: §6 · p. 9

A.1 ADAPTIVE KL CONTROL

  • The KL coefficient rises when mean student-weighted policy shift is positive and falls when it is negative, subject to bounds of 0.5 and 2.5.The controller starts at 2.5, changes by 1% per update, and uses the updated coefficient in the subsequent actor update.Source: §A.1 · p. 13

A.2 DETAILED LOSS AGGREGATION

  • S2D-OPD restricts Direct-OPD’s positionwise average to selected positions while leaving each local update unchanged.The KL penalty and policy-shift terms apply to the same selected positions; sampled prefixes, candidate sets, mask, and KL coefficient stay fixed during differentiation.Source: §A.2 · pp. 13–14

B PROOF OF THE MASS-INVARIANCE CONSTRUCTION

The construction holds the student and candidate-set conditionals fixed while shrinking the shared probability mass assigned to candidate tokens, separating Direct-OPD reward and gradient from checkpoint divergence.

B PROOF OF THE MASS-INVARIANCE CONSTRUCTION

  • The construction varies each checkpoint with ε while holding the student, candidate set, and conditional distributions fixed; ε is the total mass assigned to student candidates.The student has positive probabilities, and the candidate set is a strict subset of the vocabulary with K ≥ 2.Source: B PROOF OF THE MASS-INVARIANCE… · p. 14
  • For each candidate token, the Direct-OPD reward is invariant to ε because the common factor cancels in the ratio.The local reward gradient at the student checkpoint is also independent of ε and nonzero with respect to the student logits.Source: B PROOF OF THE MASS-INVARIANCE… · p. 14
  • Teacher–reference JSD and both KL directions scale with ε, so all three divergences vanish as ε approaches zero.The shared tail contributes zero to each divergence, while ε factors out of the candidate-set sums.Source: B PROOF OF THE MASS-INVARIANCE… · p. 14
  • As shared candidate mass vanishes, the Direct-OPD reward and its gradient remain fixed and nonzero, while JSD and both KL divergences vanish.The construction assumes distinct, strictly positive teacher and reference distributions on the candidate set and a common positive distribution on its complement.Source: B PROOF OF THE MASS-INVARIANCE… · p. 14

C TOKEN OVERLAP ACROSS DIVERGENCE RANKS

The overlap analysis compares how training on JSD-ranked states changes the student’s top-16 token overlap with both the post-RL teacher and pre-RL reference.

C TOKEN OVERLAP ACROSS DIVERGENCE RANKS

  • Higher-JSD bins generally attain greater overlap with the post-RL teacher, broadly tracking the performance ordering across bins.The panels track the student against the post-RL teacher and pre-RL reference; overlap is the fraction of shared tokens between top-16 sets.Source: C TOKEN OVERLAP ACROSS DIVERGENCE RANKS · p. 15, Figure 7 · p. 15
  • The highest-JSD bin increases teacher overlap while keeping reference overlap above its initial level.This shows that the strongest transfer in this comparison does not require a trade-off between overlap with the two checkpoints.Source: C TOKEN OVERLAP ACROSS DIVERGENCE RANKS · p. 15
  • Reference overlap alone does not follow the performance ordering: the 10–20 bin recovers close to its initial reference overlap despite much smaller gains in teacher overlap.The overlap experiment also includes a uniformly sampled 10% control.Source: C TOKEN OVERLAP ACROSS DIVERGENCE RANKS · pp. 14–15

D SENSITIVITY TO RETENTION RATIO

For Qwen3-1.7B under the JustRL teacher pair, retaining different fractions of high-JSD positions produced broadly similar learning curves, while low-JSD-only training could degrade transfer.

D SENSITIVITY TO RETENTION RATIO

  • Retaining 5%, 10%, 15%, or 20% of high-JSD positions yielded broadly similar learning curves and substantial gains over the base model on AIME 2024 and AIME 2025.No retention ratio consistently dominated across both benchmarks and training checkpoints.Source: D SENSITIVITY TO RETENTION RATIO · p. 15, Figure 8 · p. 16
  • Low-JSD-only training can substantially degrade transfer, while varying the retained fraction within the high-JSD region has a comparatively small effect.The most selective setting did not consistently outperform broader subsets.Source: D SENSITIVITY TO RETENTION RATIO · pp. 15–16
  • JSD is used to screen supervision rather than to assume that larger divergence always implies greater transfer utility.The similar performance across retention ratios supports 10% as a common default in this setting.Source: D SENSITIVITY TO RETENTION RATIO · pp. 15–16

E CASE STUDY: WHAT POLICY CHANGES DOES SELECTION PRESERVE?

The case study finds that JSD retains states with substantial teacher–reference changes, including numerical corrections, continuation preferences, and local errors. Selection depends on the state’s distributional change, not the sampled token or whether the change is correct.

E CASE STUDY: WHAT POLICY CHANGES DOES SELECTION PRESERVE?

  • At the same sampled token, JSD is 0.0001 for a valid calculation and 0.1735 for an incorrect boundary, producing different mask decisions.For the incorrect boundary, RL moves probability from 3 to 4 (∆p ≈−0.55 and +0.56).Source: E CASE STUDY: WHAT POLICY CHANGES DOES… · p. 16
  • At retained states, RL lowers the probability of Alternatively; at a later state it favors Let, Given, and Since instead.The later decrease for Alternatively is ∆p ≈−0.71.Source: E CASE STUDY: WHAT POLICY CHANGES DOES… · p. 17, Figure 10 · p. 17
  • JSD can retain an adverse shift: the student’s odd judgment receives increased probability, and a later reconsideration preference does not undo that earlier shift.These four cases were selected by hand for interpretability and do not estimate how often each pattern occurs.Source: E CASE STUDY: WHAT POLICY CHANGES DOES… · pp. 16–17, Figure 11 · p. 17
  • JSD retains states with numerical corrections, continuation preferences, and local errors, while Direct-OPD’s log-ratio determines which tokens are encouraged at those states.The score concentrates supervision where RL changed the teacher most, without filtering shifts by correctness.Source: E CASE STUDY: WHAT POLICY CHANGES DOES… · p. 17

F TRAINING DETAILS

Experiments use the math subset of Skywork-OR1-RL-Data and a step-by-step response prompt requiring a final answer line. Runs use verl on eight NVIDIA H200 GPUs, with one training run per configuration.

F TRAINING DETAILS

  • All experiments use the math subset of Skywork-OR1-RL-Data and prompt step-by-step solutions with the answer on its own final line.The prompt specifies the form “Answer: $Answer” and asks that the answer appear on its own line.Source: F TRAINING DETAILS · p. 18
  • Experiments run with verl on 8 NVIDIA H200 GPUs, and each configuration is trained with a single run.The confidence interval resamples held-out problems and does not reflect variation across training seeds.Source: F TRAINING DETAILS · p. 18

G LIMITATIONS

The limitations concern what JSD selects, the transfer setting evaluated, and the strength of per-setting evidence. In particular, JSD measures shift magnitude rather than correctness or student need, and experiments use one run per configuration.

G LIMITATIONS

  • JSD can retain large shifts toward errors as readily as corrections, so retained-supervision quality still depends on the teacher’s RL.The authors suggest selectors incorporating direction or correctness, such as verifier or outcome signals, as a possible extension.Source: G LIMITATIONS · p. 18
  • JSD depends on the student only through its top-K candidate set, so it may spend budget where the student already follows the shift and discard shifts the student lacks.Combining teacher–reference divergence with student-side signals is proposed as a natural extension.Source: G LIMITATIONS · p. 18
  • Evidence covers transfer from 1.5B teachers to students up to 8B on mathematical reasoning with verifiable rewards; code and agentic tasks remain untested.The authors note that weak-to-strong transfer is most valuable for much larger students, where RL is most costly.Source: G LIMITATIONS · pp. 18–19
  • Each configuration is trained once, so reported confidence intervals capture held-out-problem variation rather than variation across training seeds.The authors say repeated runs would strengthen per-setting comparisons.Source: G LIMITATIONS · p. 19
Loading 2609.29142v1…