图解混合专家模型 (MoE) :50 张图带你深入了解 LLM 的核心组件

混合专家模型(MoE)通过多个专家模型和路由器提升LLM性能,实现稀疏计算,提高效率。

原文标题:混合专家模型 (MoE) 到底是什么?看完这篇,你也能讲明白!

原文作者:图灵编辑部

冷月清谈:

混合专家模型 (MoE) 是一种提升大语言模型 (LLM) 质量的关键技术。它由多个专家模型和一个路由器组成。每个专家模型都是一个前馈神经网络 (FFNN),专门处理特定类型的输入。路由器负责将输入分配给最合适的专家。

MoE 的核心思想是稀疏性,即每次只激活部分专家,从而提高效率并降低计算成本。与传统的密集模型相比,MoE 能够处理更复杂的关系,并更好地捕捉数据中的细微差别。

专家并非特定领域的专家,而是更偏向于语法层面的专家。它们擅长处理特定上下文中的特定标记。路由器通过学习输入的特征,将输入分配给最合适的专家。

负载均衡是 MoE 中一个重要的概念,它旨在确保每个专家都得到充分利用,并防止模型过拟合到某些专家。常用的负载均衡方法包括 KeepTopK 和辅助损失函数。

Switch Transformer 是一个基于 Transformer 的 MoE 模型,它简化了架构和训练过程,并提高了训练稳定性。它使用 Top-1 路由,每次只选择一个专家来处理输入。

MoE 不仅用于语言模型,也用于视觉模型,例如 Vision-MoE (V-MoE)。V-MoE 将图像分割成图像块,并将这些图像块作为 token 输入到 MoE 模型中。

MoE 模型的一个重要特点是其稀疏参数和活跃参数的区别。稀疏参数是指模型中所有专家的参数,而活跃参数是指在推理过程中实际使用的参数。由于只使用部分专家,MoE 模型在推理时运行速度更快,但需要更多的内存来加载所有专家。

怜星夜思:

1、文章提到专家并非领域专家,而是语法专家。那么,如果我想让模型在特定领域表现更好,应该如何利用MoE呢?
2、负载均衡在 MoE 中很重要,但如何选择合适的负载均衡策略呢?KeepTopK 和辅助损失函数各有什么优缺点?
3、MoE 的稀疏性带来了计算效率的提升,但也引入了新的挑战,例如 token 溢出。如何有效地减少 token 溢出,并提高模型的鲁棒性?

原文内容

当我们查看最新发布的大语言模型(LLM)时,常常会看到标题中出现 MoE 这一术语。那么,这个 MoE 代表什么?为什么这么多 LLM 都在使用它呢?
如果你还不知道什么是 MoE ,那就码住今天的文章吧!我们将图解 MoE,用 50 张图细致地探索这一重要组件——混合专家模型(MoE)。
在这篇文章中,我们将深入探讨 MoE 的两个主要组件——专家和路由器,并了解它们在典型 LLM 架构中的应用。

什么是混合专家(MoE)?

混合专家(Mixture of Experts,简称 MoE)模型,是一种利用多个不同的子模型(或“专家”)来提升大语言模型(LLM)质量的技术。

MoE 架构主要由两个核心组件构成:

1.专家模块:每个前馈神经网络(FFNN)层包含一组“专家”,每次可从中选择一个子集。这些“专家”本身通常是 FFNN。  
2.路由或者网关网络,决定将哪些标记(token)分配给哪些专家。

在采用 MoE 架构的大语言模型的每一层中,我们会发现这些专家模块往往具有这样的特点:

要知道,所谓的“专家”并不是专门针对某个领域,如“心理学”或“生物学”。最多,它只是语法层面的专家:

更具体来说,它们的专长是在特定的上下文中处理特定的标记。
路由器(网关网络)来选择最适合给定输入的专家(们):
每个专家并不是一个完整的 LLM,而是 LLM 架构中的子模型部分。

