大语言模型的简易可扩展增量预训练策略

前言

• 原论文：Simple and Scalable Strategies to Continually Pre-train Large Language Models
• 翻译文件已整理至Github项目Some-Paper-CN，欢迎大家Star！

摘要

大语言模型（LLMs）通常需要在数十亿个tokens上进行预训练，一旦有了新数据，又要重新开始训练。更有效的解决方案是增量预训练（Continue PreTraining）这些模型，这与重新训练相比可以节省大量的计算资源。然而，新数据带来的分布偏移通常会导致在旧数据上的性能下降或者对新数据的适应性不佳。在本文中，我们展示了一个简单并且可扩展的结合学习率（LR）重预热（re-warming）、学习率重衰减（re-decaying）以及添加以前的数据，足以达到与从零开始在所有可用数据上重新训练相同的性能水平，这一结论通过最终损失和在多个LLM基准（benchmarks）评估上的平均得分来衡量。。具体来说，我们展示了在微弱但切合实际的分布偏移：训练语料为常用英文预训练语料（英文->英文）和强分布偏移：训练语料为德语（英文->德语）情况下，使用大量数据集（100B tokens）和上述策略训练405M参数量的模型是有效的。在更大规模的实验中，我们选择了微弱但切合实际的分布偏移数据集，并发现我们的增量预训练策略在10B参数量的LLM效果上足以与重新训练的模型媲美。实验结果表明，通过简单、可扩展的增量预训练策略，可以成功更新LLM，只需要使用一小部分计算量就能与重新训练的模型媲美。最后，受之前工作的启发，我们提出了余轩学习率策略的代替方案，有助于避免LR重预热（re-warming）引起的遗忘，并且并且不受固定tokens预算的限制。

引言

在过去几年中，大型预训练模型在语言建模、视觉理解、文生图和文生视频方面实现了巨大的性能提示。大型语言模型是所有这些改进的核心，为人类通过语言与机器学习算法进行交互提供了一种直观的手段。

虽然LLM是当前生成式人工智能技术的基石，但它们的训练和更新成本高得令人望而却步。然而，随着新的、更高质量的数据集不断出现，企业需要更新其模型，以便在竞争中保持领先。目前，LLM 需要结合新旧数据进行重新训练。现有工作旨在通过实现低成本超参数优化或提供在给定计算预算下最大化性能的指南来降低这些训练成本。然而，这些工作都假设模型将从随机初始化开始训练，这就提出了以下问题： 从业者是否应该总是结合现有数据集，从随机初始化开始训练，以获得最佳性能？对模型的每次更新都这样做很快就会变得过于昂贵。

为了避免完全的重新训练，我们探索了简单且可扩展的增量学习策略，以便在大量新数据（200B+ token）上继续预训练 LLM（多达 10B 参数）。我们将我们的设置称为 “增量预训练（continual pre-training）”，并强调它有别于文献中的现有设置，因为我们考虑了大量的输入数据。在这项工作中，我们并不打算改进在所有可用数据上随机初始化训练出来的模型的性能。相反，我们将在现有数据集的基础上训练出来的模型视为基线，我们希望通过大规模持续学习策略的组合来匹配这些模型的性能。

