Stacking Your Transformers: A Closer Look at Model Growth for Efficient LLM Pre-Training

作者/机构: Wenyu Du (The University of Hong Kong), Tongxu Luo (HKUST, USTB), Zihan Qiu (Tsinghua University), Zeyu Huang (University of Edinburgh), Yikang Shen (MIT-IBM Watson AI Lab), Reynold Cheng (The University of Hong Kong), Yike Guo (HKUST), Jie Fu (HKUST)

A1 主要贡献

本文旨在解决大语言模型（LLM）预训练成本高昂的问题，通过一种名为“模型增长”的方法来加速训练过程。模型增长的核心思想是利用已训练好的小模型来初始化并加速大模型的训练。然而，现有模型增长方法在LLM预训练中的可行性尚未得到充分探索。作者指出了阻碍其应用的三个关键障碍（Obstacles）：
* O1: 缺乏全面评估：现有方法大多在BERT等编码器模型上进行评估，缺乏在Llama等主流解码器LLM上的标准比较，或比较不充分。
* O2: 未经检验的可扩展性：现有方法是否能在大模型尺寸和海量训练数据上持续有效，尚不明确。
* O3: 缺乏经验性指导：不像缩放定律（Scaling Laws）那样有明确的指导原则，模型增长缺乏如何选择关键超参数（如增长时机和增长因子）的实践指南。

为了克服这些障碍，本文做出了以下主要贡献：
1. 系统性评估模型增长技术：将现有的模型增长方法归纳为四种原子增长算子（G_direct, G_learn, G_zero, G_random），并设计了一个标准的LLM预训练测试平台，在宽度和深度两个维度上对它们进行了全面评估。
2. 发现并验证Gstack的优越性：实验发现，一种简单的深度方向堆叠算子——G_direct的深度版本（后文称Gstack），在训练损失和下游任务性能上均显著优于其他方法和从头训练的基线。例如，在训练一个7B模型时，Gstack相比从头训练在达到相同损失时节省了54.6%的训练tokens。

A3 背景知识

2 相关工作 - 用于高效预训练的模型增长

神经网络增长思想的起源。神经网络的增长思想可以追溯到上世纪90年代【24–26】。而开创性的工作Net2Net【10, Net2net: Accelerating learning via knowledge transfer, 2015, arXiv preprint arXiv:1511.05641】标志着一个里程碑，它是深度学习时代首次研究模型增长的尝试。Net2Net通过随机分割旧神经元和注入新的恒等层来扩展模型的宽度和深度，同时保持原有功能（即函数保持）。

基于扩展现有参数的方法。Net2Net的宽度分割方法代表了一系列旨在将现有神经元“扩展”到所需更大尺寸的工作。Bert2Bert【11, bert2bert: Towards reusable pretrained language models, 2021, arXiv preprint arXiv:2110.07143】是基于BERT的Net2Net宽度扩展。StagedGrow【13, Staged training for transformer language models, 2022, International Conference on Machine Learning】通过拼接两个相同的层来使宽度加倍，并通过将最终损失减半来保持函数不变。Lemon【12, Lemon: Lossless model expansion, 2023】建议在Bert2Bert的神经元分割中集成一个参数，旨在打破权重对称性。在深度方面，StackedBert【14, Efficient training of bert by progressively stacking, 2019, International conference on machine learning】简单地堆叠重复的层来形成一个更深的模型。与上述直接复制/分割的方法相反，LiGO【15, Learning to grow pretrained models for efficient transformer training, 2023, arXiv preprint arXiv:2303.00980】提出了一种基于学习的方法，通过学习一个从较小模型参数到较大模型参数的线性映射来初始化较大模型的参数。

基于独立初始化新参数的方法。除了扩展现有参数的方法外，还有一些工作是在不依赖现有参数的情况下初始化新参数。例如，MSG【17, Masked structural growth for 2x faster language model pre-training, 2024】提出了一种多阶段增长策略，逐步扩展Transformer组件，其中新增长的神经元使用掩码机制进行随机初始化，以确保函数保持。此外，一些工作将特定值（如零）分配给新初始化的神经元，以消除其影响【16, Gradmax: Growing neural networks using gradient information, 2022, arXiv preprint arXiv:2201.05125】、【12, Lemon: Lossless model expansion, 2023】。

