Stable training improves on-policy distillation in language models
TV-Regulated OPD: Direction Matters in On-Policy Distillation
Machine Learning
Summary
Training large language models can be unstable because of noisy feedback during learning. The authors found that just knowing whether parts are better or worse, not how much better, helps keep performance strong. They created a new method that smooths this feedback, making training steadier and results more reliable. This approach improves how expert knowledge is passed to smaller models after initial training.
What this means in practice
- •For language model engineers: Improve the reliability of fine-tuning large models by reducing training instability and variance in performance.
- •For ai service providers: Deliver more consistent updates to deployed language models by stabilizing knowledge distillation processes post-training.
Authors
Han Xiao, Yifan Niu, Dongyi Liu, Chang Luo, Jia Li
Abstract
On-Policy Distillation (OPD) facilitates the transfer of knowledge from domain expert to student in the post-training phase of Large Language Models (LLMs). However, the supervision signals in mainstream OPD methods suffer from high variance and noise which is generally instable during training. In this work, we systematically investigated what really matters to the performance and the fundamental mechanisms behind the instability during training. We found that retaining only the sign of token-level advantages is sufficient to achieve the performance comparable to standard OPD. Meanwhile, smoother and bounded advantages can stabilize the training process without sacrificing its performance. These motivated us to shape the advantages using the Total Variation (TV) and propose a robust TV regulated On-Policy Distillation (TV-OPD) method. Benefiting from the bounded and diminished advantages, TV-OPD exhibits stable training dynamics and steady late-stage performance. We conducted comprehensive experiments and found that, across various settings, TV-OPD consistently achieved better performance and lower variance in the late-stage of training.