专家(Experts)

为了探讨专家(MoE中的E)的代表意义和它们的工作方式,我们首先需要了解 MoE 原本要替代的是什么——密集层。
密集层 (Dense Layer)
混合专家(MoE)起始于 LLM 相对基础的功能,即前馈神经网络(FFNN)。

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

一个 FFNN 允许模型使用由注意力机制创建的上下文信息,并通过进一步转换这些信息来捕捉数据中更为复杂的关系。
然而,FFNN 的规模会迅速增长。为了学习这些复杂的关系,它通常会扩展它接收到的输入:
稀疏层 (Sparse Layers)

传统 Transformer 中的前馈神经网络(FFNN)被称为密集模型,因为所有参数(权重和偏置)都会被激活。没有任何参数被遗漏,所有参数都参与输出计算。

如果我们更仔细地观察这种密集模型,可以发现输入会在某种程度上激活所有参数:

与此相对,稀疏模型仅激活一部分参数,并且与混合专家(MoE)紧密相关。
为了说明这一点,我们可以将密集模型切割成多个部分(所谓的专家),重新训练它,并且在给定时间内只激活一组专家:
其底层思想是每个专家在训练过程中学习不同的信息。然后,在进行推理时,仅使用与给定任务最相关的特定专家。
当被问到问题时,我们可以选择最适合执行给定任务的专家:

专家(Experts)学到什么?

如前所述,专家学习的内容比整个领域更为细致。因此,称它们为“专家”有时被认为是误导性的。
ST-MoE 论文中编码器模型的专家专业化
然而,解码器模型中的专家似乎并没有同样类型的专业化。不过,这并不意味着所有专家都是平等的。
一个很好的例子可以在 Mixtral 8x7B 论文中找到,其中每个标记都以第一个专家的选择为颜色。
这个视觉图展示了专家倾向于关注语法而非特定领域。
因此,尽管解码器专家似乎没有特定的专长,但它们确实会被一致地用于某些类型的标记。

专家的架构

尽管将专家视为密集模型中的隐藏层并将其切分成块是很形象的做法,但它们通常本身就是完整的 FFNN:
由于大多数 LLM 具有多个解码器块,因此给定文本将在多个专家之间传递,直到文本生成:
所选的专家很可能在标记之间有所不同,这就导致了不同的“路径”被选择:
如果我们更新解码器块的可视化,它现在会包含更多的 FFNN(每个专家一个):
解码器块现在有多个 FFNN(每个都是一个“专家”),可以在推理时使用。

路由机制

现在我们有了一组专家,那么模型是如何知道应该使用哪些专家的呢?
在专家之前,添加了一个路由器(也叫网关网络),它经过训练来选择哪个专家适合给定的标记。
路由器
路由器(或网关网络)也是一种前馈神经网络(FFNN),它根据特定的输入来选择专家。它输出概率,并利用这些概率来选择最匹配的专家:
专家层返回所选专家的输出,并乘以门值(选择概率)。
路由器与专家(其中只有少数几个被选中)一起构成了 MoE 层:
给定的 MoE 层有两种大小,分别是稀疏的混合专家或密集的混合专家模型。
这两者都使用路由器来选择专家,但稀疏 MoE 只选择几个,而密集 MoE 则选择所有专家,但可能以不同的分布进行选择。
例如,给定一组标记,MoE 会将标记分配给所有专家,而稀疏 MoE 只选择少数专家。
在当前的 LLM 状态下,当你看到“MoE”时,它通常是一个稀疏 MoE,因为它允许你使用专家子集。这在计算上更便宜,这是 LLM 的一个重要特性。
专家的选择
门控网络可以说是任何 MoE 模型中最重要的组成部分,因为它不仅决定了在推理过程中选择哪些专家,还决定了在训练过程中选择哪些专家。
在最基本的形式中,我们将输入(x)与路由器权重矩阵(W)相乘:
然后,我们对输出应用 SoftMax 操作,为每个专家创建概率分布 G(x):
路由器使用这种概率分布来选择与给定输入最匹配的专家。
最后,我们将每个路由器的输出与每个选定的专家相乘,并对结果进行求和。
让我们把所有内容放在一起,看看输入是如何通过路由器和专家流动的:

