引言

在当今的人工智能领域，缩放定律 (Scaling Laws) 占据了统治地位: 如果你想要一个更智能的模型，就把通过做大。模型参数量已经从数百万膨胀到数十亿，现在更是达到了万亿级别。然而，我们正在撞上一堵墙。对于许多研究人员和应用来说，运行这些庞大的稠密模型所需的纯粹计算成本正变得不可持续。

混合专家 (Mixture-of-Experts, MoE) 架构应运而生。MoE 承诺了两全其美: 它将模型的总大小 (容量) 与其计算成本 (推理延迟) 解耦。通过在处理每个 token 时仅激活网络的一小部分，MoE 模型可以拥有巨型模型的知识容量，同时保持小得多的模型的运行速度。

但这其中有个陷阱。从头开始训练一个 MoE 模型是出了名的困难。它受到不稳定性、负载平衡问题的困扰，并且需要海量的数据才能正确收敛。

如果我们不必从头开始呢？如果我们能拿一个强大的、预训练好的稠密模型——比如 LLaMA-2——并像做手术一样将其转化为一个稀疏的混合专家模型呢？

这正是论文 “LLaMA-MoE: Building Mixture-of-Experts from LLaMA with Continual Pre-training” (LLaMA-MoE: 通过持续预训练基于 LLaMA 构建混合专家模型) 所探讨的内容。研究人员提出了一个新的框架，将稠密的 LLaMA-2 7B 模型“升级改造 (upcycle) ”为高效的 LLaMA-MoE 模型。在这篇文章中，我们将拆解他们的方法，探索他们如何将神经网络拆分为专家，并分析这个“弗兰肯斯坦式”的拼接模型是否真的比其稠密原型表现更好。

背景: 稠密 vs. 稀疏

在深入探讨方法之前，让我们简要阐明一下架构上的转变。

在传统的 稠密 Transformer (如原始的 LLaMA 或 GPT-3) 中，前馈网络 (FFN) 层中的每一个参数都用于处理每一个 token。如果你向模型输入单词“apple”，整个大脑都会被点亮。

在 稀疏混合专家 (MoE) 中，FFN 层被一组“专家 (Experts) ”所取代。对于任何给定的 token，一个“路由器 (Router) ” (或门控) 决定哪些专家最适合处理它。如果 token 是“apple”，路由器可能会将其发送给专注于食物或物体的专家，而忽略专注于代码或语法的专家。

该论文的目标就是利用 LLaMA-2 已经学到的权重，将前者转化为后者。

核心方法: 对语言模型进行手术

从 LLaMA 到 LLaMA-MoE 的转变涉及两个关键阶段: 专家构建 (Expert Construction) 和 持续预训练 (Continual Pre-training) 。

图 1: 构建 LLaMA-MoE 模型的主要框架。(a) LLaMA 中的原始 FFN 被分割成不同的专家。(b) 在转换后的 LLaMA-MoE 中，隐藏状态由部分选择的专家处理，而不是所有专家。

如图 1 所示，该过程首先获取原始 FFN (a) 并将其拆分为多个独立的专家 (b)。一旦结构发生变化，模型就需要进行进一步的训练以适应其新的稀疏特性。

第一阶段: 专家构建

这项研究最具挑战性的部分是决定如何切割现有的矩阵。LLaMA 模型在其 FFN 中使用 SwiGLU 激活函数。在数学上，FFN 层的输出 \(y\) 定义为:

定义带有 SwiGLU 激活函数的 FFN 输出的公式。

这里，\(W_{up}\)、\(W_{gate}\) 和 \(W_{down}\) 是承载模型知识的巨大权重矩阵。为了创建一个 MoE，研究人员将这些矩阵分割成 \(N\) 个更小的部分。如果原始隐藏维度大小为 \(d_h\)，并且我们想要 \(N\) 个专家，那么每个专家的隐藏大小将为 \(m = d_h / N\)。