现有方法在LLM时代的局限性。以上所有方法主要在BERT或LLM预训练的早期阶段进行了探索。而本文的目标是在LLM时代对模型增长技术进行首次系统性回顾。据作者所知，FLM-101B【20, Flm-101b: An open llm and how to train it with $100k budget, 2023】是唯一使用增长方法【17, Masked structural growth for 2x faster language model pre-training, 2024】来加速十亿级LLM预训练的现有LLM。然而，这项工作缺乏一个从头开始训练的基线模型，因此难以评估模型增长技术的有效性。相比之下，本文旨在通过建立一个标准化的测试平台，在LLM预训练中比较从头训练的LLM和使用各种增长方法的LLM，从而提供一项全面的研究。 ## A2 方法细节 ### 3 系统性评估LLM预训练的模型增长 **现有模型增长方法的局限性**。现有的模型增长方法【14, Efficient training of bert by progressively stacking, 2019, International conference on machine learning】、【11, bert2bert: Towards reusable pretrained language models, 2021, arXiv preprint arXiv:2110.07143】、【18, Progressively stacking 2.0: A multi-stage layerwise training method for bert training speedup, 2020, arXiv preprint arXiv:2011.13635】、【15, Learning to grow pretrained models for efficient transformer training, 2023, arXiv preprint arXiv:2303.00980】、【12, Lemon: Lossless model expansion, 2023】、【13, Staged training for transformer language models, 2022, International Conference on Machine Learning】、【16, Gradmax: Growing neural networks using gradient information, 2022, arXiv preprint arXiv:2201.05125】、【17, Masked structural growth for 2x faster language model pre-training, 2024】主要在BERT【27, Bert: Pre-training of deep bidirectional transformers for language understanding, 2019】上进行评估，对像Llama【21, Llama 2: Open foundation and fine-tuned chat models, 2023】这样的仅解码器大规模语言模型的关注有限。此外，由于训练设置不同，这些增长方法通常不具可比性【14, Efficient training of bert by progressively stacking, 2019, International conference on machine learning】、【11, bert2bert: Towards reusable pretrained language models, 2021, arXiv preprint arXiv:2110.07143】、【17, Masked structural growth for 2x faster language model pre-training, 2024】、【12, Lemon: Lossless model expansion, 2023】。即使一些增长型LLM实验被评估过，其结果也常常不完整【20, Flm-101b: An open llm and how to train it with $100k budget, 2023】、【15, Learning to grow pretrained models for efficient transformer training, 2023, arXiv preprint arXiv:2303.00980】。为了克服这些限制，本文首先将现有工作总结为四种原子增长算子来代表这些增长技术。然后，我们建立一个标准化的LLM训练测试平台，用四种增长算子在深度和宽度方向上预训练LLM，并使用训练损失和Harness【23, A framework for few-shot language model evaluation, 2023】中的八个评估指标对结果进行评估。

3.1 使用增长算子增长LLM

四种原子增长算子。近年来，研究人员通过利用较小的预先存在的模型来提高训练大模型的效率【10, Net2net: Accelerating learning via knowledge transfer, 2015, arXiv preprint arXiv:1511.05641】、【11, bert2bert: Towards reusable pretrained language models, 2021, arXiv preprint arXiv:2110.07143】、【14, Efficient training of bert by progressively stacking, 2019, International conference on machine learning】、【18, Progressively stacking 2.0: A multi-stage layerwise training method for bert training speedup, 2020, arXiv preprint arXiv:2011.13635】、【15, Learning to grow pretrained models for efficient transformer training, 2023, arXiv preprint arXiv:2303.00980】、【12, Lemon: Lossless model expansion, 2023】、【13, Staged training for transformer language models, 2022, International Conference on Machine Learning】、【16, Gradmax: Growing neural networks using gradient information, 2022, arXiv preprint arXiv:2201.05125】、【17, Masked structural growth for 2x faster language model pre-training, 2024】。这些最先进的方法可以分为两类：第一类侧重于从现有神经元派生出新神经元【10, Net2net: Accelerating learning via knowledge transfer, 2015, arXiv preprint arXiv:1511.05641】、【11, bert2bert: Towards reusable pretrained language models, 2021, arXiv preprint arXiv:2110.07143】、【14, Efficient training of bert by progressively stacking, 2019, International conference on machine learning】、【12, Lemon: Lossless model expansion, 2023】、【15, Learning to grow pretrained models for efficient transformer training, 2023, arXiv preprint arXiv:2303.00980】，第二类侧重于独立初始化新参数【18, Progressively stacking 2.0: A multi-stage layerwise training method for bert training speedup, 2020, arXiv preprint arXiv:2011.13635】、【13, Staged training for transformer language models, 2022, International Conference on Machine Learning】、【16, Gradmax: Growing neural networks using gradient information, 2022, arXiv preprint arXiv:2201.05125】、【17, Masked structural growth for 2x faster language model pre-training, 2024】。基于这两条研究路线，本文总结了四种原子增长算子：
* (A) G_direct：直接复制和堆叠旧层（深度方式）或分割同一层内的神经元（宽度方式）。
* (B) G_learn：使用一个可学习的映射矩阵从现有参数生成扩展参数。
* (C) G_zero：将新参数设置为零。
* (D) G_random：随机初始化新参数。
这四种算子的图示见图2。G_direct和G_learn从现有神经元产生新神经元，而G_zero和G_random则独立初始化新参数。作者关注简单、普适的增长算子，因为LLM预训练计算成本高昂。算子的正式定义和与现有增长方法在设计上的差异，请参阅附录A。

