新智元报道
编辑：编辑部
MoE 已然成为 AI 界的主流架构，不论是开源 Grok，还是闭源 GPT-4，皆是其拥趸。然而，这些模型的专家，最大数量仅有 32 个。最近，谷歌 DeepMind 提出了全新的策略 PEER，可将 MoE 扩展到百万个专家，还不会增加计算成本。
如果你熟悉当前 LLM 的主流架构，混合专家（MoE）技术想必是老朋友之一。有人甚至会说，MoE 是使大模型崛起的关键因素之一。
开源的 Mixtral、DBRX、Grok 等模型都使用了 MoE，而且根据 Soumith Chintala 等大佬的推测，GPT-4 也是一个规模为8×220B 的 MoE 模型。
类似 GPT-4，多数 MoE 模型都会将专家数量限制在较少数量，一般不会超过 16 或 32。
然而，DeepMind 研究科学家 Xu Owen He 最近就在一篇独立发表的论文中，提出了一种全新的方法——PEER（参数高效专家检索，Parameter Efficient Expert Retrieval），可以将专家数量扩展到百万数量级。
论文地址：https://arxiv.org/abs/2407.04153
这究竟是如何做到的？参数量不会爆炸吗？不会造成收益递减吗？如何实现能在百万个专家中实现高效检索？
背景与介绍
Transformer 架构中，每个块内都包含注意力层和前馈层（FFW），注意力层用于计算序列中 token 之间的关系，FFW 网络则负责存储模型知识。
我们当然希望 LLM 能在参数中隐式存储更多知识，但 FFW 的计算成本和激活内存会随之线性增加。稠密模型中，FFW 层已经占据了总参数量的2/3，是扩展的主要瓶颈之一。
MoE 模型虽然参数量也很大，但每次推理时不会动用整个模型的能力，而是将数据路由到小型且专门的「专家模块」，因此能在 LLM 参数增加的同时，让推理所需的计算成本基本不变。
那么专家数量（即 MoE 模型的「粒度」）是不是越多越好？
这要考虑多个因素，包括模型参数总量、训练 token 数量和算力的预算。
2022 年的一项研究认为，模型总参数量不变时，存在一个能达到最优性能的「最佳粒度」。专家数量超过这个与之后，模型性能就会进入「平台期」。
Unified Scaling Laws for Routed Language Models
论文地址：https://arxiv.org/pdf/2202.01169
然而，今年年初 Krajewski 等人发表的一篇论文反驳了这个观点。他们发现，如果同时增加训练所用的 token 数量，那么更高粒度可以提高性能。
论文地址：https://arxiv.org/pdf/2402.07871
受到这种细粒度 MoE Scaling Law 的启发，作者推断，模型容量的持续改进将带来具有高粒度的 LLM，即包括大量微型专家的模型。
除了能带来更高效的扩展，增加专家数量还有另外一层好处——终身学习。
之前有研究表明，通过简单地添加新专家并进行适当正则化，MoE 模型就可以适应连续的数据流。冻结旧专家、仅让新专家权重更新，就可以在保持可塑性的同时防止灾难性遗忘。
在终身学习环境中，数据流可能达到无限长度，甚至永无止境，因此论文探索的专家数量的扩大就显得十分重要。
百万 MoE 所系
PEER 层设计
形式上，PEER 可以表达为函数，包含如下三部分：
- N 个专家组成的专家池，其中每个专家 ei 和f有相同的函数签名
- 一组相应的N个乘积键
- 用于将输入向量x映射到查询向量q(x)的查询网络
用表示 top-k 运算符，给定输入x，从N个专家中检索到k个专家的过程可以表示为：
之后使用 softmax 或 sigmoid 等激活函数应用于 top-k 个专家的查询-键内积，获得路由分数：
最后，计算路由分数加权的专家分数之和，作为 PEER 层输出。
上述的 PEER 层可以插入到 Transformer 架构骨干的中间，或者代替 FFW 层，整体计算流程如图 2 所示。
PEER 层示意图
乘积键检索
由于要使用非常多的专家（N≥106），直接计算公式 1 中的前k个指标可能会非常耗费资源。
为此，研究人员提出了乘积键检索技术——通过连接来自两个独立d/2 维的子键C和C′（C, C′ ⊂ Rd/2）的向量来进行创建：
这种K的笛卡尔积结构，能够高效地找到前k位专家。
也就是，先将查询向量q(x)拆分为两个子查询 q1 和 q2，并分别对子查询与子键之间的内积进行前k个操作：
然后得到一组 k2 候选键集合， ，在数学上，这可以保证从K中与q(x)最相似的k个键在这个候选集合中。
此外，候选键与q(x)之间的内积只是子键与子查询之间内积的总和： 。
因此，可以再次将前k个操作应用于这k的平方个内积，以从原始的乘积键集合K中获得前k个匹配键。
最终，公式 1 中通过穷举搜索进行的 top-k 专家检索的复杂度，从O(N·d)降到了O((√N + k2)d)。
参数高效专家和多头检索
通常的 MoE 架构中会将专家隐藏层设计为 FFW 层相同大小，但 PEER 中的每个专家 ei 则小得多，仅仅是包含一个神经元、一个隐藏层的单例 MLP：
其中 vi、ui 都不是参数矩阵，而是与x维度相同的向量，表示 ReLU 或 GELU 等激活函数。
每个专家只有如此少的参数，怎么可能有强大的表达能力？
此处，作者借鉴了多头注意力的做法，使用了「多头检索」（multi-head retrieval）机制。
也就是说，查询过程并不是仅有一次，而是使用h个独立的查询网络，分别计算自己的查询向量并检索出一组 top-k 专家，但他们共享相同的专家池和每个专家的乘积键。
h 个「检索头」的输出相加后得到最终检索结果：
可以证明，当每个头只检索一名专家时（k=1），使用含h个检索头的 PEER 层，等效于使用含有h个神经元的专家。
这就意味着，这种多头检索的过程相当于动态组装出含有h个神经元的专家 MLP 网络。
与现有 MoE 中多个专家之间完全隔离的状态相比，这种设计允许专家间共享隐藏神经元，从而提升了参数效率和知识迁移的能力。
算法 1 描述了 PEER 前向计算的简化实现。要想达到高效的实现，可能还需要专门的硬件内核来加速 embedding 查找，以及与 einsum 操作的融合。
为什么用大量的小专家
论文的背景介绍中，我们从直觉和经验层面推导出使用大量专家的优点，但作者也在此处用公式推导的方法证明了这一点。
首先，我们可以用 3 个超参数表征 MoE 层的设置：参数总量P、每个 token 激活的参数 Pactive 以及单个专家的大小 Pexpert。
之前提及的，Krajewski 论文中提出的「细粒度 MoE Scaling Law」就可以用这三个参数表示为如下形式：
其中，L表示损失值，D是训练 token 总量，G是激活专家的数量，a、b、g、α、β、γ等字母都表示常量。
为了降低损失值、提高模型性能，我们需要提升P、D、G的值，同时也需要限制 Pactive 的大小，因为计算和内存成本主要由激活参数量决定。
值得注意的是，Pactive 对应的内存占用也和 batch 中的 token 数量有关，但与 batch 大小或序列长度无关，因为在处理 batch 或序列的每个 token 时只需要存储模型的一份副本。
由于专家大小 Pexpert=Pactive/G，因此专家数量N=P/Pexpert=P·G/Pactive。
这就意味着，如果想要在增加P、G的同时保持 Pactive 不变，就应该减小专家大小 Pexpert，并增加专家数量N。
这就解释了，我们为什么需要大量的小专家。
一般来说，对于只有单个隐藏层的专家模型，其参数大小为：Pexpert=(2dmodel+1)dexpert，相应地，Pactive=(2dmodel+1)dactive。
其中 dmodel、dexpert、dactive 分别表示表示 Transformer 的隐藏维度、单个专家的隐藏神经元数量以及每个 token 激活的参数总量。
PEER 中使用了尽可能小的专家大小，即 dexpert=1，因此最大程度地降低了 dactive，大小仅为检索头数量乘以每次检索的专家数量（h·k）。
实验
预训练 isoFLOP 分析
作者使用 isoFLOP 分析将 PEER 与各种基线方法进行了比较。
这里，他选择了固定的 FLOP 预算（6e18 和 2e19），并同时改变模型大小和来自 C4 数据集的训练 token 数，以获得 isoFLOP 曲线。
isoFLOP 曲线上的每个点，都具有相同的计算成本，作者根据其模型大小和在 C4 上的最终验证困惑度（perplexity），来绘制这些点。
对于密集基线模型，通过改变层数、注意力头数和模型维度来改变它们的大小。
对于 MoE、PKM 和 PEER 方法，作者选取了每个考虑的密集模型，并分别用 MoE、PKM 和 PEER 层，替换中间块中的 FFN 层（例如，在 12 块 Transformer 中，替换第 6 块中的 FFN）。
在 MoE 中，作者使用了专家选择路由算法。该算法有效地解决了专家负载不平衡问题，并且通常优于 token 选择 MoE。
每个专家的大小与相应密集模型中原始 MLP 的大小相同，作者使用了 128 个专家来覆盖与 PEER 模型相同的模型大小范围。
这种类型的 MoE 代表了标准的粗粒度方法，由少量大型专家组成。
在 PKM 中，作者使用了 1024×2 个储存器，8 个头，每个头选择 32 个存储器（top k=32）。
此外，研究者还应用了查询批量归一化，这是原始 PKM 论文中推荐的，用于提高记忆使用率。
在 PEER 中，同样使用了 10242 个专家，8 个头，每个头选择 16 个存储器（top k=16）。
默认情况下，作者也启用了查询批归一化（BatchNorm）来增加专家使用率。与专家选择 MoE 基线不同，PEER 代表了一种细粒度方法，使用了大量小型专家。
在所有模型大小和方法中，保持了一致的批大小（128）和序列长度（2048）。通过将总计算预算除以每个训练步骤的 FLOP 来计算训练步骤数。
图 1 展示的是 isoFLOP 曲线。
与密集 FFW 基线相比，稀疏替代方案使 isoFLOP 曲线向下和向右移动，因为其引入了更多的总参数P，但使用了更少或相同数量的活跃参数 Pactive。
另外，在相同的计算预算下，PEER 模型达到了最低的计算优化困惑度。
语言建模数据集评估
根据 isoFLOP 曲线确定每种方法的计算最优模型后，研究人员在几个流行的语言建模数据集上评估了这些预训练模型的性能，包括 Curation Corpus、Lambada、the Pile、Wikitext 和预训练数据集 C4。
评估结果如表 1 所示，作者根据模型在训练期间使用的 FLOP 预算将其分组。
消融实验
改变专家总数
图 1 中 isoFLOP 图所示的模型都有超过一百万（10242）个专家。
接下来，研究者进行了一项消融研究，研究专家数量N的影响，N决定了等式 9 中的总参数数量P。
作者选择了 isoFLOP 最优模型，并改变 PEER 层中的专家数量（N = 1282, 2562, 5122, 10242），同时活跃专家的数量不变（h = 8, k = 16）。
实验结果如图3(a)所示，可以看出，isoFLOP 曲线在 10242 个专家的 PEER 模型和相应的密集骨干模型（未将中间块的 FFW 层替换为 PEER 层）之间进行插值。
这表明，仅仅增加专家数量，就可以提高模型性能。
改变活跃专家的数量
另外，作者还对活跃专家 hk 数量进行了消融研究，hk 等同于等式 9 中的粒度G。
他系统地改变了活跃专家的数量（hk = 32, 64, 128, 256, 512），同时保持专家总数不变（N = 10242）。
此外，对于给定的 hk，同时改变h和k以找到最优组合。结果中的 isoFLOP 曲线，以头（h）数为横轴绘制，如图3（b）所示。
结果表明，在考虑的数值范围内，更高的h·k值通常会带来更好的性能。
值得注意的是，随着h·k的增加，最优的h也在增加。然而，性能逐渐趋于饱和，增加活跃专家的数量也会增加设备内存消耗，可能需要额外的加速器设备。
因此，在实际应用中，应该根据性能、设备数量和计算资源需求之间的权衡来选择适当的h·k值。
专家使用和查询批归一化
由于 PEER 层中有超过一百万个专家，对于许多人来说，一定好奇其在推理过程中实际选择了多少专家？它们的使用是否均匀分布？
为了分析这点，研究者在 C4 验证集的所有 token 中，为每个专家 ei 保留一个累积的路由分数（router score），表示为： 。
其中，gi (x)是指当 token x 作为输入时用于聚合专家输出的路由分数。如果专家 ei 未被选中，则 gi (x)=0。
从这些累积的路由分数中，可以获得一个经验概率分布向量，表示为 。
代表 C4 验证集上所有专家的分布。
然后，作者计算了 Lample 等人提出的以下指标来评估专家的使用情况和分布：
- 专家使用：推理过程中检索到的专家比例：
- 不均匀性：z与均匀分布之间的 KL 散度
其中，N是专家总数。
默认设置下，作者在查询网络上添加了一个批归一化（BN）层，目的是增加训练过程中的专家使用率。
对此，他研究添加这个 BN 层对上述指标的影响。
表 2 展示了不同专家数量下的专家使用和不均匀性，包括使用和不使用 BN 的情况。
可以看到，即使对于 100 万个专家，专家使用率也接近 100%，而使用 BN 可以导致专家的使用更加平衡，并降低困惑度。
这些发现证明了，PEER 模型在利用大量专家方面的有效性。
最后，作者还比较了有和没有批归一化（BN）的 isoFLOP 曲线。
如下图 4 所示，使用 BN 的 PEER 模型通常能达到更低的困惑度。虽然差异不是很显著，但在 isoFLOP 最有区域附近最为明显。
作者介绍
Owen He 是 Google DeepMind 的研究科学家，专注于持续学习和基础模型。
在加入 DeepMind 之前，Owen 是在格罗宁根大学和不来梅雅各布大学（康斯特大学）攻读博士学位，导师是 Herbert Jaeger。
在此期间，Owen 访问过 Mila、马克斯·普朗克研究所、伯尔尼大学和苏黎世联邦理工学院，并在 Google DeepMind 进行过实习。
他曾组织过 ICML 2021 持续学习研讨会，目前担任 Conference on Lifelong Learning Agents 的出版主席。
参考资料：
https://arxiv.org/abs/2407.04153
https://venturebeat.com/ai/deepminds-peer-scales-language-models-with-millions-of-tiny-experts/
