用强化学习动态调度训练数据，让大模型预训练更高效、效果更好。
Holistic Data Scheduler for LLM Pre-training via Multi-Objective Reinforcement Learning
- 设计多目标奖励机制，从数据质量、损失变化、权重规律三方面优化数据混合策略。
- 在The Pile上减少44%训练迭代次数，达到当前最优模型的验证困惑度。
- 适用于追求训练效率与模型性能的开发者，尤其适合大规模语言模型训练。
训练数据的构成，由数据源的多样性及其混合策略决定，是大语言模型（LLM）预训练的核心。在线数据混合（ODM）通过在训练过程中自适应调整数据组合，成为提升训练效率的有前景方向。然而，现有方法受限于单一优化视角，未能充分考虑复杂预训练中多维度动态数据组合的需求。为此，我们提出整体数据调度器（HDS），将数据调度建模为连续控制空间中的强化学习问题，并采用Soft Actor-Critic（SAC）算法以保证在高维策略空间中的稳定性和样本效率。HDS的核心是一个新型多目标、全局性奖励函数，融合了数据驱动的质量奖励、损失驱动的跨域影响奖励以及基于权重范数的模型驱动奖励。我们在不同规模的LLM上进行了系统实验。在The Pile基准上，HDS以44%更少的训练迭代次数达到次优方法的最终验证困惑度；同时，在MMLU零样本任务上提升7.2%，并在其他基准上保持一致优势，展现出显著提升训练效率和最终模型能力的能力。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。