路由的复杂性

然而,这个简单的函数常常导致路由器会选择相同的专家,因为某些专家可能比其他专家学习得更快:
不仅会导致选择的专家分布不均,还有一些专家几乎不会被训练到,这导致在训练和推理过程中都会出现问题。
因此,我们希望在训练和推理过程中,专家之间具有相等的重要性,这就是我们所说的负载均衡。从某种意义上说,这是为了防止对同一专家的过拟合

负载均衡

为了平衡专家的重要性,我们需要关注路由器,因为它是决定在特定时刻选择哪些专家的主要组成部分。
KeepTopK
负载均衡路由器的一种方法是通过一个简单的扩展,叫做 KeepTopK。通过引入可训练的(高斯)噪声,我们可以防止总是选择相同的专家:
然后,除了你想激活的前 k 个专家(例如 2 个)之外,其余专家的权重将被设置为 -∞:
通过将这些权重设置为 -∞,这些权重上的 SoftMax 输出将导致概率为 0:
KeepTopK 策略是许多大型语言模型仍然使用的一种方法,尽管有很多有前景的替代方案。请注意,KeepTopK 也可以在没有添加额外噪声的情况下使用。
Token 选择
KeepTopK 策略是将每个 token 路由发送到少数几个选定的专家。这种方法称为 Token Choice,允许将给定的 token 发送到一个专家(top-1 路由):
或发送到多个专家(top-k 路由):
一个主要的好处是,它允许对专家的贡献进行加权和整合。
辅助损失
为了在训练期间获得更均匀的专家分布,辅助损失(也称为负载均衡损失)被添加到网络的常规损失中。
它添加了一个约束,强制要求专家具有相等的重要性。
这个辅助损失的第一个组成部分是对整个批次中每个专家的路由器值进行求和:
这给出了每位专家的重要性分数,表示无论如何输入,某位专家被选择的可能性。
我们可以使用这个来计算变异系数(CV),它告诉我们不同专家之间的重要性分数差异有多大。
例如,如果重要性分数差异很大,CV 就会很高:
相反,如果所有专家的权重分数相似,CV 就会很低(这是我们希望的):
使用这个 CV 分数,我们可以在训练过程中更新辅助损失,目的是尽可能降低 CV 分数(从而给每个专家相等的重要性):
最后,辅助损失作为一个独立的损失加入到训练中进行优化。

专家容量

不平衡不仅存在于被选择的专家中,也存在于分配给专家的 token 分布中。
例如,如果输入 token 不成比例地分配给某个专家,而忽视了其他专家,那么也可能导致训练不足:
这里的问题不仅是哪些专家被使用,而是它们被使用的频率。
解决这个问题的一种方法是限制每个专家可以处理的 token 数量,即专家容量。当一个专家的容量达到上限时,剩余的 token 会被发送到下一个专家:
如果两个专家的容量都已满,那么 token 将不会被任何专家处理,而是被发送到下一层,这就是 token 溢出。

通过 Switch Transformer 简化 MoE

