Training a smaller AI model to mimic a larger one usually requires comparing their predictions over the entire vocabulary at every step, which is computationally expensive. This process, called distillation, is typically slow and memory-intensive.
The researchers tested whether using only the top 0.01% of the teacher model's vocabulary per step, its most likely next-word predictions, could maintain quality. They ranked the teacher's token probabilities, kept only the highest ones, and renormalized them. This was tested on several model pairs of different sizes and families.
Using just 0.01% of the vocabulary caused no drop in accuracy on the MMLU-Pro benchmark. This reduced training time by 38% and peak memory by 24% by caching only those few scores. The approach works, but the safe cutoff depends on the model family.