导读

最近半年大模型的价格战杀疯了， 其中一个重要的原因，就是很多公司的大模型都采用了MOE架构， 在保证模型参数不变的情况下， 降低推理计算的参数量， 不仅能保证效果还能大幅降低费用。

而DeepMind的研究人员则更为夸张， 常识专家数量达到百万级别的MOME， 下文即为该模型架构的介绍。

谷歌DeepMind的新研究，旨在利用百万专家的混合来构建大模型LLM的推理效果

深入探讨百万专家混合（MoME）架构的发展，该架构在性能和计算效率方面超越了传统的LLM

目前所有大型科技公司都在争先恐后地开发比现有模型性能更好的LLM。

增加模型大小、数据集大小，以及计算量，然后，就拥有了比以前更好的模型。

根据Scale Law，Google DeepMind的研究人员发现，以某种方式调整模型架构也会显著提高其性能和训练效率。

他们的洞察力来自Transformer架构，作为LLM核心，在密集的前馈（FFW）层中参数存储了大部分事实知识。

这些层占据了Transformer中所有参数的2/3。

这其实是一个比较大的问题

它使传统Transformer的计算占用量大大增加，因为它与其包含的参数数量成正比。

为了解决这个问题，并从专家混合架构（MoE）中汲取灵感，研究人员开发了PEER（参数高效专家检索），这是一个新增加的Layer，可以适应现有的Transformer架构，并且可以像以前从未有过的那样改进它。

之后这个层可以与由超过一百万个微小专家组成的专家混合LLM一起使用！

在ArXiv的论文中展示了百万专家混合（MoME）架构如何在计算效率方面超越了传统的LLM和以前的MoE模型（专家更少但更大）！    

以下是一个深入探讨这种新颖架构如何工作、如何开发以及与以前的LLM架构相比表现如何的故事。

混合专家模型

将计算成本与模型参数数量解耦，这种操作并非最近才出现。在N年前推荐系统中就已经开始使用类似的架构。为什么是推荐系统呢？一个可能前些年推荐系统比较火， 另外可能应用的场景也比较能赚钱。

Mixture-of-Experts (MoE) model 在2017年被引入并广泛流行于推荐场景。

该类模型结合了多个称为专家的神经网络，这些专家在不同的技能集上有所专长。

一个门控网络负责根据给定的查询选择性地激活这些专家，并组合它们的加权输出以产生结果。

由于这项研究，人们发现使用稀疏激活的专家模块而不是密集的前馈网络（FFWs）更有效。

然后，在2024年初，这些MoE模型的缩放定律被引入大模型

数学定义

其中：

•L(N,D,G) 是MoE模型的最终测试损失

•N 是非嵌入参数的总数

•D 是训练标记的数量

•G 是粒度或活跃专家的数量

•a,b,g,γ,α,β 是常数

根据这些定律，增加模型大小、更多的训练标记和更高的粒度（活跃专家的数量）会影响MoE模型，导致更低的损失并使其更有效。

该公式揭示了一个重要的结论

应该使用许多小的专家来扩展MoE架构，而不是像以前那样使用少量大的专家！

数百万呢？

但是将路由到数百万专家也很困难

将输入查询路由到数百万专家是一项具有挑战性的任务。

为了为给定的查询路由或选择相关的专家，流程如下：

首先，计算输入查询向量和每个专家的关键向量之间的内积。

这些分数决定了专家对给定输入查询的相关性。

然后对这些分数进行排序，并选择前k个分数进行路由。

这个过程在计算上是昂贵的，并且其复杂性随着专家的总数（N）和关键和查询向量的维度（d）的增长而增加。

这个问题被研究人员解决了，他们引入了一个新的路由层架构。

Parameter Efficient Expert Retrieval (PEER) Layer
参数高效专家检索（PEER）层

Google DeepMind的研究人员通过构建一个名为Parameter Efficient Expert Retrieval (PEER)的新层架构解决了低效路由的问题。

PEER层可以插入到Transformer主干中，也可以用来替换其前馈层。

