1% of Tokens Can Be Enough: On Gradient Estimation in On-Policy Distillation

Most token positions in a student rollout are not worth supervising — not because they're unimportant, but because a single sampled next token produces a gradient estimate so noisy it can actively mislead training. That's the core observation here, and it's more precise than the usual "not all tokens matter" heuristic.

The paper introduces the information-efficiency ratio (IER): a signal-to-noise measure derived from Fisher geometry rather than Euclidean distance. At a fixed prefix, the reverse-KL gradient is an expectation over the full vocabulary, but in practice you estimate it from one sampled token. IER quantifies how reliably that single sample captures the true gradient direction, under the optimal scalar baseline that minimizes variance. The distribution of IER scores turns out to be extremely heavy-tailed — fewer than 0.1% of tokens score above 1, meaning estimated noise exceeds signal for the vast majority of positions. Supervising everything, as standard on-policy distillation does, is therefore not just wasteful but potentially counterproductive.

Computing IER exactly requires full-vocabulary statistics at every prefix, which is infeasible. The practical approximation uses a top-K candidate set drawn from both student and teacher logits plus the sampled token, with small values assigned for missing overlap. This keeps overhead modest: adding IER-based selection increases mean step time by roughly 2% to 2.5% relative to full on-policy distillation.

The experimental results are striking at the extremes. At a 0.1% token budget — roughly one token per trajectory in the medical task — IER alone achieves a HealthBench overall score of 45.25 against full distillation's 45.77. On mathematical benchmarks, IER at 0.1% approaches or exceeds full distillation on several settings. When combined with existing usefulness scores via soft OR and AND operators, the combined selectors match or beat full distillation at budgets between 0.1% and 1% across both math and medical reasoning tasks.

The gains are not universal. Entropy-based selectors don't reliably improve when augmented with IER, and performance is non-monotonic with budget — adding more tokens sometimes hurts. The paper is candid that both usefulness scores and IER are approximations with no theoretical guarantee of improvement. Thinking-on mode is notably harder: usefulness-alone selectors can underperform the base student model, while adding IER generally recovers performance.

One practical caveat: sparse supervision here doesn't mean shorter training runs. Complete trajectories are still generated and scored; only the gradient update is restricted to selected tokens. The computational savings from sparsity aren't realized yet, though the authors suggest trajectory truncation once the IER budget is met as a direction for future work.

The framing — usefulness and gradient reliability as two orthogonal axes — is the real contribution. Prior work optimized along one dimension; this paper shows the other dimension has comparable leverage at small budgets.