Large language models learn which tokens to distill for better efficiency
Learning What to Distill: Bilevel Top-K Token Selection for Self-Distillation in Large Language Models
Machine Learning
Summary
Large language models are powerful but expensive to run, so smaller models try to learn from them through a process called knowledge distillation. The authors study improving this process when the smaller model teaches itself, by learning which parts (tokens) of the text are most important to focus on rather than treating all parts equally. They use a two-level optimization method that lets the model adaptively pick the key tokens during training. Their method improves mathematical reasoning performance with only a small increase in computation.
What this means in practice
- •For mobile developers: Build smaller language models that run efficiently on phones by focusing distillation on the most important text parts to preserve reasoning skills.
- •For cloud service engineers: Reduce inference costs by training compact language models with adaptive token selection to maintain accuracy in complex reasoning tasks.
Authors
Heng Liang, Xinwen Zhang, Hongchang Gao
Abstract
Large language models have shown strong reasoning capabilities, but their high inference costs make knowledge distillation an important approach for transferring such capabilities to compact models in resource-constrained scenarios. On-policy self-distillation further reduces the reliance on external large teacher models while improving the reasoning ability of compact language models. However, existing methods typically either distill all token positions uniformly or select tokens using fixed heuristic criteria, assigning the same distillation strength to the selected positions rather than adaptively learning which tokens are most beneficial for distillation. To address these limitations, we propose BiToK-SD (Bilevel Top-K Token Selection for Self-Distillation), a bilevel-optimization-based token selection method that learns where distillation should be applied during on-policy self-distillation. Specifically, BiToK-SD is formulated as a bilevel optimization problem, where the lower-level problem models Top-K token selection as a differentiable threshold-based relaxation, allowing the selected positions to adapt as the student policy evolves, while the upper-level problem performs knowledge distillation on the selected positions. Experiments on mathematical reasoning benchmarks show that BiToK-SD achieves the best average performance among all compared methods while requiring only lightweight additional computation.