本文基于论文《Patch-Level Training for Large Language Models》，论文作者全部来自微信翻译团队。
论文: https://arxiv.org/pdf/2407.12665
代码: https://github.com/shaochenze/PatchTrain
导语近年来，大语言模型（Large Language Models，LLMs）以其卓越的语言理解和生成能力引领着人工智能技术的前沿。这些模型之所以能够实现如此令人瞩目的性能，很大程度上归功于它们庞大的模型参数量和对应的训练数据规模。然而，随之而来的是模型训练成本的急剧上升。据报道，OpenAI今年的模型训练开支预计将高达30亿美元。Anthropic的首席执行官更是预测，在未来三年内，大模型的训练成本将上升到100亿美元至1000亿美元之间。这一趋势清晰地表明，在下一代大模型的研发中，其高昂的训练成本将成为亟待解决的主要瓶颈。
具体而言，训练大模型所需的计算量（FLOPs）与模型参数量N以及训练数据中的文本单元（即token）数量D近似成正比，可由下式表示：
因此，训练成本的优化可以从减小模型参数量N和文本单元数D两方面入手。模型成长法（Model Growth）便是一类从参数量着手的方法，这类方法不是在一开始就为模型配置全部参数，而是通过在训练过程中逐步扩大模型的参数规模来减少整体训练成本。然而，从根本上说，模型的性能依赖于充足的参数量来存储丰富的知识和提升复杂推理的能力。当参数量不足时，模型所能习得的知识与能力也是有限的，因此很难达到与全参数训练相似的效果。
那么，对于训练数据中包含的文本单元数量，我们是否有优化的空间呢？在传统的训练框架下，模型通过预测每个token并更新对应的参数来从训练数据中学习。这种方法虽然简单直观，但对于参数量庞大的LLMs而言，其效率却不尽人意。训练数据蕴含的所有知识稀疏地分布在数以万亿计的token中，而每个token所能关联到的相关参数仅占模型总参数的一小部分。这种稀疏性意味着，尽管模型在每次学习过程中投入了大量的计算资源，但实际上只有一小部分参数得到了有效更新，从而导致了训练效率的低下。
这一问题的核心在于，现有的训练方法是基于token单元的，每个token的处理都需要模型的完整参与。因此，信息在token上的稀疏分布严重影响了训练的效率。如果能够找到一种方式，增大模型在每次学习过程中接触和处理的信息量，即增大每个文本单元的信息密度，那么就有可能在不牺牲模型性能的前提下，显著提高训练的效率。这也正是本文将要介绍的方法：将多个token聚合成一个更高信息密度的单元——我们称之为“patch”，并以此作为训练语言模型的基本单元，最后使用少量训练数据将模型调整回token级别。如下图所示，这一方法的损失函数为橙色曲线，其训练成本仅为传统方法的一半，同时还能达到无损甚至更高的性能水平。
方法介绍我们将多个token聚合成为一个更高信息密度的单元——我们称之为“patch”——来作为训练语言模型的基本单元，训练流程如下图所示。取patch大小为固定值K，我们将连续的K个token聚合为一个patch，patch embedding直接取为这些token embedding的均值。模型的训练目标是预测下一个patch，即后续的K个token，因此模型的输出需要同时与这K个token都计算交叉熵损失。
上述流程没有引入任何额外的模型参数，与token级别的训练流程在本质上保持一致，区别仅在于文本的基本单元从token变为了patch，这种设计旨在减小模型过渡到token级别时受到的阻力。我们在比例为λ的训练数据上进行patch级别的训练，随后在余下1-λ比例的数据上将模型调整回token级别，整体的训练成本为传统训练方式的λ/K+1-λ倍。在实践中，我们通常设置K=4、λ=2/3，此时训练成本将相比传统的token级别训练方法降低一半。上述方法的具体实现可以参考https://github.com/shaochenze/PatchTrain，核心代码仅如下10行，易于在各种主流训练框架中实现与应用。
实验效果我们选定训练数据为包含360B tokens的pile数据集，训练4种不同参数规模的Transformer模型，在下表中对比了两种训练方式的表现。结果显示，patch级别训练带来的效率提升几乎是无损的，在将训练成本减半的同时，模型在困惑度上与传统训练方法持平，在准确率上甚至实现了稳定的性能提升。我们认为，这种性能提升可能源自两个方面：首先，patch级别的训练引入了一定程度的正则化效果；其次，通过压缩序列长度，模型可能更加擅长捕捉长距离的依赖关系。
相比我们的实验设置，业界主流的大模型在模型和训练数据的规模上都远远更大。因此，我们进一步对所提方法在在不同训练规模下的性能变化趋势进行观察，以评估其在更大训练规模下的可扩展性。下面两表分别展现了当模型和数据规模扩大时，两种训练方法的性能变化对比。结果显示，与传统训练方法相比，patch级训练的性能随模型参数量的增长略慢，但随训练数据量的增长却显著更快。这可以从知识迁移的角度解释：模型参数量越大，则模型从patch级迁移到token级的过程越为复杂；而当训练数据量增加时，可以利用更多数据来完成模型的调整，从而降低了迁移难度。
最后，我们重新审视最初提出的问题 —— 将K个token聚合成一个更高信息密度的patch单元，能否克服信息稀疏性，让更多的模型参数得到有效更新？下图展示了在训练过程中，模型各层神经元的激活比例受patch大小K的影响。如图所示，神经元的激活数确实随K值增大而显著增加，在模型的浅层部分尤为突出，这是模型在patch级别训练效率更高的一个有力佐证。
结语训练数据蕴含的知识稀疏地分布在数以万亿计的token中，而训练的目标便是将这些知识嵌入到模型参数中。现行的训练机制需要逐个处理所有token，在整个模型中寻找与之相关的部分进行参数更新，这一过程因token中信息的稀疏性而显得相对低效。本文提出将多个token聚合成一个更高信息密度的patch单元来训练语言模型，这种训练方式带来的效率提升几乎是无损的，在将训练成本降低一半的同时，仍能保持与传统训练方法持平甚至更优的性能表现。从理想角度看，将数据中的知识搬运到模型中应是一个近乎线性复杂度的过程，我们相信这其中仍存在可观的优化空间。
