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

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.

Large language models are powerful but expensive to run, too big for a phone, a laptop, or the compute budget most university labs work with. A common fix is distillation: you take a large, capable model (the "teacher") and train a much smaller one (the "student") to imitate it, so the student retains much of the teacher's ability at a fraction of the cost. The catch is in how the imitation is scored. The standard method rewards the student mainly for matching the teacher's most confident guesses: the handful of words it thinks are most likely to come next. But a model's prediction is really a long ranked list of possibilities, and a surprising amount of useful information sits in the "tail": the many less-likely words the teacher still considers. By fixating on the top guesses, ordinary distillation throws away much of that. Our method, Tail-Aware Distillation (TAD), rebalances this by separating out the teacher's top predictions from the rest and deliberately gives the tail a stronger voice, so the student learns the full shape of the teacher's reasoning rather than just its loudest opinions, all at no extra computational cost. The payoff is that strong models can be trained on a shoestring. Using a single GPU for about a week, we distilled small models that match or beat much more expensive recent methods, with especially large gains in mathematical reasoning: a roughly 1-billion-parameter model became a genuinely capable maths solver, and a 7-billion one rivaled the best systems built on the same base, at a fraction of their training cost. The broader aim is to lower the barrier to entry, so that producing a good small model doesn't require a tech giant's resources, making it more practical for academic groups, smaller labs, and researchers in resource-constrained settings to build capable models, and easier to run them on everyday devices.