进一步扩大释放 transformer 还能保持计算效率的潜力。
标准 Transformer 体系结构中的前馈（FFW）计算成本和激活内存的线性增加会随着隐藏层宽度的增加而增加。大语言模型（LLM）如今，随着体积的增加，稀疏混合专家（MoE）架构已成为解决这一问题的可行方法，它将模型尺寸与计算成本分开。很多新兴的 MoE 所有模型都能达到相同的体积，更好的性能和更强的性能。最近发现的细粒度 MoE 扩展定律表明，更高的粒度可以带来更好的性能。然而，由于计算和优化的挑战，现有的 MoE 模型仅限于低数量专家。本周二，Google DeepMind 新研究引入了参数高效的专家检索机制，利用乘积密钥技术从100万微型专家中进行稀疏检索。
链接：https://arxiv.org/abs/2407.04153该方法试图有效地将路由学习索引结构连接到大量的小专家中，从而将计算成本与参数计数分开。与密集的 FFW、粗粒度 MoE 以及产品密钥存储器 (PKM) 与层相比，表现出卓越的效率。本工作介绍了参数高效专家检索（PEER) 架构（parameter efficient expert retrieval），使用产品密钥（product key）对大量专家进行检索，将计算成本与参数量分开。该设计在实验中展示了优异的计算性能水平，并将其定位为扩展基础模型的密集型 FFW 层的竞争替代方案。这项工作的主要贡献是：极端 MoE 探索设置:与以前相比 MoE 研究中对少数大型专家的关注不同，这项工作研究了许多小型专家没有充分探索的情况。学习索引结构用于路由：首次证明学习索引结构可以有效地路由到100多万专家。新的层设计：将产品密钥路由与单神经元专家相结合，我们介绍了 PEER 层，它可以在没有大量计算费用的情况下扩展层容量。实证结果表明，与密度 FFW、粗粒度 MoE 以及产品密钥内存 (PKM) 其效率高于层。综合消融研究：我们研究了 PEER 语言建模任务的不同设计选择(如专家数量、活动参数、头数、查询批量标准化)。在本节中，研究人员详细介绍了高效参数专家的搜索 (PEER) 层，它是一种混合专家架构，使用路由中的产品钥匙和单神经元 MLP 作为专家。下图 2 展示了 PEER 层内的计算过程。
PEER 层概览。从形式上看，PEER 层是函数 f : R^n → R^m，它由三个部分组成：一个由三个部分组成： N 由专家组成的池 E := {e_i}^N_i=1，每一位专家 e_i : R^n → R^m 与 f 共享相同的签名；一组相应的 N 个产品密钥 K := {k_i}^N_i=1 ? R^d ；以及查询网络 q : R^n → R^d，它将输入向量 x ∈ R^n 映射到查询向量 q (x)。令 T_k 表示 top-k 运算符。给定输入 x，首先检索 k 这些专家的子集，相应的产品键和查询 q (x) 内积最高。
然后非线性激活(例如(例如) softmax 或 sigmoid）应用于前 k 为了获得路由分数，专家查询键内积。
最后，通过线性组合加权路由分数的专家输出来计算输出。
产品密钥检索。因为研究人员计划使用大量专家（N ≥ 10^6)，简单计算公式 1 中的前 k 因此，产品密钥检索技术应用于索引的成本可能很高。它们不使用 N 个独立的 d 以维向量为密钥 k_i，相反，通过连接两个独立的连接 d/2 维子密钥集(即 C, C ′ ? R d/2） 创建它们的中向量：
高效参数专家和多头检索专家。与其他 MoE 这些架构通常将每个专家的隐藏层设置为不同的架构 FFW 层的大小相同。而在 PEER 中，每个专家 e_i 都是单例 MLP，换句话说，它只有一个带有单个神经元的隐藏层:
研究人员没有改变单个专家的规模，而是使用多头检索来调整 PEER 这类似于层的表达能力 transformer 多头注意力机制和 PKM 多头记忆。具体来说，它们是用的 h 一个独立的查询网络，每个网络计算自己的查询并检索一组单独的查询网络 k 专家。然而，不同的头共享相同的专家池，并拥有相同的产品密钥。这个 h 个头输出简要总结如下：
为什么要有大量的小专家？给定的 MoE 层可以用三个超参数来表示：参数总数 P、每个 token 活动参数的数量 P_active 以及单个专家的规模 P_expert。Krajewski 等人 (2024) 表明，MoE 模型的 scaling law 以下形式：
对于 PEER，研究人员通过设置 d_expert = 1 使用尽可能小的专家规模，激活神经元的数量乘以每个头检索的专家数量：d_active = hk。因此，PEER 粒度总是 G = P_active/P_expert = d_active/d_expert = hk。
在根据 isoFLOP 在确定了每种方法的最佳计算模型后，研究人员在以下流行语言建模数据集中评估了这些预训练模型的性能：
下表 1 评估结果显示出来。研究人员根据训练期间使用的内容 FLOP 对模型进行预算分组。可见，PEER 对这些语言建模数据集的困惑程度最低。
在消融实验中，研究人员改变了专家总数。下图 1 中 isoFLOP 曲线上有100多万个模型(1024^2) ?）专家。
研究人员选择了 isoFLOP 模型的最佳位置，并改变了 PEER 层次专家数量（N = 128^2、256^2、512^2、同时保持活跃专家数量不变，1024^2)（h = 8，k = 16)。结果如下图所示 3 (a) 所示。可以看到，isoFLOP 曲线在具有 1024^2 个专家的 PEER 插值在模型和相应的密集主干之间，无需使用 PEER 中间块中的层替换 FFW 层。这说明只要增加专家数量，模型性能就可以提高。与此同时，研究人员改变了活跃专家的数量。他们系统地改变了活跃专家的数量（hk = 32、64、128、256、512)同时保持专家总数不变（N = 1024^2)。给定的 hk，另一方面，研究人员联合改变 h 和 k 确定最佳组合。下图 3 (b) 绘制了关于头数的数量 (h) 的 isoFLOP 曲线。
下表 2 列出使用和不使用 BN 不同数量专家的专家利用率和不均匀性。即使对 1M 个别专家的利用率也接近 并且使用100% BN 它可以使专家的利用率更加平衡，更加混乱。这些发现证明了这一点 PEER 模型利用了大量专家的有效性。
研究人员也比较了有无 BN 的 isoFLOP 曲线。下图 4 显示，有 BN 的 PEER 模型通常可以达到较低的混乱。虽然差异不明显，但在 isoFLOP 最佳区域附近的差异最为明显。
PEER 研究只有一个作者 Xu He（Owen），他是 Google DeepMind 2017年研究科学家 荷兰格罗宁根大学博士毕业。以上是一篇作者论文。谷歌提出了数百万专家的细节，超越了密集的前馈和稀疏的Moe。请关注其他相关文章！
p
