RegMix-用回归任务解决大模型数据混合问题

大型语言模型在预训练阶段，怎么选数据、怎么配比例，一直是个让人头疼的问题。数据混合比例的好坏，直接关系到能不能用更少的token、更低的loss来加速预训练进程。手动去调各个数据源的配比，不仅难以扩展，而且大概率不是最优解。

今天我们来聊聊一个挺巧妙的方法——

。它的核心思路，说白了就是把数据混合问题转化成回归任务。具体来说，先按不同的数据混合比例训练一批小模型，拿到结果；然后用这些结果训练一个回归模型；接着遍历所有可能的比例，让回归模型帮忙找到那个最优解；最后再用这个最佳配比去训练真正的大模型。

Paper: https://arxiv.org/abs/2407.01492
Github: https://github.com/sail-sg/regmix

他们用这个方法，训练了512个参数只有1M的小模型，拟合出回归模型后，筛选出top64的混合比例，然后用这些比例去训练一个1B参数的模型。结果很直观：最优配比下训练的模型，验证集loss就是最低的。

方法

整个流程其实不复杂。首先，生成一堆随机的数据混合比例，用这些比例采样数据来训练小模型。然后，把这些混合比例当作特征，模型训练的目标值当作标签，拿去拟合一个回归模型。接下来，在更大的数据混合比例空间中，用这个回归模型去预测哪个比例能拿到最佳目标值。最后一步，就是用这个模拟出来的最佳比例，去训练真正的大模型。

训练小模型当然是多多益善，但考虑到成本，得控制次数。所以初始化比例的时候，既要保证多样性，也要确保每个数据领域都有极端值出现。具体采样时，他们是基于token分布，采用狄利克雷分布来生成混合比例的。

详见：mixture_config/synthesize_mixture.py

至于回归模型的选择，他们试了

和两种，实际效果后面会看到差别。

结果

实验用的是Pile数据集中不涉及版权的17个子集，具体组成如下表。

一个关键发现是：用512个1M小模型（在1B token上训练）拟合出的回归模型，与直接在25B token数据上训练出来的1B模型进行排序对比，相关性高达97.12%。这意味着小模型的排序结果，已经能非常准确地预示大模型的表现趋势。

这里有个细节值得注意：训练

本身，比训练的总token数更重要，对回归模型的效果影响更大。而且，在建模能力上，LightGBM明显优于线性回归。

另一个值得关注的点是，数据混合比例对下游任务的影响相当显著。比如在Lambada数据集上，最好和最差的配比，效果能差出14.6%。

有意思的是，实验结果还挑战了一个传统认知。过去我们总觉得维基百科数据质量高，是评估LLM最具代表性的数据集。但实验发现，

。比如Pile-CC这个数据集，用它做验证集时得出的损失值，与下游任务的相关性要强得多。

最后，RegMix还能帮我们发现不同领域数据之间复杂的交互作用。这些相互作用，光靠人类固有的经验，是很难直接看出来的。