作为Mixture-of-a-Million-Experts (MoME)架构的基础，该层有三个组成部分：

1. 专家池

2. 一组产品密钥（每个专家对应一个密钥），其中每个密钥确定了对于给定输入查询每个专家的相关性

3. 一个查询网络，用于将输入向量映射到查询向量（用于与产品密钥进行比较以选择最相关的专家）   

给定一个输入向量x，PEER层的操作如下。

1. 输入向量x传递到一个查询网络，产生一个查询向量q(x)

2. 计算该查询向量与产品密钥（两者的维度均为d）的内积

3. 选择具有最高内积的前k个专家

对这些前k个专家的内积应用非线性激活函数（Softmax或Sigmoid）。这产生了决定每个选定专家在最终输出中重要性的路由器分数。

每个专家的输出乘以其对应的路由器分数作为权重，这些值线性组合以产生PEER层的最终输出。

现在，看看第3步的方程，我们计算前k个索引。              
这就是所有低效性所在。

这种方法的复杂度是O(Nd)，其中d是向量的维度，N是专家的总数。    

该问题在2019年发表的名为'Large Memory Layers with Product Keys'的ArXiv研究中得到解决，研究人员设计了一个结构化内存，可以适应基于Transformer的架构，以提高其性能并减少推断时间。

该研究引入了高效产品密钥检索的技术，这些技术被借用来构建PEER层。

它的工作原理如下：

不再计算N个不同的d维键向量，而是通过连接来自两组子键的向量来创建这些键。

这两组子键用C和C'表示，每个子键的维度为d/2，包含√N个子键。

这导致了一种产品密钥K的笛卡尔积结构。

同样，查询向量q(x)被分成两个子查询向量q(1)和q(2)，每个子向量的维度为d/2。

接下来，通过计算这些子键和子查询向量之间的内积来进行前k个专家的选择。

