REGMIX: 作为语言模型预训练的回归数据配比 原创

一、结论写在前面

论文标题：RegMix: Data Mixture as Regression for Language Model Pre-training

论文链接：​​https://arxiv.org/pdf/2407.01492​​

代码：​​https://github.com/sail-sg/regmix​​

大型语言模型预训练中的数据配比对性能有显著影响，但如何确定有效的配比仍不明确。论文提出REGMIX，通过将数据配比问题形式化为回归任务，自动识别高性能的数据配比。REGMIX包括使用多种数据配比训练一组小型模型，并拟合一个回归模型来预测给定各自配比的模型的性能。利用拟合的回归模型，论文模拟排名最高的配比，并使用它来训练一个大规模模型，其计算量是之前的几个数量级。

为了实证验证REGMIX，论文训练了512个具有1M参数的模型，用于1B个不同配比的token，以拟合回归模型并找到最佳配比。使用这种配比，论文训练了一个1B参数的模型，用于25B个token（即比之前大1000倍，长25倍），论文发现它在64个候选的1B参数模型中表现最佳，这些模型使用了其他配比。

为了验证REGMIX，论文使用不同的数据配比训练了1M和1B参数的模型。通过在1Btoken上训练512个1M参数的模型，论文发现REGMIX优化的数据配比在64个1B参数模型（比1M参数大1000倍）和训练时间长25倍（25Btoken）的模型中表现更优，如图1所示。此外，使用REGMIX优化的数据配比产生的模型性能优于人工选择，并且与旗舰DoReMi方法[64]的性能相当，尽管它需要的总计算量更少，并允许并行训练。论文还发现：(1) 数据配比显著影响下游性能，导致单任务性能差异高达14.6%；(2) 通用网络语料库（如CommonCrawl），而非维基百科，与下游任务性能提升的正相关性最强；(3) 领域间的交互复杂且常常与直觉相悖，突显了像REGMIX这样的自动化方法的必要性；(4) 数据配比效应超越了缩放定律，REGMIX通过综合考虑所有领域来捕捉其复杂性。    

二、论文的简单介绍

2.1 论文的背景

大规模公开数据集的可用性已成为创建大型语言模型（LLMs）的关键因素。大多数数据可在互联网上获取，包括学术论文（例如arXiv）、书籍（例如古腾堡计划）和代码配比(例如GitHub)。在创建首批LLMs之一，GPT-3 时，作者已经认识到选择最佳训练数据的重要性，因此他们决定增加维基百科的采样比例，因其被认为具有高质量。然而，这种手动数据选择不具备可扩展性，可能导致次优选择。随着用于LLM预训练的数据规模和多样性不断增长，确定最优数据配比变得愈发具有挑战性。这引发了一个关键的研究问题：论文如何以可扩展且高效的方式选择最优数据配比？

先前的工作[64, 16, 2]采用小规模模型（“代理模型”）来预测大规模语言模型的领域权重。这些工作通过训练具有大量token(例如，1000)，有时甚至与用于训练LLM的数量相同，并通过监控训练动态来动态调整分配策略。然而，随着用于预训练LLM的训练数据不断增长，这些方法变得效率低下。为当前模型（如Llama-3）训练代理模型，需要使用多达15T token，这在当前方法下可能过于昂贵且过于缓慢，不值得。    

图1：左：论文假设数据配比在不同模型大小和训练token数量下的秩不变性。利用这一假设，论文使用在较少token上训练的小模型来预测用于训练具有更多token的大型模型的有效数据配比。右：通过训练512 x 1M模型，论文的方法在训练64 x 1B模型之前识别出最佳数据配比。预测的最佳数据配比，用红色星号表示，实现了最低的验证损失。

论文认为在有限数量的token上训练小型模型足以预测用于大型语言模型（LLM）训练的有效数据配比。论文的关键假设是数据配比的秩不变性，即在不同模型大小和训练token数量下，数据配比对模型性能影响的相对排名是一致的。在这一假设下，关键挑战在于从近乎无限数量的潜在数据配比中发掘出排名最高的数据配比。为此，论文将数据配比选择视为回归任务。论文不是通过每一种可能的配比来穷尽地训练小型模型，而是仅训练一组小型模型，每个模型具有独特的数据配比。基于这些模型的性能及其配比，论文拟合一个回归模型来预测其他数据配比的性能。论文的方法比先前的工作更具可扩展性，因为它允许并行训练小型代理模型，而不是长时间训练单个模型。此外，回归模型提供了关于领域交互的洞察，有助于理解和数据整理。