标准化增长流程。为了公平比较这四种增长算子在LLM预训练中的效果，本文定义了一个标准的“单跳”增长过程，该过程包括两个训练阶段：增长前的小模型训练和增长后的大模型训练。首先，使用d个tokens训练小LLM，然后，使用算子G将其增长到目标LLM，非嵌入参数的增长因子为g，接着对大LLM进行D个tokens的持续预训练。该过程中有两个关键因素值得注意：增长因子g和用于基础模型训练的数据量d（可解释为“增长时机”）。本文进一步通过分别检查深度增长G↑和宽度增长G→来评估每个增长算子。具体来说，我们从在d = 10B tokens上训练的基础模型（400M LLM）开始，在两个方向上应用这四种算子，将它们扩展到1.1B的目标大小（增长因子g约等于4），然后继续训练额外的D = 97.5B tokens。LLM的架构配置和训练细节见附录B。

3.2 预训练1.1B LLM

实验结果与发现。本文报告了训练损失、八个标准Harness NLP基准测试以及平均准确率和加速比的结果，如图3所示。关键发现是，深度增长G↑相比于宽度增长G→和从头训练模型都表现出显著的加速效果，而令人惊讶的是，G→并未提供任何显著优势。在深度增长算子中，G↑_direct、G↑_learn和G↑_zero都优于基线和G↑_random。G↑_random在本研究中表现不佳，可能是因为它设计用于渐进式的“小步”增长【17, Masked structural growth for 2x faster language model pre-training, 2024】，而本文的统一方法使用单步增长。最值得注意的是，深度堆叠G↑_direct在所有增长算子中脱颖而出，在加速比、训练损失和几乎所有Harness评估指标上都超越了竞争对手。例如，与从头训练100B tokens的模型相比，G↑_direct实现了49.1%的显著效率提升。加速比的计算请参阅附录B.2。附录C提供了这些算子的更多实验，包括它们的损失训练和评估图。

4 深入研究深度堆叠（Gstack）

聚焦Gstack。经验证据表明，某些增长算子，特别是G↑_direct，与从头开始训练的基线方法相比，在LLM预训练中表现出令人印象深刻的加速效果。现在我们将注意力转向对G↑_direct进行更深入的审视。为了方便，我们此后将G↑_direct称为Gstack。Gstack算子的定义如下：
M_target = M_base ◦ M_base ◦ · · · ◦ M_base (g次)，其中M_base是一个用d个tokens训练的小基础模型，M_target是目标模型，g是增长因子。

本节旨在解决引言中概述的两个主要挑战（O2和O3）：1）在扩展场景下评估Gstack的性能，即更大的模型尺寸和更多的训练tokens；2）确定使用Gstack时的超参数，即增长时机d和增长因子g。

4.1 Gstack的扩展性

扩展Gstack的模型尺寸。我们的扩展实验涉及两种更大的模型尺寸：3B和7B。我们首先训练层数是我们目标层数四分之一的小模型（增长因子g = 4），使用10B tokens（d = 10B）。然后，我们对这两种尺寸的堆叠模型使用超过300B tokens（D = 300B）进行训练。图4和图5分别显示了训练3B和7B LLM 300B tokens时的损失，以及使用Harness评估器评估的NLP基准平均准确率。Gstack的加速效果在两种模型和所有评估指标上都是一致的。例如，对于3B模型，图4显示，与使用240B tokens训练的基线3B模型相比，Gstack在预训练中实现了54.5%的加速，NLP基准平均准确率提高了2.1。

Gstack在不同模型尺寸下的表现。当比较1B、3B和7B模型时，很明显Gstack的优势并未随着模型尺寸的增加而减弱，这意味着即使在更大的模型上，其加速效应也可以被利用。评估结果的详细信息，包括使用指令微调的评估，可以在附录D中找到。附录E将我们的基线与开源LLM Pythia和tinyLlama进行了比较。

扩展Gstack的训练Tokens数量。接下来，我们在另一个维度上评估堆叠算子的可扩展性——使用更多的tokens进行训练。这一点尤其重要，因为最近关于高效训练算法有效性的讨论引发了争论【22, No train no gain: Revisiting efficient training algorithms for transformer-based language models, 2023】，即某些策略可能初期学习更快，但在给予更多训练数据时，最终表现可能与普通训练方法相似或更差。因此，我们的目标是使用Gstack在大量的训练tokens上预训练一个LLM。具体来说，我们在一个410M的LLM上进行了使用750B tokens的实验。遵循上一节的实验设置，我们设置增长率g = 4和增长时机d = 10B，并在目标410M LLM上进行750B tokens的持续预训练。相比于chinchilla推荐的410M模型的8B tokens【4, Training compute-optimal large language models, 2022】，我们的实验设置也超过了这个值近100倍，达到了750B tokens。