然而，简单的继续在新数据上训练模型，往往会导致性能远远低于在所有可用数据上重新训练的结果，这通常是由于：（1）适应性差（未能优化新数据集）或 （2）灾难性遗忘（在以前的数据集上损失了大量能力）。首先，适应性问题是我们的核心问题，因为在大型数据集上进行训练的成本很高。人们大概不会选择花费大量计算资源在新数据集上进行训练，却只为适应新数据集。然而，大多数性能优异的开源 LLM 在训练结束时都会将学习率衰减到一个很小的值。因此，我们假设，在新数据集上进行训练时，学习率必须重新增加（re-increased）和重新衰减（re-decayed），以提高单位计算量的适应性。我们注意到，增量训练的相关文献尚未对此进行深入研究。其次，如果要充分发挥持续预训练的潜力，灾难性遗忘是一个需要克服的关键困难。适应数千亿个新token固然重要，但绝不能以抹去LLM中的大部分现有知识为代价。最近的研究表明，在LLM微调设置中，添加以前的数据（只要1%）就足以在很大程度上减轻遗忘。虽然持续对大量新数据进行预训练几乎肯定会比微调导致更多遗忘，但我们假设，适量的添加原数据可以减轻遗忘——即使在我们的实验中也是如此。此外，最近的研究表明，预训练和扩大模型规模都有助于减少遗忘的影响。因此，我们预计语言模型容量和预训练数据集规模同步增长的趋势将使模型的持续学习能力越来越强，这表明我们的实验结果只会随着模型规模的扩大而改善。

鉴于增量训练在大大降低重新训练模型的成本方面具有巨大潜力，而且LLMs有可能成为强大的增量学习者，我们提出了以下问题：当应用简单且可扩展的增量训练技术时，相较于在所有可用数据上从随机初始化开始预训练的LLMs，增量预训练的LLMs的性能差异有多大？为了回答这个问题，我们对用于LLM预训练的增量预训练技术进行了大规模实证研究。我们的实证评估涵盖了大型（10B参数量）和小型（405M参数量）模型，以及弱（英语 → 英语）和强（英语 → 德语）分布偏移。我们的主要贡献可归纳如下：

• 我们确定了使用余弦学习策略对预训练模型的学习率重预热（re-warming）和重衰减（re-decaying）的效果，表明在增量预训练期间，重预热（re-warming）和重衰减（re-decaying）对适应新数据是必要的。
• 我们确定了添加原数据的效果，同时在两种分布转移和多次改变添加原数据的百分比。 我们发现，即使在对数千亿个新token进行模型更新时，也可以通过适量添加原数据来显著减少遗忘。
• 我们通过两个模型的参数大小和分布变化证明，将 LR重预热（re-warming）、LR重衰减（re-decaying）和添加原数据简单且可扩展地结合在一起，可使增量预训练模型的平均性能达到与在所有数据基础上重新训练的模型类似的水平，同时大大减少了计算量。
• 我们为LLM的增量预训练提出了无限学习率策略（infinite learning rate schedules）（允许跨数据集平滑过渡的策略），这是避免与学习率重预热（re-warming）相关的优化困难的一种可行方法。

主要结论和启示

我们的实验结果假定，增量预训练LLM依次经历两个或多个预训练阶段。也就是说，我们的结果适用于随机初始化一个增量预训练LLM，并在按照数据集$D_0$，$D_1$，$D_{N-1}$的顺序进行预训练的情况，其中$N\ge 2$且$tokens(D)\ge 100B$。我们注意到，这包括了以下两种情况：一种是涉及到已经在$D_0$上进行过预训练的开源模型的情况，另一种是企业可能希望训练一个初始的LLM，并打算在新的数据上不断地对其进行预训练的情况。新数据可能与之前的数据相似，对应于弱分布偏移（比如来自于网络上不同领域的最新信息），也可能与之前的数据截然不同，对应于强分布偏移（比如来自全新语言的数据）。我们的实验评估考虑到了这些困难，发现适当应用LR重预热（re-warming）、LR 重衰减（re-decaying）和添加原数据足以在弱和强分布偏移以及两种模型大小中媲美重新训练的性能（如图 1）。为了使我们的研究成果尽可能为社会所了解，我们现在提供应用我们研究成果的经验法则：

• 注意——以下指南是基于我们当前的知识编写的。

