3.1 核心假设与整体框架#

如 图 1 所示，左图中，作者提出了一个叫排序不变性的假设：如果数据配比 A 在小模型、少量 Token 的训练效果下优于配比 B，那么在扩大模型参数和增加训练 Token 数量后，配比 A 依然会优于配比 B。在这个前提下，我们无需训练大的代理模型，只需要在极少量 Token 上训练大量 Tiny Models，就足以预测出大型 LLM 的最佳数据配比。

右图中，作者先训练 512 个 1M 参数的模型，拟合出一个预测模型，然后，他们用这个预测模型，对 64 种不同数据配比进行了性能预测，并真实训练了 64 个 1B 模型来验证是否这个假设准确。

其中：

• X 轴代表预测排名，从 0 到 64，越靠左边代表 RegMix 这个方法认为这个配比的效果越好

• Y 轴代表真实的验证集损失，越靠下表示大模型训练出来的真实效果越好

• 蓝色圆点代表 64 个真实的 1B 模型，可以看到明显的左下到右上的趋势（强正相关）即预测排名越靠后的模型，真实的验证集损失约越高

• 位于最左下角的红色五角星，这个是 RegMix 方法预测排名第一的配方，实验也证明了，用这个配方训练出来的 1B 模型，其真实的验证集损失也是 64 个模型中最低的。

3.2 RegMix 方法概览#

RegMix 方法将寻找最佳数据配比的优化问题转化成了一个标注的回归预测任务，包含：

1. 训练小型代理模型 Group，生产多种随机数据配比（采用狄利克雷分布来生成训练配比，一个是要满足和为 1，另一个是要覆盖极端比例情况）

2. 拟合回归模型，数据配比权重作为特征输入（X），将验证集 Loss 作为预测目标（Y），本文测试了线性回归和基于树的 LightGBM 模型。

3. 模拟与预测，随机生成海量的数据配比，输入到训练好的回归模型中，预测出它们的验证集Loss，挑选出 Loss 最低的 Top 配比（实际上 Top 100 取平均以增加鲁棒性）

4. 训练大规模模型，将预测出的最佳配比直接应用到大型 LLM 的预训练中。

3.3 狄利克雷分布#

为了要充分覆盖数据配比空间，尤其是一些极端配比，采用常规的随机方法很难均匀覆盖，所以作者采用了狄利克雷分布来生成多种不同的训练配比。

狄利克雷分布满足以下性质：

1. 天然符合数据配比的数学约束（归一化与非负），从狄利克雷分布中采样的任何一组随机向量 $x = \{ x_1, x_2, ..., x_k \}$ 均满足，每一个维度的值都是非负的（$x_i \geq 0$），且所有维度的值相加永远为 1（$\Sigma_{i=1}^kx_i = 1$）。

2. 其分布由一组浓度参数 $\alpha = (\alpha_1,\alpha_2,...,\alpha_k)$ 决定，采样结果的期望分布与 $\alpha_i$ 的相对大小直接相关，从数学期望上来看，采样得到的第 $i$ 个数据域权重期望值为：$E[x_i]=\frac{\alpha_i}{\Sigma_{j=1}^k \alpha_j}$。这样利用这些参数来注入物理先验，即可以把真实各个Domain的 Token 分布情况作为基础参数来进行控制，确保生产的配比不会违背现实数据的量级（比如用极其稀少的数据去占满80%的数据量预算）。

3. 可以通过缩放参数来灵活控制样本性质，狄利克雷分布不仅能控制期望中心的位置，还能通过缩放参数 $\alpha$ 的整体绝对量级，来控制配比性质：

   • 生成极端配比（稀疏性探索）：当整体缩小参数（论文中乘以 0.1，使得 $\alpha_i < 1$）时，概率密度会向边缘集中，此时生成的配比会极度倾斜，大部分权重会落在某一种数据上，其余接近 0，这样，几乎只会采样到真实数据量大的几个 domain，而一些稀缺的 domain 几乎直接归 0。

   • 生成均匀配比（密集型探索）：当整体放大参数（论文中乘以 0.5，使得 $\alpha_i > 1$）时，概率密度会向中心收缩，此时生成的配比非常均衡，各领域数据雨露均沾。

所以，如果直接使用简单的随机方法：

1. 很可能生成大量非法的配比（如某个 domain 占比巨大，但数据量不足）

2. 根据统计学规律，随机生成数字并求百分比时，这些数字最终的占比会强烈地向平均值。（比如抛硬币，32次全部都是正面的这种极端例子，其概率几乎为0）

3. 无法调控生成的这组数字的宏观状态，没有探索能力。

4.1 评估指标#

论文用了三种评估指标，来评估回归模型能不能预测正确：