第 \(j\) 个专家的参数是基于一组索引 \(S_j\) 从原始矩阵中切分出来的:

显示如何为特定专家 j 切分权重的公式。

关键问题是: 哪些神经元去哪个专家那里?

作者研究了四种不同的策略来划分这些参数。这是神经网络的一个典型“先天与后天”的实验——精心的聚类重要吗，还是随机就足够了？

1. 神经元独立策略 (Neuron-Independent)

这些方法将划分视为一个简单的集合分割问题。

  • 独立 (随机) : 中间神经元的索引 \(\{1, 2, ..., d_h\}\) 被打乱并随机分成 \(N\) 个大小相等的集合。这是最简单的方法。
  • 独立 (聚类) : 研究人员对 \(W_{up}\) 矩阵的行向量使用平衡 k-means 聚类。其想法是将具有相似激活模式的神经元归类到同一个专家中。

2. 神经元共享策略 (Neuron-Sharing)

受网络剪枝的启发，这些方法假设某些神经元比其他神经元更“重要”，也许应该在专家之间共享或区别对待。他们根据神经元被剪枝时的损失变化 (泰勒展开) 来计算神经元的重要性。

基于损失梯度更新神经元重要性向量 v 的公式。

  • 共享 (内部) : 神经元被复制。如果一个神经元对多个潜在的聚类都非常重要，它会被复制到多个专家中。
  • 共享 (外部/间) : 一组最关键的神经元被留作“共享残差块”，始终处于激活状态，而其余不太关键的神经元则分配给各专家。

门控机制

一旦专家建立起来，就需要一种机制将 token 路由给他们。该模型使用 Top-K 门控网络。对于每个输入 \(x\)，门控计算一个分数并选择最好的 \(k\) 个专家。

所选专家的 MoE 输出求和公式。

第二阶段: 持续预训练

经过“手术”后，模型实际上处于脑损伤状态。连接被切断并重新组织。虽然权重是从 LLaMA 初始化的，但新的门控网络是随机初始化的。为了恢复模型的语言能力，研究人员必须进行持续预训练 。

他们使用了 SlimPajama 数据集 (RedPajama 的清洗版本) ，包含 6270 亿个 token。他们研究的一个主要部分集中在数据采样上——即决定给模型输入多少来自不同领域 (维基百科、GitHub、书籍等) 的数据。

他们实验了两种主要方法:

  1. 静态权重: 在整个训练过程中固定数据比例 (例如，遵循原始 LLaMA 配方或“Sheared-LLaMA”配方) 。
  2. 动态权重: 动态调整数据混合比例。如果模型在编码任务上表现挣扎，采样器会自动增加 GitHub 数据的比例。

图 5: 四个领域的数据采样权重变化。动态策略随时间变化，而静态策略保持平坦。

如图 5 所示，动态策略 (绿线和青线) 在训练期间急剧改变数据消耗，而静态策略 (蓝线和红线) 保持恒定。

实验与结果

研究人员训练了几个变体，主要关注 LLaMA-MoE-3.5B 。 需要注意的是，这里的“3.5B”指的是激活参数 (推理成本) ，而总参数量仍然高得多 (与原始 7B 相似) 。

他们将自己的模型与其他具有相似激活参数量的开源模型进行了比较，例如 Open-LLaMA-3B 和 Sheared-LLaMA-2.7B。

哪种构建方法胜出了？

这可能是论文中最令人惊讶的发现。尽管对相似神经元进行聚类或共享重要神经元背后的逻辑很复杂，但 随机划分 (独立-随机) 的表现最好。

图 3: 不同专家构建方法的模型性能。独立 (随机) 获得了最好的结果。

在图 3 中，看那条蓝线 (独立随机) 。与聚类 (绿色) 或共享 (红色/青色) 相比，它始终能达到更高的准确率 (左图) 和更低的损失 (中图) 。

为什么? 作者认为，由于门控和专家在持续预训练阶段是同时训练的，复杂的划分可能会在构建时引入偏差，使得门控更难学习正确的路由。随机性提供了一个平衡的起点，允许模型更有效地自组织。

