Improving small AI models by selectively learning from larger ones
Recursive Offline Distillation via Entropy Gain Sampling and Answer Mixing
Machine Learning
Summary
Smaller AI language models often learn by copying reasoning steps from larger models, but this can be hard without full access to the larger model’s details. The authors propose a new way to teach small models better by picking the most useful examples based on uncertainty and mixing in correct answers to help guide learning. Their method works better than simply copying everything and learns faster. They also studied different ways to create and adjust these learning traces for better training.
What this means in practice
- •For natural language processing developers: Train smaller language models more efficiently by selectively using reasoning steps from larger models guided by uncertainty and mixing in correct answers.
- •For ai application engineers: Build compact AI systems that improve reasoning accuracy without needing full access to large model internals using this selective distillation approach.
Authors
Andrey Goncharov, Daniil Vyazhev, Simon Epanov, Alexey Zaytsev
Abstract
Distillation of reasoning chains from Large Language Models (LLMs) into their smaller siblings is a standard training approach. Frequently, researchers do not have full access to a teacher model and its internal representations, which limits the resulting performance. We propose a recursive offline distillation method combining two components: entropy-based sampling for selective distillation and answer mixing. Specifically, for three open models (3-4B), at each epoch, we sample reasoning chains weighted by entropy difference with an auxiliary medium model (70-72B) and, next, supplement thinking traces with single-token ground truth answers. Our approach consistently outperforms naive distillation with comparable data budgets ($57.15 \pm 1.01$ vs 51.61 on average) and offers a much steeper learning curve. Additionally, we present a comparative analysis of methods to generate and normalize distillation traces.