图2：使用Hacker News、GitHub和Philpapers作为训练域，以StackExchange域上的损失作为目标（，展示论文的方法。通过使用小规模代理模型训练日志拟合回归模型，并用于预测模拟空间内最佳数据配比，从而直接预测大规模语言模型预训练的数据配比

2.2 REGMIX：数据配比作为回归

如图2所示，论文的方法涉及四个关键步骤：（1）生成随机数据配比并在此基础上训练小规模代理模型。（2）使用表1：Pile数据集概览，其中因版权问题不再可用的数据集以灰色标记。在论文的实验中，论文使用可获得的17个领域来研究语言模型预训练的数据配比作为特征，目标值作为标签，拟合一个线性回归模型。（3）在更大规模上模拟数据配比空间，并利用回归模型识别目标值的最佳配比。（4）使用模拟的最佳数据配比训练一个大规模模型。    

2.2.1 训练小规模代理模型

第一步是在多个不同的数据配比上训练一组小规模代理模型。为了减少所需的运行次数，论文的目标是选择一系列涵盖每个领域从0%到100%极端权重的多样化数据配比。论文通过基于词频分布的Dirichlet分布来实现这一点，这使论文能够采样广泛的范围值，并将回归模型暴露于各种极端情况。同时，基于词频分布确保了整体数据配比在统计上反映了数据的可用性。例如，这可以防止任何词频低于1%的领域被过度强调，这在大型训练中是不可行的，因为该领域没有足够的可用词频。在实践中，论文将词频分布乘以0.1到5.0之间的值，以构建各种稀疏和近似均匀的分布，然后将这些分布向量作为Dirichlet分布的超参数alpha。

在训练小规模代理模型几步之后，论文可以获得几个训练良好的小模型。例如，在论文的主要实验中，每个代理模型包含1M参数，并在1B词频上进行训练。然后，论文可以选择在领域或基准上评估这些训练好的模型，以获得论文想要优化的目标值。通常，目标值可以是某个领域的损失，如图2所示的StackExchange领域。一旦论文获得了这些目标值，论文就可以使用数据配比作为特征，目标值作为标签来拟合一个回归模型。

2.2.2 拟合回归模型

第二步是使用数据配比作为特征，目标值作为标签来拟合一个回归模型。回归任务是一个传统的监督学习任务，涉及基于输入特征X=( x_1, x_2, ..., x_n )预测连续目标变量y。目标是找到一个函数f，该函数最好地将输入特征映射到目标变量，使得y=f(X)+epsilon，其中epsilon表示数据中的误差或噪声。输入特征X对应于数据配比的领域权重，而目标变量y是论文想要优化的值。使用这些数据，论文训练回归模型，学习一个函数来预测基于任意数据配比的目标值，而无需进一步训练。    

线性回归。线性回归模型在回归分析中应用广泛。它假设输入特征与目标变量之间存在线性关系，可以表示为：

其中，omega_0 是截距，omega=(omega_1, ..., omega_n) 是与相应输入特征 x_1, ..., x_n 相关的系数。系数 omega 通常使用普通最小二乘法等技术进行估计，旨在最小化预测值与实际值之间的残差平方和。

表2展示了论文基于512x 1M模型在1Btoken上训练的结果来拟合回归模型，并在未见过的数据配比上评估1M、60M和1B参数模型的表现。皮尔逊相关系数 r 和均方误差（MSE）衡量损失预测性能，而斯皮尔曼相关系数 rho 比较预测值与实际值的排名。

LightGBM回归。LightGBM 是一种强大的梯度提升算法，可用于回归和分类任务。在回归的背景下，LightGBM 学习一组决策树来预测目标变量。该过程由基于梯度的优化算法指导，该算法最小化指定的损失函数（例如均方误差)。此外，LightGBM 被设计为高效且可扩展，适合大型数据集。

