# Don't Ignore the Tail: Decoupling top-K Probabilities for Efficient Language Model Distillation

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

## Key facts

- Cited by 7 papers in the map (data through 14 Oct 2026).
- Builds on 37 of its 49 references that are in the map.
- First posted on arXiv 24 Feb 2026.

## Builds on (most cited first, 30 of 37)

- [Proximal Policy Optimization Algorithms](https://mlsanity.com/paper/1707.06347/) (2017)
- [Llama 2: Open Foundation and Fine-Tuned Chat Models](https://mlsanity.com/paper/2307.09288/) (2023)
- [Chain-of-Thought Prompting Elicits Reasoning in Large Language Models](https://mlsanity.com/paper/2201.11903/) (2022)
- [Distilling the Knowledge in a Neural Network](https://mlsanity.com/paper/1503.02531/) (2015)
- [Training Verifiers to Solve Math Word Problems](https://mlsanity.com/paper/2110.14168/) (2021)
- [On Calibration of Modern Neural Networks](https://mlsanity.com/paper/1706.04599/) (2017)
- [Qwen2.5 Technical Report](https://mlsanity.com/paper/2412.15115/) (2024)
- [DistilBERT, a distilled version of BERT: smaller, faster, cheaper and lighter](https://mlsanity.com/paper/1910.01108/) (2019)
- [FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness](https://mlsanity.com/paper/2205.14135/) (2022)
- [OPT: Open Pre-trained Transformer Language Models](https://mlsanity.com/paper/2205.01068/) (2022)
- [The Pile: An 800GB Dataset of Diverse Text for Language Modeling](https://mlsanity.com/paper/2101.00027/) (2020)
- [HuggingFace's Transformers: State-of-the-art Natural Language Processing](https://mlsanity.com/paper/1910.03771/) (2019)
- [Gemma 2: Improving Open Language Models at a Practical Size](https://mlsanity.com/paper/2408.00118/) (2024)
- [Phi-3 Technical Report: A Highly Capable Language Model Locally on Your Phone](https://mlsanity.com/paper/2404.14219/) (2024)
- [MiniLM: Deep Self-Attention Distillation for Task-Agnostic Compression of Pre-Trained Transformers](https://mlsanity.com/paper/2002.10957/) (2020)
- [Solving Quantitative Reasoning Problems with Language Models](https://mlsanity.com/paper/2206.14858/) (2022)
- [On-Policy Distillation of Language Models: Learning from Self-Generated Mistakes](https://mlsanity.com/paper/2306.13649/) (2023)
- [Sequence-Level Knowledge Distillation](https://mlsanity.com/paper/1606.07947/) (2016)
- [Galactica: A Large Language Model for Science](https://mlsanity.com/paper/2211.09085/) (2022)
- [WizardMath: Empowering Mathematical Reasoning for Large Language Models via Reinforced Evol-Instruct](https://mlsanity.com/paper/2308.09583/) (2023)
- [TinyLlama: An Open-Source Small Language Model](https://mlsanity.com/paper/2401.02385/) (2024)
- [Llemma: An Open Language Model For Mathematics](https://mlsanity.com/paper/2310.10631/) (2023)
- [Multilingual Translation with Extensible Multilingual Pretraining and Finetuning](https://mlsanity.com/paper/2008.00401/) (2020)
- [Prioritized Training on Points that are Learnable, Worth Learning, and Not Yet Learnt](https://mlsanity.com/paper/2206.07137/) (2022)
- [Stochastic Beams and Where to Find Them: The Gumbel-Top-k Trick for Sampling Sequences Without Replacement](https://mlsanity.com/paper/1903.06059/) (2019)
- [RegMix: Data Mixture as Regression for Language Model Pre-training](https://mlsanity.com/paper/2407.01492/) (2024)
- [OpenWebMath: An Open Dataset of High-Quality Mathematical Web Text](https://mlsanity.com/paper/2310.06786/) (2023)
- [Orca-Math: Unlocking the potential of SLMs in Grade School Math](https://mlsanity.com/paper/2402.14830/) (2024)
- [Compact Language Models via Pruning and Knowledge Distillation](https://mlsanity.com/paper/2407.14679/) (2024)
- [Rho-1: Not All Tokens Are What You Need](https://mlsanity.com/paper/2404.07965/) (2024)

## Built on by (most cited first, 7 of 7)

- [Tail-Aware Top-k On-Policy Distillation](https://mlsanity.com/paper/2608.14728/) (2026)
- [When Top-K Misses the Decision: Tool-Call Drift in Multi-Teacher On-Policy Distillation](https://mlsanity.com/paper/2607.07050/) (2026)
- [When Mean CE Fails: Median CE Can Better Track Language Model Quality](https://mlsanity.com/paper/2605.24667/) (2026)
- [Speculative Decoding and the Curse of Multilinguality](https://mlsanity.com/paper/2605.30580/) (2026)
- [Let the Data Decide: Supervision Analysis, Capability Trade-offs, and Adaptive Objective Routing in Continued Pre-Training via Off-Policy Distillation](https://mlsanity.com/paper/2607.16246/) (2026)
- [The Other Half of the Memory Wall: Serving 35B MoEs from SSD with Trained Routing Prediction](https://mlsanity.com/paper/2609.18063/) (2026)
- [ReTaCo: Residual-Target Control for On-Policy Distillation](https://mlsanity.com/paper/2609.39275/) (2026)

## Abstract (arXiv)

The core learning signal used in language model distillation is the standard Kullback-Leibler (KL) divergence between the student and teacher distributions. Traditional KL divergence tends to be dominated by the next tokens with the highest probabilities, i.e., the teacher's modes, thereby diminishing the influence of less probable yet potentially informative components of the output distribution. We propose a new tail-aware divergence that decouples the contribution of the teacher model's top-K predicted probabilities from that of lower-probability predictions, while maintaining the same computational profile as the KL Divergence. Our decoupled approach reduces the impact of the teacher modes and, consequently, increases the contribution of the tail of the distribution. Experimental results demonstrate that our modified distillation method yields competitive performance in both pre-training and supervised distillation of decoder models across various datasets. Furthermore, the distillation process is efficient and can be performed with a modest academic budget for large datasets, eliminating the need for industry-scale computing.