• Spearman 秩相关系数：它衡量预测的排序结果与真实排序有多一致，在不考虑并列的情况下，一般如下：

$\rho = 1 - \frac{6\Sigma_id_i^2}{n(n^2-1)}$

其中， $d_i$ 是第 $i$ 个样本预测名次和真实名次的差。

• Pearson 相关系数：它衡量两个连续变量之间的相关程度：

$r = \frac{\sum_i (x_i - \bar{x})(y_i - \bar{y})}{\sqrt{\sum_i (x_i - \bar{x})^2} \sqrt{\sum_i (y_i - \bar{y})^2}}$

• 均方误差（MSE）：衡量预测 Loss 与真实 Loss 是否线形一致，$\text{MSE} = \frac{1}{n} \sum_i (\hat{y}_i - y_i)^2$

以及用了 accuracy 或者 normalized accuracy 来衡量模型在下游任务上的指标：

• 0-shot：不给模型示例，直接让模型回答题目

• 5-shot：在提示词中给模型 5 个示例，再让模型回答题目

4.2 结果分析#

如 表 2 所示，可以看到，在 1M 模型上，LightGBM 的 Spearman 秩相关系数达到了 98.45，在 60M 的模型上，有 98.64，在 1B 模型上，仍高达 97.12，对比之下，线性回归模型的性能整体更弱。

这个实验证明了：

1. 排序不变性是有一定依据的。

2. 且领域相关性应该是非线性的，所以 LightGBM 会比 Linear 模型性能更好。

如 图 3 所示，作者探讨了，是多训练一些 Token 有效，还是多训练一些代理模型更有效。图中绘制了 Spearman 秩相关系数随训练 Token 数量变化的曲线，分别比较了 64、128、512 个代理模型，可以看到：

1. 增加代理模型的数量，比单纯增加 Token 收益更明显；

2. 增加单个代理模型的训练 Token，收益很快就会饱和，大约在 0.25B tokens 左右这个拐点之后就呈现明显的边际递减。

如 表 3 所示，作者讨论了不同的数据配比是否会显著影响下游任务。表中显示，在 64 个 1B 模型上，只修改数据配比单不改模型规模，很多任务的分数差距都非常明显，比如 Lambda 最差是 18.9，最好是33.5，差了 14.6分。这也说明了，数据配方优化会显著改变模型的最终能力。

如 图 4 所示，可以看到，这个热力图横轴是各个下游任务，纵轴上不同验证域的 Loss 与下游表现的相关性，作者这里做了取反操作（越红代表该领域 Loss 下降时，下游任务可能会变得更好），作者原本预期 Wikipedia 会最有代表性，但结果发现 Pile-CC 的相关性最强，甚至对 HellaSwag 几乎相关性为 1.0，这也说明了，高质量数据不一定最有助于下游泛化，覆盖广、贴近真实语言使用场景的网页语料，可能更加重要。

如 表 4 所示，这里对比了 RegMix和不同方法的性能对比：

• Human：The Pile 的人工配比（46.6）

• DoReMi：已有的自动配比 Baseline 方法（48.6）

• Pile-CC Only：只使用 Pile-CC 作为 Baseline（48.5）

• RegMix：本文方法（48.6）

同时，在计算量上：DoReMi: $3.7 \times 10^{19}$ FLOPS，而RegMix：$3.5 \times 10^{18}$ FLOPS，差不多只用了前者 10% 的算力。

此外，这里不仅说明了，使用 RegMix 很强也可以节约算力，也说明了直接使用 Pile-CC 这个 Baseline 就已经能做到很强了。

如 图 5 所示，这里左图中训练语料包含 Pile-CC，而右图中完全不包含，但两者仍然都是最小化 Pile-CC 验证损失，可以看到，即使目标域不出现在训练语料中的情况下，RegMix 的效果仍然优于 Baseline，这样说明其确实学到了一种多目标域的配比规律。

如 图 6 所示，左图是 Pile 数据域之间的关系，右图是代码域之间的关系，颜色越深，说明某个训练域权重增加时，对某个目标域 Loss 的影响越明显。比如，哲学文本（PhilPapers）这个领域，在他们的线性分析中，发现其对很多别的域都有正面作用，这也说明了，配比优化不是哪个域好加多喂哪个域的数据的这种线性关系，而是要考虑全局交互的非线性关系。

如 图 7 所示，这里把 1M 模型的训练日志可视化了，横轴上某个域在数据混合中的权重，纵轴是该领域的对数验证损失，如果存在很简单的 scaling law，这些点回大致落在一条清晰的直线上，但结果是，DM Mathematics 这种 domain看起来还是比较接近 log-log 线性，但 Pile-CC 等很多域则明显更复杂，分布很散。作者认为，单独看某个 domain 自己的权重，不足以预测它的 Loss，必须把整数据配比一起建模。