Training methods compared under a distribution matching view
A Distribution-Matching Perspective on LLM Training Algorithms
Machine LearningArtificial Intelligence
Summary
Training large language models involves teaching a computer to predict words like humans do. This work shows that many teaching methods can be seen as trying to make the model's guesses match a target pattern. The researchers found one method works better when the model is already close to the target but worse when it's far. They also spotted a problem that makes some methods worsen over time and proposed a fix that keeps the model’s variety while still improving. This helps clarify why some existing training methods behave as they do and guides better approaches.
What this means in practice
- •For machine learning engineers: Choose or design language model training algorithms based on how close a current model is to its target distribution to improve training efficiency and stability.
- •For reinforcement learning practitioners: Use conditional KL regularization to reduce long-term degradation of learned policies when the target depends on the current model, enhancing training robustness.
A theory result. No direct application yet.
Authors
Yanxing Huang, Bingzhang Yan, Zihan Tang, Wenzhe Li, Peng Li, Yang Liu
Abstract
In this work, we propose a unified framework for training algorithms across different stages of modern large language model development, based on the perspective of distribution matching. We show that a broad class of language model training algorithms can be interpreted as minimizing a certain divergence between the current model and a target distribution, and that characterizing the properties of this target distribution provides substantial insight into the behavior of the corresponding training algorithms. By analyzing specific forms of the target distribution, we establish that On-Policy Distillation (OPD) enjoys an advantage over classical Knowledge Distillation (KD) when the current model distribution is close to the target distribution, while exhibiting a disadvantage when the two distributions are far apart. Furthermore, when the target distribution itself depends on the model being optimized, we show that the accumulation of estimation variance can induce a deterministic degradation of the target distribution. These results provide a theoretical explanation, at least in part, for the long-term behavior of existing Reinforcement Learning (RL) algorithms. With this intuition we further designed Conditional KL (CKL) regularization for mitigating the degradation of the target distribution. CKL can effectively preserve mode diversity while relaxing the constraint imposed by classical KL regularization on reward improvements. We believe this theoretical perspective can bring some new understanding of language model training algorithms and provide useful guidance for future research.