长程训练下的Gstack表现。图6a的训练动态表明Gstack在这种情况下仍然有效。具有相似发现的评估结果细节可以在附录D.3中找到。基于LLM预训练的卓越稳定性【37, A survey of large language models, 2023】、【38, Megascale: Scaling large language model training to more than 10,000 gpus, 2024】，我们估算了损失改进并将其绘制在图6b中。拟合曲线表明，即使在8T tokens之后，Gstack仍将继续展现加速效应，这比推荐的token数量【4, Training compute-optimal large language models, 2022】长了1000多倍。值得注意的是，在8T训练后的这种损失改进对于LLM预训练并非微不足道，因为先前的研究【39, Understanding emergent abilities of language models from the loss perspective, 2024】表明，即使在后期阶段的微小改进也可能对下游性能产生相对实质性的影响。从LLM实践者的角度来看，考虑到“过训练”（即使用比缩放定律【3–5】推荐的更多数据来训练LLM）已成为普遍做法，这一点也至关重要。一个显著的例子是使用15T tokens训练LLama 3-8B，这几乎是chinchilla缩放定律【4, Training compute-optimal large language models, 2022】推荐的token数量的100倍。因此，这一发现为Gstack在整个实际LLM预训练过程中持续的卓越加速提供了信心。

估算缩放定律。为了进一步探索我们的发现，我们将四个模型（410M、1.1B、3B和7B）绘制在同一张图上，并尝试使用Gstack算子揭示我们的“缩放定律”。遵循【3, Scaling laws for neural language models, 2020】和【4, Training compute-optimal large language models, 2022】，我们使用方程L(C) = a * C^b来定义缩放幂律，其中a和b是我们需要拟合的常数，C代表FLOPs，L(C)表示在此FLOP下的模型最终损失。我们使用SciPy【40, Scipy 1.0: fundamental algorithms for scientific computing in python, 2020, Nature Methods】中的curve_fit函数来拟合从头训练模型和Gstack模型，并在图7中展示了估算的缩放定律。该图显示，我们的Gstack缩放定律相比于从基线LLM估算的缩放定律表现出更高的效率，在达到相同的目标损失时需要少得多的计算资源。然而，鉴于其他缩放定律研究【3, Scaling laws for neural language models, 2020】、【4, Training compute-optimal large language models, 2022】投入的大量计算资源，我们承认我们的Gstack缩放定律是受计算限制的初步估计，全面的研究留待未来进行。

4.2 确定使用Gstack的增长时机和增长因子

问题定义。我们在第4.1节全面验证了Gstack相比于从头训练的有效性。然而，要将Gstack纳入LLM的预训练过程，我们需要确定两个关键的超参数：增长时机（d）和增长因子（g）。在我们之前的实验中，我们依赖于对这些参数的临时选择，因此在使用Gstack时缺乏系统性的确定方法。虽然有研究调查增长时机【41, When to grow? a fitting risk-aware policy for layer growing in deep neural networks, 2024】，但其设置与LLM预训练大相径庭。因此，本节为希望在LLM预训练过程中优化使用Gstack算子的实践者提供明确的指导。

Gstack策略下的计算预算。我们首先给出一个正式的定义。在给定的计算预算C下，已建立的缩放幂律【3, Scaling laws for neural language models, 2020】、【4, Training compute-optimal large language models, 2022】可以指导非嵌入参数N和训练tokens数量D的选择，以在从头训练的情况下实现最低的模型损失。然而，当固定的预算C被分配用于寻找使用Gstack算子的最优模型训练策略时，调整超参数变得更加复杂，因为这涉及到两个训练阶段。因此，总计算预算C可以表示为两个部分的和：C = C1 + C2。这里，C1和C2分别代表训练初始小模型所需的flops C1 = FLOPs(n, d)，和训练大模型所需的flops C2 = FLOPs(N, D)，其中n和d表示小模型的参数和训练tokens，N和D表示大模型的参数和训练tokens。由于大模型是通过增长因子g增长的，使得N = gn，我们有 C = C1 + C2 = FLOPs(g, N, d) + FLOPs(N, D) = FLOPs(g, N, d, D)。

优化目标。因此，在给定预算C的情况下，我们的目标是找到优化的D, N, d, g值，以最小化损失L(D, N, d, g)。然而，同时优化上述四个变量的计算成本可能很高。因此，我们不搜索全局最优解，而是分别确定与Gstack密切相关的两个因素：小模型的训练tokens（增长时机）d和增长因子g。

确定增长时机 d。我们首先探究增长时机的影响，即小模型的训练token d。具体来说，我们将Gstack算子应用于一系列用d = 0B, 1B, 5B, 10B, 20B, 50B tokens训练的小模型。随后，我们以增长因子g = 4将它们堆叠到目标层数，并为一组固定的计算FLOPs进行训练。我们使用三种目标模型尺寸N = 410M, 1.1B, 3B重复上述实验，并将每组IsoFLOP点绘制在图8a、8b和8c中。令人惊讶的是，即使是仅用1B tokens训练的小模型，与直接堆叠随机初始化的小模型（表示为“0B”）相比，也显示出显著的加速。而0B的性能与从头训练的模型相似，这表明堆叠本身并不能作为一种有效的初始化方法。此外，通过应用平滑技术将模型IsoFLOP曲线建模为抛物线，我们确定了为每个FLOP计数最小化损失的最优d值，这引导我们假设存在一个涉及N、C和d的对数方程。