Switch Transformer 是第一个基于 transformer 的 MoE 模型,它处理了 MoE 中的训练不稳定性问题(如负载均衡)。它简化了架构和训练过程,同时提高了训练稳定性。
切换层
Switch Transformer 是一个 T5 模型(编码器-解码器),它用切换层替代了传统的 FFNN 层。切换层是一个稀疏 MoE 层,为每个 token 选择一个专家(Top-1 路由)。
路由器不进行任何特殊计算,它只是计算输入与专家权重的乘积,并取其 softmax(与之前做的操作相同)。
这种架构(Top-1 路由)假设每次路由只需要一个专家来学习如何路由输入。这与我们之前看到的有所不同,后者假设 token 应该路由到多个专家(Top-k 路由)来学习路由行为。
容量因子
容量因子是一个重要值,它决定了每个专家可以处理多少 token。Switch Transformer 扩展了这一点,引入了一个直接影响专家容量的容量因子。
专家容量的组成部分非常直接:
如果我们增加容量因子,每个专家就能处理更多的 token。
然而,如果容量因子过大,我们就浪费计算资源。相反,如果容量因子过小,由于 token 溢出,模型性能将下降。
辅助损失
为了进一步防止丢失 token,引入了辅助损失的简化版本
与其计算变异系数,这个简化的损失衡量了每个专家的 token 分派比例与路由器概率分布之间的比率:
由于目标是让 N 个专家之间的路由统一,我们希望向量 P 和 f 的值为 1/N。
α 是一个超参数,我们可以用它来微调训练过程中这个损失函数的权重。过高的值会超越主损失函数,过低的值对负载均衡几乎没有作用。

视觉模型中的混合专家(MoE)

MoE 不仅是语言模型独有的技术。视觉模型(如 ViT)利用基于 transformer 的架构,因此也有可能使用 MoE。
简要回顾一下,ViT(视觉 transformer)是一种将图像分割为图像块并像 token 一样处理的架构。
这些图像块(或 token)随后被映射为嵌入(并添加了额外的位置嵌入),然后再输入到常规编码器中:
当这些图像块进入编码器时,它们就像 token 一样被处理,这使得该架构能够很好地利用 MoE。
Vision-MoE
Vision-MoE(V-MoE)是 MoE 在图像模型中的首次实现之一。它以我们之前看到的 ViT 为基础,并将编码器中的密集 FFNN 替换为稀疏 MoE。
这使得通常比语言模型小的 ViT 模型可以通过添加专家来进行大规模扩展。
为了减少硬件约束,每个专家使用了小的预定义专家容量,因为图像通常有很多图像块。然而,低容量往往会导致图像块丢失(类似于 token 溢出)。
为了保持容量低,网络会为图像块分配重要性分数,并优先处理那些重要性较高的图像块,这叫做批处理优先路由。
因此,如果 token 的百分比下降,我们仍然应该看到重要的图像块被路由。
优先路由使得我们能集中处理最重要的图像块,从而减少处理的图像块数量。
从稀疏混合专家到软混合专家
在 V-MoE 中,优先级评分器帮助区分更重要和不太重要的图像块。然而,图像块被分配给每个专家后,未处理的图像块中的信息会丢失。
软 MoE 旨在通过混合图像块,从离散到软图像块(token)分配。
在第一步中,我们将输入 x(图像块嵌入)与可学习的矩阵 Φ 相乘。这为我们提供了路由信息,告诉我们某个 token 与给定专家的相关性。
然后,通过对路由信息矩阵(列)进行 softmax 操作,我们更新每个图像块的嵌入。
更新后的图像块嵌入实际上是所有图像块嵌入的加权平均值。
从视觉上看,就好像所有图像块都被混合了。这些合成的图像块然后被发送到每个专家。生成输出后,它们再次与路由矩阵相乘。
路由矩阵在 token 层面影响输入,在专家层面影响输出。
最终,我们得到的是“软”图像块/token,而不是离散输入。

Mixtral 8x7B 的稀疏参数与活跃参数

