Auto Byte

专注未来出行及智能汽车科技

微信扫一扫获取更多资讯

Science AI

关注人工智能与其他前沿技术、基础学科的交叉研究与融合发展

微信扫一扫获取更多资讯

单一作者论文,谷歌提出百万专家Mixture,超越密集前馈、稀疏MoE

释放进一步扩展 Transformer 的潜力,同时还可以保持计算效率。


标准 Transformer 架构中的前馈(FFW)层会随着隐藏层宽度的增加而导致计算成本和激活内存的线性增加。在大语言模型(LLM)体量不断增大的现在,稀疏混合专家(MoE)架构已成为解决此问题的可行方法,它将模型大小与计算成本分离开来。很多新兴的 MoE 模型都可以实现相同体量之上,更好的性能与更强大的表现。

最近发现的细粒度 MoE 扩展定律表明,更高的粒度可带来更好的性能。然而由于计算和优化方面的挑战,现有的 MoE 模型仅限于低数量专家。

本周二,Google DeepMind 的新研究引入了一种参数高效的专家检索机制,其利用乘积密钥技术从一百万个微型专家中进行稀疏检索

图片

链接:https://arxiv.org/abs/2407.04153

该方法尝试通过用于路由的学习索引结构有效地串联到大量微小专家,从而将计算成本与参数计数分离。与密集的 FFW、粗粒度 MoE 和产品密钥存储器 (PKM) 层相比,表现出卓越的效率。

这项工作引入了参数高效专家检索(PEER) 架构(parameter efficient expert retrieval),利用产品密钥(product key)检索高效地路由到大量专家,将计算成本与参数量分离。这种设计在实验中展示了卓越的计算性能水平,将其定位为用于扩展基础模型的密集 FFW 层的竞争性替代方案。这项工作的主要贡献是:

极端 MoE 设置的探索:与以前的 MoE 研究中对少数大型专家的关注不同,这项工作研究了众多小型专家的未充分探索的情况。

用于路由的学习索引结构:首次证明学习索引结构可以有效地路由到超过一百万个专家。

新的层设计:将产品密钥路由与单神经元专家相结合,我们引入了 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 曲线确定每种方法的计算最优模型后,研究者在以下几个流行语言建模数据集上评估了这些预训练模型的性能:

  • Curation Corpus
  • Lambada
  • Pile
  • Wikitext
  • 预训练数据集 C4

下表 1 展示了评估结果。研究者根据训练期间使用的 FLOP 预算对模型进行了分组。可以看到,PEER 在这些语言建模数据集上的困惑度最低。

图片

