PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Run Yang, Runpeng Dai, Jie Sun, Jielei Zhang, Fan Zhou, Hongtu Zhu, Peiyi Li, Longwen Gao

arXiv:2608.29846v1cs.CLcs.LG

TL;DR

Sampled-token OPD is efficient but often improves pass@1 without transferring the teacher’s diversity to pass@k. This paper introduces an entropy-influence diagnostic and IDA-OPD, which selectively shrinks entropy-contracting updates; experiments show improved pass@k at lower information cost while broadly maintaining pass@1.

  • Problem

    Sampled-token OPD frequently suffers diversity distillation failure: pass@1 improves while pass@k plateaus, limiting transfer of the teacher’s diversity.

  • Method

    IDA-OPD uses First-Order Local Entropy Influence to preserve entropy-expanding updates and apply divergence-adaptive shrinkage to entropy-contracting advantages using sampled-token information.

  • Results

    IDA-OPD consistently improves pass@k over same-budget methods, matches the strongest teacher-informed methods at lower cost, and broadly maintains vanilla OPD’s pass@1.

  • Takeaways & Limitations

    The method recovers generation diversity through sampled-token distillation without full-vocabulary teacher information.

  • Takeaways & Limitations

    Hard-masking every entropy-contracting update sacrifices knowledge signals and produces the worst pass@1 drop among the evaluated variants.

Abstract

from arXiv · show

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.

1 Introduction

Sampled-token OPD reduces teacher-information costs but often improves pass@1 without improving pass@k, creating a diversity distillation failure. The paper attributes this to entropy-contracting updates and proposes IDA-OPD, which selectively shrinks low-discrepancy contracting updates while preserving entropy-expanding updates.

  • Sampled-token OPD frequently improves pass@1 while pass@k plateaus, failing to transfer the teacher’s generation diversity.
  • IDA-OPD preserves entropy-expanding updates and applies divergence-adaptive advantage shrinkage to entropy-contracting updates, strongest at small teacher–student discrepancies.
  • IDA-OPD improves pass@k under the same sampled-token budget, matches strong teacher-informed methods at lower information cost, and broadly maintains pass@1.
  • Entropy-contracting updates, including many low-discrepancy positions, explain the diversity distillation failure.
  • First-Order Local Entropy Influence separates each update’s entropy effect into the teacher–student log-probability gap and the student’s local probability structure.

2 Preliminaries

Standard OPD matches teacher and student distributions on student-induced prefixes, but exact reverse-KL computation requires costly full-vocabulary teacher logits. Sampled-token OPD uses a one-token estimator requiring only the teacher log-probability of the sampled token, though this can reduce diversity and cause entropy collapse.

  • OPD optimizes reverse KL between student and teacher next-token distributions on student-induced prefixes.
  • Exact reverse-KL computation requires full-vocabulary teacher logits at every prefix, making it prohibitively expensive.
  • Sampled-token OPD estimates reverse KL by drawing one token from the student distribution and querying only that token’s teacher log-probability.
  • Although locally unbiased, the one-token estimator can reduce diversity and cause premature entropy collapse.

3 Understanding Diversity Distillation Failure

Sampled-token OPD updates affect entropy through the interaction between the advantage and the student’s local probability distribution, not through advantage alone. First-Order Local Entropy Influence captures this interaction and predicts whether updates contract or expand entropy.

  • 3.1 Advantage Is Insufficient: On real training positions, advantage alone has a weak relationship with measured entropy change because identical advantages span both signs.This motivates modeling the advantage jointly with the student’s current distribution.
  • 3.1 Advantage Is Insufficient: The same positive advantage can decrease or increase entropy depending on whether it reinforces a high- or low-probability token.High-probability reinforcement sharpens the distribution, whereas low-probability reinforcement spreads probability mass.
  • 3.2 First-Order Local Entropy Influence: Figure 1 shows that IH(y) tracks measured one-step entropy change and predicts its sign more closely than the advantage.This connects the analytical decomposition to real training positions.
  • 3.2 First-Order Local Entropy Influence: Under a local logit-space step, the entropy change satisfies ∆H = η · IH(y) + O(η2), with IH(y) = AyDy.The influence’s sign provides a token-wise signal for identifying entropy-contracting and entropy-expanding updates.
  • 3.3 Where Entropy Loss Originates: Cumulative entropy contraction peaks near δy ≈0, where many individually weak updates accumulate into the dominant entropy drain.A secondary entropy-loss cluster remains at the high-divergence negative tail.
  • 3.2 First-Order Local Entropy Influence: First-Order Local Entropy Influence decomposes the leading entropy effect into the teacher–student log-probability gap and the student’s local probability structure.The advantage scales the update, while the student distribution determines its direction-dependent effect.

4 Influence-Directed Adaptive On-Policy Distillation

