Influence-Directed Distillation: Solving the Diversity Bottleneck in Sampled-Token On-Policy Distillation

AI summary

Overview

Research area: Large language model post-training, specifically on-policy distillation (OPD) of reasoning and code capabilities from a teacher model to a student model.

Technical level: Advanced — the paper requires familiarity with knowledge distillation, reverse KL estimation, policy-gradient advantages, and policy entropy.

Scope: The paper diagnoses why sampled-token on-policy distillation improves a student's pass@1 while failing to inherit the teacher's pass@k diversity, derives a first-order entropy proxy for this effect, and proposes a divergence-adaptive shrinkage method (IDA-OPD) that restores diversity using only the teacher's sampled-token log-probability.

What This Paper Is About

Sampled-token on-policy distillation is cheap: it trains a student on its own generated tokens and only needs the teacher's probability for that one sampled token, avoiding full-vocabulary teacher logits. However, it suffers from diversity distillation failure — the student's pass@1 goes up while its pass@k plateaus, meaning it never inherits the teacher's range of solutions. The authors set out to explain where this entropy loss comes from at the level of individual updates, and to fix it without paying the cost of full-vocabulary or top-K teacher distributions.

Key Contributions

1. Root-cause diagnosis of diversity distillation failure. The authors trace the pass@1-up / pass@k-flat gap to entropy-contracting updates, and show empirically that a substantial share of cumulative entropy loss occurs at low-divergence positions where teacher and student already agree.

2. First-Order Local Entropy Influence. They formulate a signed metric, I_H(y) = A_y · D_y, that decouples each update's entropy effect into the teacher–student log-probability gap A_y and a student-side local probability structure term D_y. Its sign is shown to control the entropy trajectory under real training.

3. IDA-OPD (Influence-Directed Adaptive On-Policy Distillation). A method that leaves entropy-expanding updates intact and reweights entropy-contracting advantages by a divergence-adaptive shrinkage weight w_y = |q_y − p_y| / (q_y + p_y), using no teacher information beyond the sampled-token log-probability.

4. Empirical validation across math and code. IDA-OPD is evaluated on two model scales for mathematics and one for code, improving pass@k over same-budget methods and matching the strongest teacher-informed baselines at strictly lower information cost.

Main Findings

• Advantage alone does not determine entropy. Measured one-step entropy change plotted against A_y forms a "broad two-sided fan": updates with similar advantages produce entropy changes of opposite sign. The authors illustrate that the same positive advantage sharpens the distribution when reinforcing a high-probability mode but flattens it when promoting a low-probability alternative. In contrast, I_H(y) "correlates tightly with the measured entropy change and predicts its sign."

• Entropy contraction peaks at low discrepancy, not high. Binning updates by normalized discrepancy δ_y ∈ [−1, 1], the authors find the cumulative entropy loss "overwhelmingly peaks" at δ_y ≈ 0, driven by a large volume of highly aligned tokens. A secondary cluster appears at the high-divergence negative tail (δ_y ≈ −1.0). The token-count histogram exposes the massive population of low-discrepancy tokens that accumulate into the dominant drain.

• Theoretical guarantee of the shrinkage shape. Proposition 1 states that w_y A_y = O(A_y²) as A_y → 0, and w_y A_y → A_y as |A_y| → ∞ — near-quadratic attenuation of low-discrepancy updates and near-lossless preservation of high-divergence corrections.

• Pass@16 gains across all four math benchmarks at both scales. In the 4B setting, IDA-OPD improves over standard OPD by +8.3 (HMMT Feb), +4.6 (AIME 2024), and +4.3 (AIME 2025). In the 8B setting it yields +7.2, +4.2, and +6.0 respectively. IDA-OPD attains the highest pass@16 across all four benchmarks at both scales.

• Diversity can exceed the teacher. IDA-OPD matches or surpasses the teacher's pass@16 on several benchmarks: 83.3% vs. 80.0% on 4B AIME 2024, and 76.7% vs. 73.3% on 8B AIME 2025.

• Pass@1 is broadly maintained or improved. The Δ vs OPD rows in Table 1 show pass@1 changes of +1.6, +2.1, +0.1, +1.5 (8B) and +2.1, +1.5, −0.2, +1.1 (4B) across the four benchmarks. The single negative value is −0.2 on HMMT Feb at the 4B scale.

