深入探讨百万专家混合（MoME）架构的开发，该架构在性能和计算效率上前所未有地超越了传统大型语言模型（LLM）

我们正身处一场大型语言模型（LLM）的竞争之中。

这可能并不显而易见，但所有大型科技公司都在争相开发超越现有模型的更优秀LLM。

增加模型规模、数据集大小以及计算量，你就能得到比以往更优秀的模型。

基于这一规模化定律，谷歌DeepMind的研究人员发现，以特定方式调整模型架构也能显著提升其性能和训练效率。

他们的洞察源于这样一个事实：作为LLM核心的Transformer架构，其大部分事实知识存储在密集的前馈（FFW）层中。

这些层占据了Transformer中三分之二的参数。

而这正是一个问题。

它大幅增加了传统Transformer的计算足迹，因为其计算量与所含参数数量成线性正比。

为解决这一问题并受专家混合架构（MoE）启发，研究人员开发了PEER（参数高效专家检索），这是一种能融入现有Transformer架构并以前所未有的方式对其进行改进的新型层。

这一层随后可用于由超过百万个微型专家组成的专家混合LLM中！

他们在ArXiv上的论文展示了这种百万专家混合（MoME）架构如何在保持计算效率的同时，超越传统LLM及先前参数较少但规模较大的MoE模型。

以下是深入探讨这一新型架构如何运作、其开发过程以及与先前LLM架构相比性能表现的故事。

让我们从混合专家模型开始

将计算成本与模型参数数量解耦的探索并非新鲜事。

2017年，混合专家（Mixture-of-Experts，MoE）模型被引入并普及。

该模型结合了多个被称为专家的神经网络，这些专家专长于不同的技能集。

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

得益于这项研究，发现使用稀疏激活的专家模块而非密集的前馈网络（FFWs）更为高效。

随后，在2024年初，针对这些MoE模型的缩放法则被提出。

它们在数学上描述如下——

其中：

L(N,D,G) 是MoE模型的最终测试损失
N 是总的非嵌入参数数量
D 是训练令牌的数量
G 是粒度，即活跃专家的数量
a, b, g, γ, α, β 是常数

根据这些法则，增加模型大小、更多训练令牌以及更高的粒度（活跃专家数量） 影响MoE模型，导致损失降低并使其更高效。

这揭示了一个重要洞见——

我们需要使用众多小型专家来扩展MoE架构，而非像之前那样使用少量大型专家。

那数百万个呢？

向数百万专家路由的难题

将输入查询路由到百万级别的专家是一项艰巨的任务。

路由或从总数为 N 的专家中选择与给定查询相关的专家的过程如下：

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

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

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

这一过程在计算上开销巨大，其复杂度随着专家总数 (N) 以及关键和查询向量的维度 (d) 的增长而增加。

研究人员通过引入一种新的路由层架构解决了这一问题。

接下来让我们详细探讨这一架构。

参数高效专家检索（PEER）层

Google DeepMind的研究人员通过构建一种名为参数高效专家检索（PEER）的新层架构，解决了路由效率低下的问题。

PEER层可以插入到Transformer主干架构中，或用于替换其前馈层。

这一构成百万专家混合（MoME）架构基础的层包含三个组成部分：

专家池
一组产品键（每个专家对应一个键），每个键决定每个专家对给定输入查询的相关性
一个查询网络，功能是将输入向量映射到查询向量（用于与产品键比较，选择最相关的专家）

给定一个输入向量x，PEER层的运作方式如下：

输入向量x传递到一个查询网络，生成查询向量q(x)
计算此查询向量与产品键（维度均为d）的内积
选择内积最高的k个专家

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

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

现在，看步骤3的方程，我们计算top-k索引。这就是所有低效所在。

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

这一问题在2019年发表于ArXiv的研究《具有产品键的大型记忆层》中得到解决，研究人员设计了一种结构化记忆，可融入基于transformer的架构，提高其性能并减少推理时间。

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

其工作原理如下：

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

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

这导致产品键集合K形成笛卡尔积结构。

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

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

这产生两组top-k子键：I(C)和I(C')。

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

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

因此，可以再次对这些内积应用top-k操作符，从原始产品键集合(K)中获得top-k匹配键。

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

由于现在只需与√N个子键而不是N个完整键比较，这显著将top-k选择过程的复杂度从O(Nd)降低到O((√N + k²)d)。

扩展PEER层至多头部

在传统的MoE（专家混合）模型中，每个专家的隐藏层大小与其他前馈层相同。

而百万专家混合（MoME）架构则采用了不同的方法。

在这里，每个专家是一个单一隐藏层的MLP，即它仅有一个隐藏层和一个神经元。

为了提高粒度（记住MoE的扩展定律），MoME将单个专家的大小保持在最小，并通过多头部检索来增加其表达能力。

这与Transformer中的多头注意力机制类似。

不同于仅使用一个查询网络，多个独立的查询网络（头部）从共享的专家池（具有相同的产品键集合）中计算各自的查询，并检索一组不同的top-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 )
    
    # 使用乘积键检索匹配度最高的专家的权重
    # indices和scores的形状为'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）进行初始训练的阶段。

通过此分析，研究人员在等量计算资源（因此称为 iso-FLOP ）下，使用 C4 数据集比较了 MoME 预训练性能与其他基准的表现。

比较的四种不同模型包括：

带有密集前馈层（FFW）的 LLM
带有 128 个专家和专家选择路由的混合专家模型（MoE）（专家选择 top-k 令牌而非令牌选择 top-k 专家）
带有产品键记忆层（PKM）的 LLM ，具有 1024² 记忆
带有 PEER 层的 LLM（MoME 架构），具有 1024²（即 1,048,576）个专家

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

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

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

这种向右下方的移动起初可能看似违反直觉，因为人们可能预期增加更多参数会增加计算成本（或将曲线向上和向右推）。

然而，事实并非如此，因为尽管这些稀疏模型引入了更多的总参数，它们仍然激活并利用其中的一小部分。

结果显示在下方的图中，可以看出**在相同的计算预算下，带有 PEER 层的 MoME 架构达到了最低的计算最优困惑度 **。

语言建模数据集上的性能表现

在根据isoFLOP分析选出计算最优模型后，这些模型的性能在多个流行的语言建模数据集上进行了评估。

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

MoME架构变体的性能表现

研究者接下来探讨了总专家数和活跃专家数的变化对困惑度的影响。

研究发现，单纯增加总专家数能提升性能，且无需额外成本。

然而，在保持总专家数不变的情况下改变活跃专家数，性能会在一定程度内提升，随后达到饱和并增加内存消耗。

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

研究者最终探讨了查询批量归一化（即在查询向量使用前对一批输入进行归一化处理）对模型困惑度的影响。

本分析中使用了两个额外指标：

专家使用率：指推理过程中活跃专家的数量。
KL散度** / 不均匀性**：衡量给定概率分布与均匀分布的偏离程度。

结果显示，使用查询批量归一化能更均衡地利用专家，并获得更低的困惑度。

值得注意的是，即使专家数量超过百万，模型的专家使用率仍接近100%。

此外，isoFLOP曲线表明，模型在使用查询批量归一化时，通常能实现更低的困惑度（尤其是在isoFLOP最优区域附近）。

这篇研究论文的结果非常出色，我对这将如何积极改变LLM的未来感到非常乐观。

对此你有何看法？欢迎在下方评论区留言。

Barry's Home

谷歌DeepMind的新研究利用数百万专家混合构建大规模LLMs