基于多目标强化学习的大语言模型预训练全息数据调度器
阅读原文· arxiv.org现有大语言模型预训练的在线数据混合方法仅从单一角度优化。新框架全息数据调度器(HDS)将数据调度建模为连续控制空间中的强化学习问题,采用Soft Actor-Critic算法。其核心是多目标综合奖励函数,整合数据驱动质量奖励、损失驱动域间影响奖励和模型驱动权重范数奖励。在The Pile基准上,HDS达到次优方法最终验证困惑度所需训练迭代减少44%,并在MMLU 0-shot任务上提升7.2%。
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.