当查看大型语言模型LLMs的最新发布时你经常会看到标题中带有“MoE”。这个“MoE”代表什么为什么这么多LLMs都在使用它在这份视觉指南中我们将花时间探索这个重要组件——专家混合MoE通过超过 50 个可视化进行深入探讨在此视觉指南中我们将介绍 MoE 的两个主要组件即专家和路由器以及它们在典型的基于LLM架构中的应用。查看与LLMs相关的更多可视化内容并支持此通讯请参阅我写的关于大型语言模型的书籍什么是混合专家MoE混合专家MoE是一种使用许多不同的子模型或“专家”来提高LLMs质量的技术。两个主要组件定义了 MoE专家 - 每个 FFNN 层现在都有一个“专家”集合其中可以选择子集。这些“专家”通常是 FFNN 本身。路由器或网关网络 - 确定哪些令牌发送给哪些专家。在具有 MoE 的LLM的每一层中我们发现某种程度上是专业的专家“专家”并非在特定领域如“心理学”或“生物学”中专业化。最多它只是词法方面的专家更具体地说他们的专长在于处理特定语境中的特定标记。路由器网关网络选择最适合给定输入的专家们每位专家不是一个完整的LLM而是LLM架构的子模型部分。专家Expert, MoE中的E为了探索专家所代表的内容以及他们的工作方式让我们首先考察 MoE 应该取代的内容密集层。密集层Dence layer混合专家MoE都始于相对基本的功能LLMs即前馈神经网络FFNN。请记住标准解码器 Transformer 架构在层归一化之后应用 FFNN一个 FFNN 允许模型使用由注意力机制创建的上下文信息进一步将其转化为捕捉数据中更复杂关系的能力。FFNN 然而其规模却迅速增长。为了学习这些复杂关系它通常会在接收到的输入基础上进行扩展稀疏层在传统的 Transformer 中FFNN 被称为密集模型因为所有参数其权重和偏差都被激活。没有任何东西被留下所有内容都用于计算输出。如果我们仔细观察密集模型注意输入如何在一定程度上激活所有参数相比之下稀疏模型仅激活其总参数的一部分并且与专家混合模型密切相关。为了说明我们可以将我们的密集模型切割成多个部分所谓专家重新训练它并在给定时间只激活一组专家底层理念是每个专家在训练期间学习不同的信息。然后在运行推理时仅使用与特定任务最相关的特定专家。当被问及问题时我们可以选择最适合执行特定任务的专家专家学到了什么如我们所见专家比整个领域学习到更细粒度的信息。1 因此有时将他们称为“专家”被视为误导。编码器模型在 ST-MoE 论文中的专家专业化。解码器模型专家似乎没有相同的类型专业化。但这并不意味着所有专家都是平等的。一个很好的例子可以在 Mixtral 8x7B 论文中找到其中每个标记都涂上了第一个专家的选择。此视觉图还表明专家往往关注语法而非特定领域。因此尽管解码器专家似乎没有特定的专长但他们似乎一直被用于某些类型的标记。专家架构尽管将专家视为密集模型中切割成块的隐藏层听起来很酷但他们本身往往是完整的全连接神经网络FFNN由于大多数LLMs具有多个解码器块给定文本在生成之前将通过多个专家所选专家可能在标记之间有所不同从而导致采取不同的“路径”如果我们更新解码器块的可视化它现在将包含更多的 FFNN每个专家一个而不是解码器块现在具有多个 FFNN每个都是一个“专家”它在推理期间可以使用。路由机制现在我们有一批专家模型是如何知道使用哪些专家的在专家之前添加了一个路由器也称为网关网络该路由器经过训练能够选择为给定令牌选择哪个专家。路由器路由器或网关网络也是一种前馈神经网络FFNN用于根据特定输入选择专家。它输出概率并使用这些概率来选择最佳匹配的专家专家层返回所选专家的输出乘以门值选择概率。路由器与专家其中只有少数被选中共同构成 MoE 层给定 MoE 层有两种大小要么是稀疏的专家混合要么是密集的专家混合。两者都使用路由器来选择专家但稀疏 MoE 只选择少数而密集 MoE 则选择所有但可能在不同的分布中。例如给定一组标记MoE 会将标记分配给所有专家而稀疏 MoE 只会选择少数专家。在当前状态LLMs下当你看到“MoE”时它通常是一个稀疏 MoE因为它允许你使用专家子集。这在计算上更便宜这是LLMs的一个重要特性。专家选择门控网络可以说是任何 MoE 最重要的组成部分因为它不仅决定了推理期间选择哪些专家还包括训练阶段。在其最基本的形式中我们将输入x乘以路由权重矩阵W然后我们对输出应用 SoftMax为每位专家创建概率分布 G(x)路由器使用这种概率分布来选择给定输入的最佳匹配专家。最后我们将每个路由器的输出与每个选定的专家相乘并将结果相加。让我们把所有内容整合起来探索输入如何通过路由器和专家流动路由的复杂性然而这个简单的函数往往导致路由器选择相同的专家因为某些专家可能比其他专家学习得更快不仅选定的专家分布不均而且一些专家几乎都没有经过培训。这导致在训练和推理过程中都存在问题。相反我们希望在训练和推理过程中专家之间保持同等重要性我们称之为负载均衡。从某种意义上讲这是为了防止对同一专家的过度拟合。负载均衡为了平衡专家的重要性我们需要将路由器视为主要组件以决定在特定时间选择哪些专家。保持顶部 K一种路由器负载均衡的方法是通过一个简单的扩展称为 KeepTopK 2。通过引入可训练高斯噪声我们可以防止总是选择相同的专家然后除了您想要激活的前 k 个专家例如 2 个之外其余专家的权重将被设置为-∞通过将这些权重设置为 -∞这些权重上的 SoftMax 输出将导致概率为 0The KeepTopK 策略是许多 LLMs 尽管有许多有希望的替代方案仍在使用的策略。请注意KeepTopK 也可以在不添加额外噪声的情况下使用。Token 选择The KeepTopK 策略将每个令牌路由到几位选定的专家。这种方法称为 Token Choice 3允许将特定令牌发送给一位专家top-1 路由或超过一个专家top-k 路由一个主要好处是它允许权衡和整合专家各自的贡献。辅助损失为了在训练期间获得更均匀的专家分布网络常规损失中添加了辅助损失也称为负载均衡损失。它增加了一个约束迫使专家们具有同等重要性。该辅助损失的第一个组成部分是对整个批次中每个专家的路由器值进行求和这为我们提供了每位专家的重要性得分它代表了无论输入如何某位专家被选中的可能性。我们可以用这个来计算变异系数CV它告诉我们专家之间的重要性分数差异有多大。例如如果重要性得分差异很大CV 将会很高相比之下如果所有专家的重要性评分都相似则 CV 值将较低这正是我们追求的目标使用此 CV 得分我们可以在训练过程中更新辅助损失使其尽可能降低 CV 得分从而给予每个专家同等的重要性最后将辅助损失作为一个单独的损失项用于训练优化。专家能力不平衡不仅存在于被选中的专家中也存在于发送给专家的代币分布中。例如如果输入标记不成比例地发送给某个专家而不是另一个这也可能导致欠训练这里不仅仅是关于使用哪些专家而是他们被使用了多少。该问题的解决方案是限制特定专家可以处理的令牌数量即专家容量 4。当专家达到容量时产生的令牌将被发送到下一个专家如果两位专家都已达到其容量则令牌不会被任何专家处理而是发送到下一层。这被称为令牌溢出。简化 MoE 与 Switch Transformer之一是处理 MoE如负载均衡训练不稳定问题的第一个基于 transformer 的 MoE 模型——Switch Transformer。5 它简化了大部分架构和训练过程同时提高了训练稳定性。交换层开关变换器是一种 T5 模型编码器-解码器它用切换层替换了传统的 FFNN 层。切换层是一个稀疏 MoE 层为每个标记选择一个专家Top-1 路由。路由器在计算选择哪个专家时没有特殊技巧只是将输入乘以专家的权重后取 softmax与我们之前做的一样。这种架构top-1 路由假设路由器只需要 1 位专家就能学习如何路由输入。这与我们之前所见的假设不同我们之前假设应该将令牌路由到多个专家top-k 路由以学习路由行为。容量系数容量系数是一个重要值因为它决定了专家可以处理多少个标记。Switch Transformer 通过引入直接影响专家容量的容量系数来扩展这一点。专家能力组成部分简单明了如果我们提高容量系数每位专家将能够处理更多的标记。然而如果容量系数过大我们会浪费计算资源。相比之下如果容量系数过小由于标记溢出模型性能将下降。辅助损失为进一步防止丢失标记引入了辅助损失的简化版本。而不是计算变异系数这种简化的损失衡量了分配的标记分数与每个专家的路由概率分数之比由于目标是让 N 个专家之间具有统一的令牌路由我们希望向量 P 和 f 的值为 1/N。α 是一个超参数我们可以用它来微调训练过程中这个损失函数的重要性。值过高会超过主损失函数值过低对负载均衡的帮助很小。专家混合模型MoE 并非仅限于语言模型的技巧。视觉模型如 ViT利用基于 transformer 的架构因此有潜力使用 MoE。作为快速回顾ViT视觉 Transformer是一种将图像分割成类似标记处理的补丁的架构。6这些补丁或令牌随后被投影到嵌入带有额外的位置嵌入中然后再输入到常规编码器中这些补丁进入编码器的那一刻它们就像标记一样被处理这使得这种架构非常适合 MoE。Vision-MoEVision-MoEV-MoE是图像模型中 MoE 的第一个实现之一。7 它采用我们之前看到的 ViT并将编码器中的密集 FFNN 替换为稀疏 MoE。这允许 ViT 模型通常比语言模型尺寸更小通过添加专家进行大规模扩展。每个专家使用了一个小的预定义专家容量以减少硬件限制因为图像通常有许多块。然而低容量往往会造成块被丢弃类似于令牌溢出。为了保持容量低网络将重要性分数分配给补丁并首先处理这些补丁以便溢出的补丁通常不太重要。这被称为批量优先路由。因此即使令牌百分比下降我们仍然应该看到重要的补丁被路由。优先路由通过关注最重要的部分允许处理更少的补丁。在 V-MoE 中优先级评分器帮助区分重要程度不同的补丁。然而补丁被分配给每个专家未处理补丁中的信息丢失。软 MoE 旨在通过混合补丁从离散补丁标记分配过渡到软补丁标记分配。8在第一步中我们将输入 x补丁嵌入与可学习的矩阵Φ相乘。这为我们提供了路由信息告诉我们某个标记与特定专家的相关程度。通过随后对路由信息矩阵按列进行 softmax 操作我们更新每个补丁的嵌入。更新的补丁嵌入实际上是所有补丁嵌入的加权平均值。视觉上就像所有的补丁都被混合在一起了。这些组合补丁随后被发送给每位专家。在生成输出后它们再次与路由矩阵相乘。路由矩阵影响令牌级别的输入和专家级别的输出。因此我们得到“软”补丁/令牌而不是处理离散输入。活跃参数与稀疏参数在 Mixtral 8x7B 中莫埃MoE引人入胜的一大原因是其计算需求。由于在特定时间只使用专家子集我们能够访问比实际使用的更多参数。尽管给定的 MoE 有更多参数需要加载稀疏参数但由于我们只在推理时使用一些专家激活参数因此激活的参数较少。换句话说我们仍然需要在您的设备上加载整个模型包括所有专家稀疏参数但在进行推理时我们只需要使用子集活动参数。MoE 模型需要更多的 VRAM 来加载所有专家但在推理时运行更快。让我们通过一个例子来探讨稀疏参数与活跃参数的数量Mixtral 8x7B。9这里我们可以看到每个专家的大小是 5.6B而不是 7B尽管有 8 个专家。我们不得不加载 8x5.6B46.7B个参数包括所有共享参数但推理时我们只需要使用 2x5.6B12.8B个参数。结论这标志着我们与专家混合模型的旅程结束希望这篇帖子能让你更好地理解这种有趣技术的潜力。现在几乎所有模型集都至少包含一个 MoE 变体感觉它将在这里停留。