2.2.3 模拟与预测

一旦论文训练了回归模型，论文就可以高效地探索整个可能的数据配比空间。通过使用训练好的模型预测每个潜在数据配比的目标值，论文可以快速识别产生最佳目标值的输入。这种基于模拟的优化相对廉价，因为模拟和回归预测在计算上都很快。例如，对1,000,000个数据配比进行预测仅需不到10个CPU秒。    

2.2.4 大规模模型训练

在通过模拟识别出最佳数据配比后，论文将排名靠前的数据配比推广到包含更多token的大规模模型训练中。如图 2所示，论文直接使用最佳数据配比来训练更大的模型。实际上，为了提高回归预测的稳健性，论文选择前100个配比并将其平均作为大规模训练的数据配比。

2.3 回归预测评估

论文评估REGMIX预测未见数据配比效果的能力。首先，论文使用小型（即1M参数）模型的训练工件拟合回归模型，并评估小型模型上的损失预测性能。然后，为了验证论文的排名不变性假设，论文测试学习到的回归模型在不同模型大小和token数量上预测排名的能力。

2.3.1 实验设置

数据集和模型。论文使用表1中描绘的Pile数据集的domain进行实验。由于版权问题，论文使用HuggingFace上不违反版权问题的17个子集。论文考虑线性回归和LightGBM回归模型，其中目标变量y设置为Pile-CC领域的验证损失。

训练与评估。该回归模型利用512个512x 1M模型的训练成果，基于10亿个token进行拟合，并在256个未见过的数据配比体上对1M、60M模型（各自以10亿个token训练）以及64个未见过的数据配比体上对1B模型（各自以250亿个token训练）进行评估。

评估指标。论文采用三种不同指标来基准测试论文的回归模型：(1) 斯皮尔曼等级相关系数（ρ）是一种非参数度量方法，用于评估两个排序变量之间关联的强度和方向。(2) 皮尔逊相关系数（I）则是衡量两个变量间线性关系的相关性。（3）均方误差（MSE）是一种常用指标，用于通过测量预测值与实际值之间的平均平方差来评估回归模型。

图 3：Spearman 秩相关系数 rho 在不同训练token数量和不同代理模型数量下，线性回归（左）和 LightGBM 回归（右）预测秩与真实秩之间的相关性图。如图所示，增加代理模型数量显著提升 rho，而增加更多训练token则收益递减。          

2.3.2 实验结果

模型大小间的强相关性。如表 2 所示，LightGBM 模型在所有三个指标上均优于线性回归模型，并且在对具有更多训练token的大型模型进行评估时，其优势变得更加明显。同时，使用 1B token训练的 1M 模型在未见过的 1B 模型与 25B token配比上达到 97.12% 的高相关性，直接验证了论文的秩不变性假设。

代理模型数量优于训练token数量。在相同的小规模训练 FLOPs 预算下，论文可以增加token数量（即训练token数量）或代理模型数量。因此，论文研究哪种方法会产生更好的性能。如图 3 所示，增加代理模型的训练token在约 0.25B token后饱和。相比之下，增加代理模型数量持续提升性能，尤其是对于 LightGBM 模型。值得注意的是，在 0.2B token上训练的 512 个模型的性能超过了在 0.8B token上训练的 128 个模型，表明增加代理模型数量比在超过某个token阈值后增加训练token数量更为有效。

2.4 在下游任务上的评估

论文将方法应用于现实下游任务，以展示其有效性。在评估时，论文排除了根据先前工作 [36] 和论文的观察报告的性能跟踪中表现出较大性能差异 (例如，RIL) 的特定基准。最终，论文选择了以下基准作为下游任务：Social IQA 、HellaSwag 、PiQA 、OpenBookQA 、Lambada 、SciQ 、ARC Easy 、COPA 、RACE、LogiQA 、QQP、WinoGrande 和 MultiRC。这些基准涵盖了多样化的任务，能够全面评估 REGMIX 在现实世界中的影响。对于每个基准，如果 lm-eval-anrnss 提供归一化准确率，则使用该指标，否则使用常规准确率。

2.4.1 数据配比显著影响下游性能

