Reverse-KL distillation shows different data needs for online versus offline training
Reverse-KL Policy Distillation: Statistical Guarantees and Limits
Machine Learning
Summary
Training AI models to copy smarter models can be done in different ways, and some ways need more examples than others to work well. This paper explains when it helps to see the full range of a teacher model’s choices versus just the choice made. It also shows that learning directly from the AI’s own actions requires more careful methods than learning from examples chosen by the teacher. The authors provide new mathematical limits and algorithms to get the best learning results in these settings.
What this means in practice
- •For machine learning engineers: Design training methods to reduce data needs when teaching language models by carefully choosing online versus offline distillation approaches.
- •For reinforcement learning developers: Improve policy learning efficiency by applying importance weighting and trajectory-loss techniques informed by the paper’s statistical guarantees.
A theory result. No direct application yet.
Authors
Haolin Liu, Chen-Yu Wei, Julian Zimmert
Abstract
Reverse Kullback-Leibler (KL) distillation has shown empirical promise for training language models, but its statistical guarantees and limits remain poorly understood. In particular, the effects of two design choices remain unclear: Does access to the teacher's full action distribution provide an advantage over observing only the probability it assigns to the sampled action? Must training trajectories come from the student (online distillation), or can teacher-generated trajectories (offline distillation) suffice? Our work takes a first step toward understanding the theory underlying these questions. We study the sample complexity of online reverse-KL distillation under two feedback models: the teacher reveals either its full action distribution or only the probability of the student-generated action. We establish minimax bounds for the former and show that substantially more data may be required for the latter. In contrast, offline reverse-KL distillation can attain the same minimax rate whether the teacher reveals its full distribution or only its probability of the teacher-generated action. The key distinction is that teacher-generated actions themselves reveal the teacher's preferences, while queries at student-selected actions can leave those preferences unobserved. This highlights an additional statistical challenge for online distillation compared to offline distillation. These statistical guarantees do not follow simply from optimizing standard distillation objectives: our lower bounds show that common online objectives can fail when candidates are evaluated only on past student-generated trajectories. We develop trajectory-loss and importance-ratio truncation techniques together with policy-space learning to achieve optimal online distillation guarantees.