文|庞德公
编辑|郭嘉
随着隐藏层宽度的增加,标准Transformer架构中的前馈 (FFW) 层会导致计算成本和激活内存的线性增加。在阅读本篇文章之前建议阅读稀疏混合专家 (MoE) 架构。MoE已经成为解决此问题的可行方法,它将模型大小与计算成本分离。
最近发现的MoE缩放定律表明,更高G可以带来更好的性能。由于计算和优化挑战,现有的 MoE 模型仅限于少数专家。Google提出的PEER(参数高效专家检索),这是一种新颖的层设计,它利用Product Key Memory从大量微型专家(超过一百万)中进行稀疏检索。在语言建模任务上的实验表明,在性能计算方面,PEER层优于密集FFW和粗粒度 MoE。通过高效利用大量专家,PEER 释放了进一步扩展Transformer模型的潜力,同时保持了计算效率。
知识点1:假设P为参数总数、每个代币 𝑃active为每个Token被激活的参数,𝑃expert为单个专家的大小。Krajewski(2024)表明,MoE模型的标度定律具有以下形式:ℒ 是最终的测试损失, 𝑎,𝑏,𝑔,𝛾,𝛼,𝛽是常数, 𝐷是训练Token的总数,𝐺是活动专家的粒度:
无论何种规模参数,L肯定是越小越好。为了提高模型性能,那么按照公式就需要扩大𝑃、𝐷和𝐺规模。一方面要让G和P越大越好,另一方面𝑃active 要越小越好,这样才能体现MoE的价值,毕竟内存成本主要由训练和推理期间的𝑃active 决定。
因为专家模型的大小expert=active
于是就诞生了“参数高效专家检索”(PEER),N为100万,Pexpert为1。即这种技术使用了超过一百万个微小的“专家”——只有一个神经元的小型神经网络——而不是传统变压器模型中使用的大型前馈层。
PEER层可以插入到 transformer主干的中间,也可以用于替换 FFW层。给定上一层的状态向量x,查询Q将其映射到查询向量 q(x),然后将其与Product Key进行比较以计算路由器得分并检索前 k 个专家e1, ..., ek。检索到的专家做出预测ei(x)后,使用softmax归一化的路由器得分作为权重,对其输出进行线性组合。
专家数量太多,如何高效地提升Product Key检索十分关键,例如如何在
为了提高效率,将Key拆分成两个独立的
对于这些候选Key,再计算查询向量q(x) 与它们的内积,并选出前k个最终匹配。这种方法将原始的计算复杂度大幅地降低,显著提高了检索效率。这种技术通过减少计算量,实现了在处理大量专家时的高效匹配。
在语言建模实验中,PEER在效率方面优于传统的Transformer模型和以前的MoE方法。在相同的计算能力下,PEER在各种基准测试中表现更好。
缩放定律
这个研究利用了Loss函数反向的设计PEER,也解释了PEER的成功原因。它们从数学上描述MoE模型的性能如何随着其大小和训练数据量的增加而增加。大量的小型专家可以在不大幅增加计算成本的情况下提高模型的整体容量。在同等算力的基础上,这种技术还是表现不俗。
研究人员认为,PEER的另一个优势是“终身学习”的可能性。由于可以很容易地添加新的专家,因此PEER模型理论上可以不断吸收新信息,而不会忘记它已经学到的知识。