论文略读：RegMix: Data Mixture as Regression for Language Model Pre-training

iclr 2025 688

大型语言模型（LLM）的预训练数据混合比例对模型性能有显著影响，但如何确定一个有效的数据混合策略仍不清楚。为此，我们提出了 RegMix，一种将数据混合选择建模为回归任务、用于自动识别高性能数据混合比例的方法。

RegMix 首先在多种数据混合上训练大量小模型，通过回归模型预测未见数据混合的性能，最后将预测效果最好的混合比例应用于大规模模型的训练，以此实现计算资源的大规模放大利用。

在实证验证中，我们训练了 512 个拥有 100 万参数的小模型，每个模型在 10 亿 tokens 上训练，用于拟合回归模型并预测最佳混合策略。随后，我们使用该预测的最佳混合比例训练了一个 10 亿参数模型，训练数据为 250 亿 tokens（即模型大小扩大 1000 倍，训练数据量扩大 25 倍）。实验证明，该模型在 64 个候选混合策略中表现最佳。

更进一步，RegMix 在最多 70 亿参数、训练量达 1000 亿 tokens 的实验中，始终优于人类选择的数据混合方案，并在仅使用 10% 计算资源的情况下匹敌或超过 DoReMi。

我们的实验还揭示了以下关键发现：

1. 数据混合比例对模型性能具有重大影响；

2. 与常见认知相反，网页数据集相比于高质量数据（如 Wikipedia）对下游性能有更强正相关性；

3. 不同领域之间的交互十分复杂，常识性直觉经常无效，因此需要 RegMix 这样的自动方法；

4. 数据混合效应超越了传统的 scaling laws（扩展规律）。

代码开源地址：https://github.com/sail-sg/regmix。