论文
通过多目标强化学习的 LLM 预训练 的整体数据调度器
Holistic Data Scheduler for LLM Pre-training via Multi-Objective Reinforcement Learning
摘要
训练数据的组成由来源的多样性及其混合策略决定,是 大语言模型 (LLM) 预训练 的基石。在线数据混合(ODM)是一种在训练过程中自适应调整数据混合的技术,已成为提高效率的一个有前途的方向。然而,现有方法受到对单一优化视角的依赖的限制,这从根本上忽略了复杂的 LLM 预训练 从多个维度考虑动态数据组合的需要。为了克服这个限制,我们引入了整体数据调度器(HDS),一种新颖的在线数据混合框架。 HDS 将数据调度挑战表述为连续控制空间中的强化学习问题,并利用 Soft Actor-Critic (SAC) 算法在探索高维策略空间时的稳定性和样本效率。 HDS 的核心在于一种新颖的多目标整体奖励函数,它集成了三个关键视角:数据驱动的质量奖励、捕获域间影响的损失驱动的奖励以及基于权重规范的模型驱动的奖励。为了验证我们的设计并确定其最佳配置,我们对各种尺寸的LLM进行了系统实验。在 The Pile 基准测试中,HDS 达到了次佳方法的最终验证困惑度,训练迭代次数减少了 44%。此外,它在 MMLU 0-shot 任务上实现了 7.2% 的改进,并且在其他基准测试中取得了一致的收益,展示了其提高训练效率和最终模型能力的能力。