Don’t Ignore the Tail: Decoupled Distillation Produces Top Maths Students on an Academic Budget

Sayantan Dasgupta, Trevor Cohn, Timothy Baldwin
Proceedings of the 43rd International Conference on Machine Learning, PMLR 306:23031-23044, 2026.

Abstract

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 teacher’s highest-probability 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 teacher modes and, consequently, increases the contribution of the distribution’s tail. Experimental results demonstrate that our modified distillation method yields competitive performance in both pre-training and supervised distillation for mathematical reasoning of decoder models across various datasets. Furthermore, the distillation process is efficient and can be performed on modest academic budgets for large datasets, drastically reducing the computational costs typically associated with large-scale distillation.

Cite this Paper

BibTeX
@InProceedings{pmlr-v306-dasgupta26a, title = {Don’t Ignore the Tail: Decoupled Distillation Produces Top Maths Students on an Academic Budget}, author = {Dasgupta, Sayantan and Cohn, Trevor and Baldwin, Timothy}, booktitle = {Proceedings of the 43rd International Conference on Machine Learning}, pages = {23031--23044}, year = {2026}, editor = {Zhang, Tong and Dudik, Miroslav and Jaggi, Martin and Agarwal, Alekh and Li, Sharon and Schuurmans, Dale and Zhu, Jerry and Berkenkamp, Felix and Dong, Hanze and Bietti, Alberto}, volume = {306}, series = {Proceedings of Machine Learning Research}, month = {06--11 Jul}, publisher = {PMLR}, pdf = {https://raw.githubusercontent.com/mlresearch/v306/main/assets/dasgupta26a/dasgupta26a.pdf}, url = {https://proceedings.mlr.press/v306/dasgupta26a.html}, abstract = {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 teacher’s highest-probability 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 teacher modes and, consequently, increases the contribution of the distribution’s tail. Experimental results demonstrate that our modified distillation method yields competitive performance in both pre-training and supervised distillation for mathematical reasoning of decoder models across various datasets. Furthermore, the distillation process is efficient and can be performed on modest academic budgets for large datasets, drastically reducing the computational costs typically associated with large-scale distillation.} }
Endnote
%0 Conference Paper %T Don’t Ignore the Tail: Decoupled Distillation Produces Top Maths Students on an Academic Budget %A Sayantan Dasgupta %A Trevor Cohn %A Timothy Baldwin %B Proceedings of the 43rd International Conference on Machine Learning %C Proceedings of Machine Learning Research %D 2026 %E Tong Zhang %E Miroslav Dudik %E Martin Jaggi %E Alekh Agarwal %E Sharon Li %E Dale Schuurmans %E Jerry Zhu %E Felix Berkenkamp %E Hanze Dong %E Alberto Bietti %F pmlr-v306-dasgupta26a %I PMLR %P 23031--23044 %U https://proceedings.mlr.press/v306/dasgupta26a.html %V 306 %X 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 teacher’s highest-probability 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 teacher modes and, consequently, increases the contribution of the distribution’s tail. Experimental results demonstrate that our modified distillation method yields competitive performance in both pre-training and supervised distillation for mathematical reasoning of decoder models across various datasets. Furthermore, the distillation process is efficient and can be performed on modest academic budgets for large datasets, drastically reducing the computational costs typically associated with large-scale distillation.
APA
Dasgupta, S., Cohn, T. & Baldwin, T.. (2026). Don’t Ignore the Tail: Decoupled Distillation Produces Top Maths Students on an Academic Budget. Proceedings of the 43rd International Conference on Machine Learning, in Proceedings of Machine Learning Research 306:23031-23044 Available from https://proceedings.mlr.press/v306/dasgupta26a.html.

Related Material