拟合d的预测公式。拟合后，我们得到a = 0.88, b = 163.27 和 c = -5.74，并在图9中绘制了等高线图。可以看出，我们估计的曲线与实际最优点吻合得很好。

确定增长因子 g。我们要确定的另一个因素是增长因子g。由于3B和7B参数的模型具有相同的深度，我们使用两种模型尺寸进行实验：1.1B（24层）和3B（32层）。具体来说，我们为1.1B模型改变堆叠因子为g = 2, 4, 8, 24，为3B模型改变为g = 4, 8, 16, 32，同时保持基础模型用d = 10B tokens进行训练。平滑的IsoFLOP曲线绘制在图10中。有趣的是，即使使用一个相对较浅的2层基础模型和g = 16的增长因子，我们也观察到与基线3B模型（g = 1）相比有显著的改进。然而，当使用1层基础模型时，Gstack的表现不及基线。我们的曲线表明，最优增长因子g位于2和4之间。然而，与确定训练token d不同，由于计算限制，我们无法生成足够的数据来估计N、C和g之间的关系。因此，这项工作建议使用一个恒定的增长因子g = 4。我们也在附录F中包含了我们初步估计的g的方程和等高线图。第4.2节的所有评估结果列在附录G中。

5 消融与讨论

目的。为了进一步深入了解在LLM预训练中采用模型增长技术，我们对Gstack的变体进行了消融实验，并讨论了一般模型增长技术中的函数保持问题。

5.1 消融实验：如何堆叠？

“单跳”堆叠 vs. 渐进式堆叠。值得注意的是，Gstack与StackedBERT【14, Efficient training of bert by progressively stacking, 2019, International conference on machine learning】中提出的算法不同，后者采用的是渐进式堆叠策略。因此，我们比较了我们的“单跳”Gstack和他们的渐进式堆叠方法。遵循StackBERT中介绍的方法，我们采用两步堆叠策略。给定我们1.1B、24层的目标模型尺寸，我们从一个6层模型开始。随后，我们在10B tokens上训练它，并通过堆叠将模型深度加倍，重复此步骤两次（训练-堆叠-训练-堆叠）以达到期望的规模。我们的实验表明，Gstack在损失和下游评估上均优于渐进式堆叠方法。例如，评估结果显示，在预训练大模型100B tokens时，Gstack比渐进式堆叠的平均准确率高2.4，Wikitext PPL好0.6。结果见附录H.1。我们进一步比较了其他堆叠变体，例如通过插值进行堆叠和部分堆叠某些层，这些方法也分别在LlamaPro【42, Llama pro: Progressive llama with block expansion, 2024, arXiv preprint arXiv:2401.02415】和Solar【43, Solar 10.7 b: Scaling large language models with simple yet effective depth up-scaling, 2023, arXiv preprint arXiv:2312.15166】中被采用，详细发现见附录H.2和H.3。

5.2 讨论：为什么函数保持会失败？

函数保持（FP）的定义。函数保持（FP）是大多数模型增长方法【10–12, 17】背后的一个关键概念。这个想法很直观，即一个更大的模型应该初始化能够表示与较小模型中相同函数的参数，即对所有输入x，f(x; Θ^(s)) = f(x; Θ^(l)_init)。我们在附录I.1中给出了FP的数学定义。

破坏函数保持的实验。我们发现，我们的Gstack方法违反了FP，却在本研究中成为最有效的方法，这很有趣。为了进一步调查，我们进行了一个简单的消融研究，通过在严格遵守FP的算子G→_direct上引入噪声来破坏FP。我们通过两组参数的加权组合来初始化新神经元：一组来自G→_direct，另一组来自随机初始化。权重因子由一个噪声比率控制。我们的发现很有趣。经过40B tokens的训练后，添加20%的噪声在Wikitext PPL上比原始的G→_direct高出0.27，在平均准确率得分上高出0.41。

在Gstack上添加噪声。我们还为Gstack添加了噪声。当我们添加20%的噪声时，我们的LLM表现略好于无噪声模型。然而，当噪声水平超过20%时，性能显著恶化。这些结果表明，函数保持可能不是模型增长的唯一决定因素。换句话说，探索加速大模型训练的方法和在增长过程中严格保持函数可能代表了两个重叠但又不同的研究方向。实验细节在附录I.2中提供。

A4 实验环境

• 数据集: 使用开源的 Slimpajama-627B 数据集【47, SlimPajama: A 627B token cleaned and deduplicated version of RedPajama, 2023, https://www.cerebras.net/blog/ slimpajama-a-627b-token-cleaned-and-deduplicated-version-of-redpajama】进行预训练。
• 模型架构: 采用类似Llama的Decoder-only Transformer架构。实验涉及的模型规模包括410M, 1.1B, 3B, 7B。具体的超参数（如层数、隐藏层维度、注意力头数）见下表。
  Table 1: Hyperparameters
  

  

