RegMix-D：通过代理训练轨迹实现动态数据混合

背景

在大型语言模型（LLM）的预训练阶段，数据混合（Data Mixture）的选择至关重要。不同的数据组合直接影响模型的最终性能。现有的主流方法，如 RegMix，通常采用一种静态策略：通过在少量小规模代理运行（proxy runs）上拟合回归模型，来选定一个固定的最佳数据混合比例。

然而，这种静态方法忽略了训练过程中的动态变化。随着训练的深入，模型对不同数据子集的受益程度可能会发生变化。现有的静态混合方案无法适应这种变化，可能导致训练后期的效率低下或性能瓶颈。

核心内容

本文提出了 RegMix-D，这是对现有 RegMix 方法的一个简单而有效的扩展，旨在实现动态数据混合。

核心洞察：利用完整的损失轨迹

RegMix-D 的关键观察在于：代理运行（proxy runs）不仅产生最终的损失值（endpoint losses），还产生了完整的损失轨迹（full loss trajectories）。

传统的 RegMix 仅利用最终的损失值来训练回归模型，从而预测静态的最佳混合比例。而 RegMix-D 利用这些完整的损失轨迹来训练回归模型。这意味着模型可以学习数据混合比例与训练进度（如步数、epoch）之间的关系，从而预测在多个训练阶段的最优混合比例。

两种部署模式

RegMix-D 支持两种灵活的部署模式，以适应不同的计算资源和训练需求：

1. 离线变体（Offline Variant）： 在目标模型训练之前，先生成一个完整的数据混合调度表（mixture schedule）。该调度表定义了在整个训练过程中，不同阶段应使用的数据混合比例。

2. 在线变体（Online Variant）： 在训练过程中，根据实时观察到的损失值动态调整数据混合比例。这种方法更加灵活，能够实时响应模型在训练过程中的表现。

实验验证

作者在 1B 参数量的目标模型上，使用了 Pile 数据集中的 25B tokens 进行了实验。实验结果展示了 RegMix-D 的显著优势：

• 性能提升：在 13 个下游任务中，RegMix-D 的性能 consistently（一致地）优于 RegMix 和 DoReMi 等基线方法。
• 代理效率：RegMix-D 保持了极高的代理计算效率。即使仅使用 128 个代理模型（仅为 RegMix 所需代理计算预算的 25%），RegMix-D 依然能够超越 RegMix 的性能。

这表明，通过利用更丰富的训练轨迹信息，RegMix-D 能够在更少的计算资源下实现更优的数据混合策略。

关键要点

• 从静态到动态：RegMix-D 将数据混合从单一的静态选择扩展为动态调整，能够适应训练不同阶段的需求。
• 利用完整轨迹：核心创新在于利用代理运行产生的完整损失轨迹，而不仅仅是最终损失值，从而训练出能预测多阶段最优混合比例的回归模型。
• 灵活部署：提供离线（预生成调度表）和在线（实时调整）两种模式，兼顾计划性与灵活性。
• 高效能比：在仅使用 25% 代理计算预算的情况下，RegMix-D 仍能超越原 RegMix 方法，证明了其数据利用效率的提升。
• 广泛适用性：在 1B 参数模型和 Pile 数据集上的实验显示，该方法在多个下游任务中均能带来稳定的性能提升。

意义与影响

RegMix-D 的提出标志着大语言模型预训练数据策略的一个重要进步。它证明了在数据混合优化中，时间维度的信息（即训练轨迹）具有巨大的挖掘潜力。

1. 资源优化：对于计算资源有限的研究团队或企业，RegMix-D 提供了一种在较少代理计算成本下获得更优数据混合方案的路径，降低了预训练的经济门槛。
2. 训练效率：动态混合策略允许模型在训练早期和后期分别关注不同类型的数据，可能加速收敛并提升最终模型的泛化能力。
3. 方法论启示：该工作鼓励研究者重新审视已有的代理实验数据。许多现有的静态混合方法可能无意中丢弃了大量有价值的动态信息，未来研究可进一步探索如何利用这些轨迹信息优化其他预训练超参数或架构设计。

总之，RegMix-D 通过简单而深刻的洞察，将数据混合从“一次设定，全程使用”转变为“动态调整，全程优化”，为高效、高性能的大模型预训练提供了新的工具和方法论支持。