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

# RegMix-D: 通过代理训练轨迹的动态数据混合
来源：https://arxiv.org/html/2606.18663
Kaiyan Zhao1,2, Zhongtao Miao1, Akiko Aizawa2, Yoshimasa Tsuruoka1 1东京大学,2国立情报学研究所 \{kaiyan1006, miao, tsuruoka\}@logos\.t\.u\-tokyo\.ac\.jp, aizawa@nii\.ac\.jp

###### 摘要

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

RegMix-D: 通过代理训练轨迹的动态数据混合

Kaiyan Zhao1,2, Zhongtao Miao1, Akiko Aizawa2, Yoshimasa Tsuruoka1 1东京大学,2国立情报学研究所\{kaiyan1006, miao, tsuruoka\}@logos\.t\.u\-tokyo\.ac\.jp, aizawa@nii\.ac\.jp

## 1 引言

预训练大型语言模型（LLM）需要从不同领域组装大规模语料库[Belenki et al., 2025 (https://arxiv.org/html/2606.18663#bib.bib1); Chen et al., 2026 (https://arxiv.org/html/2606.18663#bib.bib4)]。这些领域之间的混合比例会显著影响下游性能，使得数据混合选择成为关键的设计选择[Feng et al., 2024 (https://arxiv.org/html/2606.18663#bib.bib7)]。近期研究提出了自动化方法来优化这些比例，包括 DoReMi[Xie et al., 2023 (https://arxiv.org/html/2606.18663#bib.bib29)]，它使用参考模型上的分组 DRO[Oren et al., 2019 (https://arxiv.org/html/2606.18663#bib.bib21)]，以及 RegMix[Liu et al., 2025 (https://arxiv.org/html/2606.18663#bib.bib16)]，它在小型代理模型集合上训练回归模型，以预测损失与混合比例的关系。

虽然这些方法有效，但它们共有一个假设：存在一个单一的最优混合，并且可以在整个预训练过程中使用。然而，多个领域的最优组成可能在训练过程中发生变化[Mo et al., 2025 (https://arxiv.org/html/2606.18663#bib.bib20)]。在这种观点下，静态混合必然会在训练过程中妥协，导致性能次优。

近期一些方法采用动态混合来克服这一局限[Chen et al., 2025 (https://arxiv.org/html/2606.18663#bib.bib3); Wang et al., 2025 (https://arxiv.org/html/2606.18663#bib.bib27); Ma et al., 2025 (https://arxiv.org/html/2606.18663#bib.bib18)]。然而，这些方法在标准训练之上引入了额外开销。参见第2节 (https://arxiv.org/html/2606.18663#S2) 的讨论。

为此，我们提出 RegMix-D，它是 RegMix 的简单扩展，支持动态混合。我们的关键观察是，代理运行期间收集的损失轨迹已经包含动态预测所需的信息：每次代理运行产生完整的损失曲线，而不仅仅是终点。我们训练回归模型，该模型接受当前训练状态：\(t, m, \ell\)，其中 \(t\) 是代理步数，\(m\) 是当前混合，\(\ell\) 是当前损失，来预测下一个采样步的损失。该模型可以两种模式部署（图1 (https://arxiv.org/html/2606.18663#S1.F1)）：一种**离线**模式，在目标训练开始前递归生成混合调度；另一种**在线**模式，在训练期间查询回归模型并原地调整混合。两种模式都不引入额外的优化机制。

参见图注图 1：RegMix-D 概述。我们在代理损失轨迹上训练回归模型 \(f\)（左），然后以两种模式部署 \(f\)（右）：**离线**在目标训练前递归生成完整混合调度；**在线**在目标训练期间使用观察到的损失查询 \(f\)，以原地调整混合。我们在 25B 个 tokens 的 Pile 数据集[Gao et al., 2020 (https://arxiv.org/html/2606.18663#bib.bib8)]上，使用 1B 参数目标模型评估 RegMix-D，遵循标准的 RegMix 设置。我们的贡献如下：

- •我们提出了 RegMix-D，这是 RegMix 的动态扩展，具有离线和在线变体，除现有回归框架外无需额外训练机制。
- •在所有测试的粒度下，RegMix-D 在验证损失上持续优于 RegMix，并在 13 项下游任务上取得优越性能。
- •RegMix-D 数据高效：128 个代理模型（RegMix 计算预算的 25%）足以超越使用 512 个代理模型的 RegMix。

## 2 相关工作

#### 静态数据混合。

传统方法为整个训练选择单一混合。DoReMi[Xie et al., 2023 (https://arxiv.org/html/2606.18663#bib.bib29)] 使用参考模型的分组 DRO 来提高高过剩损失领域的权重。DoGE[Fan et al., 2024 (https://arxiv.org/html/2606.18663#bib.bib6)] 基于梯度层面的泛化贡献重新加权领域。RegMix[Liu et al., 2025 (https://arxiv.org/html/2606.18663#bib.bib16)] 使用不同混合训练小型代理模型，并拟合一个回归模型，该模型根据候选混合预测验证损失，然后通过搜索选择最优混合。另一条工作线路通过加权模型合并而非直接代理训练来选择混合[Li et al., 2026 (https://arxiv.org/html/2606.18663#bib.bib14); Wang et al., 2026 (https://arxiv.org/html/2606.18663#bib.bib26)]。RegMix-D 直接建立在 RegMix 的回归框架之上，将单混合选择泛化为时变调度。

#### 动态数据混合

近期一些方法采用时变混合。Aioli[Chen et al., 2025 (https://arxiv.org/html/2606.18663#bib.bib3)] 在目标训练期间在线估计混合律参数，并通过指数梯度下降更新比例。TiKMiX[Wang et al., 2025 (https://arxiv.org/html/2606.18663#bib.bib27)] 使用目标分布上的组梯度影响重新加权领域。AC-ODM[Ma et al., 2025 (https://arxiv.org/html/2606.18663#bib.bib18)] 在代理模型上训练演员-评论家网络，并将其迁移到目标训练。RegMix-D 在三个方面与它们不同：\(i\) **目标训练期间无额外机制**：现有动态方法在目标训练之上引入了在线优化、梯度影响计算或强化学习查询，而 RegMix-D 将所有优化限制在代理阶段。\(ii\) **重用 RegMix 的信号**：我们从 RegMix 流程中已经产生的代理损失轨迹进行预测，无需新的仪器。\(iii\) **两种部署模式**：单个回归模型同时支持离线和在线自适应。

## 3 RegMix-D

### 3.1 预备知识

RegMix[Liu et al., 2025 (https://arxiv.org/html/2606.18663#bib.bib16)] 通过训练 \(M\) 个小型**代理**模型（每个采样的混合一个）进行固定步数 \(T_p\)，记录最后一步的验证损失，并拟合一个将混合映射到预测验证损失的回归模型，来选择单一静态数据混合。然后通过搜索候选混合并选择预测损失最低的混合来识别最优混合。**目标**模型随后使用该单一混合从头开始训练整个运行。这丢弃了代理运行中的两类信息：\(i\) 每条代理轨迹的中间损失，以及 \(ii\) 最优混合在训练过程中如何演变。

### 3.2 轨迹条件回归

我们在训练过程中以多个步区间记录验证损失，而不仅仅是终点损失。给定所需的切换点数量 \(N\)（在这些点混合将发生变化），我们从代理运行中选择 \(N+1\) 个区间 \(\{t_1, t_2, \ldots, t_{N+1}\}\)。对于混合为 \(m^{(i)}\) 的每次代理运行，这会产生验证损失的轨迹 \(\{\ell_1^{(i)}, \ldots, \ell_{N+1}^{(i)}\}\)。

我们为新的回归模型 \(f_N\) 在连续区间上构建训练对，如图 1 (https://arxiv.org/html/2606.18663#S1.F1) 第 1 部分所示。对于每个代理 \(i\) 和每个相邻对 \((t_\ell, t_{\ell+1})\)，我们形成一个训练示例：

\[
\underbrace{\big(t_\ell,\; m^{(i)},\; \ell_\ell^{(i)}\big)}_{\text{输入}} \;\longrightarrow\; \underbrace{\ell_{\ell+1}^{(i)}}_{\text{目标}}.
\tag{1}
\]

也就是说，\(f_N\) 学习预测下一个步区间的验证损失，条件是 (a) 当前代理步 \(t_\ell\)，(b) 当前混合 \(m^{(i)}\)，以及 (c) 当前观察到的损失 \(\ell_\ell^{(i)}\)。这总共产生 \(M \times N\) 个训练对。关键的是，\(f_N\) 预测的是**局部**转移 \((t_\ell, m, \ell_\ell) \rightarrow \ell_{\ell+1}\)，而不是全局混合到终点的映射。这种局部性正是实现动态调度的关键：在任意训练状态下，我们可以询问 \(f_N\) 哪种混合有望在下一个区间内将损失驱动到最低。注意，\(f_N\) 是特定于给定 \(N\) 的，因为训练对依赖于区间编号。

### 3.3 离线混合调度 (算法1 (https://arxiv.org/html/2606.18663#alg1))

给定在所需 \(N\) 处训练的回归模型 \(f_N\)，我们可以生成具有 \(N\) 个切换点 \(\{s_1, s_2, \ldots, s_N\}\) 的混合调度，对应于用于训练 \(f_N\) 的步区间。调度将目标训练划分为 \(N+1\) 个段。第一个段使用人工策划的 Pile 先验 \(m_0^{\text{human}}\)，反映了训练开始时关于目标语料库没有先验知识的假设，并且每个后续段使用回归预测的混合。

从初始损失 \(\ell_0\) 开始，我们遍历切换点 \(s_1, \ldots, s_N\)。在每个切换点 \(s_j\)，我们搜索混合 \(m_j^*\)，该混合在 Dirichlet 抽样的候选上最小化 \(f_N(s_j, m, \ell_{j-1})\)，记录 \((s_j, m_j^*)\) 作为调度条目，并更新 \(\ell_j \leftarrow f_N(s_j, m_j^*, \ell_{j-1})\) 作为下一次迭代的输入损失。新颖之处在于如何在 \(N\) 个切换点上查询 \(f_N\)：预测是**递归**进行的，每次预测的输出损失作为输入馈入下一次预测。最终的调度 \(\{(s_j, m_j^*)\}_{j=1}^N\) 通过 \(s_j^{\text{target}} = s_j \cdot (T_{\text{target}} / T_p)\) 映射到目标训练，其中 \(T_{\text{target}}\) 是目标模型的训练步数。

### 3.4 在线动态混合 (算法2 (https://arxiv.org/html/2606.18663#alg2))

在离线模式中递归使用预测可能会迭代累积误差。在线变体在每个切换点用**观察到的**目标模型损失替代预测损失，从而在目标训练期间将后续混合决策建立在实际情况基础上。

具体来说，在代理轨迹上训练的同一个回归模型 \(f_N\) 在目标训练开始时被加载为冻结的预测器。在每个切换点 \(S_j = s_j \cdot (T_{\text{target}} / T_p)\)，我们使用在目标模型上测量的验证损失 \(\hat{\ell}_{S_j}\) 查询 \(f_N\)，然后通过与离线变体相同的 Dirichlet 抽样和 top-\(k\) 平均选择下一段的混合。

#### 跨尺度损失校正。

由于 \(f_N\) 是在代理规模损失（代理模型大小 \(P_p\)）上训练的，而目标损失来自更大的模型（大小 \(P_t\)），两者不能直接比较。我们在调用 \(f_N\) 之前应用幂律校正将目标损失映射到代理规模[Kaplan et al., 2020 (https://arxiv.org/html/2606.18663#bib.bib10)]：

\[
\tilde{\ell}_{S_j} = \hat{\ell}_{S_j} \cdot \bigl(P_t / P_p\bigr)^{\beta},
\tag{2}
\]

其中 \(\beta\) 是一个小的正常数。在我们的设置中，\(P_t / P_p = 1000\)，我们使用 \(\beta = 0.05\)，选自 \(\beta \in \{0.01, 0.02, 0.05, 0.10\}\) 的扫参；在此范围内最终下游性能变化小于 \(0.25\)（表3 (https://arxiv.org/html/2606.18663#A1.T3)），表明在线自适应对 \(\beta\) 的精确选择是鲁棒的。因此，在切换点 \(S_j\) 的查询变为 \(f_N(s_j, m, \tilde{\ell}_{S_j})\)。

## 4 实验

表 1：13 项任务的下游性能，在 0-5 shot 设置上平均。所有方法在 25B 个 tokens 上训练 1B 目标模型。RegMix-D 变体使用 \(N=5\) 个切换点。**粗体**和<u>下划线</u>分别表示每行最好/第二好的结果。RegMix-D (128) 使用 128 个代理模型，约为 RegMix 512 代理计算预算的 25%。

### 4.1 实验设置

遵循 Liu 等人 (2025 (https://arxiv.org/html/2606.18663#bib.bib16))，我们使用在 25B Pile[Gao et al., 2020 (https://arxiv.org/html/2606.18663#bib.bib8)] 数据集（17 个领域）上训练的 1B TinyLlama[Zhang et al., 2024 (https://arxiv.org/html/2606.18663#bib.bib31)] 作为目标模型，并使用 1M 参数的 TinyLlama 变体作为主要实验中的代理模型。我们与三个基线进行比较：Human、DoReMi[Xie et al., 2023 (https://arxiv.org/html/2606.18663#bib.bib29)] 和 RegMix[Liu et al., 2025 (https://arxiv.org/html/2606.18663#bib.bib16)]。我们报告了 RegMix-D (128) 和 RegMix-D (512)，以及各自的两种部署变体：**离线**和 **在线**，主要实验中使用 \(N=5\) 个切换点。报告了来自 lm-eval-harness [1](https://github.com/EleutherAI/lm-evaluation-harness)[Gao et al., 2024 (https://arxiv.org/html/2606.18663#bib.bib9)] 的 13 项任务的平均结果。更多细节见附录 A.1 (https://arxiv.org/html/2606.18663#A1.SS1)。

### 4.2 主要结果

表 1 (https://arxiv.org/html/2606.18663#S4.T1) 报告了下游性能，RegMix-D 在两种代理预算下都持续优于基线。值得注意的是，即使仅使用 128 个代理（RegMix 使用的代理计算的 25%），RegMix-D 的两个变体在平均得分上仍超过 RegMix。图 2 (https://arxiv.org/html/2606.18663#A1.F2) 显示了训练过程中 pile-cc 验证损失：所有四个 RegMix-D 变体都实现了持续低于 RegMix 的损失。相应的动态混合与 RegMix 的静态混合有显著差异，我们在附录 A.4 (https://arxiv.org/html/2606.18663#A1.SS4) 的图 3 (https://arxiv.org/html/2606.18663#A1.F3) 中可视化了 pile-cc 权重轨迹。

在线通常优于离线。在两种代理预算下，在线变体都实现了比离线变体更高的平均准确率。我们将其归因于代理与目标动态之间的差距：离线调度完全从代理损失轨迹生成，而在线预测则基于目标模型实际观察到的损失。