最初，论文训练了64个模型，每个模型具有10亿参数，使用不同的数据配比。每个模型都在Pile数据集[18]的250亿个词元上进行训练，词元的分配基于其相应领域的权重。表3展示了每个下游任务中最差和最佳模型的性能。报告的性能是从零样本到五样本评估的平均值，使用lm-eval-harness评估框架[19, 4]进行评分。论文发现数据配比显著影响下游性能，其中最大的性能\triangle在Lambada任务上达到14.6。这突显了研究最优数据配比的重要性。    

表3：论文实验了64个模型，每个模型具有10亿参数，训练于不同的数据配比，并在各种基准上评估它们的性能。每个任务的报告性能是从零样本到五样本设置的平均分数，遵循Muennighoff等人[41]。

图4：64个10亿参数模型在各领域验证损失与下游性能的相关性。注意，在计算相关性时，论文取损失值的负值，因为这使得可视化更直观。这同样适用于图6。    

2.4.2 网络语料库对下游任务性能的提升最为显著

论文在图4（a）中展示了论文64个1B模型在不同领域的验证损失与其在各种下游任务上的表现之间的相关性。在可视化之前，论文假设维基百科（en）子集上的验证损失会与大多数下游任务显示出强烈的关联，因为它是一个高质量的数据集，并且许多下游任务源自维基百科文本。同样，先前的工作通常将WikiText [38]作为标准基准来指示语言模型的性能。

然而，出乎意料的是，Pile-CC数据集上的验证损失与大多数下游任务显示出最强的相关性。例如，HellaSwag任务与Pile-CC验证损失之间的相关系数非常接近1.0。这一意外结果挑战了传统的假设，即WikiText是评估深度分布式模型（DDMs）最具代表性的数据集。此外，这一结果与先前研究[17, 24]的发现相符，这些研究发现网络数据集上的验证损失与下游性能密切相关。

此外，论文分析了模型在C4100Domain验证集上的损失之间的相关性，该验证集取自C4数据集[47]，并被认为与Pile-CC共享相似的分布，因为它们都源自CommonCrawl语料库。由于CommonCrawl是一个多样化领域的集合，论文原本预期每个领域的损失与下游任务之间的相关性会有所不同。然而，令人惊讶的是，超过85%的领域都与Pile-CC表现出非常强的相关性（完整的相关性图表见附录D）。这一点在www.ign.com域名上得到了很好的例证，如图4（b）所示，它与Pile-CC的整体相关性图表非常相似。这也表明，Pile-CC与下游任务性能之间的高相关性可能归因于它对各种主题和领域的广泛覆盖。

表4 不同数据选择方法的性能比较。Human指的是The Pile[18]中提出的权重，Pile-CC Only表示仅在Pile-CC组件上进行训练，而DoReMi则是指Xie et al.[64]中的权重。每个任务报告的性能是五个不同运行中从O-shot到5-shot设置的平均分数及标准差。论文估算了达到训练数据配比所需计算量（以FLOPs衡量）。对于每个任务，显著优于Human基线的分数以粗体突出显示，显著性由Cohen's确定。为了全面评估，论文还报告了使用LightEval的评估结果，遵循Penedo et al.[44]在附录C.4中的设置。LightEval结果显示，REGMIX略优于DoReMi和Pile-CC Only。    

2.4.3 通过REGMIX的数据配比提升下游性能

先前的工作表明，数据配比方法能够通过使用较少的训练token实现更小的验证损失（或困惑度）来加速LLM预训练[64]。然而，一个关键问题是应该优化哪个验证损失？最直观的方法，也是先前工作所采用的方法，是尽量减少所有领域的损失。然而，基于论文对100万条训练日志的研究，论文发现这在实践中几乎是不可能实现的。没有任何数据配比能够同时超越所有领域验证损失上的人工选择。这表明，简单地最小化所有领域的损失很可能是不可行的。因此，论文选择优化Pile-CC验证损失，以期在下游任务上实现普遍性能提升，因为它与下游性能显示出最高的相关性。

论文实施了两种方法来确定数据配比。第一种方法依赖于人类的直觉。由于Pile-CC与其自身分布应该是最接近的匹配，论文假设仅在Pile-CC上进行预训练可能会比基线产生更好的性能。第二种方法利用REGMIX，以Pile-CC验证损失作为目标变量。论文采用LightGBM来预测能够最小化Pile-CC验证损失的数据配比。