MoE 的一个重要特点是其计算需求。由于在给定时间只使用专家子集,因此我们可以访问比我们实际使用更多的参数。
尽管一个 MoE 有更多的参数需要加载(稀疏参数),但因为我们只在推理过程中使用部分专家,所以激活的参数较少。
换句话说,我们仍然需要将整个模型(包括所有专家)加载到设备上(稀疏参数),但在推理时,我们只需要使用其中的一部分(活跃参数)。MoE 模型需要更多的 VRAM 来加载所有专家,但在推理时运行速度更快。
让我们通过一个例子——Mixtral 8x7B,来探索稀疏参数与活跃参数的数量。
这里,我们可以看到每个专家的大小是 5.6B,而不是 7B(尽管只有 8 个专家)。
我们不得不加载 8x5.6B(46.7B)个参数(以及所有共享参数),但我们只需要使用 2x5.6B(12.8B)个参数进行推理。

结论

到这里关于混合专家(MoE)的探讨就结束了!希望这篇文章能让你对这一有趣技术的潜力有更深的理解。如今,几乎所有的模型集合中都包含了至少一个 MoE 变体,感觉它已经成为了长期存在的技术。
本文作者 Maarten Grootendorst,他也是“袋鼠书”Hands-On Large Language Models 的作者如果你想查看更多关于大语言模型可视化的内容,相信“袋鼠书”能够帮到你,这本书的中文版《图解大模型》,将于 2025 年 4 月上市,小伙伴们敬请期待~

英文版封面,中文版封面

《图解大模型》

Jay Alammar,Maarten Grootendorst | 著

李博杰 | 译

美亚 4.7 星评,畅销书 Hands-on 系列新作。只要具备 Python 基础,就可以通过本书学习大语言模型,并将大语言模型的能力应用到真正的 AI 实践中。

本书将为 Python 开发人员提供使用大模型的实用工具和概念,帮助大家掌握实际应用场景。你将学习如何利用预训练的大型语言模型进行文案撰写、文本摘要、语义搜索等任务,构建超越关键词匹配的智能系统。


扫码进群,一起学习~👇


初审:邢璐

复审:董冲
终审:郑梦鸽

我感觉可以根据实际情况选择,如果模型对负载均衡的要求不高,可以用 KeepTopK;如果对负载均衡的要求比较高,可以用辅助损失函数,或者两者结合使用。也可以参考一些相关的论文或开源代码,看看别人是怎么做的。

顺着楼上的思路,我觉得可以更激进一点,直接用特定领域的模型来初始化MoE的专家,而不是从头训练。这样相当于给专家注入了领域知识,应该能更快地达到理想的性能。

补充一下,KeepTopK 适合专家数量较少的情况,而辅助损失函数更适合专家数量较多的情况。另外,辅助损失函数的超参数调整比较 tricky,需要一些经验和技巧。

另外,我觉得 token 溢出本身也提供了一些信息,可以利用这些信息来改进模型。比如,可以把溢出的 token 作为一种特殊的输入,训练一个专门处理溢出 token 的专家。

领域知识的融入是个难题,我觉得可以尝试结合知识图谱,将领域知识注入到专家的训练数据中,或者直接用领域知识图谱来指导路由器的选择,让路由器能更好地理解输入的领域上下文,从而选择更合适的专家。

针对你的问题,我觉得可以考虑预训练阶段对特定领域的语料进行强化训练,并在MoE的专家选择上做一些调整,比如增加特定领域相关的专家数量,或者修改路由机制,使其更倾向于选择特定领域的专家。这样应该能提升模型在该领域的性能。

关于如何减少 token 溢出,我觉得可以从两方面入手,一是增加专家容量,二是改进路由机制。增加专家容量可以减少溢出的概率,但会增加计算成本;改进路由机制可以更有效地分配 token,避免某些专家过载。

还可以考虑动态调整专家容量,根据输入的复杂度动态调整每个专家的容量,这样可以更好地平衡计算效率和模型性能。

关于负载均衡策略的选择,KeepTopK 的优点是简单易实现,计算开销小。但缺点是可能导致某些专家利用率不足。辅助损失函数的优点是可以更精细地控制专家选择,但缺点是实现起来比较复杂,需要调整额外的超参数。