可视化图解MOE大模型的7个核心问题:专家、路由、负载均衡及其用于视觉模态


今天是2025年3月2日,星期日,北京,小雨。

更具象化地理解一个技术,是很有益的。最近看到一个工作,可视化图解moe讲的挺不错,转载过来并做翻译,当作记录,值得收藏慢慢看。

原文在:https://newsletter.maartengrootendorst.com/p/a-visual-guide-to-mixture-of-experts。主要包括7个问题:

1、什么是专家混合模型(Mixture of Experts, MoE)?
2、什么是专家?
3、路由机制如何工作?
4、如何进行负载均衡?
5、如何通过Switch Transformer简化MoE?
6、专家混合模型在视觉模型中如何应用?
7、Mixtral 8x7B中的活跃参数与稀疏参数?

专题化,体系化,会有更多深度思考。大家一起加油。


一、什么是专家混合模型(Mixture of Experts, MoE)?

专家混合模型(Mixture of Experts, MoE)是一种通过使用许多不同的子模型(或称为“专家”)来提升大语言模型(LLMs)质量的技术。

两个主要组成部分定义了专家混合模型:

  1. 专家(Experts)—— 每一个前馈神经网络(FFNN)层现在都有一组“专家”,可以选择其中的一部分。这些“专家”通常是前馈神经网络本身。
  2. 路由器或门控网络(Router or gate network)—— 决定哪些token(tokens)被发送到哪些专家。

在带有专家混合模型的每一层大语言模型中,都能找到(相对专业化的)专家。

了解“专家”并不是指在特定领域(如“心理学”或“生物学”)的专业化。最多,它只是在词汇层面上学习语法信息:

更具体地说,它们的专长在于处理特定上下文中的特定token。

路由器(门控网络)会选择最适合给定输入的专家:

每个专家并不是一个完整的语言模型,而是语言模型架构的一部分子模型。

二、什么是专家?

为了探究专家的含义以及它们是如何工作的,首先需要考察专家混合模型(MoE)所替代的对象:密集层。

1、密集层

专家混合模型(MoE)从语言模型的一个基本功能出发,即前馈神经网络(FFNN)。

请记住,标准的仅解码器的Transformer架构在层归一化之后应用了FFNN:

FFNN允许模型利用由注意力机制生成的上下文信息,并将其进一步转换以捕捉数据中更复杂的关系。

然而,FFNN的规模会迅速增长。为了学习这些复杂的关系,它通常会扩展其接收到的输入:

2、稀疏层

传统Transformer中的FFNN被称为密集模型,因为其所有参数(权重和偏置)都被激活。没有任何部分被遗漏,所有参数都被用来计算输出。

如果仔细观察密集模型,会发现输入在一定程度上激活了所有参数:

相比之下,稀疏模型只激活其总参数的一部分,并且与专家混合模型密切相关。

为了说明这一点,可以将密集模型分割成若干部分(所谓的“专家”),重新训练它,并且在给定时间只激活其中一部分专家:

其核心思想是,每个专家在训练过程中学习不同的信息。然后,在推理时,只使用那些与给定任务最相关的特定专家。

当被问及一个问题时,可以选择最适合该任务的专家:

3、专家学习了什么?

正如之前所见,专家学习的信息比整个领域更加细致。因此,有时将它们称为“专家”被认为具有误导性。

在ST-MoE论文中,编码器模型的专家专业化如下图所示: 然而,解码器模型中的专家似乎并没有这种类型的专业化。但这并不意味着所有专家都是相同的。

一个很好的例子可以在Mixtral 8x7B论文中找到,其中每个token都用第一个专家选择的颜色token。

这个可视化还表明,专家倾向于关注语法而不是特定领域。

因此,尽管解码器专家似乎没有特定的专业化,但它们似乎始终被用于某些类型的token。

4、专家的架构

尽管将专家视为被分割的密集模型的隐藏层很直观,但它们通常是完整的前馈神经网络:

由于大多数语言模型有多个解码器模块,给定文本会在生成文本之前通过多个专家:

所选专家可能因token而异,从而导致不同的“路径”被采用:

如果更新解码器模块的可视化,它现在将包含更多的前馈神经网络(每个专家一个):

解码器模块现在有多个前馈神经网络(每个都是“专家”),在推理过程中可以使用。

三、路由机制如何工作?

既然已经有一组专家,模型是如何知道该使用哪些专家的呢?

在专家之前,添加了一个路由器(也称为门控网络),它被训练用来为给定token选择专家。

1、路由器

路由器(或门控网络)也是一个前馈神经网络,用于根据特定输入选择专家。它输出概率值,用这些概率值来选择最匹配的专家:

