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

AI summary

Overview

Research area: Efficient post-training of large language models, specifically on-policy distillation (OPD) and weak-to-strong transfer of reinforcement-learning (RL) improvements from a small model to a larger one.

Technical level: Intermediate. The paper assumes familiarity with policy gradients, KL divergence, Jensen–Shannon divergence (JSD), and top-K token distributions, but its central argument is stated in accessible probabilistic terms.

Scope: The paper identifies a probability-mass blind spot in Direct-OPD's log-ratio reward, proves it with an exact construction, and proposes a JSD-based state-selection mask that retains only the top 10% of states per response.

What This Paper Is About

Direct-OPD transfers RL gains from a small teacher to a larger student by rewarding the student at every response position with the teacher–reference log-ratio of the sampled token. That reward measures only relative change, so it can stay large even when both the post-RL teacher and its pre-RL reference assign almost no probability mass to the student's candidate tokens. The paper asks whether every state's policy shift is actually worth distilling, and answers "no" by masking supervision at low-divergence states.

Key Contributions

1. A probability-mass mismatch in Direct-OPD. Through an exact construction, the authors show that the log-ratio reward and its local gradient on the student can remain fixed while the probability mass behind the policy shift vanishes — and with it the teacher–reference JSD and both KL directions between the checkpoints. They further show that JSD bounds how much the teacher's behavior can change at every state.

2. Stable and effective selective transfer (S²D-OPD). A method that ranks student-sampled states by teacher–reference JSD and masks Direct-OPD supervision at low-divergence states, retaining the top 10% of states per response. It requires no extra forward passes because it reuses teacher and reference probabilities Direct-OPD already computes.

3. Understanding the gains from selective transfer. Controlled experiments show that performance broadly rises with JSD percentile, that the top bin outperforms a uniformly sampled 10% subset, and that the lowest bin degrades the student below its initialization and eventually collapses.

4. Robustness evidence. Gains persist across divergence measures (JSD, forward KL, reverse KL), selection scope (response-level vs. batch-level), and retention rates from 5% to 20%.

Main Findings

• Held-out improvement in seven of eight settings. Across two teacher pairs (R1-Distill-1.5B → JustRL-1.5B and Nemotron-1.5B → QuestA-1.5B) and four students (Qwen3-1.7B, Qwen3-4B, Qwen3-8B, R1-Distill-7B), S²D-OPD beats dense Direct-OPD on held-out Test Avg. in seven of eight teacher–student settings and ties in the eighth (R1-Distill-7B under the JustRL pair, where both reach 40.5).

• Overall gain is statistically significant. Mean held-out accuracy improves from 46.36% to 47.31%, a gain of 0.95 points (95% CI: 0.40–1.54; p = 6.4 × 10⁻⁴), assessed over 93 held-out problems using a paired problem bootstrap and an exact one-sided sign-flip test.

• Early learning is not slowed. Under the JustRL teacher pair, S²D-OPD matches or exceeds dense Direct-OPD accuracy over the first 60 steps for all four students, despite discarding 90% of states. Later, it is more stable on Qwen3-4B and Qwen3-8B; dense Direct-OPD drops below the initial student on Qwen3-4B by step 300.

• Dense supervision is largely redundant. A uniformly sampled 10% control reaches a peak validation accuracy of 51.2, on par with dense Direct-OPD (51.3), supporting the view that OPD is "data-overfed but algorithm-starved" (attributed to Fu et al., 2026c).

• Which states are retained matters, not how many. At the same 10% budget, the top JSD bin reaches 52.9, above both the random control (51.2) and dense Direct-OPD (51.3). Across all eleven subsets, the gain in peak validation accuracy over the initial student is approximately linear in the logarithm of the subset's mean training JSD (R² = 0.97).

• Low-divergence supervision can be actively harmful. Trained in isolation, the three lowest JSD bins end below the initial student, and the two lowest collapse. Top-K overlap analysis suggests the top bin moves the student toward the teacher while preserving overlap with the reference, whereas lower bins move it away from the reference without a commensurate gain in teacher alignment.

• JSD and reward magnitude rank states differently. In one Qwen3-1.7B rollout under the JustRL pair, RL lowers the sampled token's probability from 0.12 to 2 × 10⁻⁵, giving that state the largest Direct-OPD reward magnitude among the response's 8,192 states, even though the teacher's distribution moves only 0.12 in total variation. Its JSD (0.044) falls below the retention cutoff (0.067), so S²D-OPD masks it while retaining a state with a smaller log-ratio but a larger redistribution of mass (JSD 0.262).

