arXiv：cs.LG· Heng Liang, Xinwen Zhang, Hongchang Gao·· 4 小时前AI 评分31
BiToK-SD：面向大语言模型自蒸馏的双层 Top-K Token 选择方法
Learning What to Distill: Bilevel Top-K Token Selection for Self-Distillation in Large Language Models
AI 导读
BiToK-SD 提出一种基于双层优化的 Top-K Token 选择方法，用于大语言模型的自蒸馏，通过学习蒸馏应施加在哪些 Token 位置来替代均匀蒸馏与固定启发式选择。该方法将 Top-K Token 选择建模为可微的阈值松弛，使选中位置随学生策略演变而自适应调整，上层则在选中位置执行知识蒸馏。在数学推理基准上，BiToK-SD 取得对比方法中最佳的平均性能，且仅需轻量额外计算。
正文
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.
| Subjects: | Machine Learning (cs.LG) |
| Cite as: | arXiv:2610.07247 [cs.LG] |
| (or arXiv:2610.07247v1 [cs.LG] for this version) | |
| https://doi.org/10.48550/arXiv.2610.07247 arXiv-issued DOI via DataCite (pending registration) |
Submission history
From: Heng Liang [view email]
[v1]
Mon, 5 Oct 2026 18:48:00 UTC (505 KB)
来源：arXiv：cs.LG · arxiv.org