专家层返回所选专家的输出乘以门控值(选择概率)。

路由器与专家(其中只选择少数)共同构成了MoE层:

给定的MoE层有两种形式,要么是稀疏的专家混合,要么是密集的专家混合。

两者都使用路由器来选择专家,但稀疏MoE只选择少数专家,而密集MoE选择所有专家,尽管可能以不同的分布方式。

例如,对于一组token,MoE会将token分配给所有专家,而稀疏MoE只会选择少数专家。

在当前的语言模型中,当看到“MoE”时,它通常是指稀疏MoE,因为它允许使用专家的一个子集。这种计算成本更低,对于语言模型来说是一个重要的特性。

2、专家如何选择

门控网络可以说是任何MoE中最重要的组成部分,因为它不仅决定了在推理过程中选择哪些专家,还在训练过程中发挥关键作用。

在其最基本的形式中,将输入(x)乘以路由器权重矩阵(W):

然后,在输出上应用SoftMax,为每个专家创建一个概率分布G(x):

路由器使用这个概率分布来选择给定输入的最佳匹配专家。

最后,将每个路由器的输出与每个选定的专家的输出相乘,并将结果相加。

将所有内容放在一起,探索输入是如何通过路由器和专家流动的:

3、路由的复杂性

然而,这种简单的函数通常会导致路由器选择相同的专家,因为某些专家可能比其他专家学得更快:

不仅会选择的专家分布不均,而且有些专家可能根本得不到训练。这在训练和推理过程中都会导致问题。

相反,希望在训练和推理过程中,专家的重要性是平衡的,这被称为负载均衡。在某种程度上,这是为了防止对相同的专家过度拟合。

四、如何进行负载均衡?

为了平衡专家的重要性,需要关注路由器,因为它是在给定时间决定选择哪些专家的主要组件。

1、KeepTopK

一种负载均衡路由器的方法是通过一个简单的扩展,称为KeepTopK。通过引入可训练的(高斯)噪声,可以防止总是选择相同的专家:

然后,除了你想要激活的顶部k个专家(例如2个)之外,所有其他专家的权重将被设置为-∞:

通过将这些权重设置为-∞,SoftMax对这些权重的输出将导致概率为0:

KeepTopK策略是许多语言模型仍然使用的策略,尽管有许多有希望的替代方案。注意,KeepTopK也可以在没有额外噪声的情况下使用。

2、token选择

KeepTopK策略将每个token路由到少数选定的专家。这种方法被称为token选择,允许将给定token发送到一个专家(top-1路由):

或者发送到多个专家(top-k路由):

一个主要好处是,它允许对专家的贡献进行加权和整合。

3、辅助损失

为了在训练过程中实现专家的更均匀分布,辅助损失(也称为负载均衡损失)被添加到网络的常规损失中。

它增加了一个约束,强制专家具有同等重要性。

这个辅助损失的第一个组成部分是将每个专家的路由器值在整个批次中求和:

这为提供了每个专家的重要性分数,它代表了无论输入是什么,给定专家被选择的可能性。

可以用它来计算变异系数(CV),它告诉专家之间的重要性分数有多不同。

例如,如果重要性分数的差异很大,CV将很高:

相比之下,如果所有专家的重要性分数相似,CV将很低(这是期望的):

使用这个CV分数,可以在训练过程中更新辅助损失,使其尽可能降低CV分数(从而给予每个专家同等重要性):

最后,辅助损失作为单独的损失在训练过程中进行优化。

4、专家容量

不平衡不仅出现在被选择的专家中,还出现在发送到专家的token分布中。

例如,如果输入token不成比例地发送给一个专家而不是另一个专家,那么这也可能导致训练不足:

在这里,问题不仅在于使用了哪些专家,还在于它们被使用的程度。

解决这个问题的一个方法是限制一个给定专家可以处理的token数量,即专家容量。当一个专家达到容量时,随后的token将被发送到下一个专家:

如果两个专家都达到了容量,token将不会被任何专家处理,而是被发送到下一层。这被称为token溢出。

五、如何通过Switch Transformer简化MoE?

Switch Transformer是第一个处理MoE训练不稳定性问题(如负载均衡)的基于Transformer的MoE模型。它简化了架构和训练过程,同时提高了训练稳定性。

1、切换层

Switch Transformer是一个T5模型(编码器-解码器),它用切换层替换了传统的FFNN层。切换层是一个稀疏的MoE层,为每个token选择一个专家(Top-1路由)。

路由器在计算选择哪个专家时没有使用任何特殊技巧,而是对输入乘以专家权重后的结果取SoftMax(与之前所做的相同)。