• 硬件配置: 实验在NVIDIA Hopper集群上进行。7B模型的训练速度约为每天100B tokens。
• 软件配置:
  • 代码库: 基于TinyLlama【45, Tinyllama: An open-source small language model, 2024】的代码库进行预训练。
  • 并行与加速: 使用FSDP（Fully Sharded Data Parallel）进行分布式训练，并集成了FlashAttention 2.0【46, Flashattention: Fast and memory-efficient exact attention with io-awareness, 2022】等加速技术。

A5 实验结果

原子增长算子评估 (1.1B模型)

• 实验内容: 对四种原子算子（G_direct, G_learn, G_zero, G_random）在宽度（G→）和深度（G↑）两个增长方向上进行评估，将一个400M的基础模型（训练了10B tokens）增长到1.1B。
• 实验结果:
  • 深度增长（G↑）显著优于宽度增长（G→）和从头训练基线。
  • 在深度增长中，G↑_direct (即Gstack) 表现最佳，在训练损失和几乎所有8个NLP基准测试中都取得了最好的结果。
  • 与从头训练100B tokens相比，Gstack 实现了49.1%的训练速度提升 (图3)。
• 结论: 简单的深度堆叠 (Gstack) 是最高效的模型增长策略。

Gstack可扩展性评估

• 模型尺寸扩展 (3B & 7B):
  • 实验内容: 将Gstack应用于3B和7B模型的预训练中，基础模型层数为目标模型的1/4，在10B tokens上训练后堆叠，然后继续训练至300B tokens。
  • 实验结果: Gstack的加速效果在更大模型上保持稳定。对于3B模型，在240B tokens时加速54.5% (图4)；对于7B模型，在220B tokens时加速55.3% (图5)。在损失和下游任务准确率上均显著优于从头训练。
  • 结论: Gstack的优势不会随着模型规模增大而减弱。
• 训练数据扩展 (410M模型, 750B tokens):
  • 实验内容: 将一个410M模型使用Gstack训练超过750B tokens，远超Chinchilla定律建议的8B tokens。
  • 实验结果: Gstack的优势在长程训练中持续存在。在400B tokens时加速53.1%，在700B tokens时仍有31.0%的加速 (图6a)。通过对损失差异进行对数拟合，预测其优势可以持续到8T tokens以上 (图6b)。
  • 结论: Gstack在“过训练”场景下依然有效，符合当前LLM训练的实践趋势。
• 缩放定律估算:
  • 实验内容: 基于410M、1.1B、3B、7B四个模型的实验结果，分别拟合了从头训练和使用Gstack的缩放定律。
  • 实验结果: Gstack的缩放定律曲线显示，达到相同的目标损失所需的计算量（FLOPs）远低于从头训练 (图7)。
  • 结论: Gstack提供了一种更高效的计算扩展路径。

Gstack实践指南研究

• 增长时机 (d):
  • 实验内容: 在410M、1.1B和3B模型上，测试了不同增长时机 d (0B, 1B, 5B, 10B, 20B, 50B tokens) 对最终模型性能的影响。
  • 实验结果: 存在一个最优的增长时机d，过早或过晚增长都会影响效果。即使只用1B tokens训练基础模型也比不训练（d=0）要好得多。作者拟合出一个公式 d = a * log(N) * log(C) + c 来预测给定模型大小N和计算预算C下的最优d (图8, 图9)。
  • 结论: 增长时机d是一个关键超参数，可以通过公式进行估计。
• 增长因子 (g):
  • 实验内容: 在1.1B和3B模型上，测试了不同增长因子g (从2到32) 的影响。
  • 实验结果: 最优的增长因子g在2到4之间。过大的增长因子（如从1层或2层模型直接堆叠到32层）虽然仍优于基线，但不是最优选择 (图10)。
  • 结论: 建议使用 g=4 作为一个稳健的默认值。

消融实验与讨论

• 堆叠方式:
  • 实验内容: 比较了“单跳”堆叠、渐进式堆叠、插值堆叠和部分层堆叠。
  • 实验结果: “单跳”堆叠优于渐进式堆叠。保持原始层间连接（高连接率Rc）的堆叠方式（如整体堆叠）优于打乱连接的方式（如插值堆叠）。堆叠所有层（all）> 堆叠中间层（middle）≈ 堆叠后部层（back）≫ 堆叠前部层（front）。
  • 结论: 简单地将整个小模型复制堆叠是最佳策略。
• 函数保持 (Function Preserving):
  • 实验内容: Gstack本身不满足函数保持。作者通过在满足函数保持的G→_direct算子上添加噪声来主动破坏它。
  • 实验结果: 对G→_direct添加20%的噪声后，性能反而优于原始的严格函数保持版本。对Gstack添加少量噪声也有轻微提升。
  • 结论: 严格的函数保持可能不是模型增长成功的必要条件，甚至可能限制了模型的学习潜力。

A6 结论

