基于多目标强化学习的大语言模型预训练整体数据调度器 Holistic Data Scheduler for LLM Pre-training via Multi-Objective Reinforcement Learning
使用SAC算法+多目标奖励函数优化LLM预训练数据混合策略
前置知识
在线数据混合
Online Data Mixing (ODM) 是一种在 LLM 预训练过程中动态调整各数据域采样权重的技术。与离线静态方法不同，ODM 在每个训练步骤根据模型当前状态（如损失、梯度等）自适应地修改数据混合比例，从而加速收敛和提升最终性能。传统 ODM 方法通常将此问题建模为多臂老虎机问题，直接使用域损失作为奖励信号。
论文的核心就是改进 ODM 方法，理解 ODM 的工作原理和局限性是理解 HDS 创新点的基础。
Soft Actor-Critic (SAC)
SAC 是一种离策略、最大熵的深度强化学习算法。它由演员网络（策略）和两个评论家网络（价值函数）组成，通过优化最大熵目标来探索更稳定的策略。SAC 的优势在于其样本效率和训练稳定性，特别适合探索高维连续控制空间。论文使用 SAC 来学习动态数据混合策略，其中动作空间是各数据域的采样概率分布。
SAC 是 HDS 框架的算法基础，理解其工作机制（特别是演员-评论家架构和熵正则化）对理解 HDS 如何在连续控制空间中优化数据调度至关重要。
MTLD (Measure of Textual Lexical Diversity)
MTLD 是衡量文本词汇多样性的指标，计算文本中词汇重复模式的平均长度。MTLD 值越高表示文本词汇越丰富、句式越复杂，MTLD 值较低则表示文本词汇重复度高、结构相对简单。论文使用 MTLD 作为文本复杂度的代理指标，将各数据域的 MTLD 归一化到 [0,1]，然后根据训练进度设计奖励函数，引导模型从简单文本逐步过渡到复杂文本。
MTLD 是 HDS 多目标奖励函数中「计划词汇多样性奖励」的核心组件，理解这个指标才能理解 HDS 如何实现「简单到复杂」的课程学习策略。
研究动机
现有在线数据混合方法存在根本性局限：它们都从单一优化角度来调整数据混合比例。早期 ODM 方法将问题建模为多臂老虎机，直接使用每个域的训练损失作为奖励，这种方法过于短视，无法捕捉域间的潜在协同效应。后续的 AC-ODM 虽然引入了基于梯度对齐的奖励来考虑跨域影响，但仍然只关注训练动态这一单一维度。更近期的 AutoScale 和 Aioli 等方法虽然整合了更全面的框架，但它们的优化视角仍然狭窄——要么关注数据质量，要么关注训练损失，要么关注跨域影响——而忽略了需要从多个维度全面考虑复杂 LLM 预训练过程的动态数据组成需求。在 The Pile 数据集的 22 个不同域上训练 LLM 时，每个域的独特知识贡献和它们之间的复杂相互作用被现有方法忽略。
本文的目标是本文提出整体数据调度器（HDS），目标是设计一个能够在多个维度上同时优化数据混合策略的在线调度框架。具体目标包括：同时考虑数据质量、训练动态和模型稳定性；通过多目标强化学习在连续控制空间中优化数据调度；在提升训练效率的同时改善最终模型性能。论文在 Pythia-1B 模型上验证，目标是使用更少的训练迭代达到同等或更好的最终性能，并在下游任务（如 MMLU）上实现准确率提升。
与已有工作不同的是，本文的独特切入角度是将 LLM 预训练的数据调度问题重新定义为连续控制空间中的多目标强化学习问题，而不是像现有方法那样从单一维度优化。核心创新是设计了多目标整体奖励函数，同时整合三个不同视角：基于跨域梯度对齐的数据驱动奖励、基于计划词汇多样性的课程学习奖励、以及基于权重范数变化的模型稳定性奖励。这与之前只关注一个方面的方法形成鲜明对比，例如 AC-ODM 只关注梯度对齐，ODM 只关注域损失。此外，论文使用 Transformer 架构的轻量级演员和评论家网络来有效建模高维状态向量中复杂的相互依赖关系，这也是与之前方法使用简单 MLP 的区别。
核心方法
HDS 将动态数据调度形式化为马尔可夫决策过程（MDP），使用 Soft Actor-Critic 算法在连续控制空间中优化数据混合策略。整体框架包括：状态空间捕获模型当前性能、学习速度和内部稳定性；动作空间是各数据域的采样概率向量；奖励函数综合三个目标的加权组合。在每一步训练中，SAC 代理观察状态、输出动作、采样数据批次、更新 LLM 参数、计算奖励、存储转换到经验回放缓冲区，然后从缓冲区采样小批量更新演员和评论家网络。这种循环使 HDS 能够在整个预训练过程中动态优化数据混合策略。
HDS 的核心创新是多目标整体奖励函数，整合三个不同视角：跨域影响奖励 $r_{align}$ 衡量一个域的训练对所有其他域的积极影响，通过该域梯度与所有其他域梯量的内积计算；计划词汇多样性奖励 $r_{diversity}$ 基于 MTLD 指标引导模型在早期训练阶段学习简单文本（低 MTLD），后期过渡到复杂文本（高 MTLD），实现课程学习；模型稳定性奖励 $r_{stability}$ 基于选定层权重的 L2 范数变化，惩罚剧烈的参数波动以促进平滑收敛。最终奖励为 $\mathbf{r}_t = w_{align}\mathbf{r}_t^{align} + w_{diversity}\mathbf{r}_t^{diversity} + w_{stability}r_t^{stability} \cdot \mathbf{1}$。这与之前只使用域损失或梯度对齐单一信号的方法有本质区别。
方法步骤详情
HDS 算法流程：第一步，初始化演员网络 $\pi_{\theta_A}$、两个评论家网络 $Q_{\theta_{C1}}$ 和 $Q_{\theta_{C2}}$、目标评论家网络、经验回放缓冲区、LLM 模型参数 $\theta_M$。第二步，对于每个训练步骤 $t = 1$ 到 $T$：观察状态 $\mathbf{s}_t$（包括样本数向量、训练步数、域损失、损失变化、权重范数、权重范数变化）；从策略采样动作 $\mathbf{a}_t$；根据 $P_{\mathbf{a}_{t+1}}$ 采样数据批次 $B_t$；更新 LLM 参数 $\theta_{t+1}^M \leftarrow \theta_t^M - \eta_M \sum_{i=1}^K a_{t+1}^i \nabla L(\theta_t^M, B_t^i)$；使用公式 (1)-(5) 计算奖励向量 $\mathbf{r}_t$；观察下一状态 $\mathbf{s}_{t+1}$；将转换 $(\mathbf{s}_t, \mathbf{a}_t, \mathbf{r}_t, \mathbf{s}_{t+1})$ 存入缓冲区。第三步，对于 $n_{ac} = 1$ 到 $N_{ac}$：从缓冲区采样小批量；计算每个转换的目标值；更新评论家网络；更新演员网络；更新熵温度；使用 Polyak 平均更新目标网络。
技术新颖性
HDS 的技术新颖性体现在多个方面：首次将 LLM 预训练数据调度形式化为多目标强化学习问题，而之前的方法要么使用多臂老虎机（离散动作），要么只优化单一目标；设计了新颖的计划词汇多样性奖励，基于 MTLD 指标实现从简单到复杂的课程学习，这是数据混合文献中首次引入文本复杂度的显式课程；使用 Transformer 架构的轻量级演员和评论家网络（8 个编码器块，隐藏维度 512，约 500 万参数），利用自注意力机制建模高维状态向量中复杂的相互依赖关系；多目标奖励函数的权重经过系统超参数优化（$w_{align}=1, w_{diversity}=10, w_{stability}=10$），实现三个目标的平衡。在技术实现上，论文提出了实用的扩展指南：状态表示使用层采样策略（偶数层），对齐奖励使用模型后半部分的层，代理网络规模约为 LLM 参数的 0.5%。
实验结果
HDS 在 Pythia-1B 模型和 The Pile 数据集上展现了显著的训练效率和最终性能提升。训练 500 亿 tokens（41,667 步）后，HDS 的最终验证困惑度为 13.24，比静态 TPW 基线（15.31）降低 13.6%，比 ODM（14.52）降低 8.9%，比 AC-ODM（13.97）降低 5.3%。更重要的是效率提升：HDS 达到 AC-ODM 最终困惑度只需约 18,000 步，比 AC-ODM 的 21,330 步减少 44%；达到 TPW 最终困惑度只需 17,917 步，比 TPW 的 41,667 步减少 57%。在 MMLU 基准上，HDS 的 0-shot 准确率为 0.26915，比 AC-ODM（0.25146）提升 7.2%，5-shot 准确率为 0.31064，比 AC-ODM（0.29868）提升 4.0%。在 COPA、PIQA、WinoGrande 等任务上也持续超越基线。计算开销分析显示，引入 SAC 代理增加了 26.5M 参数（约 2.6%），每步时间从 2.47 秒增加到 2.49 秒（增加 <1%），但实现了 2.21x 的 wall-clock 加速比。扩展实验在 Pythia-12B 模型上验证了可扩展性，在 20,832 步时 HDS 困惑度为 4.89，比 ODM（7.32）低 33% 以上。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| The Pile 验证困惑度 | 验证困惑度（越低越好） | 13.24 | AC-ODM: 13.97 | 5.3% 降低 |
| MMLU 0-shot | 准确率 | 0.26915 | AC-ODM: 0.25146 | 7.2% 提升 |
| MMLU 5-shot | 准确率 | 0.31064 | AC-ODM: 0.29868 | 4.0% 提升 |
| 训练效率 | 达到 TPW 基线困惑度所需步数 | 17,917 步 | TPW: 41,667 步 | 57% 减少 |
| COPA | 0-shot 准确率 | 0.71000 | AC-ODM: 0.69800 | 1.7% 提升 |
| PIQA | 0-shot 准确率 | 0.69840 | AC-ODM: 0.69670 | 0.2% 提升 |
| Pythia-12B 验证困惑度 | 20,832 步验证困惑度 | 4.89 | ODM: 7.32 | 33.2% 降低 |
局限与改进
作者在论文中没有明确讨论局限性，但可以观察到几个潜在限制：首先，HDS 引入了额外的超参数（三个奖励权重、层选择策略、代理网络规模等），虽然论文提供了系统的超参数分析，但这些参数在不同模型规模和数据集上的最优配置可能需要重新调优，增加了实际应用的复杂度。其次，HDS 的计算开销虽然相对较小（每步时间增加 <1%），但仍然需要额外的内存来存储经验回放缓冲区和维护两个评论家网络，在资源受限的环境中可能成为瓶颈。第三，论文的实验主要在 The Pile 数据集上进行，该数据集的域结构相对清晰（22 个域），但在更复杂或域边界模糊的数据集上，基于梯度对齐和 MTLD 的奖励机制可能需要调整。第四，论文没有讨论在更长时间训练（如 Chinchilla 缩放定律建议的更多 tokens）下的表现，课程学习策略是否仍然最优有待验证。最后，虽然论文验证了在 12B 模型上的可扩展性，但对于更大规模的模型（如 70B 或更大），代理网络的架构和状态表示可能需要进一步优化。
独立分析的弱点
HDS 的主要弱点之一是奖励函数中的 MTLD 指标可能无法全面反映文本复杂度。MTLD 主要衡量词汇多样性，但忽略了语义复杂度、逻辑结构、推理难度等其他重要维度。改进方向是引入更全面的复杂度度量，如句子长度、依存句法复杂度、语义密度等多维指标，构建更鲁棒的复杂度表征。第二个弱点是梯度对齐奖励的计算需要存储每个域的梯度，这在数据域数量很多时可能成为内存瓶颈。改进方向是使用梯度投影或低秩近似来减少内存开销，或者使用分布式梯度计算。第三个弱点是 HDS 的训练过程包含两个相互耦合的学习系统（LLM 和 SAC 代理），这使得训练稳定性变得更加复杂。改进方向是引入更鲁棒的训练策略，如自适应学习率调整、梯度裁剪、或预训练代理网络。第四个弱点是论文的实验主要在英文数据集上进行，在其他语言（特别是低资源语言）上的效果未知。改进方向是在多语言数据集上验证和调整 HDS，考虑不同语言特有的复杂度指标。
未来方向
作者提出的未来研究方向包括：将 HDS 扩展到更大规模的模型（如 70B 或更大）并验证性能；探索在其他预训练数据集（如 CommonCrawl、Wikipedia 等）上的适用性；研究如何将 HDS 与其他训练优化技术（如混合专家、序列并行等）结合。基于论文成果可以延伸的方向包括：将多目标奖励框架扩展到更多维度，如域覆盖度、知识互补性、下游任务对齐等；探索将 HDS 应用于微调阶段，根据下游任务动态调整数据混合；研究如何将 HDS 与人类反馈（RLHF）结合，在预训练阶段就考虑对齐目标；开发自适应的超参数调整机制，使 HDS 能够根据训练进度自动调整奖励权重；探索在分布式训练环境中的实现和优化，减少通信开销。
复现评估
论文提供了良好的复现支持。代码已公开在 Zenodo（DOI: 10.5281/zenodo.18123749），使用 The Pile 开源数据集（825GB，22 个域）和基于 GPT-NeoX 的修改版 Pythia 套件。实验硬件使用 Intel Xeon Platinum 8468 CPU 和八张 NVIDIA H800 80GB GPU，全局批大小 1152（每 GPU 微批大小 8，梯度累积 18 步），序列长度 1024。论文详细提供了模型配置（16 层 Transformer，隐藏维度 2048，16 个注意力头）和训练超参数（学习率从 2.5e-5 到 2.5e-4 的线性预热然后余弦衰减，总步数 41,667）。HDS 代理使用 8 个 Transformer 编码器块，隐藏维度 512，约 500 万参数。奖励权重为 $w_{align}=1, w_{diversity}=10, w_{stability}=10$。论文还提供了详细的消融研究、超参数敏感度分析和扩展指南。综合评估复现难度为中等，主要挑战是需要大量计算资源（8 张 H800 GPU）和长训练时间。
论文图表
