MoE++: 颜水成团队与北大袁粒团队推出新一代MoE架构,专家推理速度倍增,性能全面提升
发表于 2024-10-22 10:55:20

颜水成领衔的昆仑万维2050研究院与北大袁粒团队联合提出了MoE++,相较于传统的MoE,具有更快的推理速度和更高的性能。

MoE++是一个通用框架,能够与任何现有的MoE模型无缝集成。其核心创新在于引入了“零计算量专家”,这一设计带来了三大优势:

(1) 降低计算成本:MoE++允许每个Token使用可变数量的FFN专家,甚至可以完全跳过当前的MoE层。

(2) 提升性能:通过减少简单Token所需的FFN专家数量,MoE++使更多专家能够专注于复杂的Token,释放出比传统MoE更大的性能潜力。

(3) 易于部署:零计算量专家的参数极小,可以在每个GPU上同时部署所有的零计算量专家,避免了分布式FFN专家部署带来的通信开销和专家负载不均的问题。

此外,MoE++还让每个Token在选择专家时参考前一层的路由路径。

实验结果表明,在0.6B到7B参数规模的LLMs上,MoE++在相同模型大小的情况下,相比传统MoE,性能更优,同时实现了1.1到2.1倍的专家吞吐速度。该模型权重已开源。

image.png

图1丨MoE++和普通MoE的对比

MoE++是如何做到的?

现有的大多数混合专家(MoE)方法通常为所有Token激活固定数量的FFN专家。

在许多研究中,每个Token会选择Top-2 FFN专家,并将它们的输出加权合成为下一层的输入。然而,并非所有Token的预测难度都是相同的。

例如,像逗号等简单符号可能只需要一个FFN专家来处理。甚至在某些特殊情况下,某些Token如果与当前MoE层的专家不匹配,绕过该层而不选择Top-2 FFN专家反而可能更为高效。

基于这一见解,研究团队认为,现有MoE使用的固定混合机制可能导致训练和推理效率下降,同时限制模型性能。

为了在提升速度的同时增强性能,研究团队提出了一种通用的异构MoE框架,称为MoE++。

具体来说,团队引入了三种零计算量专家:

Zero专家,输出空向量;

Copy专家,将输入直接作为输出;

Constant专家,用可训练的向量替代输入作为输出。

如图1所示,与传统MoE方法不同,MoE++允许每个Token使用可变数量的FFN专家,接受恒定向量的替换,甚至完全跳过当前的MoE++层。

这种异构结构通过扩大网络的组合空间,提升了模型的拟合能力,并显著降低了计算成本。此外,研究团队还将前一层的路由分数整合到当前层的专家选择中,使Token在选择专家时能够参考其先前的路由路径,从而实现更稳定的专家分配。

研究团队认为,新设计的MoE架构应满足以下标准:

(1) 设计应尽量简化,以高效处理简单的Token;

(2) 为了确保公平比较,新增参数应保持在可忽略的范围。

在这些原则的指导下,研究团队引入了零计算量专家,每个专家仅执行最基本的操作。如图2(a)所示,团队设计了三种类型的零计算量专家:Zero专家、Copy专家和Constant专家,分别对应丢弃、跳过和替换操作。

image.png

图2丨MoE++的核心组成部分

Zero专家。最简单的零计算量专家是丢弃当前输入的Zero专家。本质上,Zero专家的存在可以将Top-2 MoE++层降级为Top-1 MoE++层。具体来说,当Zero专家被激活时,Top-2 MoE++层的输出将等同于另一个专家的单独输出。这样,Zero专家的引入提升了模型在处理简单Token和复杂Token时的灵活性。

Copy专家。Copy专家直接将输入作为输出,直观上相当于跳过当前MoE++层。具体而言,当输入Token与现有专家的匹配较差时,选择绕过MoE++层可能更为有利。

Constant专家Constant专家通过可训练向量替换输入Token。然而,完全替换会导致输入Token信息的丢失。为此,研究团队引入了可训练的权重矩阵,用于动态预测替换的比例。由于Constant专家的计算开销极小,因此仍被归类为零计算量专家。

路由分数残差。由于MoE++包含异构专家,因此与普通MoE相比,路由器的设计变得更加关键。为此,如图2(b)所示,研究团队提出了一种路径感知路由器,它在选择合适的专家时考虑了前一层所采用的路径。具体来说,MoE++将前一层的路由分数通过一个可学习的转换矩阵合并到当前层的专家选择中。这些路由分数残差使每个Token在选择专家时可以考虑其先前的路由路径。

为什么MoE++比MoE更好(Why)