在消融实验中,研究者改变了专家总数量。下图 1 中 isoFLOP 曲线所示的模型都有超过一百万个(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 年博士毕业于荷兰格罗宁根大学。

图片

工程专家检索Google DeepMind
相关数据
DeepMind机构

DeepMind是一家英国的人工智能公司。公司创建于2010年,最初名称是DeepMind科技(DeepMind Technologies Limited),在2014年被谷歌收购。在2010年由杰米斯·哈萨比斯,谢恩·列格和穆斯塔法·苏莱曼成立创业公司。继AlphaGo之后,Google DeepMind首席执行官杰米斯·哈萨比斯表示将研究用人工智能与人类玩其他游戏,例如即时战略游戏《星际争霸II》(StarCraft II)。深度AI如果能直接使用在其他各种不同领域,除了未来能玩不同的游戏外,例如自动驾驶、投资顾问、音乐评论、甚至司法判决等等目前需要人脑才能处理的工作,基本上也可以直接使用相同的神经网上去学而习得与人类相同的思考力。

https://deepmind.com/
参数技术

在数学和统计学裡,参数(英语:parameter)是使用通用变量来建立函数和变量之间关系(当这种关系很难用方程来阐述时)的一个数量。

超参数技术

在机器学习中,超参数是在学习过程开始之前设置其值的参数。 相反,其他参数的值是通过训练得出的。 不同的模型训练算法需要不同的超参数,一些简单的算法(如普通最小二乘回归)不需要。 给定这些超参数,训练算法从数据中学习参数。相同种类的机器学习模型可能需要不同的超参数来适应不同的数据模式,并且必须对其进行调整以便模型能够最优地解决机器学习问题。 在实际应用中一般需要对超参数进行优化,以找到一个超参数元组(tuple),由这些超参数元组形成一个最优化模型,该模型可以将在给定的独立数据上预定义的损失函数最小化。

注意力机制技术

我们可以粗略地把神经注意机制类比成一个可以专注于输入内容的某一子集(或特征)的神经网络. 注意力机制最早是由 DeepMind 为图像分类提出的,这让「神经网络在执行预测任务时可以更多关注输入中的相关部分,更少关注不相关的部分」。当解码器生成一个用于构成目标句子的词时,源句子中仅有少部分是相关的;因此,可以应用一个基于内容的注意力机制来根据源句子动态地生成一个(加权的)语境向量(context vector), 然后网络会根据这个语境向量而不是某个固定长度的向量来预测词。

映射技术

映射指的是具有某种特殊结构的函数,或泛指类函数思想的范畴论中的态射。 逻辑和图论中也有一些不太常规的用法。其数学定义为:两个非空集合A与B间存在着对应关系f,而且对于A中的每一个元素x,B中总有有唯一的一个元素y与它对应,就这种对应为从A到B的映射,记作f:A→B。其中,y称为元素x在映射f下的象,记作:y=f(x)。x称为y关于映射f的原象*。*集合A中所有元素的象的集合称为映射f的值域,记作f(A)。同样的,在机器学习中,映射就是输入与输出之间的对应关系。

插值技术

数学的数值分析领域中,内插或称插值(英语:interpolation)是一种通过已知的、离散的数据点,在范围内推求新数据点的过程或方法。求解科学和工程的问题时,通常有许多数据点借由采样、实验等方法获得,这些数据可能代表了有限个数值函数,其中自变量的值。而根据这些数据,我们往往希望得到一个连续的函数(也就是曲线);或者更密集的离散方程与已知数据互相吻合,这个过程叫做拟合。

规范化技术

规范化:将属性数据按比例缩放,使之落入一个小的特定区间,如-1.0 到1.0 或0.0 到1.0。 通过将属性数据按比例缩放,使之落入一个小的特定区间,如0.0到1.0,对属性规范化。对于距离度量分类算法,如涉及神经网络或诸如最临近分类和聚类的分类算法,规范化特别有用。如果使用神经网络后向传播算法进行分类挖掘,对于训练样本属性输入值规范化将有助于加快学习阶段的速度。对于基于距离的方法,规范化可以帮助防止具有较大初始值域的属性与具有较小初始值域的属相相比,权重过大。有许多数据规范化的方法,包括最小-最大规范化、z-score规范化和按小数定标规范化。

神经元技术

(人工)神经元是一个类比于生物神经元的数学计算模型,是神经网络的基本组成单元。 对于生物神经网络,每个神经元与其他神经元相连,当它“兴奋”时会向相连的神经元发送化学物质,从而改变这些神经元的电位;神经元的“兴奋”由其电位决定,当它的电位超过一个“阈值”(threshold)便会被激活,亦即“兴奋”。 目前最常见的神经元模型是基于1943年 Warren McCulloch 和 Walter Pitts提出的“M-P 神经元模型”。 在这个模型中,神经元通过带权重的连接接处理来自n个其他神经元的输入信号,其总输入值将与神经元的阈值进行比较,最后通过“激活函数”(activation function)产生神经元的输出。

查询技术

一般来说,查询是询问的一种形式。它在不同的学科里涵义有所不同。在信息检索领域,查询指的是数据库和信息系统对信息检索的精确要求

语言模型技术

语言模型经常使用在许多自然语言处理方面的应用,如语音识别,机器翻译,词性标注,句法分析和资讯检索。由于字词与句子都是任意组合的长度,因此在训练过的语言模型中会出现未曾出现的字串(资料稀疏的问题),也使得在语料库中估算字串的机率变得很困难,这也是要使用近似的平滑n元语法(N-gram)模型之原因。

推荐文章
暂无评论
暂无评论~