这种架构(top-1路由)假设路由器只需要1个专家来学习如何路由输入。这与之前看到的情况相反,之前假设token应该被路由到多个专家(top-k路由)以学习路由行为。

2、容量因子

容量因子是一个重要的值,因为它决定了一个专家可以处理多少token。Switch Transformer通过引入直接影响专家容量的容量因子来扩展这一概念。

专家容量的组成部分很简单:

如果增加容量因子,每个专家将能够处理更多的token。

然而,如果容量因子过大,将浪费计算资源。相反,如果容量因子过小,由于token溢出,模型性能将下降。

3、辅助损失

为了进一步防止token丢失,引入了一个简化的辅助损失版本。

与其计算变异系数,这个简化的损失对每个专家的token分配比例与路由器概率比例进行加权:

由于目标是让token在N个专家之间均匀路由,希望向量P和f的值为1/N。

α是一个超参数,可以在训练过程中用它来微调这个损失的重要性。如果值过高,它将取代主要损失函数;如果值过低,它对负载均衡的作用将很小。

六、专家混合模型在视觉模型中如何应用?

MoE并非仅限于语言模型。视觉模型(如ViT)利用基于Transformer的架构,因此也有潜力使用MoE。

简单回顾一下,ViT(视觉Transformer)是一种架构,它将图像分割成小块,这些小块的处理方式类似于token。

这些小块(或token)在被送入常规编码器之前,会被投影到嵌入空间(并添加位置嵌入):

当这些小块进入编码器时,它们被当作token处理,这使得这种架构非常适合使用MoE。

1、Vision-MoE

Vision-MoE(V-MoE)是图像模型中MoE的第一个实现之一。它采用了之前看到的ViT,并用稀疏MoE替换了编码器中的密集FFNN。

这使得通常比语言模型更小的ViT模型能够通过添加专家进行大规模扩展。

为了减少硬件限制,每个专家都使用了一个小的预定义专家容量,因为图像通常有许多小块。然而,低容量往往会导致小块被丢弃(类似于token溢出)。

为了保持低容量,网络会为小块分配重要性分数,并优先处理这些小块,这样溢出的小块通常不那么重要。这被称为批量优先级路由。

因此,即使token的比例减少,仍然应该看到重要的小块被路由。

优先级路由通过专注于最重要的小块,允许处理更少的小块。

2、从稀疏到软MoE

在V-MoE中,优先级评分器有助于区分重要和不重要的小块。然而,小块被分配给每个专家,未处理的小块中的信息会丢失。

软MoE旨在从小块(token)的离散分配转变为软分配,通过混合小块来实现。

首先,将输入x(小块嵌入)乘以一个可学习的矩阵Φ。这为提供了路由器信息,它告诉某个token与给定专家的相关性。

然后,对路由器信息矩阵(在列上)取SoftMax,更新每个小块的嵌入。

更新后的小块嵌入实际上是所有小块嵌入的加权平均值。

从视觉上看,就好像所有小块都被混合了。这些混合的小块然后被发送到每个专家。在生成输出后,它们再次与路由器矩阵相乘。

路由器矩阵在token级别影响输入,在专家级别影响输出。

因此,得到的是被处理的“软”小块/token,而不是离散的输入。

七、Mixtral 8x7B中的活跃参数与稀疏参数?

MoE之所以有趣,一个很大的原因是它的计算需求。由于在给定时间只使用一部分专家,拥有比实际使用的更多的参数。

尽管一个给定的MoE有更多的参数需要加载(稀疏参数),但在推理过程中,由于只使用一些专家,激活的参数更少。

换句话说,仍然需要将整个模型(包括所有专家)加载到你的设备上(稀疏参数),但在推理时,只需要使用其中的一部分(活跃参数)。MoE模型需要更多的显存来加载所有专家,但在推理时运行速度更快。

以Mixtral 8x7B为例,探讨稀疏参数与活跃参数的数量。

在这里,可以看到每个专家的大小为56亿参数,而不是70亿(尽管有8个专家)。

需要加载8×56亿参数(467亿参数,加上所有共享参数),但在推理时,只需要使用2×56亿参数(128亿参数)。

参考文献

1、https://newsletter.maartengrootendorst.com/p/a-visual-guide-to-mixture-of-experts

关于

老刘,NLP开源爱好者与践行者,主页:https://liuhuanyong.github.io。

对大模型&知识图谱&RAG&文档理解感兴趣,并对每日早报、老刘说NLP历史线上分享、心得交流等感兴趣的,欢迎加入社区,社区持续纳新。

加入会员方式:关注公众号,在后台菜单栏中点击会员社区->会员入群加入



(文:老刘说NLP)

欢迎分享

发表评论