# Learning What to Distill: Bilevel Top-K Token Selection for Self-Distillation in Large Language Models

arXiv: https://arxiv.org/abs/2610.07247 · HTML: https://mlsanity.com/paper/2610.07247/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Placed in Test-time compute and efficient thinking, one of 17 directions in LLM reasoning.
- Builds on 19 of its 37 references that are in the map.
- First posted on arXiv 5 Oct 2026.

## Placed in

- LLM reasoning / [Test-time compute and efficient thinking](https://mlsanity.com/lr/test-time-compute-and-efficient-thinking/)
- Knowledge distillation / [Knowledge distillation for LLMs](https://mlsanity.com/distill/knowledge-distillation-for-llms/)

## Builds on (most cited first, 19 of 19)

- [Distilling the Knowledge in a Neural Network](https://mlsanity.com/paper/1503.02531/) (2015)
- [DeepSeekMath: Pushing the Limits of Mathematical Reasoning in Open Language Models](https://mlsanity.com/paper/2402.03300/) (2024)
- [DeepSeek-R1: Incentivizing Reasoning Capability in LLMs via Reinforcement Learning](https://mlsanity.com/paper/2501.12948/) (2025)
- [Self-Distilled Reasoner: On-Policy Self-Distillation for Large Language Models](https://mlsanity.com/paper/2601.18734/) (2026)
- [Reinforcement Learning via Self-Distillation](https://mlsanity.com/paper/2601.20802/) (2026)
- [OpenThoughts: Data Recipes for Reasoning Models](https://mlsanity.com/paper/2506.04178/) (2025)
- [A Survey on Knowledge Distillation of Large Language Models](https://mlsanity.com/paper/2402.13116/) (2024)
- [Self-Distillation Enables Continual Learning](https://mlsanity.com/paper/2601.19897/) (2026)
- [On-Policy Context Distillation for Language Models](https://mlsanity.com/paper/2602.12275/) (2026)
- [TIP: Token Importance in On-Policy Distillation](https://mlsanity.com/paper/2604.14084/) (2026)
- [Scaling Reasoning Efficiently via Relaxed On-Policy Distillation](https://mlsanity.com/paper/2603.11137/) (2026)
- [Black-Box On-Policy Distillation of Large Language Models](https://mlsanity.com/paper/2511.10643/) (2025)
- [CRISP: Compressed Reasoning via Iterative Self-Policy Distillation](https://mlsanity.com/paper/2603.05433/) (2026)
- [NeuralNDCG: Direct Optimisation of a Ranking Metric via Differentiable Relaxation of Sorting](https://mlsanity.com/paper/2102.07831/) (2021)
- [Large-scale Stochastic Optimization of NDCG Surrogates for Deep Learning with Provable Convergence](https://mlsanity.com/paper/2202.12183/) (2022)
- [SelecTKD: Selective Token-Weighted Knowledge Distillation for LLMs](https://mlsanity.com/paper/2510.24021/) (2025)
- [Rethinking Selective Knowledge Distillation](https://mlsanity.com/paper/2602.01395/) (2026)
- [Understanding Reasoning in LLMs through Strategic Information Allocation under Uncertainty](https://mlsanity.com/paper/2603.15500/) (2026)
- [Algorithmic Foundations of Empirical X-risk Minimization](https://mlsanity.com/paper/2206.00439/) (2022)

## Abstract (arXiv)

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.