单一作者论文,谷歌提出百万专家,超越密集前馈稀疏
机器之心报道
编辑:泽南、杜伟
释放进一步扩展Transformer的潜力,同时还可以保持计算效率。
标准Transformer架构中的前馈(FFW)层会随着隐藏层宽度的增加而导致计算成本和激活内存的线性增加。在大语言模型(LLM)体量不断增大的现在,稀疏混合专家(MoE)架构已成为解决此问题的可行方法,它将模型大小与计算成本分离开来。很多新兴的MoE模型都可以实现相同体量之上,更好的性能与更强大的表现。
最近发现的细粒度MoE扩展定律表明,更高的粒度可带来更好的性能。然而由于计算和优化方面的挑战,现有的MoE模型仅限于低数量专家。
本周二,GoogleDeepMind的新研究引入了一种参数高效的专家检索机制,其利用乘积密钥技术从一百万个微型专家中进行稀疏检索。
链接:https://arxiv.org/abs/2407.04153
该方法尝试通过用于路由的学习索引结构有效地串联到大量微小专家,从而将计算成本与参数计数分离。与密集的FFW、粗粒度MoE和产品密钥存储器(PKM)层相比,表现出卓越的效率。
这项工作引入了参数高效专家检索(PEER)架构(parameterefficientexpertretrieval),利用产品密钥(productkey)检索高效地路由到大量专家,将计算成本与参数量分离。这种设计在实验中展示了卓越的计算性能水平,将其定位为用于扩展基础模型的密集FFW层的竞争性替代方案。这项工作的主要贡献是:
极端MoE设置的探索:与以前的MoE研究中对少数大型专家的关注不同,这项工作研究了众多小型专家的未充分探索的情况。
用于路由的学习索引结构:首次证明学习索引结构可以有效地路由到超过一百万个专家。
新的层设计:将产品密钥路由与单神经元专家相结合,我们引入了PEER层,它可以扩展层容量而无需大量计算开销。实证结果表明,与密集FFW、粗粒度MoE和产品密钥内存(PKM)层相比,其效率更高。

综合消融研究:我们研究了PEER的不同设计选择(例如专家数量、活动参数、头数量和查询批量规范化)对语言建模任务的影响。
方法介绍
本节中,研究者详解了参数高效专家检索(PEER)层,它一种混合专家架构,使用路由中的产品密钥和单神经元MLP作为专家。下图2展示了PEER层内的计算过程。
实验结果
我们首先来看在语言建模数据集上的评估结果。
在根据isoFLOP曲线确定每种方法的计算最优模型后,研究者在以下几个流行语言建模数据集上评估了这些预训练模型的性能:
CurationCorpus
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研究只有一位作者XuHe(Owen),他是GoogleDeepMind研究科学家,单一作者论文,谷歌提出百万专家,超越密集前馈稀疏2017年博士毕业于荷兰格罗宁根大学。