• Masking helps more under the lower-divergence teacher pair. Mean JSD is higher for JustRL than for QuestA across all states and the retained top 10%, yet held-out gains over dense Direct-OPD are larger under QuestA: 1.0–1.7 points versus 0.0–1.2 for JustRL across the four students.

• Robustness across design choices. Forward KL and reverse KL selection both beat dense Direct-OPD from step 160 onward on all three Qwen3 students and avoid the Qwen3-4B late decline (final accuracy 71.6–74.1 versus 67.2). Batch-level selection matches response-level selection on Qwen3-1.7B (52.3 versus 52.9 under JustRL; 51.7 versus 51.7 under QuestA). Varying retention from 5% to 20% on Qwen3-1.7B changes peak validation accuracy by less than one point (52.2–53.0).

Methodology in Plain English

The authors first build a mathematical counterexample. They construct a teacher and reference distribution that agree on the shape of their preferences over the student's top-K candidates but assign them an arbitrarily small total probability mass, parameterized by ε. As ε goes to zero, every per-candidate log-ratio — and therefore the Direct-OPD reward and its gradient — stays exactly the same, while the JSD and both KL directions shrink proportionally to ε and vanish. They also prove a general bound: the difference in probability that any two distributions assign to any event is at most the square root of twice their JSD, so a small JSD means the teacher's behavior provably barely changed anywhere.

From this they design S²D-OPD. At each state visited during the student's own rollouts, they compute teacher–reference JSD over the student's top-K candidates plus a residual token that collects all remaining probability mass. Because the distributions are compressed this way, the score is a lower bound on the full-vocabulary JSD by the data-processing inequality. Within each response, they keep only the top 10% of states by this score and apply the unchanged Direct-OPD objective — both the reward term and the student-anchor KL term — only at those positions. The adaptive KL coefficient is updated using all valid positions, with α₀ = 2.5, ε = 0.01, and range [0.5, 2.5].

Experiments use two publicly released teacher pairs, so no RL is run by the authors. Students are trained on Skywork-OR1-RL-Data, checkpoints are selected on AIME 2024 and AIME 2025, and evaluation is on three held-out benchmarks — AIME 2026, HMMT November 2025, and HMMT February 2026 — all of which postdate the release of every student model. Results are reported as Avg@32, with Test Avg. averaging over 93 held-out problems.

Why This Matters

Impact on research. The paper reframes OPD state selection as a question about the teacher's behavior change rather than the size of its log-ratio, and it supplies both a formal impossibility result for the existing reward and a cheap, drop-in fix. It also connects OPD to the observation that RL-induced policy shifts are sparse and concentrated in a small fraction of tokens.

Real-world applications (implications drawn from the paper's setting):

• Improving large reasoning models without paying the compute cost of running RL at their scale, by running RL on a small model and transferring the result.
• Post-training pipelines where only released checkpoints are available rather than live RL infrastructure.
• Cost-sensitive training runs that cannot afford dense per-token supervision over every rollout position.
• Model-upgrade workflows where a larger student already outperforms the post-RL teacher, a case Direct-OPD targets and ordinary OPD cannot handle.

Industry relevance. Discarding 90% of supervised positions without extra forward passes directly reduces training supervision cost while improving held-out accuracy, which matters for teams whose bottleneck is rollout and training compute rather than data. The "data-overfed but algorithm-starved" finding also suggests that adding more distillation data has limited returns relative to selecting where supervision is applied.

Future Directions

• Extending state selection to adaptive retention rates, rather than the fixed top-10% (ρ = 0.1) used by default.
• Applying selective supervision to domains beyond mathematical reasoning, which the paper names as a natural next step.
• Explaining mechanistically why low-divergence supervision harms the student; the top-K overlap analysis is offered as a mechanism but the paper notes the selection results do not themselves prove that removing low-divergence states improves transfer.
• Testing whether the divergence-based masking transfers to other policy-shift methods such as Proxy-OPD and ExOPD, and whether the divergence measure should be chosen per teacher pair, given that forward KL tracked JSD closely while reverse KL was weaker on Qwen3-4B.

Target Audience

Researchers and engineers working on LLM post-training, reinforcement learning from verifiable rewards, and knowledge distillation, particularly those interested in weak-to-strong transfer, on-policy distillation objectives, or token-level supervision design. It is most useful to readers comfortable with policy-gradient objectives and divergence measures; the theoretical construction is compact enough to follow without a heavy mathematics background.

Authors’ abstract

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.