• The gap widens with k. On AIME 2024 and AIME 2025 for k = 1 to 64, IDA-OPD and OPD are close at k = 1 (~+2 points), but the gap widens to +4 to +5 points by k = 8–16 and is sustained through k = 64.

• Results transfer to code. On MBPP+ and LiveCodeBench, OPD leaves pass@16 flat (72.9 vs. the student's 73.5 on MBPP+), while IDA-OPD improves +2.1/+0.9 on MBPP+ and +1.3/+1.0 on LiveCodeBench (pass@1/pass@16). IDA-OPD reaches 71.6/73.8 on MBPP+ and 28.1/55.2 on LCB.

• Both design choices are necessary. Ablations on Qwen3-4B-Non-Thinking-RL-Math → Qwen3-4B: full IDA-OPD scores 56.7/83.3 (AIME24) and 53.3/70.0 (AIME25). Removing the I_H gate drops to 53.8/81.5 and 51.2/68.4; hard masking drops to 52.1/82.0 and 50.4/68.9; using sign(A_y) as the gate gives the worst pass@16 at 76.8/65.2.

• Linear shrinkage beats other shapes. Holding the I_H gate fixed, comparing Constant (54.6/79.2), Sqrt (55.9/80.4), linear "Ours" (56.7/83.3), and Square (54.1/81.6) on AIME24. Constant is worst for leaving low-divergence drain uncurbed; Square over-suppresses moderate-divergence corrections and posts the lowest pass@1.

• The proxy holds under real training dynamics. In Figure 4(a), masking entropy-expanding updates drives entropy below vanilla OPD, masking entropy-contracting updates lifts it highest, and IDA-OPD's soft reweighting sits between the two while staying well above OPD. The authors note the theorem assumes idealized logit-space single-step updates, but the global trajectory shifts indicate I_H remains a robust directional proxy under parameter-space optimization with shared weights and AdamW dynamics.

• Token-level entropy moves toward the teacher. On AIME 2024, in the lowest-entropy bin (0.0–0.1) OPD piles up ~74% of mass versus the teacher's ~57%, while IDA-OPD stays close at ~55% and retains more mass in the high-entropy region (≥ 1.0) where OPD is depleted.

• Efficiency is preserved. IDA-OPD transfers only one scalar per position (O(L)); teacher-informed methods (EOPD, AOPD) transfer a truncated top-K teacher distribution at selected positions (O(L_sel K), K >> 1). The only extra cost, the D_y entropy term, reuses student logits already materialized in the forward pass.

Methodology in Plain English

The authors start from a standard setup: a teacher model generates a target distribution, and a student model is trained on prefixes it produced itself. To avoid computing the teacher's full next-token distribution, they use the cheap "sampled-token" variant — draw one token from the student, ask the teacher only for that token's log-probability, and use the difference as an advantage signal.

Their first move is diagnostic. They ask: does the advantage alone tell you whether an update will raise or lower the student's entropy? They plot measured one-step entropy change against the advantage on real training positions and find no clean relationship. So they do a first-order mathematical expansion of the entropy change under a single logit-space gradient step. The expansion factors neatly into two pieces: the advantage, and a term that depends only on the student's current probability distribution. The product of the two, which they call First-Order Local Entropy Influence, predicts the sign of the entropy change.

Their second move is to find out where the entropy actually disappears. They bin updates by a normalized teacher–student discrepancy measure and, in each bin, measure cumulative entropy loss across genuine optimizer steps (not the proxy). The loss concentrates where teacher and student already agree — a huge population of low-impact tokens that collectively drain entropy.

Their third move is the intervention. For each update, they compute the sign of the influence. If it is non-negative (entropy-expanding), the advantage is left untouched. If it is negative (entropy-contracting), the advantage is multiplied by a symmetric relative disagreement weight that is near zero when teacher and student agree and approaches one as they diverge. This preserves genuine teacher corrections while softening the low-value updates responsible for the drain.

They then evaluate on competition mathematics (AIME 2024, AIME 2025, HMMT 2025 Feb, HMMT 2025 Nov) and on code (MBPP+, LiveCodeBench), using two teacher→student settings for math (8B and 4B Qwen3-Non-Thinking distillations) and one for code. Mathematics training uses DeepMath103K filtered to difficulty level 6. Evaluation uses the unbiased pass@k estimator of Chen et al. (2021) with n = 128 sampled responses per problem at temperature 1.0 and top-p = 1.0, reporting pass@1 and pass@16. Baselines are standard OPD, REOPOLD, EOPD, AOPD, an entropy bonus, and advantage shaping. They also run ablations that disable the influence gate or the shrinkage individually, compare four shrinkage shapes, and track policy entropy and token-level entropy histograms during and after training.

Why This Matters

Impact on research. The paper reframes diversity loss in distillation as a per-update phenomenon rather than a global entropy knob. Its central diagnostic is teacher-free — it needs only signals the student already produces — which challenges the assumption that protecting diversity requires reintroducing expensive teacher distributions. It also supplies a mechanistic explanation (low-discrepancy, high-volume updates) for a failure mode that prior work observed but had not attributed to specific positions.

Real-world applications:

• Reasoning assistants that must solve competition-style mathematics and are deployed with repeated sampling or self-consistency voting, where pass@k matters more than single-shot accuracy.
• Code generation tools that benefit from returning several distinct candidate solutions rather than several near-identical ones.
• Post-training pipelines for smaller, cheaper student models that must be distilled from larger teachers under compute or bandwidth constraints.
• Any deployment that needs to sample many candidate outputs and select among them, where collapsing to a few dominant modes reduces the benefit of test-time scaling.

Industry relevance. The paper notes that OPD is integrated into the post-training pipelines of recent large-scale systems including GLM-5, KiMi, and Qwen3. For teams that already use sampled-token OPD for its efficiency, IDA-OPD is attractive because it adds no teacher query beyond the sampled token and reuses logits already computed in the forward pass, so it fits inside existing infrastructure rather than requiring a new top-K teacher distribution channel.

Future Directions

• Beyond reasoning and code. The authors demonstrate the diagnosis and intervention on mathematics and code domains; whether the same low-discrepancy entropy drain dominates in other settings (long-form writing, dialogue, agentic tasks) is not established by the reported experiments.

• Theory-to-practice gap. Theorem 1 assumes idealized logit-space, single-step gradient updates. The authors present Figure 4(a) as empirical evidence that the proxy remains directional under parameter-space optimization with shared weights and AdamW dynamics, but a fuller theoretical account of that regime is left open.

• Choosing the shrinkage shape in principle. The linear map performed best among the four shapes tested, but the paper selects it empirically rather than deriving the optimal f(w_y) from the analysis. A principled derivation of the shrinkage form is a natural extension.

• Where the boundary lies for hard interventions. The ablation shows that masking all entropy-contracting updates protects diversity but causes the largest pass@1 damage, while full IDA-OPD balances the two. Quantifying how much knowledge signal can be sacrificed before pass@1 degrades materially remains an open tuning question.

Target Audience

Researchers and engineers working on LLM post-training, knowledge distillation, and reinforcement-learning-based fine-tuning — particularly those already using sampled-token on-policy distillation and observing that pass@1 improves while pass@k stalls. It will also interest readers studying policy entropy, exploration, and test-time scaling, and practitioners who need diversity-preserving distillation without paying for full-vocabulary or top-K teacher signals.

Authors’ abstract

Sampled-token on-policy distillation (OPD) efficiently transfers capabilities from teacher to student using student-generated tokens, requiring teacher probabilities only for sampled tokens. Yet it frequently suffers from diversity distillation failure: the student's pass@1 improves while its pass@$k$ plateaus, failing to inherit the teacher's diversity. To explain this, we introduce First-Order Local Entropy Influence, a signed first-order proxy that decouples each update's entropy effect into the teacher--student log-probability gap and the student's local probability structure, and empirically links entropy contraction to negative-influence positions. Motivated by this, we propose Influence-Directed Adaptive On-Policy Distillation (IDA-OPD): rather than relying on costly full-vocabulary Forward-KL objectives, it preserves entropy-expanding updates while replacing entropy-contracting ones with divergence-adaptive advantage shrinkage, using only the teacher's sampled-token log-probability. Experiments on reasoning-oriented distillation show IDA-OPD consistently improves pass@$k$, inheriting the teacher's diversity through distillation, matches the strongest teacher-informed methods at strictly lower cost, and broadly maintains vanilla OPD's pass@1, all without full-vocabulary teacher information.