• 学习率策略：

  • 如果学习率在初始数据集的预训练过程中从一个较大的值${\eta }_{max}$ 以余弦衰减到一个较小的值 ${\eta }_{min}$，那么以下指南可以帮助您进行增量预训练模型：
  • 将学习率从$\mathcal{O}({\eta }_{max})$重新升温并再次衰减到$O(\mathcal{O}({\eta }_{min})$可以改善模型对新数据集的适应性。
  • 降低学习率策略中的最大学习率可以帮助减少遗忘，而增大最大学习率可以改善模型对新数据的适应性。
  • 无限学习率计划（Infinite LR schedules）是余弦衰减策略有前景的替代方案。它们在任务之间转变为较高的恒定学习率，从而通过避免在任务之间重新预热（re-warming）学习率，帮助防止与优化相关的遗忘。它们还避免了对特定的tokens预算的限制，因为在训练过程中可以随时使用最终的指数衰减来将模型训练至收敛。

• 添加原数据：

  • 我们建议添加 5%的原数据，作为默认比例。 对于较强的分布偏移，应该使用更大的比例，而对于较弱的分布偏移，可以使用 1%的原数据。
    

图 1：增量预训练可降低更新模型的计算成本，同时保持相似的最终验证平均性能。我们展示了 在两个数据集的基础上训练的Pile ∪ SlimPajama(SP)/German(Ger.) 基线模型的结果，我们认为这是性能的上限。我们还报告了两个增量预训练模型的性能。“在Pile上预训练”（PT on Pile）从一个预训练的Pile检查点开始，只使用学习率重预热（re-warming）和重衰减（re-decaying），而“Replay（在Pile上预训练）”则使用学习率重预热（re-warming）和重衰减（re-decaying），并对SlimPajama添加5%的原始数据，对German添加25%的原始数据。我们观察到，学习率重预热（re-warming）和重衰减（re-decaying）和添加原数据的组合使我们的增量预训练模型在需要显著更少计算量的情况下，能够达到与基线模型相似的平均性能。我们注意到，这种设置假定有一个预训练模型可用。

相关工作

增量学习（训练）（Continual learning）

增量训练（CL）方法旨在从不断变化的数据分布中学习，在适应新数据的同时，保留通过先前训练收集的知识。增量训练的关键挑战在于，既要避免遗忘过去的信息，又要适应新的信息。这种权衡被称为 “刚性-可塑性困境”（the rigidity-plasticity dilemma）。

即使在小规模场景中，CL方法也很方便，可以避免从头开始重新训练，或解决数据可用性问题。然而，在大规模情况下，CL不仅仅是一种便利，可能还需要处理大量持续收集的数据。最近训练规模的扩大，尤其是LLM的扩大，为CL提供了新的机会，可以降低重新训练的成本，提高内存、计算和存储的效率。正如联合学习可以实现在空间共址的不同代理之间共享计算和数据一样，增量训练可以通过时间逐步实现计算和数据共享，成为大规模训练的有用工具。

最近的研究表明，SGD和Adam等优化器在DNN中具有有趣的知识保留特性，这可能有利于大规模CL，而且只需添加少量原数据就足以促进知识积累。在这项工作中，我们希望在大型语言模型预训练的背景下从这些方法的效率中获益，并通过正确的学习率策略和添加原数据（replay）策略来提升它们。

预训练，模型规模和增量学习（训练）

现有几项研究评估了预训练和模型规模对增量训练的影响。Cossu 等人（2022 年）研究了语言和视觉的预训练方案。他们发现，无监督和自我监督的预训练在减轻遗忘方面起着根本性的作用，而监督则会损害性能。同样，Mehta 等人（2023 年）发现，预训练模型的遗忘比随机初始化模型少，这是因为它们的权重位于损失空间较平坦的区域。他们还发现，较大的模型遗忘较少，这与Ramasesh等人（2022 年）和Mirzadeh等人（2022 年）的研究结果有关。前者发现预训练模型随着规模的扩大遗忘较少，这可能是由于隐藏表征随着规模的扩大更加正交。后者发现，与参数相当的深度神经网络相比，广度神经网络的遗忘更少。Hernandez等人（2021 年）建立了迁移的比例定律：可以预测神经网络在新任务上的表现的方程，是其参数数量和预训练数据集大小的函数。作者发现，随着参数数量的增加，这种正迁移的效果会持续改善。最后，Scialom等人（2022 年）表明，自回归LLM具有很强的增量训练能力，他们假设这与其预训练目标有关。

领域自适应增量预训练 （DACPT）

现有的工作考虑了领域自适应增量预训练（DACPT），在这种情况下，LM可以依次获得一系列未标记的领域，而实践者希望以自我监督的方式对每个领域进行训练，同时保持每个领域的性能。虽然目标与我们的类似，但我们考虑的是混合多个领域的通用预训练数据集，而不是特定领域的数据集。Ke等人（2022 年）假设在新领域进行训练时，以前领域的数据不可用，并针对这种情况开发了一种新技术，其中包括对所有以前任务的参数进行重要度掩码，以防止在以掩码语言建模（MLM）为目标进行预训练时发生遗忘。Gururangan等人（2020 年）研究了RoBERTa（也是 MLM）的领域和任务自适应预训练，并为高效增量预训练贡献了一种样本选择策略。同样，Xie 等人（2023 年）也提出了一种数据选择策略，降低了增量预训练的计算成本（自回归 LM）。Qin 等人（2023 年）研究了重新循环利用以前基础LM的微调适配器层，作为新适配器的初始化，以根据特定任务适配不断更新的基础LM版本。最近，Wu 等人（2024 年）提出了LLaMA Pro，这是一种对 LLM 进行增量预训练的方法，可以在不遗忘先前知识的情况下学习新任务。不过，与我们考虑调整所有现有权重的工作不同，LLaMA Pro需要在每次新更新时扩大模型的规模，并只调整新权重。

应用于垂直领域的LMs增量训练

一些相关工作将增量预训练应用于特定任务和领域（Sun 等人，2020；Jang 等人，2022a;b；Gong 等人，2022；Zan 等人，2022；Yadav 等人，2023a；Ma 等人，2023；Yang 等人，2024）。虽然这些工作也利用了增量预训练技术，但与我们的工作不同的是，它们专注于特定领域，而不是一般的预训练技术，专注于规模较小的数据集（小于10B tokens）与较小的模型。现有研究中唯一接近我们数据集规模的是（Gogoulou 等人，2023 年），该研究在英语、丹麦语、冰岛语和挪威语数据集（各为 730 亿）中探索了连续自回归语言建模。虽然他们没有添加原始数据（replay），但对学习率进行了重预热（re-warming）和重衰减（re-decaying）。目前唯一接近我们模型规模的研究是（Yang 等人，2024 年）。他们在小规模的植物科学学术数据上不断对LLaMA2进行增量预训练和指令微调。这项同时进行的工作采用了与我们提出的非常相似的增量训练设置：混合原始数据（replay）、重预热（re-warming）和重衰减（re-decaying）。虽然与我们的工作不同，他们没有建立一个受控实验框架来系统评估这些方法在增量预训练中的有效性，但很高兴看到进一步的实验证据验证了我们的方法。

学习率策略

一些研究考察了不同学习率（LR）策略对神经网络训练稳定性和最终性能的影响。Goyal 等人（2018）发现，在训练初期逐步预热 LR 有助于克服优化难题，尤其是在mini-batch规模较大的情况下。此外，Popel & Bojar（2018）还强调了训练后 LN Transformers 时热身阶段的重要性。另一方面，Xiong 等人（2020 年）发现，Pre-LN Transformers 更为稳定，可能不需要热身阶段。You 等人（2019）探索了 LR 衰减的作用，发现较大的初始 LR 会阻止网络记忆嘈杂数据，而较小的 LR 则有助于学习复杂模式。卡普兰等人（2020）探索了预训练大型语言模型（LLM）的 LR 策略，发现策略选择对性能没有显著影响。Hoffmann 等人（2022 年）纠正了这一错误发现，他们发现 LR 策略确实起着重要作用。Hoffmann 等人（2022 年）和 Rae 等人（2021 年）确立了在预训练LLM时使用余弦策略的最佳实践，并已被广泛采用。相比之下，Raffel 等人（2023 年）和 Zhai 等人（2022 年）探索了大规模预训练时遵循反平方根衰减的LR策略。Raffel 等人（2023 年）利用反平方根衰减（Inverse Square Root Schedule）训练 LLM，从而可以灵活调整训练步数。在 Zhai 等人（2022 年）的研究中，作者利用这些被称为"无限学习率策略（infinite learning rate schedules） "的策略来训练视觉变换器。这些策略可实现无限训练，并在一次运行中评估多个训练持续时间。我们提出的LLM无限学习率策略（第 7.4 节）就是受这一思想的启发。

背景与方法

在本节中，我们将提供与LLMs增量预训练相关的背景和方法。

线性预热和余弦衰减策略

Hoffmann等人（2022 年）和 Rae 等人（2021 年）确立了在预训练LLM时使用余弦策略的最佳实践。具体来说，他们建议从线性预热阶段开始，然后将学习率衰减到其最大值的 10 倍，这样余弦周期的终点就会与tokens数量相匹配。虽然线性预热的持续时间各不相同，但大多数著名的作品都将0.1%到0.5%的预热作为训练步骤（Zhao 等人，2023 年）。鉴于许多流行的开源模型（Touvron 等人，2023b;a; Almazrouei 等人，2023）都遵循这一增量训练策略，因此了解其细微差别对于增量预训练此类模型至关重要。该策略首先在$T_{warmup}$个时间步长（timesteps）上线性增加学习率，或者直到某个时间步$t_{ann}=T_{warmup}$：
$${\eta }_t={\eta }_{max}\cdot \frac{t}{T_{warmup}}$$

其中，${\eta }_t$是迭代到$t$时的学习率值，${\eta }_{max}$是最大学习率。然后，策略会在$T_{ann}$个时间步上过渡到余弦退火阶段，等同于直到某个时间步$t_{end}=T_{ann}+t_{ann}$：
$${\eta }_t={\eta }_{min}+\frac{{\eta }_{max}-{\eta }_{min}}{2}\cdot \left(cos\left(\pi \cdot \frac{t-t_{ann}}{t_{end}-t_{ann}}\right)+1\right)$$

其中${\eta }_{max}$是最大学习率，${\eta }_{min}$是最小学习率。

保持Replay后计算量相等

在许多实验中，我们将使用添加了原始数据（Replay）训练的模型与不使用原始数据（Replay）训练的模型进行比较。在进行这种比较时，我们会保持训练两种模型的计算量不变。也就是说，我们会相应减少从新数据集中看到的token数量，以适应从原始数据（Replay）缓冲区看到的额外token。我们将这种添加原始数据的方式称为Compute-equivalent Replay。例如，假设数据集$D_0$和$D_1$各包含 100B tokens。我们希望在$D0$和$D1$上顺序训练的模型 (a) 与在$D_0$和$D_1$上顺序训练并加入5%原始数据的模型 (b) 进行比较。模型 (a) 将看到来自两个数据集的所有tokens，共计200B个tokens。模型（b）将看到 $D_0$的100B tokens和$D_1$的95B tokens，再加上$D_0$的5B原始tokens，共计200B tokens 。这样，两个模型的计算量相同。

例如，在我们的实验中，只有两个数据集$(D_0、D_1)$，我们在训练$D_1$时使用$D_0$的Replay数据。我们按照在$D_0$上进行预训练时的顺序Replay数据，因为在初步实验中，我们在Replay数据时没有观察到明显的差异。Replay样本的选择方法将留待今后研究。我们将使用重放数据的模型称为 “$D_{1}x\%Replay$”，其中$x$