1灵活的计算量分配。MoE++通过为简单Token分配较少的FFN专家,优化了计算资源的分配,从而使更多FFN专家能够专注于处理更具挑战性的Token。

正如图3所示,研究团队发现,在MoE++中,动词激活的FFN专家数量最多,其次是名词,而拆分后的词片激活的FFN数量最少。这表明,MoE++能够让语义较少的Token使用更少的FFN专家,从而释放更多专家来处理语义更丰富的Token。因此,MoE++不仅减少了计算开销,还提升了整体性能。

image.png

图3丨在MoE++中不同Token所需要的平均FFN专家数量

2稳定的路由。MoE++将前一层的路由分数合并到当前层的专家选择中。这些路由分数残差使每个Token在选择专家时考虑其先前的路由路径。

如图4所示,路由分数残差有效地建立了不同MoE++层之间的联系,减小了路由分数的方差。同时,路由分数残差不改变路由分数的均值和取值范围。因此,路由分数残差有助于在MoE++中实现异构专家架构的稳定路由。

image.png

图4丨路由分数残差对路由分数分布的影响

3)更低的计算复杂度。如下表所示,MoE++具有比普通MoE更低的理论计算复杂度。

image.png

实验结果

从0.6B的参数量逐渐扩展到7B参数量的大量实验结果表明,MoE++方法明显优于普通MoE方法。

与相同大小的普通MoE模型相比,MoE++的专家吞吐量提高了约15% ~ 111%,同时具有更高的性能。

image.png

image.png

现有的LLMs模型通常需要大量的训练预算,比如OpenMoE-8B/32E使用1.1T Tokens,TinyLlama-1.1B使用3T Tokens。

研究人员也将MoE++模型的训练预算扩展到1T Tokens。研究人员发现MoE++模的性能与具有2到3倍激活参数的稠密模型相当。

值得注意的是,MoE++优于OpenMoE-8B/32E,这是一个从零开始训练的更大的MoE模型,使用更多的训练Tokens。这些结果表明,MoE++框架是一种很有前途的LLMs框架方案。

image.png

任务级专家负载分布的可视化

研究人员探索了MoE++模型中跨不同任务的专家负载分。这些可视化揭示了几个有趣的发现:

(1)专家负载在不同层之间存在相关性,特别是在相邻层之间。例如,当第j层激活很大比例的FFN专家时,第j + 1层很可能也会以同样大的比例激活FFN专家。

(2)与中间层相比,浅层和最后一层的专家分配模式在不同任务之间的差异更大。这表明该模型主要通过其浅层和最终层而不是中间层来适应不同的任务。未来的工作可以集中在这些层中设计更复杂的结构,以增强模型对不同任务的适应性。

(3)不同任务中每个Token激活的FFN专家数量存在显著差异,但并不一定是更简单的任务激活更少的FFN专家。例如,ARC Challenge任务通常比ARC Easy任务激活更多的FFN专家。这些结果表明,MoE++模型根据知识内容和Token级别的复杂性来分配专家,而不是根据整体任务难度来分配专家。

(4)在所有专家类型中,Zero专家的平均激活次数最高,更简单的任务显示出更高的平均激活次数。例如,ARC Easy任务比ARC Challenge任务激活更多的零专家。这表明Zero专家的激活水平可能可以作为模型任务难度的一个指标。

(5)在MoE++模型的所有层中,不同任务主题的专家分配差异显著,这表明MoE++模型通过采用不同的专家分配模式来处理不同主题的任务。

image.png

image.png

image.png

image.png

image.png

「免责声明」:以上页面展示信息由第三方发布,目的在于传播更多信息,与本网站立场无关。我们不保证该信息(包括但不限于文字、数据及图表)全部或者部分内容的准确性、真实性、完整性、有效性、及时性、原创性等。相关信息并未经过本网站证实,不对您构成任何投资建议,据此操作,风险自担,以上网页呈现的图片均为自发上传,如发生图片侵权行为与我们无关,如有请直接微信联系g1002718958。

CSDN官方微信
扫描二维码,向CSDN吐槽
微信号:CSDNnews
微博关注
【免责声明:CSDN本栏目发布信息,目的在于传播更多信息,丰富网络文化,稿件仅代表作者个人观点,与CSDN无关。其原创性以及文中陈述文字和文字内容未经本网证实,对本文以及其中全部或者部分内容、文字的真实性、完整性、及时性本网不做任何保证或者承诺,请读者仅作参考,并请自行核实相关内容。您若对该稿件有任何怀疑或质疑,请立即与CSDN联系,我们将迅速给您回应并做处理。】