论文比较了论文提出的方法与强基线的性能，包括由人类为Pile [18]进行的筛选，以及DoReMi [64]。对于DoReMi，论文直接从其报告的最佳域权重中获取数据配比，并重新归一化所有可用的\textbfI /域。这可能导致DoReMi的性能相对于原始报告结果有所下降。如表4所示，Pile-CC Only和REGMIX都显示出与基线相比的强劲性能。在广泛使用的HellaSwag基准测试中，REGMIX比人类选择提高了6.8。此外，REGMIX在8项任务性能中击败了所有其他三种方法。    

图5：左：预训练语料库中包含Pile-CC的不同方法在Pile-CC上的验证损失。右：预训练中排除Pile-CC的情况下，在Pile-CC上的验证损失。

图 6：使用线性回归模型对不同目标域验证损失与训练域权重之间相关性的可视化。左侧基于 Pile 数据集，右侧基于 Stack 数据集。高度相关表明增加训练域权重对降低目标域验证损失有积极影响。

在 14 个案例中，Pile-CC Only 表现突出，平均得分最高。Pile-CC Only 的惊人强劲表现进一步强化了论文上一节的结论：网络语料库对下游性能有益。最终，REGMIX 超越了表 3 中的最佳模型，证明论文的自动数据配比方法比随机搜索更高效。

尽管 Pile-C 验证损失是下游性能的有益指标，但它可能不适用于所有感兴趣的任务。有时论文无法假设验证集与训练集数据分布相似，而是面临分布外场景。为验证论文的方法在分布外场景中的有效性，论文完全排除 Pile-CC 域，并利用剩余域找到最小化 Pile-CC 验证损失的最优数据配比。如图 5（右）所示，论文提出的方法仍优于基线方法。这表明 REGMIX 无论目标域是否分布内或分布外都具有鲁棒性。在此设置下，论文还提供了回归评估结果，如图 5 所示。

2.4.4 域间交互对人类理解具有挑战性

为了理解不同域之间的影响，论文在图 6 中可视化了线性回归模型的系数 ( \omega)。该可视化揭示了各数据域如何相互贡献，展示了它们之间的复杂交互。论文还展示了基于 Stack 数据集 [28] 训练的每个 1M 代码模型的代码相关性图。令人惊讶的是，无论是域交互可视化还是代码相关性图，都显示了复杂的关系，这些关系对于人类专家来说难以完全理解。例如，Pile 数据集中的 PhilPapers 域似乎在线性回归建模下为所有其他域带来了增益，这是一个挑战直觉的人类理解的非显而易见发现。这些可视化突显了确定最优数据配比的内在复杂性，强调了论文自动化 REGMIX 方法在高效识别高性能配比方面的价值，而非仅依赖人类直觉。    

图 7：跨多种数据配比的1M训练日志可视化。

2.4.5 数据配比效应超越缩放定律

近期研究[68, 20]证明了数据配比缩放定律的可行性。然而，论文在第5.4节的发现表明，域权重与验证损失之间的关系比缩放定律所暗示的要复杂得多。为了可视化这种复杂性，论文在图7中绘制了所有1M训练日志的实验点。如果数据配比的缩放定律成立，论文期望看到所有域之间存在明显的对数-对数线性关系。然而，论文的结果揭示了一个更为细致的图景。例如，DM Mathematics域，可能由于其与其他域相比的独特分布，显示出损失与域权重之间接近对数-对数线性关系。相比之下，大多数域如Pile-CC显示出更复杂的模式，预测验证损失变得不简单。如图所示，域间交互似乎很复杂，仅基于域在配比中的权重来预测验证损失变得具有挑战性。这些发现表明，尽管缩放定律提供了有价值的见解，但它们可能无法完全捕捉数据配比动态的复杂性。论文的方法通过将整个数据配比作为回归模型的输入来应对这一挑战，提供了一个更全面的框架，用于理解和预测验证损失，同时兼顾所有域权重。    

本文转载自 ​​AI帝国​​，作者： 无影寺