IDA-OPD selectively attenuates entropy-contracting updates according to teacher–student disagreement rather than penalizing them uniformly. This targets the low-discrepancy updates driving aggregate entropy loss while preserving high-discrepancy teacher corrections.

  • 4.1 Diagnosing the Entropy Drain: High-divergence updates contribute to entropy loss, but cumulative contraction is dominated by the low-discrepancy region near δy ≈0.The dominant effect arises from the large number of highly aligned tokens, each with minimal individual impact.
  • 4.1 Diagnosing the Entropy Drain: Uniformly penalizing all updates with IH(y) < 0 can suppress critical high-divergence teacher signals while failing to target the main entropy drain.The discrepancy distribution motivates selective rather than uniform intervention.
  • 4.2 Divergence-Adaptive Shrinkage: IDA-OPD attenuates Ay only at entropy-contracting positions, using a symmetric relative disagreement weight wy to measure remaining correction need.The method preserves entropy-expanding updates and requires only sampled-token teacher information.
  • 4.2 Divergence-Adaptive Shrinkage: The retention weight wy is small when teacher and student agree and approaches 1 as disagreement grows, without introducing signal beyond Ay.This makes attenuation strongest for low-discrepancy updates and weakest for high-discrepancy corrections.
  • 4.2 Divergence-Adaptive Shrinkage: The shrinkage is sign-preserving: it never reverses teacher corrections, vanishes near agreement, and leaves high-disagreement corrections intact.Proposition 1 formalizes near-quadratic attenuation as Ay →0 and near-lossless preservation as |Ay| →∞.

5 Experiments

Experiments across mathematics and code show that IDA-OPD recovers diversity while broadly preserving pass@1, using sampled-token information at lower cost than teacher-informed baselines. Ablations and entropy analyses support both its influence-based gating and divergence-adaptive shrinkage.

  • Main Results: IDA-OPD attains the highest pass@16 across all four mathematics benchmarks at both 4B and 8B scales.It improves over standard OPD by +8.3 on HMMT Feb, +4.6 on AIME 2024, and +4.3 on AIME 2025 in the 4B setting.
  • Main Results: IDA-OPD matches or exceeds teacher-informed AOPD and EOPD while avoiding their expensive top-K teacher distributions.It uses a teacher-free entropy diagnostic and no full-vocabulary teacher information.
  • Main Results: +2.1/+0.9 on MBPP+ and +1.3/+1.0 on LiveCodeBench: IDA-OPD improves pass@1/pass@16 in the code domain.This extends the observed diversity-recovery pattern beyond mathematics.
  • Pass@k Performance: +4 to +5 points: IDA-OPD’s pass@k advantage over OPD emerges by k = 8–16 and persists through k = 64.The methods are close at k = 1, where a single high-probability mode can suffice.
  • Ablation Study: The influence gate and divergence-adaptive shrinkage are complementary, while linear shrinkage best balances suppressing low-discrepancy drain against preserving high-discrepancy corrections.Using sign(Ay) misidentifies entropy-risk positions, whereas the linear mapping tops every reported Table 4 column.
  • Ablation Study: Hard masking protects entropy but overprotects: discarding every entropy-contracting update sacrifices knowledge signals and causes the worst pass@1 drop.The ablation shows that entropy control must preserve useful teacher corrections.
  • Token-Level Entropy Analysis: IDA-OPD stays closer to the teacher’s token-entropy distribution by avoiding OPD’s low-entropy pile-up and retaining more high-entropy tokens.OPD places approximately 74% versus the teacher’s 57% in the 0.0–0.1 bin, while IDA-OPD remains near 55%.
  • Main Results: IDA-OPD retains OPD’s cheap O(L) teacher-information footprint, adding only an entropy term that reuses student logits from the forward pass.Teacher-informed EOPD and AOPD use truncated top-K teacher distributions with O(LselK), K ≫1.

6 Related Work

Related OPD methods differ in how they select, reweight, or reshape sampled-token updates, balancing diversity protection against teacher-information cost. The strongest teacher-informed approaches restore diversity but reintroduce full or top-K teacher distributions.

  • On-Policy Distillation for Language Models: On-policy distillation trains on student-sampled prefixes, reducing the train–inference mismatch present in offline distillation.OPD variants differ in the per-token teacher information they require.
  • Token Selection and Reweighting in OPD: Token-selection methods reweight updates using teacher–student divergence, token importance, teachability, or optimization granularity.Examples include TIP, Token Teachability, and Filter-Then-Reweight.
  • Token Selection and Reweighting in OPD: Other methods decouple exploration from imitation, reshape sampled-token credits, or restore entropy with forward-KL penalties at uncertain tokens.AOPD, REOPOLD, negative-sample reinforcement, and EOPD represent these strategies.
  • Token Selection and Reweighting in OPD: EOPD and AOPD achieve strong diversity protection but reintroduce full or top-K teacher distributions, sacrificing sampled-token efficiency.These methods locate salient positions using teacher–student disagreement or uncertainty.

7 Conclusion

IDA-OPD addresses diversity distillation failure in sampled-token OPD with a cheap entropy-influence proxy and adaptive treatment of entropy-contracting updates. It improves pass@k at the sampled-token budget while broadly maintaining vanilla OPD’s pass@1.

  • 7 Conclusion: First-Order Local Entropy Influence flags whether sampled-token updates expand or contract entropy without requiring full teacher distributions.The proxy is described as cheap and teacher-free.
  • 7 Conclusion: IDA-OPD preserves entropy-expanding updates and applies divergence-adaptive advantage shrinkage to entropy-contracting updates.The intervention uses the influence signal to modify only contracting updates.
  • 7 Conclusion: IDA-OPD consistently improves pass@k over methods with the same sampled-token budget, reaches parity with the strongest teacher-informed methods at lower cost, and broadly maintains vanilla OPD’s pass@1.The reported comparisons use no full-vocabulary teacher information.
Loading 2608.29846v1…