性能分析

最终模型与竞争对手相比如何？

图 2: ARC-c 和 HellaSwag 数据集上的模型性能以及训练损失。

图 2 展示了 LLaMA-MoE 模型在 2000 亿 token 训练过程中的轨迹。我们看到在 ARC-c 和 HellaSwag 等推理任务上的准确率稳步上升。

至关重要的是, LLaMA-MoE-3.5B 显著优于同等大小的稠密模型。 例如，在 ARC-Challenge 上，它的得分为 44.2 , 击败了 Open-LLaMA-3B (40.1) 和 Sheared-LLaMA-2.7B (41.6)。这验证了核心假设: 你可以通过将 7B 模型升级为 MoE，将其智能压缩进 3.5B 模型的推理预算中。

相比从头训练的优势

这篇论文的主要卖点之一是效率。升级 LLaMA 真的比从随机初始化开始训练一个新的 MoE 模型更好吗？

图 7: LLaMA-MoE-3.5B (2/8) 与从头开始训练 MoE 的模型性能对比。

答案是毫无疑问的肯定。 图 7 将 LLaMA-MoE (蓝色) 与从头开始训练的模型 (红色) 进行了比较。LLaMA-MoE 模型在准确率和更低损失方面拥有巨大的先发优势，并且在整个训练过程中保持了这种领先地位。“从头开始”的模型很难追赶上来，这证明了原始 LLaMA 权重中嵌入的知识被成功保留并在 MoE 版本中得到了利用。

深入探究: 专家真的会专业化吗？

MoE 模型一个迷人的方面是可解释性。我们通常希望专家能够“专业化”——例如，专家 1 成为“语法专家”，专家 2 成为“编码专家”。这种情况在 LLaMA-MoE 中发生了吗？

研究人员可视化了网络不同层的路由统计数据。

图 8: 第 1、8、28 和 32 层的专家路由统计。

图 8 展示了不同数据领域 (CommonCrawl, Wikipedia, arXiv, GitHub) 的专家激活热力图。

  • 浅层 (第 1 层) : 所有领域的热力图看起来非常相似。这表明早期层处理语言的通用特征 (句法、基本构词) ，这些特征适用于从书籍到代码的所有内容。
  • 深层 (第 32 层) : 在这里，我们看到了明显的模式。对比 GitHub 列和 CommonCrawl 列。“热” (深红色) 的专家是不同的。这表明在深层，专家确实已经专业化了。一些专家变成了“程序员”，而另一些更喜欢普通的网络文本。

通过观察最后一层中领域之间的路由相似性，进一步分析了这种专业化。

图 9: 第 32 层的专家路由差异。

在图 9 中，我们可以看到像 StackExchange 和 GitHub (图表右下角) 这样的技术领域共享相似的路由模式 (由较浅的颜色/较低的距离表示) ，而它们与像 Books 这样的通用文本有显著不同。这证实了模型正在根据语义上下文智能地路由 token。

结论与启示

LLaMA-MoE 论文为混合专家模型的普及化提供了一个令人信服的蓝图。它证明了我们不需要 Google 或 OpenAI 那样的海量计算资源来从头训练 MoE。相反，我们可以通过“升级改造”现有的高质量稠密模型，站在巨人的肩膀上。

主要收获:

  1. 随机性法则: 令人惊讶的是，随机分割 FFN 是升级改造时初始化专家最有效的方法。
  2. 效率: LLaMA-MoE 在相同的推理成本下实现了比稠密模型更好的性能，并且收敛速度远快于从头训练的 MoE。
  3. 专业化: 由此产生的模型在深层表现出明显的专家专业化，能够适应代码与散文等不同的领域。

这种方法为创建更小、更块、更专业的模型打开了大门，无需昂贵的预训练成本，使得高性能大语言模型在资源受限环境中的部署变得更加触手可及。