RegMix-D: Dynamic Data Mixing via Proxy Training Trajectories

📄 arXiv: 2606.18663v1 📥 PDF

作者: Kaiyan Zhao, Zhongtao Miao, Akiko Aizawa, Yoshimasa Tsuruoka

分类: cs.CL

发布日期: 2026-06-17

备注: Work in progress

💡 一句话要点

提出RegMix-D以实现动态数据混合提升语言模型预训练效果

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 动态数据混合 语言模型 预训练 回归模型 损失轨迹 代理效率 下游任务

📋 核心要点

  1. 现有方法如RegMix仅能选择单一静态数据混合，无法适应训练过程中的动态变化。
  2. RegMix-D通过利用代理运行生成的完整损失轨迹，训练回归模型以预测多个训练阶段的最佳数据混合。
  3. 在25B个Pile数据集上进行的实验表明，RegMix-D在13个下游任务中表现优于RegMix和DoReMi，且代理效率更高。

📝 摘要（中文）

数据混合选择对于大型语言模型的预训练至关重要。现有方法如RegMix通过在小规模代理运行上拟合回归模型选择单一静态混合。我们提出RegMix-D，这是对RegMix的简单扩展，支持动态混合。我们的关键观察是，代理运行不仅产生终点损失，还生成完整的损失轨迹，这可以进一步改善数据混合。通过在这些轨迹上训练回归模型，我们可以预测多个训练阶段的最佳混合。RegMix-D支持两种部署模式：离线变体在目标训练之前生成完整的混合计划，在线变体在训练过程中根据观察到的损失动态调整混合。在对25B个Pile数据集的1B参数目标模型进行的实验中，RegMix-D在13个下游任务中始终优于RegMix和DoReMi，同时保持代理效率：即使仅使用128个代理模型（为RegMix代理计算预算的25%），也超越了RegMix。

🔬 方法详解

问题定义：现有的数据混合选择方法如RegMix只能基于小规模代理运行的静态损失进行选择，无法动态适应训练过程中的变化，导致模型性能受限。

核心思路：RegMix-D的核心思想是利用代理运行生成的完整损失轨迹，而不仅仅是终点损失，通过训练回归模型来预测在不同训练阶段的最佳数据混合，从而实现动态调整。

技术框架：RegMix-D包含两个主要模块：离线模式和在线模式。在离线模式中，系统在目标训练之前生成完整的混合计划；而在在线模式中，系统根据实时观察到的损失动态调整数据混合。

关键创新：RegMix-D的创新在于利用完整的损失轨迹进行数据混合选择，这与现有方法的静态选择方式形成了本质区别，能够更好地适应训练过程中的变化。

关键设计：在模型训练中，RegMix-D使用128个代理模型进行回归训练，损失函数设计为能够捕捉到不同训练阶段的动态变化，从而实现更高效的混合选择。具体的网络结构和参数设置在实验中进行了优化，以确保模型的高效性和准确性。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果显示，RegMix-D在13个下游任务中均优于RegMix和DoReMi，且在仅使用128个代理模型的情况下，超越了RegMix，展现出显著的性能提升，证明了其在代理效率上的优势。

🎯 应用场景

RegMix-D的研究成果在大型语言模型的预训练中具有广泛的应用潜力，能够有效提升模型在多种下游任务中的表现。其动态数据混合的能力使得模型能够更好地适应不同的训练阶段，提升训练效率和效果，未来可广泛应用于自然语言处理、机器翻译等领域。

📄 摘要（原文）

Data mixture selection is critical for Large Language Model pretraining. Existing methods such as RegMix select a single static mixture by fitting a regression model on small-scale proxy runs. We propose RegMix-D, a simple extension of RegMix to dynamic mixing. Our key observation is that proxy runs produce not only endpoint losses, but also full loss trajectories, which can be used to further improve data mixture. By training regression model on these trajectories, we can predict optimal mixtures at multiple training stages. RegMix-D supports two deployment modes: an offline variant that generates a complete mixture schedule before target training, and an online variant that adapts the mixture during training using observed loss. Experiments on 25B tokens of the Pile dataset with a 1B parameter target model show that RegMix-D consistently improves over RegMix and DoReMi across 13 downstream tasks while remaining proxy-efficient: it surpasses RegMix even with only 128 proxy models (25% of RegMix's proxy compute budget).