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
cs.CL, cs.LG
2026-08-30
Sampled-token on-policy distillation often lifts pass@1 while pass@k stalls. IDA-OPD shrinks entropy-contracting updates; on 4B AIME24, pass@16 rises from 78.7 to 83.3.
On-policy distillation is already in post-training stacks such as GLM-5, Kimi, and Qwen3. The full-vocabulary version matches the teacher's entire next-token distribution and is expensive to store on long rollouts. Sampled-token OPD queries only the log-probability of the token the student actually drew. The cheap estimator often fails at diversity: pass@1 rises, pass@k stalls, and the student does not inherit the teacher's spread of solutions.
Existing fixes split into two camps. Teacher-informed methods pull top-K distributions and add forward KL at high-disagreement tokens. Student-side entropy bonuses or advantage reshaping stay cheap but blunt, twisting entropy-raising and entropy-draining updates together.
A sampled-token update changes entropy, to first order, by IH(y)=Ay·Dy. Ay is the teacher-student log-probability gap, the usual advantage. Dy depends only on the student's current local distribution. The same sign of Ay can raise or lower entropy: reinforcing an already dominant token sharpens the distribution, promoting a rare token can flatten it. On real training positions, Ay fans out against measured entropy change; the sign of IH(y) tracks it.
The entropy drain is not mainly from large teacher-student disagreements. Binned by normalized discrepancy δy, cumulative entropy loss peaks at δy≈0. Each update is tiny; the token count is huge. The high-divergence negative tail contributes, but it is not the bulk.
IDA-OPD therefore touches only positions with IH(y)<0. Entropy-expanding updates pass through. Entropy-contracting advantages are multiplied by wy=|qy−py|/(qy+py), near 0 when teacher and student already agree and near 1 when they do not. The sign of the correction never flips. The teacher still returns one sampled-token log-probability, never a full or top-K distribution.
Math uses Qwen3-8B/4B RL-Math teachers distilled into same-size Non-Thinking students on DeepMath103K difficulty 6. Code uses a Qwen3-4B-RLCode teacher. Each problem is sampled n=128 times; pass@1 and pass@16 use the unbiased estimator.
| Setting | Method | AIME24 pass@1/16 | AIME25 pass@1/16 |
| 8B | OPD | 61.7 / 79.1 | 47.9 / 70.7 |
| 8B | IDA-OPD | 63.3 / 83.3 | 50.0 / 76.7 |
| 8B | Teacher | 60.0 / 83.3 | 53.5 / 73.3 |
| 4B | OPD | 54.6 / 78.7 | 51.8 / 65.7 |
| 4B | IDA-OPD | 56.7 / 83.3 | 53.3 / 70.0 |
| 4B | Teacher | 53.3 / 80.0 | 53.5 / 73.3 |
Versus OPD on 4B, pass@16 moves +8.3 on HMMT Feb, +4.6 on AIME24, +4.3 on AIME25. The 8B deltas are +7.2, +4.2, +6.0. Several pass@16 numbers match or exceed the teacher, including 83.3 vs 80.0 on 4B AIME24. AOPD and EOPD, which need top-K teacher distributions, recover much of the diversity; IDA-OPD still posts the highest pass@16 in the table at sampled-token cost. On MBPP+ the move is 69.5/72.9 to 71.6/73.8, and on LiveCodeBench 26.8/54.2 to 28.1/55.2: same direction, smaller gap.
Ablating the IH gate, hard-masking contracting updates, or gating on sign(Ay) drops 4B AIME24 pass@1 to 53.8, 52.1, and 54.2. A linear map from wy beats constant, square-root, and square maps.
Sampled-token efficiency is why OPD fits large post-training runs. This paper traces diversity failure to a mass of low-discrepancy, entropy-contracting updates, and gates them with student logits that the forward pass already materializes. Teams that watch pass@k lag pass@1 can plug this in without bringing back full-vocabulary teachers.
Pass@1 usually rises only slightly. The headline is diversity. If production scores a single sample, the win will feel small.
The first-order derivation assumes a small logit-space step. The authors show that the sign still steers AdamW trajectories, but IH(y) remains a proxy. Main tables distill Qwen3 Non-Thinking into Qwen3 Non-Thinking; cross-family and cross-tokenizer results are absent from the lead experiments. Code-domain pass@16 moves by about one point, so the story is carried by contest math. AIME sets have 30 problems, and pass@k is sensitive to a few items. Side effects on instruction following or safety are unreported. There is no dedicated limitations section; these bounds are visible from the experimental design.