Tail-Aware Top-k On-Policy Distillation
Not on the map. Why?
Abstract (arXiv)
On-policy distillation (OPD) has emerged as an effective paradigm for transferring knowledge between language models, where a student is trained to align its next-token distribution with the teacher's along its own trajectories. To provide dense supervision at tractable cost, many works minimize the reverse Kullback-Leibler (KL) divergence between the student and teacher's normalized distributions over the teacher's top-k tokens. However, this normalized objective discards the information about tail probability: the total probability outside the teacher's top-k tokens. As a result, the optimization can steadily increase the student's tail probability and entropy, empirically degrading downstream accuracy. To address this issue, we propose Tail-Aware Top-k OPD (TA-OPD), a novel distillation method that restores the missing tail probability signal. In particular, TA-OPD minimizes the reverse KL divergence over the top-k tokens plus a tail token that carries the tail probability. In effect, TA-OPD better aligns the student's next-token distribution with the teacher's, preventing the increase in tail probability and entropy caused by top-k normalization. Extensive experiments demonstrate the superiority of TA-OPD, improving Avg@8 by up to 8.05 points on common benchmarks. Our code is available at https://github.com/HuipengHuang/TA-OPD.
Key facts. Cited by 3 papers in the map (data through 7 Oct 2026). Builds on 41 of its 45 references that are in the map. First posted on arXiv 12 Aug 2026.
Builds on
41 of 45 references.
References · 41
- 2022Training language models to follow instructions with human feedbackcited by 15,712
- 2017Proximal Policy Optimization Algorithmscited by 14,441
- 2024The Llama 3 Herd of Modelscited by 14,192
- 2015Distilling the Knowledge in a Neural Networkcited by 10,076
- 2025Qwen3 Technical Reportcited by 8,195
- 2023Efficient Memory Management for Large Language Model Serving with PagedAttentioncited by 5,767
- 2021Measuring Mathematical Problem Solving With the MATH Datasetcited by 5,203
- 2018Think you have Solved Question Answering? Try ARC, the AI2 Reasoning Challengecited by 4,710
- 2025DAPO: An Open-Source LLM Reinforcement Learning System at Scalecited by 2,678
- 2015Trust Region Policy Optimizationcited by 2,503
- 2019Fine-Tuning Language Models from Human Preferencescited by 2,124
- 2024HybridFlow: A Flexible and Efficient RLHF Frameworkcited by 1,980
29 more
- 2024MMLU-Pro: A More Robust and Challenging Multi-Task Language Understanding Benchmarkcited by 1,758
- 2022Solving Quantitative Reasoning Problems with Language Modelscited by 1,359
- 2024OlympiadBench: A Challenging Benchmark for Promoting AGI with Olympiad-Level Bilingual Multimodal Scientific Problemscited by 1,099
- 2026DeepSeek-V4: Towards Highly Efficient Million-Token Context Intelligencecited by 877
- 2023On-Policy Distillation of Language Models: Learning from Self-Generated Mistakescited by 840
- 2015Scheduled Sampling for Sequence Prediction with Recurrent Neural Networkscited by 779
- 2016Sequence-Level Knowledge Distillationcited by 761
- 2025The Entropy Mechanism of Reinforcement Learning for Reasoning Language Modelscited by 463
- 2015Policy Distillationcited by 414
- 2025Learning to Reason under Off-Policy Guidancecited by 266
- 2026Rethinking On-Policy Distillation of Large Language Models: Phenomenology, Mechanism, and Recipecited by 262
- 2026MiMo-V2-Flash Technical Reportcited by 163
- 2026Entropy-Aware On-Policy Distillation of Language Modelscited by 151
- 2026A Survey of On-Policy Distillation for Large Language Modelscited by 148
- 2026Revisiting On-Policy Distillation: Empirical Failure Modes and Simple Fixescited by 122
- 2022Why Exposure Bias Matters: An Imitation Learning Perspective of Error Accumulation in Language Generationcited by 106
- 2024Learning Dynamics of LLM Finetuningcited by 93
- 2024Speculative Knowledge Distillation: Bridging the Teacher-Student Gap Through Interleaved Samplingcited by 85
- 2025JustRL: Scaling a 1.5B LLM with a Simple RL Recipecited by 57
- 2026The Many Faces of On-Policy Distillation: Pitfalls, Mechanisms, and Fixescited by 37
- 2025KDRL: Post-Training Reasoning LLMs via Unified Knowledge Distillation and Reinforcement Learningcited by 35
- 2026Rethinking the Trust Region in LLM Reinforcement Learningcited by 32
- 2026Trust Region On-Policy Distillationcited by 26
- 2026KL for a KL: On-Policy Distillation with Control Variate Baselinecited by 24
- 2026On the Position Bias of On-Policy Distillationcited by 23
- 2025A Comedy of Estimators: On KL Regularization in RL Training of LLMscited by 18
- 2026Don't Ignore the Tail: Decoupling top-K Probabilities for Efficient Language Model Distillationcited by 7
- 2026Characterizing, Evaluating, and Optimizing Complex Reasoningcited by 5
- 2026EMA Policy Gradient: Taming Reinforcement Learning for LLMs with EMA Anchor and Top-k KLcited by 4
Built on by
Citing papers per year: 2026: 3.
3 papers.