本文对模型增长方法在高效LLM预训练中的应用进行了实证探索。通过系统性地解决当前研究中存在的三个关键挑战，本文得出了以下结论：
1. 方法论上，本文将复杂的模型增长技术归纳为四种原子算子，并通过全面的评估发现，一种简单的深度堆叠方法Gstack在各项指标上均显著优于其他方法及从头训练的基线。
2. 可扩展性上，本文通过对模型规模（最高达7B）和训练数据量（最高达750B tokens）的扩展实验，证明了Gstack的有效性是可扩展且持续的。
3. 实践应用上，本文系统地分析了Gstack算子的两个关键超参数——增长时机和增长因子，并将其形式化为一套指导方针，为实际应用提供了具体建议。
4. 深入洞察上，通过深入的讨论和全面的消融研究，本文揭示了Gstack的一些内在机制，为理解模型增长提供了新的视角。

总的来说，这项工作为在LLM预训练中应用模型增长技术提供了坚实的经验证据和实践指导，有望推动学术界和工业界在构建更高效、更可持续的AI模型方面取得进展。

A7 附录

A 增长算子细节

A.1 四种增长算子

• G_direct：从旧参数直接派生。这是一种直观的策略，通过直接复制或分割现有神经元来扩展网络【14, Efficient training of bert by progressively stacking, 2019, International conference on machine learning】、【11, bert2bert: Towards reusable pretrained language models, 2021, arXiv preprint arXiv:2110.07143】、【12, Lemon: Lossless model expansion, 2023】。宽度扩展（如Net2Net【10, Net2net: Accelerating learning via knowledge transfer, 2015, arXiv preprint arXiv:1511.05641】）涉及将旧神经元分割为多个部分，本文采用非均匀分割以解决对称性问题。深度扩展通常是复制层，即“堆叠”【14, Efficient training of bert by progressively stacking, 2019, International conference on machine learning】，本文称此操作为Gstack。
• G_learn：通过矩阵变换生成新参数。该算子学习一个矩阵变换函数，将小模型映射到一个更大的模型【15, Learning to grow pretrained models for efficient transformer training, 2023, arXiv preprint arXiv:2303.00980】。它适用于宽度和深度扩展。优化目标是最小化语言模型损失L，其中超网络G_learn用于元学习，参数从θ映射到Θ。
  

  

• G_zero：将新参数设为0。这是一种实现函数保持的简单方法。为避免优化难题，本文采用选择性置零策略【13, Staged training for transformer language models, 2022, International Conference on Machine Learning】、【16, Gradmax: Growing neural networks using gradient information, 2022, arXiv preprint arXiv:2201.05125】、【12, Lemon: Lossless model expansion, 2023】。宽度增长时，仅将新神经元的扇出（fan-out）参数置零；深度扩展时，将新复制的Transformer块的MultiHead Attention和MLP的最终输出层置零。
• G_random：随机初始化新参数。这种方法遵循常规的随机初始化新参数的做法。为了保持函数不变性，近期研究引入了对新神经元的掩码【17, Masked structural growth for 2x faster language model pre-training, 2024】，该掩码在持续训练中逐渐移除。本文将这种新方法称为G_random算子。

A.2 本文算子与基础方法的区别

• G→_direct 与Lemon的设置相似，但因Llama架构有细微差异。
• G_learn 与LiGO方法一致，但为本文自己实现。
• G_zero 在深度上与Lemon一致，但在宽度上不同于StagedTraining。本文方法不是将宽度加倍并置零非对角线元素，而是通过置零左下角子矩阵来扩展到任意维度，更具灵活性。
• G_random 采用“单跳”增长直达目标尺寸，而非MSG的“多跳”式增长。
• G↑_direct (Gstack) 的实现与stackedBert中采用的渐进式增长技术不同，本文算子遵循更直接的方法。

A.3 G_direct 细节

• Embedding: 对于嵌入矩阵 $E \in R^{V \times d}$，扩展到 $E' \in R^{V \times D}$，G_direct通过复制某些列来实现。
  

  

• Linear: 对于线性层权重 $W \in R^{d_{out} \times d_{in}}$，扩展到 $W' \in R^{D_{out} \times D_{in}}$，G_direct 定义为：
  

  

  
  为了满足函数保持，需要确保 $\alpha + \beta = I$。
• RMSNorm: 对于RMSNorm的参数 $\mu \in R^d$，扩展到 $\mu' = \sqrt{d/D} [\mu, \mu_{0, D-d}] \in R^D$。
  

  

  
  因此，使用G_direct无法实现RMSNorm的函数保持。
• Depth (Gstack): 对于一个l层的transformer $F = f_0 \circ f_1 \circ \dots \circ f_l$，目标是扩展到L层（L是l的倍数）。直接堆叠的形式为 $F' = F \circ F \circ \dots \circ F$。

Algorithm 1 Operator Gstack
Input: Base model M_l^k with l layers trained using dataset d_k where k is iteration steps. Growth factor g.
Output: Target Model M_{gl}^0 with gl layers
M_0^0, M_1^0, ..., M_{l-1}^0 <- M_0^k, M_1^k, ..., M_{l-1}^k
for i = 1 to g-1 do
    M_{il}^0, M_{il+1}^0, ..., M_{(i+1)l-1}^0 <- M_0^k, M_1^k, ..., M_{l-1}^k