这产生了两组前k个子键：I(C)和I(C')。

然后，从这些集合中选择前k个子键，得到k²个候选键，数学上保证了与查询向量q(x)最相似的前k个键是这个候选集的一部分。

此外，候选键和查询向量q(x)之间的内积简单地是子键和子查询之间的内积之和。    

因此，可以再次将前k个运算符应用于这些内积，以从原始产品密钥集K中获得前k个匹配键。

总的来说，在这种方法中，我们使用了笛卡尔积，从√N个子键中高效生成了N个键。

由于现在我们只需要与√N个子键进行比较，而不是与N个完整键进行比较，这显著降低了前k个选择过程的复杂度，从先前方法中的O(Nd)降低到了O((√N+ k²)d)。

将 PEER 层扩展到多个头部

在传统的 MoEs 中，每个专家的隐藏层与其他前馈层的大小相同。

百万专家混合体（MoME）架构采取了不同的方式。

在这里，每个专家都是一个独立的MLP，这意味着它只有一个隐藏层和一个单独的神经元。

为了改善粒度（记住 MoEs 的缩放定律），MoME 将个体专家的大小保持在最小，并且使用多头检索来增加其表达能力。

这类似于变压器中的多头注意力。

与其只有一个，来自共享专家池的多个独立查询网络（头部）计算自己的查询并检索一组单独的前k个专家。

然后将这些头部的结果相加以产生最终输出。

实现 PEER 层前向传递的伪代码

让我们了解 原始研究论文 中描述的 PEER 层前向传递的伪代码。

这将帮助我们重新学习它的工作步骤。

def peer_forward (self , x):# 嵌入层存储所有专家的下/上投影权重self . w_down_embed = nn.Embed ( num_embeddings =self . n_experts , features =self . d_model )self . w_up_embed = nn.Embed ( num_embeddings =self . n_experts , features =self . d_model )
# 使用产品密钥检索顶部匹配专家的权重# 索引和分数的形状为 'bthk'，其中 h 是头的数量indices , scores =self . get_indices ( self . query_proj (x), self . sub_keys , top_k =self .k)w_down =self . w_down_embed ( indices )w_up =self . w_up_embed ( indices )
# 计算专家输出的加权平均值x = jnp . einsum ('btd , bthkd -> bthk ', x, w_down )x =self . activation (x)x = x * nn. softmax ( scores )x = jnp . einsum ('bthk , bthkd -> btd ', x, w_up )return x

bthkd 意味着 batch_size, time_steps, num_heads, top_k, dimension

einsum 是爱因斯坦求和

peer_forward 函数的工作方式是通过初始化嵌入层来存储所有专家的下投影和上投影权重。

（下投影权重将输入投影到较低维空间，而上投影权重将其投影回较高维空间。）

对于给定的输入，它使用查询网络找到前几个匹配专家的索引和路由器分数。

然后将前几个专家的投影权重存储到 w_down 和 w_up 中。

然后使用 einsum 操作将输入与下投影权重组合，并对结果应用非线性激活函数。    

将该结果乘以 Softmax 归一化的路由器分数，然后使用另一个 einsum（爱因斯坦求和）操作将其与上投影权重组合。

最后，这些处理过的输出被求和以产生最终结果。

最后，让我们了解一下这种架构与传统对应物相比的表现。

MoME 架构的性能

MoME 架构的性能非常出色，在各种任务中得到了证明。

预训练 IsoFLOP 分析

预训练 指的是在大量未标记文本上对 LLM 进行初始训练的阶段。

通过这项分析，研究人员将 MoME 的预训练性能与等价计算资源下的其他基线进行比较（因此，使用 C4 数据集），其中术语 iso-FLOP 被使用。

比较的四种不同模型是：

•具有稠密前馈层（FFW）的 LLM

•具有 128 个专家和专家选择路由的专家混合（MoE）（其中专家选择前 k 个标记而不是标记选择前 k 个专家）

•具有 1024² 记忆的产品键内存层（PKM）的 LLM

•具有 1024²（即 1,048,576）专家的 PEER 层（MoME 架构）的 LLM

为了创建可比较的模型，首先选择了不同层数、注意力头数和维度的密集模型。

然后，在它们的 Transformer 的中间块中，将前馈层替换为 MoE、PKM 和 PEER 层。

结果显示，与密集 FFW 基线相比，其他模型将 isoFLOP 曲线向下和向右移动。

这种向下和向右的移动可能一开始看起来有些反直觉，因为人们可能期望添加更多参数会增加计算成本（或者会将曲线向上和向右推移）。

然而，这并不是真的，因为即使这些稀疏模型引入了更多的总参数，它们仍然激活和利用其中的一个小子集。    

结果显示在下面的图中，可以看到在相同的计算预算下，具有 PEER 层的 MoME 架构实现了最低的计算优化 困惑度。

在语言建模数据集上的表现

在选择了isoFLOP分析中的计算最优模型之后，它们的性能被评估在许多流行的语言建模数据集上。

结果显示，与其他模型相比，MoME架构再次实现了最低的困惑度。

MoME 架构变体的性能

研究人员接下来研究了不同数量的总体和活跃专家对困惑度的影响。

他们发现，仅仅增加总专家的数量就能提高性能，而不需要额外的成本。

然而，当活跃专家的数量发生变化，同时保持总专家数量不变时，性能会提高，直到达到一定点后，性能就会饱和并增加内存消耗。    

使用查询批量归一化的性能差异

研究人员最终研究了查询批量归一化（在使用之前，对输入批次中的查询向量进行归一化）对模型的困惑度的作用。

在这项分析中使用了两个额外的指标。

1.专家使用：这是推理过程中活跃的专家数量。

2.KL散度/不均匀性：这衡量了给定概率分布与均匀分布的偏差程度。

结果显示，使用查询批量归一化导致专家的使用更加平衡，并获得更低的困惑度。

还应该注意的是，即使对于超过一百万个专家，模型的专家使用率也接近100%！

此外，isoFLOP曲线显示，当使用查询批量归一化时，模型通常实现更低的困惑度（特别是在isoFLOP最佳区域附近）。