Holistic Data Scheduler for LLM Pre-training via Multi-Objective Reinforcement Learning

📄 arXiv: 2606.24133v1 📥 PDF

作者: Chenhao Dang, Jing Ma, Mingjie Liao

分类: cs.LG, cs.CL

发布日期: 2026-06-23

备注: Our code is at https://github.com/DANG-ai/LLM-Training-Holistic-Data-Schedule

期刊: Proceedings of the 32nd ACM SIGKDD Conference on Knowledge Discovery and Data Mining (KDD 2026), Vol. 1, pp. 176-187, 2026

DOI: 10.1145/3770854.3780325

💡 一句话要点

提出整体数据调度器以解决LLM预训练中的数据混合问题

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 大型语言模型 在线数据混合 强化学习 数据调度 多目标优化

📋 核心要点

  1. 现有方法在LLM预训练中仅依赖单一优化视角，未能充分考虑数据组成的多维动态性，导致效率低下。
  2. 本文提出的整体数据调度器（HDS）通过强化学习框架，动态调整数据混合策略，提升训练效率。
  3. 实验结果表明，HDS在The Pile基准上减少了44%的训练迭代，并在MMLU 0-shot任务上提升了7.2%的性能。

📝 摘要（中文）

训练数据的组成，受来源多样性和混合策略的影响，是大型语言模型（LLM）预训练的基石。在线数据混合（ODM）技术通过在训练过程中自适应调整数据混合，成为提高效率的有前景方向。然而，现有方法受限于单一优化视角，忽视了LLM预训练中多维动态数据组成的复杂性。为克服这一局限，本文提出了整体数据调度器（HDS），将数据调度问题形式化为连续控制空间中的强化学习问题，并利用软演员-评论家（SAC）算法以确保在高维策略空间中的稳定性和样本效率。HDS的核心是一个新颖的多目标整体奖励函数，整合了数据驱动的质量奖励、捕捉领域间影响的损失驱动奖励以及基于权重范数的模型驱动奖励。通过对不同规模的LLM进行系统实验，HDS在The Pile基准上以44%的训练迭代减少达到了最佳验证困惑度，并在MMLU 0-shot任务上实现了7.2%的提升，展现了其提高训练效率和最终模型能力的潜力。

🔬 方法详解

问题定义：本文旨在解决现有LLM预训练方法在数据混合策略上的局限性，特别是其单一优化视角导致的效率低下问题。

核心思路：提出整体数据调度器（HDS），将数据调度视为强化学习问题，利用SAC算法在高维策略空间中实现稳定和高效的探索。

技术框架：HDS框架包括数据调度模块、奖励函数设计和策略优化三个主要部分。数据调度模块负责动态调整数据混合，奖励函数则综合考虑数据质量、领域间影响和模型权重。

关键创新：HDS的创新在于引入了多目标整体奖励函数，结合了数据驱动、损失驱动和模型驱动的奖励机制，显著提升了数据调度的效果。

关键设计：在设计中，HDS采用了SAC算法以确保样本效率，并通过系统实验确定了最优的奖励函数参数设置，确保了模型在不同规模下的有效性。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

HDS在The Pile基准上以44%的训练迭代减少达到了最佳验证困惑度，并在MMLU 0-shot任务上实现了7.2%的性能提升，显示出其在提高训练效率和模型能力方面的显著优势。

🎯 应用场景

该研究的潜在应用领域包括自然语言处理、对话系统和文本生成等，能够显著提升大型语言模型的训练效率和最终性能。未来，HDS框架有望被广泛应用于其他机器学习任务中，推动智能系统的进一步发展。

📄 摘要（原文）

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.