更具象化地理解一个技术,是很有益的。
最近看到一个工作,可视化图解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)质量的技术。
两个主要组成部分定义了专家混合模型:
-
专家(Experts)
—— 每一个前馈神经网络(FFNN)层现在都有一组“专家”,可以选择其中的一部分。这些“专家”通常是前馈神经网络本身。
-
路由器或门控网络(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不成比例地发送给一个专家而不是另一个专家,那么这也可能导致训练不足: