基于多目标强化学习的LLM预训练整体数据调度器
Holistic Data Scheduler for LLM Pre-training via Multi-Objective Reinforcement Learning
摘要
训练数据的构成取决于数据来源的多样性以及这些数据的混合方式，这是大型语言模型预训练的关键要素。在线数据混合技术是一种在训练过程中自适应调整数据混合方式的策略，它被认为是一种提高训练效率的有效方法。然而，现有方法都依赖于单一的优化思路，这忽略了复杂的大型语言模型预训练需要从多个维度考虑动态数据构成的必要性。为了克服这一限制，我们提出了整体数据调度器这一新的在线数据混合框架。该框架将数据调度问题转化为连续控制空间中的强化学习问题，并采用软演员-评论家算法来确保系统的稳定性和样本效率，从而能够探索高维策略空间。整体数据调度器的核心是一个多目标、整体性的奖励函数，它整合了三个重要的视角：基于数据的质量奖励、反映不同领域之间影响的损失驱动奖励，以及基于权重范数的模型驱动奖励。为了验证我们的设计并确定其最佳配置，我们对不同规模的大型语言模型进行了系统性的实验。在The Pile基准测试中，整体数据调度器在训练迭代次数减少44%的情况下，达到了次优方法的最终验证困惑度水平。此外，它在MMLU 0-shot任务上实现了7.2%的性能提升，在其他基准测试中也取得了一致的进步，这表明它能够同时提高训练效率和最终模型的性能。
English Abstract
The composition of training data, governed by the diversity of sources and their mixing strategy, is a cornerstone of Large Language Model (LLM) pre-training. Online Data Mixing (ODM), the technique of adaptively adjusting data mixtures during training, has emerged as a promising direction to improve efficiency. However, existing methods are constrained by their reliance on a singular optimization perspective, which fundamentally overlooks the need for complex LLM pre-training to consider the dynamic data composition from multiple dimensions. To overcome this limitation, we introduce the Holistic Data Scheduler (HDS), a novel online data mixing framework. HDS formulates the data scheduling challenge as a reinforcement learning problem in a continuous control space and leverages the Soft Actor-Critic (SAC) algorithm for its stability and sample efficiency in exploring the high-dimensional policy space. At the core of HDS lies a novel multi-objective, holistic reward function that integrates three critical perspectives: a data-driven reward for quality, a loss-driven reward capturing inter-domain influence, and a model-driven reward based on weight norms. To validate our design and determine its optimal configuration, we conducted systematic experiments on LLMs of various sizes. On The Pile benchmark, HDS reaches the final validation perplexity of the next best method with 44% fewer training iterations. Furthermore, it achieves a 7.2% improvement on the MMLU 0-shot task along with consistent gains on other benchmarks, showcasing its ability to enhance both training efficiency and final model capability.