点击下方卡片，关注 「AI视界引擎」 公众号
大型语言模型（LLM）如GPT-4和Gemini的迅速普及强调了它们在训练过程中对资源的强烈需求，由于巨大的计算和环境成本，这带来了重大挑战。为了缓解这个问题，作者提出了在LLM预训练中采用预训练权重合并的方法。这种方法利用具有共享训练轨迹的LLM预训练权重，并基于贝叶斯优化对最佳合并权重进行广泛搜索空间探索。通过各种实验，作者证明了：
作者提出的方法具有增强预训练的能力，类似于以最小的成本获得巨大利益的机会；
尽管作者的方法需要给定一个保留数据集，但仍然在预训练中展现出跨多个领域的强大泛化能力，这是预训练的一个关键方面。
1 Introduction
随着如GPT-3、GPT-4、PaLM和Gemini等大型语言模型（LLMs）的快速发展，这些模型拥有数十亿到数百亿个参数，对新LLMs的需求以及旨在提高它们能力的研究显著增加。但作者应该注意到，这些LLMs的训练要求是实质性的，不仅在计算资源、人力资源和资本资源方面，也在于能源消耗和环境影响。
例如，训练拥有2T Token 的LLaMA2 70B模型需要1,720,320 GPU小时，通过神经架构搜索开发拥有2.13亿个参数的 Transformer 可能导致的环境负担相当于五辆汽车在其整个生命周期内的CO排放量。因此，在预训练阶段降低消耗和成本已经成为该领域的一个关键挑战。
为了应对这一挑战，研究行人在LLM预训练中采用了各种策略，包括混合精度训练、零冗余优化器、持续重训练、流水线并行主义以及深度放大方法。尽管这些方法有助于提高预训练的效率和降低成本，但它们主要关注的是模型架构或优化过程，而不是直接减少预训练阶段的资源消耗。
与这些研究不同，本文关注模型融合策略，这是机器学习中的一个经典话题，以增强LLM预训练。特别是，作者利用预训练期间保存的预训练权重，并平均这些预训练权重的参数来改进预训练，而不需要大量资源，因为融合的预训练权重可以相对于单个预训练权重的输出减少组合输出的方差，同时不增加偏差。
然而，在预训练中执行预训练权重合并并非易事，因为在平均参数时可能会发现不同的局部最小值。因此，探究预训练权重合并的基本特性并明智地确定合并权重是非常重要的。
为此，作者做出了以下努力：（1）作者进行了一些试点实验，以探索预训练权重合并的特征；（2）基于试点实验中的发现，作者提出了一种基于贝叶斯优化的方法，以寻找最优或近最优的合并权重。
具体来说，作者首先通过各种试点实验探索了三个研究问题：“在预训练轨迹中应该合并哪些预训练权重？”，“应该合并多少个预训练权重？”以及“如何合并预训练权重？”。然后，基于试点实验中的发现，作者利用贝叶斯优化来优化预训练权重合并的昂贵、黑箱且无需导数的客观函数，并确定预训练权重的合并权重。
通过各项实验，作者主要发现：
- 作者提出的方法有可能增强预训练，几乎不费吹灰之力；
- 除了性能优越之外，由特定的保留数据集决定的合并，仍然在各种领域展现出强大的泛化能力，这是预训练中的关键方面。
总之，本文的贡献有三点：
- 作者 Proposal 在预训练轨迹中合并预训练权重以降低预训练成本，几乎不费吹灰之力；
- 为了找到最优的合并权重，作者利用贝叶斯优化，这擅长优化昂贵的黑盒无导数目标函数；
- 通过各种实验，作者指出作者的方法表现出卓越的性能，并且新合并的预训练权重在不同领域间保持了强大的泛化能力。
2 Pilot Experiments
本节展示了在作者对预训练权重合并的初步实验中观察到的现象。具体来说，作者主要探索以下三个研究问题：RQ1: 在预训练轨迹中应该合并哪些预训练权重？ RQ2: 应该合并多少个预训练权重？ RQ3: 如何合并预训练权重？
Experiment Setup
在初步实验中，作者使用了11个不同预训练阶段的Baichuan2预训练权重，从2000亿个标记一直到完整的26万亿个标记。
同时，作者选择了两个具有代表性的基准数据集作为测试：C-Eval，一个包含超过1万个多选题的全面评估基准，覆盖了52个不同学科和四个难度 Level ；CMMLU，一个旨在评估法学硕士的知识和推理能力的一般中文评估基准。
Which Checkpoints in the Pretraining Trajectory Should be Merged?
为了回答研究问题1（RQ1），作者检查了所有可能的成对合并情景，总共有55种组合（），并使用贪心汤策略在C-Eval和CMMLU测试集上评估了它们的性能，在这种策略中，如果预训练权重在开发数据上提高了准确度，它们会被顺序地添加到汤中。
C-Eval上的结果展示在图1中。从图中作者可以发现：
- 在C-Eval数据集中，除了Baichuan2-220B与Baichuan2-440B合并以及Baichuan2-440B与Baichuan2-660B合并的情况外，通常将相邻训练阶段的两个预训练权重合并会比单个预训练权重表现得更好，例如，合并Baichuan2-1980B与Baichuan2-2200B可以达到56.65%的准确率，而Baichuan2-2200B单独只能达到54.98%的准确率。 此外，合并Baichuan2-1980B与Baichuan2-2200B（达到56.65%的准确率）可以大幅超过最后的预训练权重（Baichuan2-2420B，达到54.82%），在测试准确率上提高了1.83%。同时，在CMMLU数据集中也呈现了相同的趋势，具体内容在附录中展示。
- 合并远距离的预训练权重可能会导致性能大幅下降，例如，在C-Eval数据集中，合并Baichuan2-220B与Baichuan2-2200B只能达到25.26%的准确率，这与训练不足的Baichuan2-220B（23.89%）的表现接近。
How Many Checkpoints Should be Merged?
基于第一次试点实验，作者进一步进行了一些实验，以研究在预训练阶段结合不同数量的预训练权重对性能的影响。具体来说，作者使用C-Eval数据集作为测试平台，采用贪心策略来合并不同预训练阶段相邻的三个或四个预训练权重。
实验结果如图3所示。从图中作者可以发现：与成对合并相比，合并三个或四个预训练权重在性能上并没有显著优势。例如，在合并两个模型时，比如将Baichuan2-1320B与Baichuan2-1540B合并，性能显著提升至53.06%（+2.14），然而，合并三个预训练权重（Baichuan2-1100B、Baichuan2-1320B和Baichuan2-1540B）仅将性能提升至51.76%（+0.84），而合并四个预训练权重（Baichuan2-880B、Baichuan2-1100B、Baichuan2-1320B和Baichuan2-1540B）性能提升更少，仅为51.01%（+0.09）。
How to Merge Checkpoints?
在预训练权重合并中，最关键的问题在于确定如何分配（合并权重）以在LLM参数空间中结合预训练权重。为了探讨这个问题，作者选择了在C-Eval数据集上具有代表性的Baichuan2-1540B与Baichuan2-1760B以及Baichuan2-2200B与Baichuan2-2420B进行合并，同时在整个合并权重分布空间([0, 1])中均匀地采样了100个点进行深入测试。
实验结果如图2所示。从图中作者可以发现：
- 如图2（a）所示，Baichuan2-1540B与Baichuan2-1760B之间存在巨大的性能差距。在这种情况下，作者观察到合并汤的性能与合并权重之间存在一定的单调趋势。随着更高的合并权重被分配给性能更好的预训练权重，合并汤的性能逐渐提高。此外，有13%的权重显示出超越优秀基模型（Baichuan2-1760B，其相比Baichuan2-1540B可以实现更高的准确率）的改进。
- 在图2（b）中，由于预训练达到后期阶段，Baichuan2-2200B和Baichuan2-2420B显示出相似的性能。在合并Baichuan2-1540B和Baichuan2-1760B时显示的单调趋势消失了。同时，76%的合并权重有潜力将合并汤的性能提升到超越更优秀的预训练权重。为了进一步验证这一发现，作者还对DeepSeek-1800B和DeepSeek-2000B进行了相同的实验，DeepSeek训练的一个70B的LLM。作者发现，实验结果与在Baichuan 2上的结果一致，详细内容见附录。
3 Method
在本节中，作者首先在子节SS3.1中阐述预训练权重合并的公式。然后，在子节SS3.2中，作者介绍了作者方法的实现，该方法能够有效且高效地确定合并权重。
Preliminary: Checkpoint Merging
在进行LLM预训练时，作者在时间已经保存了多个预训练权重，表示为。这些参数空间中多个预训练权重的线性组合被称为“预训练权重汤”(Checkpoint Soup)，可以表示为：
其中， 代表合并权重。根据初步实验的关键发现，在本文中，作者只关注成对合并，因此，方程式1可以重新表述为：
如初步实验所示，当分配了适当的时，的性能可以超过。因此，对于每个预训练权重合并迭代，主要挑战是找到最优或接近最优的合并权重。
Checkpoint Merging via BayesOpt
为了找到最优或接近最优的合并权重，作者采用了贝叶斯优化（简称“BayesOpt”）。贝叶斯优化是一种有效的全局优化策略，广泛应用于寻找那些成本高昂或难以直接评估的函数的最优解。贝叶斯优化主要包括三个组成部分：一个待优化的目标函数，一个用于建模目标函数的贝叶斯统计模型，以及一个用于决定下一步采样位置的采集函数。
预训练权重合并中的目标函数。 在预训练阶段中要优化的目标函数可以表述为：
在给定配备了参数 （如方程2中所定义）的LLM和一个特定的保留标记数据集 的条件下，目标函数 可以将合并权重 投影到LLM在 中的整体性能上，例如，在命名实体识别任务中的F1分数和在翻译任务中的BLEU分数。此外，为了减少搜索空间，合并权重的范围由超参数 控制。
高斯过程：一种贝叶斯统计模型。为了模拟方程3中的目标函数，采用了高斯过程（GP）回归。具体来说，对于一组有限的合并权重 ，它能将这些点上的目标函数值收集到一个向量中 。注意，为了方便，作者省略了中的条件。在高斯过程中，这个先验分布被认为是多元正态分布，有一个特定的均值向量和协方差矩阵，表示为。
其中 和 是具有可学习参数的平均函数和协方差函数，可以通过最大似然估计进行估计。
和 。
使用贝叶斯定理，作者有一个后验概率分布：
从上述方程中，作者发现：（1）后验平均值 是先验 与基于数据 的估计之间的加权平均值。（2）后验方差 等于先验协方差 减去一个对应于通过观察 而消除的方差的项。
获取函数。 假设作者有 和 ，目前作者找到的最佳解是 ，作者需要使用获取函数来确定在哪里采样新的合并权重（下一个评估点）。获取函数作为一个计算成本低廉的效用函数，用于指导采样决策。
本文中使用的一种采集函数是预期改进（Expected Improvement，简称EI）。具体来说，作者希望选择 以便提升LLM 的性能。由于在将 中的所有数据通过LLM传递之前， 是未知的，作者能做的就是计算改进的预期值，并选择 来最大化它。作者定义预期改进为：
在这里， 表示在观测到 在 的值后，根据后验分布（如方程5所示）取期望。然后作者在期望改进最大的点上评估合并权重。
本文中使用的另一种采集函数是上限置信区间（UCB）。与EI相比，UCB试图从乐观的角度引导搜索，其定义如下：
其中 是一个可学习的参数，用来指导探索-利用之间的权衡。考虑到目标函数 可能复杂且非凸，可能没有任何单一的选择函数在整个优化过程中都能表现最佳。因此，一种混合策略，即选择函数在每次迭代时从一个池中抽样，可能比任何单一的选择方法效果更好。
请注意，所提出的通过贝叶斯优化进行预训练权重合并的方法不仅限于成对预训练权重合并，同样也可以以相同的方式处理多个预训练权重的合并。通过贝叶斯优化进行预训练权重合并的整体过程在算法1中有所展示。
4 Experiments
Experimental Setups
数据集： 除了在初步实验中使用的C-Eval和CMMLU之外，作者还进一步选择了三个基准数据集作为测试平台：MMLU，这是一个大规模的多任务语言理解基准，包含有关学术科目的多项选择题；以及GSM8K，这是一个知名的以数学为重点的评估基准。
Baseline 方法：作者将作者的合并方法与以下强劲的 Baseline 进行了比较：
- 合并前的预训练权重 为了更好地展示模型合并后的性能变化，作者报告了合并前预训练权重的性能。
- 均匀混合方法 实施了一种公平的模型平均方法，忽略个体性能指标以确保均匀的贡献。请注意，LAWA也是均匀混合方法的一个特例。
- 贪心混合方法相反，根据预训练权重在增强模型准确性方面的表现有选择地合并，只采用那些在保留数据集上带来实质性性能改进的预训练权重。
- 费希尔加权平均（简称“费希尔”）策略性地利用费希尔信息为模型参数分配权重，从而最大化集成模型的联合后验概率，以提高效率。
- RegMean 应用回归技术减少模型间的预测方差，从而增强集成模型的总体鲁棒性和性能。
LLM预训练权重 除了在预训练百川2-7B的预训练权重外，作者还利用了训练Deepseek 7B的中间预训练权重。
主要结果
图4展示了作者方法在中后期预训练阶段在C-Eval数据集上相对于最佳基准预训练权重性能（简称"BestCKPT"）值的性能提升。同时，在表1中，作者提出了两种具有显著分数提升的合并组合。例如，Baichuan2-1980B和Baichuan2-2200B的合并相较于Baichuan2-2200B在CMMLU数据集上实现了0.59%的改进。同样，Baichuan-2200B和Baichuan2-2420B的合并也显示了显著的增益，相较于Baichuan2-2420B在C-Eval数据集上实现了0.57%的提升。
除了超过基准预训练权重外，作者的方法还在各种数据集上显著优于其他合并 Baseline 。在CMMLU数据集上，将作者提出的方法应用于将Baichuan2-1980B与Baichuan2-2200B合并，可以分别超过Uniform Soup、Greedy Soup、Fisher加权平均和RegMean 2.68%、0.59%、0.75%和1.67%。在C-Eval数据集上，也观察到了类似的趋势，作者的方法以较大的优势超过了其他方法。
为了进一步验证作者提出的方法，作者将其应用于DeepSeek 7B。表2展示了结果，测试了DeepSeek-1800B与DeepSeek-2000B在C-Eval数据集上的合并，得分为45.82，分别比Uniform Soup、Greedy Soup、Fisher加权平均和RegMean高出1.34%、1.24%、0.57%和1.87%。这些结果不仅强调了所提方法的鲁棒性和适用性，也证明了其在提升不同预训练阶段和模型性能方面的有效性。
Can the Merged Soup Generalize to Unseen Domains?
由于合并权重是在某个特定的保留标记数据集（公式3）中确定的，因此，有一个非常重要的问题需要回答：“合并后的混合模型在未见过的数据集上能否表现良好？”为了回答这个问题，作者研究了各种合并混合模型在跨领域泛化能力上的表现。在实验中，作者使用C-Eval来确定合并权重，并在CMMLU、MMLU和GSM8K上测试合并后的混合模型。
结果展示在表3中。从表中作者可以发现：
- 尽管在中国数据集上确定合并权重，但通过各种合并方法获得的合并汤（merged soups）在英语数据集上，如MMLU和GSM8K，表现一直很好，这意味着在参数空间合并预训练预训练权重并不会损害预训练权重的泛化能力。
- 与贪心汤（Greedy Soup）和费舍尔加权平均（Fisher Weighted Averaging）相比，作者提出的方法呈现出更好的结果，因为领域内设置（确定权重并使用相同领域的数据进行合并汤测试）与领域外设置之间的性能差异最小，这表明在IND设置和OOD设置上找到的合并权重很可能相同。## 5 讨论部分
The Impact of Varying the Held-out Dataset Size on Checkpoint Merging
与Wortsman等人（2022年）；Matena和Raffel（2022年）相同，作者提出的方法需要训练集来确定合并权重。因此，研究在保留数据集大小变化对结果影响是值得的。为此，作者从C-Eval验证数据中提取了一些部分作为保留数据集，并测试合并Baichuan2-2200B和Baichuan2-2420B。
结果展示在表4中。从结果来看，作者发现保留数据集的大小对作者的方法的效能似乎影响最小，即使可用的数据集有限，该方法仍然保持令人满意的表现。作者猜测原因是数据集大小对作者提出方法的影响主要是评估阶段，因此，一个小型数据集仍然可以提供足够的信息来判别权重分配的质量。以C-Eval数据集为例，即便是提取其中一小部分数据，也仍然可以有效评估模型的性能。
The Impact of Varying Merging Weight Searching Space Size
如公式3所示，超参数 控制着合并权重的搜索空间。为了研究这个超参数的影响，作者进行了一些实验。实验结果如图7所示。可以观察到，将 设置为0.5或0.7可以取得相对较好的结果，而将 设置为0.9则会导致准确度明显下降。进一步挖掘，作者发现将 设置为0.5或0.7时，都能在(0.87, 0.89)的范围内收敛到最优的合并权重。
此外，在作者的经验中，作者还发现当预训练权重之间的性能差距明显时，较窄的搜索空间是有利的，主要偏向于性能较强的预训练权重。相比之下，当合并前的预训练权重可以实现平衡的性能时，更广泛的搜索空间更为合适。
7 Conclusion
在本文中，为了减轻预训练LLM的巨大计算成本，作者提出在预训练轨迹中合并预训练权重。具体来说，作者首先进行一些初步实验来探索预训练权重合并的特点。然后，基于初步实验的发现，作者提出了一种基于贝叶斯优化的方法来寻找最优或近优的合并权重。通过各种实验，作者发现：作者提出的方法有可能增强预训练，几乎不花费额外代价。除了性能优越之外，合并后的结果在各个领域仍展现出强大的泛化能力，这意味着作者提出的方法不会损害预训练预训练权重的泛化能力。
参考
[1].Checkpoint Merging via Bayesian Optimization in LLM Pretraining.
点击上方卡片，关注 「AI视界引擎」 公众号