end
return M_{gl}^0

A.4 G_zero 细节

• Embedding: 嵌入矩阵 $E \in R^{V \times d}$ 扩展为 $E' = [E, O] \in R^{V \times D}$，其中 $O$ 是零矩阵。
• Linear: 权重矩阵 $W \in R^{d_{out} \times d_{in}}$ 扩展为 $W'$，其中新参数A和C是随机初始化的，但通过输入侧的补零，可以实现函数保持。
  

  

• RMSNorm: 参数 $\mu \in R^d$ 扩展为 $\mu' = [\alpha\mu, \xi]$，类似于G_random。
• Depth: 通过仅保留残差部分并将MHA和SwiGLU的最终线性投影初始化为零，可以实现函数保持。

A.5 G_random 细节

• Embedding: 嵌入矩阵 $E \in R^{V \times d}$ 扩展为 $E' = [E, E_{rand}] \in R^{V \times D}$，其中 $E_{rand}$ 是随机初始化的新参数。使用掩码 $c$ 来屏蔽随机初始化的部分以实现函数保持。
• Linear: 权重矩阵 $W$ 扩展时，新参数A, B, C被随机初始化。同样使用掩码来确保函数保持。
• RMSNorm: 为了实现函数保持，扩展后的参数 $\mu'$ 和缩放系数 $\alpha$ 需满足特定关系，最终 $\alpha = \sqrt{d/D}$。
  

  

• Depth: 通过仅保留残差部分并对MHA和SwiGLU层进行掩码操作，可以实现函数保持。

A.6 G_learn 细节

G_learn通过学习一系列映射矩阵 $B$ 来将小模型的参数（如 $W_Q, W_K, \dots$）变换为大模型的参数（如 $W'_Q, W'_K, \dots$）。
* 宽度扩展: 对Embedding, Attention, MLP, Output Head等层的权重矩阵应用不同的可学习变换矩阵。
* 深度扩展: 对于一个L1层的Transformer，通过可学习的参数 $D_{Q,K,V,O,ln1} \in R^{L2 \times L1}$ 和 $D_{up,down,gate,ln2} \in R^{L2 \times L1}$ 将其扩展到L2层。

H 关于“如何堆叠？”的讨论与评估结果

H.2 消融实验：f2◦f1◦f0◦f2◦f1◦f0 vs. f2◦f2◦f1◦f1◦f0◦f0 (插值)

连接率对堆叠性能的影响。为了研究层间连接是否影响堆叠性能，本文比较了两种将小模型堆叠成大模型的方法。第一种方法是将整个小模型作为一个单元直接堆叠，这能保留大部分层间连接。第二种方法是复制并交错小模型中的每一层，这几乎破坏了所有连接。为量化层间连接的保留程度，定义了连接率Rc：

实验设置与结果。实验中，将一个8层小模型堆叠成24层目标模型，增长时机d为10B tokens，增长因子g为3。Gstack的Rc为91.3%，插值堆叠（G_interpolate）的Rc为30.4%。训练损失和NLP基准平均准确率如图35所示。训练初期，插值堆叠表现与整体堆叠相当，但随着训练进行，其性能恶化。

结论。可以得出结论，堆叠的连接率越高，堆叠效果越好。附录H.3继续验证了这一结论。

H.3 消融实验：部分堆叠

部分堆叠的研究背景。部分堆叠已在LlamaPro【42, Llama pro: Progressive llama with block expansion, 2024, arXiv preprint arXiv:2401.02415】、Solar【43, Solar 10.7 b: Scaling large language models with simple yet effective depth up-scaling, 2023, arXiv preprint arXiv:2312.15166】等LLM中进行了探索。但它们的目标是堆叠一个现成的LLM（如Llama2），而本文旨在加速LLM预训练过程。

实验设计与结果。为了探究堆叠小模型的哪些层能达到最佳性能，本文进行了部分堆叠的实验。实验中，将一个6层小模型（{L1, L2, · · · , L6}）堆叠成一个24层目标模型，增长时机d=10B tokens，增长因子g=4。使用 1-234*7-56 这样的格式表示堆叠 234 层7次。

结论。表5总结了八种部分堆叠方法的Rc和堆叠部分。对于部分堆叠，结论是：堆叠所有层 > 堆叠中间层 ≈ 堆叠后部层 ≫ 堆叠前部层。同时，当堆叠部分相同时，Rc越大，性能越好。

I 函数保持细节

I.2 通过添加噪声破坏函数保持

在G→_direct上添加噪声以破坏FP。对于SwiGLU中的下投影和MultiHeadAttention中的输出投影，我们应用噪声：

在Gstack上添加噪声。既然添加噪声能提升G_direct的性能，我们也在Gstack上添加了噪声。我们将一个8层小模型堆叠到24层，然后添加α = 0.2的噪声。我们报告了训练损失和标准NLP基准平均准确率在图41中。添加噪声在训练损失上显示出优势。