Lory: 推进大型语言模型训练的新篇章

人工智能咨询培训老师叶梓转载标明出处

随着模型规模的增长，如何有效训练并利用这些模型成为了一个挑战。陈丹琦团队一项新的研究提出了一种创新的预训练方法——Lory，旨在解决大模型在混合专家（MoE）架构中的可微分性和计算效率问题。Lory通过专家合并技术，实现了模型的完全可微分，并通过因果段路由和基于相似性的数据批处理策略，显著提高了训练效率和模型性能。这项工作不仅推动了大模型的研究边界，也为未来更高效、更强大的AI系统的发展奠定了基础。

方法

Lory方法提出了一种可微分的MoE（Mixture of Experts）架构，专门为自回归语言模型设计。这一架构的核心在于专家合并技术，该技术允许模型在保持自回归特性的同时进行有效的训练。Figure 1展示了Lory的整体架构，其中包括两个关键技术：因果段路由策略和基于相似性的数据批处理方法。

因果段路由策略是Lory方法的首个关键技术。这一策略的核心思想是在处理语言模型时，不是对每个标记单独进行专家合并，而是在段级别上进行。这样做的好处是显著减少了合并操作的次数，从可能的L次减少到N次，其中N是将输入序列分割成的段的数量。这种策略不仅减少了计算量，而且通过仅使用前一段的信息来指导当前段的专家合并，有效地防止了信息泄露的问题。在Figure 1中，每个段的专家合并都是基于前一段的平均隐藏表示来计算的，这种设计确保了模型的自回归特性得以保持。

基于相似性的数据批处理方法是Lory方法的第二个关键技术。这一方法的目的是通过构建具有高相似性的段来提高专家在特定领域的专业化能力。在传统的数据批处理方法中，文档通常是随机连接的，这可能导致相邻段之间的标记来自完全不同的文档，从而影响了专家的专业化。Lory方法通过顺序连接相似的文档来构建训练实例，这种方法鼓励了相邻段之间的高相似性，使得专家能够更加专注于特定的领域或主题。在Figure 1基于相似性的数据批处理方法的示意图中，其中相似的文档被有意识地组合在一起，形成了具有高相似性的段，从而促进了专家的专业化。

通过结合这两个关键技术，Lory方法不仅提高了训练效率，还增强了模型在特定领域的专业化能力。

想要掌握如何将大模型的力量发挥到极致吗？2024年10月26日叶老师带您深入了解 Llama Factory —— 一款革命性的大模型微调工具。

留言“参加”即可来叶老师的直播间互动，1小时讲解让您轻松上手，学习如何使用 Llama Factory 微调模型。

实验

实验设置

模型：实验中训练了两种规模的解码器仅Transformer模型，分别含有0.3B和1.5B个活跃参数。对于每个FFN层，用具有相同架构的MoE层替换，其中E代表专家的数量，取值为8、16或32。模型配置和总参数计数详见附录D。所有模型使用4096-token的上下文窗口进行训练，并且采用因果段路由策略，每段长度设置为T=256。

训练细节：使用AdamW优化器进行训练，其中β1=0.9和β2=0.95，学习率设置为2e-4，并采用余弦学习率调度器。所有模型的批量大小为100万个token。采用ZeRO优化进行分布式训练。训练初期，使用与参数匹配的密集模型，并复制FFN层作为MoE模型的初始化。实验中，使用前5%的训练步骤作为预热来初始化MoE权重。此外，对学习率调度器的前5%训练步骤应用线性预热。训练过程中使用了多达64个A100 GPU。

训练数据集：从Commoncrawl数据集中随机抽取150亿个token作为训练数据。使用基于相似性的数据批处理方法构建所有训练实例。

评估数据集：在从arXiv、Books、Wikipedia、C4和Python代码（Github的Python子集）中抽取的保留评估数据集上，通过测量训练模型的困惑度来评估所有模型在语言建模任务上的性能。每个评估数据集包含1K个样本，每个样本由4096个token组成。此外，还在下游任务中评估了模型的性能，包括常识推理、阅读理解、闭卷问答和文本分类。

结果

训练效率和收敛性：Figure 2（左）展示了密集模型和具有不同模型大小的MoE模型的训练损失曲线。首先，可以发现在相同的训练token数量下，Lory模型明显比密集模型基线有更好的训练损失。对于0.3B和1.5B模型，具有32个专家的模型在训练token数量不到一半的情况下达到了相同的损失水平。这表明Lory方法在相同的训练计算量下实现了更好的性能。还观察到，使用更多的专家可以获得更多的改进。

语言建模：在语言建模任务上的评估表明，MoE模型一致性地优于密集基线，显著降低了所有领域的困惑度。例如，0.3B/32E模型在Books上的困惑度比0.3B密集模型提高了13.9%。值得注意的是，在与训练数据不同的测试领域（例如Python）中，改进最为显著，这表明了强烈的专家专业化，这在第5节中进一步探讨。

下游任务：Table 1展示了模型在下游任务上的性能。可以观察到在所有任务上都有显著的性能提升。例如，0.3B/32E模型在常识推理上平均提高了3.7%，在阅读理解上提高了3.3%，在闭卷问答上提高了1.5%，在文本分类上提高了11.1%。

消融研究和分析

在前缀路由策略中，仅根据第一个段来执行整个序列的专家合并，然后使用合并的FFN来处理序列的其余部分，而不再进行更新。Figure 3显示，仅使用前缀进行路由会导致性能大幅下降，与因果段路由相比，后者在每个段级别上提供更强大的训练信号，从而获得更好的性能。

为了研究基于相似性的数据批处理方法的重要性，比较了使用和不使用这种方法的MoE模型相对于密集模型的性能提升。Figure 4（左）展示了使用相似性批处理（sim batch）和随机批处理（rand batch）数据的密集（0.3B）和MoE模型（0.3B/8E）的训练损失。在两种设置中，MoE模型均优于密集模型。然而，使用相似性批处理时，损失改进（即密集模型和MoE模型之间的损失差异）更大，并且随着训练数据的增加，这种效果更加明显（Figure 4（右））。这些结果强烈支持了基于相似性批处理对于有效训练MoE模型的重要性。

将Lory方法与现有的Expert Choice (EC) MoE方法进行了比较，EC方法通过让每个专家根据路由权重选择top-k输入来确保训练期间的均衡负载。在推理期间，为了避免利用全局信息进行路由，每个标记被路由到top-k专家中。Figure 5展示了Lory方法与EC方法（包括段级别和标记级别路由）的训练损失曲线。Lory方法在相同的路由设置下显著优于段级别的EC，表明完全可微分的架构比使用相同路由策略的稀疏MoE更有效。与标记级别的EC模型相比，尽管Lory使用方法使用段级别路由且不需要任何高级训练技术，但Lory取得了有竞争力的结果。这些结果突出了Lory方法的显著潜力。

在分析中还探讨了专家利用率和专业化问题。利用率方面，尽管没有使用任何辅助损失来平衡负载，Lory方法仍然能够实现高专家利用率，防止MoE模型退化为密集模型。专业化方面，通过研究0.3B/8E模型在不同层的平均路由权重，发现即使在没有额外的领域级监督的情况下，训练出的MoE模型中也存在清晰的领域级专家专业化。例如，第11层的专家7专门处理arXiv领域的输入。此外，还观察到arXiv和Python代码的路由权重更为相似，而与Books和Wikipedia的相似度较低，这可能是因为LaTex代码和Python代码与自然语言的差异较大。