探秘Transformer系列之（21）--- MoE

探秘Transformer系列之（21）--- MoE


目录

• 探秘Transformer系列之（21）--- MoE
  
  
  
  • 0x00 概要
    
  • 0x01 前置知识
    
    
    
    • 1.1 MoE出现的原因
      
      
      
      • 1.1.1 神经网络的稀疏性
        
      • 1.1.2 神经网络的过载性
        
      • 1.1.3 神经元的多语义性
        
      • 1.1.4 计算资源的有限性
        
      
      
    • 1.2 MoE的核心理念
      
    
    
  • 0x02 发展历史
    
    
    
    • 2.1 重要节点
      
      
      
      • 2.1.1 Adaptive mixtures of local experts
        
      • 2.1.2 sparsely-gated mixture-of-experts layer
        
      • 2.1.3 GShard
        
      • 2.1.4 Swith Transformer
        
      
      
    • 2.2 详细时间线
      
    
    
  • 0x03 模型结构
    
    
    
    • 3.1 门控函数
      
      
      
      • 3.1.1 条件计算
        
      • 3.1.2 定义
        
      • 3.1.3 特点
        
      • 3.1.4 优化
        
        
        
        • 关键因素
          
        • 改进示例
          
        
        
      
      
    • 3.2 专家
      
      
      
      • 3.2.1 特点
        
      • 3.2.2 种类
        
      • 3.2.3 位置
        
      
      
    • 3.3 分类
      
      
      
      • 3.3.1 稠密 vs 稀疏
        
      • 3.3.2 软性（soft）门控
        
      
      
    • 3.4 比对
      
    
    
  • 0x04 计算流程
    
    
    
    • 4.1 算法
      
    • 4.2 流程
      
    • 4.3 Permutation
      
    • 4.4 实现
      
      
      
      • 4.4.1 Mistral Inference
        
      • 4.4.2 Mixtral 8x7B
        
      
      
    • 4.5 参数量
      
    • 4.6 计算量
      
    
    
  • 0x05 并行计算
    
    
    
    • 5.1 通讯需求
      
      
      
      • 5.1.1 单token
        
      • 5.1.2 多token
        
      
      
    • 5.2 专家并行
      
      
      
      • 5.2.1 定义
        
      • 5.2.2 历史
        
        
        
        • 业界鼻祖
          
        • Gshard
          
        • Switch Transformers
          
        • 小结
          
        
        
      • 5.2.3 协同
        
      • 5.2.4 如何切分
        
      • 5.2.5 优势
        
      
      
    • 5.3 All-to-All通信
      
      
      
      • 5.3.1 困境
        
      • 5.3.2 All-to-All
        
      
      
    • 5.4 分布式计算过程
      
      
      
      • 5.4.1 多种范式结合
        
      • 5.4.2 通信复杂度
        
        
        
        • TP vs EP
          
        • 通讯复杂度对比
          
        
        
      • 5.4.3 代码示例
        
        
        
        • DeepSpeed-Megatron
          
        • FastMoE
          
        
        
      
      
    
    
  • 0xFF 参考
    
  
  

0x00 概要

在足够的训练数据下，我们可以通过增加参数和计算预算来扩大语言模型规模就可以得到更强大的模型。然而，与之相关的问题是极高的计算成本。而MoE（Mixture-of-Experts/混合专家）架构通过条件计算，就可以在保持计算成本适度的情况下实现参数扩展，提供增强的模型容量和计算效率。简单理解，MoE就是将多个专家模型混合起来形成一个新的模型。但MoE不是让一个单一的神经网络处理所有任务，而是将工作分配给多个专门的“专家”，由一个门控网络决定针对每不同输入激活哪些专家。

注：全部文章列表在这里，估计最终在35篇左右，后续每发一篇文章，会修改此文章列表。
cnblogs 探秘Transformer系列之文章列表
0x01 前置知识

1.1 MoE出现的原因

MoE的出现有几个主要方面原因：神经网络的稀疏性、神经元的多语义性和计算资源的有限性。我们从中也可以看到FFN的部分劣势。
1.1.1 神经网络的稀疏性

稀疏性是指我们可以仅使用整个系统的某些特定部分执行计算。这意味着并非所有参数都会在处理每个输入时被激活或使用，而是根据输入的特定特征或需求，只有部分相关参数集合被调用和运行。
虽然Transformer构建了庞大的参数网络，但是它的某些层可能会非常稀疏，即某些神经元的激活频率会低于其他神经元。论文“MoEfication: Transformer Feed-forward Layers are Mixtures of Experts”就指出，使用ReLU等激活函数会导致大部分的激活值都是0，这导致FFNs的激活值非常稀疏。这样在每次预测过程中，对于用户当前的问题来说，FFNs 中实际只有一小部分神经元被激活并参与计算。而且模型的规模越大，其稀疏性也越强。大型模型在处理输入时激活的神经元占总体的比例更小。其实，人脑也具备类似的稀疏性。如果对于所有问题，人脑都会使用全部神经元，恐怕人脑中的“CPU”早就烧毁了。
1.1.2 神经网络的过载性

对每个输入，主流深度神经网络都会载入网络中的所有层和神经元，所有模型参数都会一同参与处理该输入数据。因为上面提到的稀疏性，我们可知这意味着在处理大量参数的过程中，需要进行大量的不必要的计算。因此，网络实际上对于它们所做的大多数预测来说都太大了。LLM 成为世界上最低效和最耗能的系统之一。
除了不受控制的消耗之外，针对每个预测运行整个模型也会对性能产生重要影响，参数数量的增加会导致训练和推理过程中计算复杂度和内存消耗的增加。在追求速度和可扩展性的实际应用中部署如此庞大的模型是一项艰巨的任务。
1.1.3 神经元的多语义性

随着应用场景的复杂化和细分化，垂直领域应用更加碎片化，人们对大模型提出了更高的要求，希望一个模型既能回答通识问题，又能解决专业领域问题。
但是，有研究人员发现，神经元具有多义性的特点。也就是说，它们不专注于一个单一的主题，而是专注于许多主题。而且重要的是，它们在语义上可能并不相关。举个例子来说，在神经网络数十亿个神经元中的一个神经元可能每次在输入主题涉及“苹果”被激活，而当输入主题涉及“手机”时，这个神经元也可能被激活。这不仅使神经网络难以解释，而且也不是一个理想的情况。因为单个神经元必须精通各种彼此几乎毫无关系的主题。想象一下，你必须同时成为神经科学和地质学的专家，这将是一项艰巨的任务。
而目前不仅仅是知识范围更加广泛，多模态带来的各自数据集都可能各自的数据特征完全不同，这导致神经元很难获取知识。更糟糕的是，学习曲线可能相互矛盾，学习一个主题的更多知识可能会影响神经元获取另一个主题知识的能力。
1.1.4 计算资源的有限性

模型规模是提升模型性能的关键因素之一。而通常来讲，模型规模的扩展会导致训练成本显著增加，因此，计算资源的限制成为了大规模密集模型训练的瓶颈。
因此，需要一种技术来拆分、消除或至少缓解这些问题。这就是MoE希望达到的目的。
1.2 MoE的核心理念

MoE 的基本思想是将模型的参数计数与其使用的计算量分离。而其背后的理念则是模型的不同组件（即"专家"）在处理数据的不同任务或特征时具有专门化的能力。这种设计灵感来源于人类社会中的专业分工。在现实生活中，如果有一个包括了多个领域知识的复杂问题，我们通常会召集一个专家团队共同解决复杂问题。每位专家都拥有独特的技能。我们先拆分这个大问题到各领域，把不同的任务先分离出来，这样才便于分发给不同领域的专家。然后让各个领域的专家先逐个解决小问题，最后再把大家集合到一起来汇总结论，攻克这个任务。
MoE正是基于上述的理念，它由两个主要部分组成：专家和门控路由机制（或者路由机制）。

• 术业有专攻。模型的不同专家（expert）拥有不同领域的专业知识，负责处理不同的计算任务或者数据。每个专家子网络专门处理输入数据的子集，共同完成一项任务。相较于深度学习网络， MoE更像是宽度学习网络。另外，MoE与集成技术的主要区别在于，对于MoE，通常只有一个或少数几个专家模型针对每个输入进行运算，是稀疏模型；而在集成技术中，所有模型都会对每个输入进行运算，然后通过某种方式来综合这些模型的输出，这是密集模型。
  
• 有条件的计算。既然不同专家负责不同领域，怎么知道要把哪个token送去哪个expert呢？因此我们就需要对神经网络实际运行的程度拥有某种“决策权”，使得针对特定输入，只有特定专家被激活并处理（在生成式大模型中，就是根据token token 来选择专家的）。这部分工作就由门控机制来完成。其实，MoE的稀疏性与dropout的原理有些类似，MoE是根据任务的具体情况选择激活一定数量的专家模型来完成这个任务，而dropout则是对神经网络中的神经元进行随机性失活。
  

通过这种范式，模型将计算与参数解耦，仅激活与特定输入相关的专家，既保持了大规模知识库的优势，又有效控制了计算成本。而且，MoE能够在远少于 Dense 模型所需的计算资源下进行有效的预训练。这意味着在相同的计算预算条件下，我们可以显著扩大模型或数据集的规模。这种可扩展且灵活的创新有效遵循了扩展规律，实现了模型容量的增长而不会导致计算需求的剧增。
0x02 发展历史

2.1 重要节点

下图是MoE发展历史上的一些重要节点。

2.1.1 Adaptive mixtures of local experts

MoE的开山之作是1991年的论文“Adaptive Mixture of Local Experts"。这篇论文引入了将复杂问题分解为子问题并分配给多个专门模型的思想。这种分而治之的策略成为了 MoE 架构的核心。
因为面对多任务学习时，多层网络的各层之间通常会有强烈的干扰效应，这会导致学习过程变慢和泛化能力差。为了解决这个问题，论文提出了一种新的监督式学习方法：由多个独立子网络（专家）组成一个系统，每个子网络独立学习整个训练数据集中的一个子集。模型使用一个门控网络（gating network）来决定每个数据应该被哪个子网络去训练，这样就可以减轻不同类型样本之间的干扰。在推理时，模型将输入同时传递给不同的子网络和门控网络，每个子网络给出自己的处理结果，门控网络会依据每个子网络的权重来决定每个子网络对当前输入的影响程度，最终给出所需的输出。
如何训练这个系统？如何让损失函数整合专家和门控网络的输出？论文作者作者提出了两种思路：鼓励竞争和鼓励合作，具体参见下图。

2.1.2 sparsely-gated mixture-of-experts layer

2017年，论文“Outrageously large neural networks: The sparsely-gated mixture-of-experts layer”首次将MoE引入自然语言处理领域，并提出了Sparse MoE（稀疏MoE）概念。和论文"Adaptive mixtures of local experts"相比，本论文 MoE 的主要区别如下：

• Sparsely-Gated：不是所有expert都会起作用，而是选择TopK的专家进行计算，即只激活部分专家对特定输入进行处理。这是一种条件计算，意味着只有部分专家被激活处理特定的输入，从而可以大大降低计算量。而且，门控网络依然为每个输入同时选取多个专家，使网络能够权衡并整合各专家的贡献，从而提升性能。这种稀疏性就是MoE可以把模型容量扩大的原因。
  
• token-level：相比于sample-level，此论文使用了在 token 级别进行处理，一个句子中不同的token使用不同的专家。由于其为每个输入令牌（token）选择对应专家的特点，该方法被称为token选择门控（token-choice gating）。
  

本论文之前的MoE中，每个专家都用于每个输入，但是每个专家的贡献由一个门控函数加权，这是一个学习到的函数，它为每个专家计算一个权重或重要性，使得所有专家的权重之和为1。由于每个专家都用于每个输入，这种方法仍然导致一个密集激活的模型，因此没有解决增加计算复杂度的问题。这种路由选择算法也叫做软性选择路由算法（也称为连续混合专家）。本论文使用的路由算法是硬选择路由算法，运行只有一部分专家用于任何给定的输入，这标志着从密集激活到稀疏模型的转变。
另外，门控网络会倾向于收敛到不均衡的状态，总是为少数专家产生较大的权重（相应的参数更新也会很不均衡）。因此，作者设计了额外的损失函数来促使所有专家具有同等的重要性，也首创了具有辅助负载平衡损失的可微分启发式方法，通过选择概率对专家输出进行加权，使门控过程可微分，从而能够通过梯度优化门控函数。这种方法随后成为MoE领域的主流研究范式。
本论文也是业界第一个实施专家并行的方案。

2.1.3 GShard

2021年的论文”GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding“是第一个将MoE的思想拓展到Transformer上的工作。具体而言，GShard 将每间隔一层的FFN层替换为 MoE 结构，MoE 中的每个专家都是一个FFN（每个专家大小相同）。每个 Token 都会通过 Gating 选择不同的专家，默认为 top2。
因为难以控制将token发给专家的概率，所以在实际操作中，可能某些expert接收到了好多token，而某些expert接收的token寥寥无几，我们管这种现象叫expert负载不均。这种情况不仅不符合我们MoE的设计初衷（术业有专攻），还影响计算效率（例如引起分布式训练中各卡通讯时的负载不均）。为了缓解“赢者通吃”问题，尽可能让不同专家处理的token数尽量均衡，Gshard提出了以下几种解决办法：

• 专家容量负载：为了确保负载平衡，GShard强制每个专家处理的token数量低于某个统一阈值，论文将其定义为专家容量。当token选择的两个专家都已经超出其容量时，该token被视为溢出token，这些token或者通过残差连接传递到下一层，或被完全丢弃。
  
• Local group dispatching（本地组调度）：将训练批次中的所有token均匀划分为 G 组，即每组包含 S = N/G 个token。所有组均独立并行处理。通过这种方式，我们可以确保专家容量仍然得到执行，并且总体负载保持平衡。
  
• Auxiliary loss（辅助损失）：添加辅助损失函数，对expert负载不均的情况做进一步惩罚。
  
• 随机路由：在Top-2 gating的设计下，GShard 始终选择排名最高的专家，但第二个专家是根据其权重比例随机选择的。直觉上认为在输出是加权平均且次要权重通常较小的情况下，次要专家的贡献可以忽略不计。
  

Gshard也提出了MoE跨设备分片的方法。当扩展到多个设备时，MoE 层在不同设备间共享，而其他所有层则在每个设备上复制。这样，整个 MoE 层的计算被分散到了多个设备上，每个设备负责处理一部分计算任务。这种架构对于大规模计算非常有效。这也解释了为什么 MoE 可以实现更大模型参数、更低训练成本。

GShard 为后续所有的 MoE 研究铺好了路：它证明了稀疏专家是有价值的，还为后来的方案指明了“容量因子”这些关键概念的重要性。
2.1.4 Swith Transformer

论文”Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity“用稀疏激活的混合专家层替换了Transformer块中的FFN，同时使用简化的门控机制使训练更加稳定，从而使 MoEs 成为语言建模应用更现实、更实用的选择。
Switch Transformer 的指导设计原则是以简单且计算高效的方式最大化 Transformer 模型的参数数量。在此原则指导下，论文做了一些有效努力，包括：简化稀疏路由、使用高效稀疏路由和增强的训练和微调技巧。Switch Transformer 作者发现仅使用一个专家也能保证模型的质量。一个专家可以让路由计算更简单，通信量也更少；一个 Token 仅对应一个专家，计算量也更少；平均每个专家对应的 batch size 至少可以减半。因此，Switch Transformer 的门控网络每次只路由到 1 个 expert，也就是每次只选取 Top1 的专家，而其他的模型都是至少 2 个专家。

2.2 详细时间线

下图是2017年之后的若干代表性的MoE模型的时序概览。时间线主要根据模型的发布日期构建。位于箭头之上的MoE模型是开源的，而位于箭头之下的则是专有闭源模型。来自不同领域的MoE模型用不同颜色标记：自然语言处理（ NLP）用绿色，计算机视觉（CV）用黄色，多模态（multimodal）用粉色，推荐系统（RecSys）用青色。

0x03 模型结构

MoE包括以下核心组件：

• 专家。在MoE架构中，专家是专门针对特定任务的子模型。专家拥有不同领域的专业知识，负责处理不同的计算任务或者特定输入子空间。形式上，每个专家网络\(f_i\)（通常是一个linear-ReLU-linear网络）由参数W来参数化，接受输入x并生成输出\(f_i(x; W_i)\)。
  
• 门控函数（也称为路由函数或路由器）：门控函数负责协调专家计算，即判定哪个输入样本应该由哪些专家处理，哪些专家将被激活并参与到当前的计算中。形式上，门控函数G（通常由linear-ReLU-linear-softmax网络组成）由参数O来参数化，接受输入x并产生输出。
  
• 聚合层（Combining Layer）：聚合层负责整合专家网络的输出，以形成最终的输出结果。很多资料并没有把聚合层单独罗列出来。
  

整个结构可以用下图来表示。门控网络输出是一个稀疏的n维向量，\(G(x)_i\)是门控网络给出的第 i 个专家的权重。\(E_i(x)\) 是第 i 个expert的输出。那么对于在当前的输入x，输出就是所有 experts 的加权和。

下图给出了MoE的总体处理流程，具体是如下五步。这里每个专家和token都有颜色编码，门控网络权重（W）有每个专家的表示（颜色匹配）。为了确定路由，路由器权重对每个token embedding（x）执行点积，以产生路由器得分（h(x)）。然后将这些分数归一化为1（p(x)）。G使用了softmax函数。

• 将输入token的embedding和门控网络权重进行点积，得到门控分数。在语言建模的上下文中，这里每一列将表示输入序列中的一个token。因此，每个token可以路由到不同的专家。
  
• 在门控分数上施加softmax将门控分数进行归一化，得到概率。此概率表示每个专家模型对该token的贡献程度，即在给定输入情境下每个专家被激活的概率。或者说此概率表明专家处理传入token的能力如何。
  
• 使用此概率分布作为权重来选择最佳匹配的专家。
  
• 专家对输入token进行处理。
  
• 专家处理之后，将每个路由器的输出与每个选定的专家相乘，并对结果求和。
  

3.1 门控函数

3.1.1 条件计算

稀疏激活是 MoE 模型的关键部分和优势之一。与所有专家或参数对输入都活跃的密集模型不同，稀疏激活确保只有一小部分专家根据输入数据被激活。这种方法在保持性能的同时减少了计算需求，因为任何时候只有最相关的专家是活跃的。
本质上，我们现在所谈论的MoE 大模型是使用条件计算来强制稀疏激活（Sparse Activation）。条件计算是探讨如何分离计算复杂性和计算量需求，并在其之间进行合理的权衡的理论。在此处的意思是动态开启/关闭神经网络的部分功能。MoE 模型条件计算的核心是学习一个计算成本低的映射函数，该函数确定网络的哪些部分——换句话说，哪些专家可以最有效地处理给定的输入。条件计算在大模型中通常使用路由网络或者门控网络来实现。它是判断选择使用哪个专家的关键，即在网络中根据输入数据有选择地激活部分单元。在当前的大语言模型中，这是通过对每个 token 进行条件判断来实现的。当模型输入一个token时，路由网络根据上下文和当前token，选择合适的专家网络来计算。这种选择性激活的直接效果是加快信息在网络中的传播速度，无论是在训练阶段还是推理阶段。通过条件计算或者稀疏性，大模型能够在增加模型规模的同时，降低计算成本，实现了一个合适的均衡。其实大模型推理中常见的早退出机制（Early-Exit）也是一种条件计算，它允许在网络的早期层级就做出决策并减少计算。
因为是条件计算，所以与具有相同参数数量的模型相比，MoE具有更快的推理速度。也因为是条件计算，所以MoE需要把所有专家系统完全加载到内存中，所以需要大量显存。
3.1.2 定义

门控网络（Gating Network）的设计和实现是Sparsely-Gated MoE 层的核心组成部分。门控网络负责为每个输入 token 选择一个稀疏的专家组合，这些专家将参与到当前的计算中。门控函数是一个可以执行一系列非线性变换的网络，该网络对概率分布进行建模，根据概率去做出相应的选择。门控网络由学习到的参数组成，并且与网络的其余部分同时进行预训练。一个典型的门控网络就是一个带有 softmax 函数的简单网络。
假定注意力层的输入数据形状是 (batch_size, seq_len, embedding_size)，则门控网络的大小是 (token_size,  expert_num)，门控网络的输入形状是 ( batch_size * seq_len, embedding_size)， 输出是  ( batch_size * seq_len, expert_num)，即每个token去向每个expert的概率。比如：

1. gates ( batch_size * seq_len = 3, expert_num = 4)：
2. [
3.     [0.2, 0.4, 0.1, 0.3],  # Token A 被分配到不同专家的概率
4.     [0.1, 0.6, 0.2, 0.1],  # Token B 被分配到不同专家的概率
5.     [0.3, 0.1, 0.5, 0.1]   # Token C 被分配到不同专家的概率
6. ]

复制代码

门控网络会学习将输入发送给哪个expert，softmax的输出作为每个专家的最终使用权重。门控函数的处理流程如下：

• 计算专家分数。门控函数接收单个token的emebdding作为输入，基于输入数据的特征进行计算，然后输出一组分数。这些分数表示每个专家模型对该token的贡献程度，或者说是表明专家处理传入token的能力如何。
  
• 计算专家的概率分布。以下图为例，门控函数使用softmax对分数进行处理，得到在给定输入情境下每个专家被激活的概率分布。这个分布反映了输入数据与各个专家相关性的大小，概率越高，表示该专家对于当前输入的预测任务越重要。下图中选择了两个专家，门控函数输出的概率是0.1和0.9，说明专家1对该token贡献是10%，专家2对该token的贡献是90%。
  
• 激活专家。门控函数将每个 Token 作为输入，并在 expert 上生成一个概率分布，以确定每个 Token 被发送给哪个 expert。根据门控输出的概率分布，一部分专家将被选中并激活。在下图中，如果使用top2 策略来选择，因为其它专家的激活概率不到0.1，因此没有被选中。专家2和专家n-1因为具有较高的激活概率，将被选中参与到后续的计算中。这意味着，只有这两个专家的参数将被用于处理当前的输入数据。假设专家2输出结果值是0.4，专家n-1输出结果值是0.5，则最终MoE返回选定专家的输出乘以门值（选择概率）是0.49。通过同时咨询多个专家对给定输入的意见，网络能够有效地权衡并整合他们的贡献，从而提升性能。
  

注：在MoE模型中，虽然动态训练门控功能是标准做法，但一些研究探索了非可训练的token选择门控机制（Non-trainable Token-Choice Gating）。这种机制的主要优势在于不需要额外的门控网络参数，通过特定的门控机制即可实现全面的负载均衡。比如，Hash Layer采用基于随机固定门控的方式，通过对输入token进行哈希，无需训练门控网络即可工作。人们还探索了其他更复杂的哈希函数，例如通过对单独预训练的Transformer模型产生的token嵌入应用k-means聚类，或者根据训练数据中token频率预计算一个哈希表，将token ID映射到专家，从而确保token到专家的分配更加平衡。
3.1.3 特点

我们接下来看看门控函数的特点。
首先，门控函数不仅决定在推理过程中选择哪些专家，还决定在训练过程中选择哪些专家。.这是因为只有让每个专家在训练期间学习到不同的信息，在推理时，才能知道哪些专家与给定的任务最相关。
其次，门控函数逐层都会选择专家。在具有 MoE 的 LLM 的每个层级中，我们都会找到（某种程度上专业的）专家。具体参加下图。

其实，在最微观层面，每个神经元就是一个专家，激活函数就起到了门控函数的作用。MoE只是把很多神经元聚类成一个专家。
另外，"OpenMoE"和Mixtral都对门控机制进行了分析，其中提到几个特点如下：

• 上下文无关的专精化（Context-independent Specialization）。MoE 倾向于简单地根据 token 级语义对 token 进行路由，即无论上下文如何，某些关键词经常被分配给同一位专家。因为路由规则和文本的语义主题无关，这说明MoE 模型中的专家各自擅长处理不同的 token，但是专家实际可能并不专门精通某一领域的知识。
  
• 路由早期习得（Early Routing Learning）。路由在预训练的早期就已建立，并且基本保持不变，因此 token 在整个训练过程中始终由相同的专家处理。这或许能启发我们设计更高效的路由机制。
  
• 序列尾部丢弃（Drop-towards-the-End）现象显著。序列后部的token因专家达到容量上限而更容易被丢弃。具体而言，在 MoE 模型中，为了保证负载均衡，通常会为每个专家设置其容量上限。当某个专家的容量达到上限时，该专家将不再接受新的 token 而将其丢弃（Drop）。如果我们从前往后为序列中的 tokens 分配专家，那么序列尾部的 tokens 将有更大的概率被丢弃，这在指令调优数据集中更为严重。
  
• 位置的局部性。相邻的token通常被路由到同一位专家，这表明token在句子中的位置会影响路由选择，会带来“高重复率”现象。这有利于减少专家负载的突发波动，但也可能导致专家被局部数据“霸占”。
  

因此，我们需要了解数据集的“局部规律”，因为一旦数据分布换了（比如从新闻文本转到代码），它原先的路由模式就可能失效。要做大规模的 MoE，就得好好考虑数据特征和专家分配之间的关系。
3.1.4 优化

关键因素

现在MoE 大模型的整体架构非常固定，而路由选择则成了关键。路由算法可以从简单（在张量的平均值上进行均匀选择或分箱）到复杂。在决定特定路由算法对问题的适用性的许多因素中，以下几个经常被讨论。

• 模型精度。MoE 模型对舍入误差很敏感，比如 softmax 中的指数操作可能会产生舍入误差，导致训练不稳定。 然而，简单地裁剪（即应用硬阈值以删除大值）门控函数输出的logit 又可能会损害模型性能。
  
• 平衡负载。我们希望尽可能让不同专家处理的token数尽量均衡。目前我们知道最简单的解决方案是根据 softmax 概率分布选择前 k 个专家。然而，这种方法会导致训练负载不平衡：训练期间，大多数token都会被分发给少数专家，因此这少数专家积累了大量的输入token，而其它专家比较空闲，这减慢了训练速度。与此同时，许多其他专家根本没有接受过足够的训练。因此需要更好的门控函数，以便在所有专家之间更均匀地分配token。
  
• 高效。如果门控函数只能串行执行，则很难实现负载均衡。假设我们有 E 个专家和N个token，则仅门函数的计算成本就至少为 O(NE)。在实际工作中，N 和 E的数量级会很大，门控函数的低效执行将使大部分计算资源（专家）在大多数时间处于闲置状态。因此，我们需要让门控函数可以高效并行实现来利用众多设备。
  

为了达成这些目标，研究人员做了不谢的努力，下图展示了MoE模型中使用的不同门控函数。包括 (a)使用top-1门控（top-1 gating）的稀疏MoE；(b)BASE层（BASE layer），(c)组合领域映射与随机门控（combining domain mapping with stochastic gating），(d)专家选择门控（expert selection gating），(e)注意力路由器（attention router），以及(f)带专家合并的软MoE（soft MoE with expert merging）。这些函数可能通过各种形式的强化学习和反向传播进行训练，做出二元或稀疏且连续、随机或确定性的门控决策。

改进示例

我们以论文”Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer“为例，看看如何对简单的 softmax 门控函数做升级，进而满足需求。
softmax 函数的问题：softmax 函数会让所有专家都会对输入进行运算，再通过门控网络的输出进行加权求和，如果experts的数量太大，就会导致计算量非常大。因此，我们需要找到一种方法能使某些专家模型的门控网络的输出为0，这样就没有必要对这个专家进行相应的计算，就可以节省计算资源。
普通” top-k 路由策略可以满足这个需求，其会根据 softmax 概率分布选择前 k 个专家。即，在Softmax函数应用于专家权重之前，执行一个KeepTopK操作，将除前k个专家之外的所有专家的权重设置为-∞。这确保了只有前k个专家在应用Softmax后权重大于0。因此，这种MoE可以帮助我们在扩大模型规模的同时保证计算量是非线性增加的（因为每个token只用过topK个expert，不需要使用全量expert），这也是我们说MoE-layer是稀疏层的原因。
普通” top-k 路由策略有一个重要缺点是，门控网络可能收敛到只激活少数专家。这是一个自我强化的问题：如果一小部分专家在早期被不成比例地选中，那么这些专家将更快地被训练，这会导致训练负载不平衡。而且，相对其它训练不足的专家，这些更快训练的专家会输出更可靠的预测，它们将继续被更多地选中。这种不平衡的负载意味着其他专家最终会成为名副其实的累赘。
带噪声的 TopK 门控 (Noisy Top-K Gating)能够缓解这个问题，其在为每个专家预测的概率值中添加了一些高斯噪声。在MoE模型中加入噪声的原因主要有以下几点：

• 提高模型的鲁棒性和泛化能力。当模型在训练或推理阶段遇到不确定或嘈杂的数据时，鲁棒性较强的模型更能保持稳定的性能。
  
• 噪声会增加一部分的随机性，减少过拟合的风险。
  
• 加入噪声可以实现不同专家之间的负载均衡。
  

该方案还为专家选择添加了两个可训练的正则化项：最小化负载均衡损失会惩罚过度依赖任何一个专家，而最小化专家多样性损失会奖励对所有专家的平等利用。我们用下图展示下上述解决思路。

• 原始处理流程如下图标号1，其中门控函数是softmax。路由策略就是将输入乘以权重矩阵并应用 softmax。
  
• 但是，这种方法并不能保证专家的选择将是稀疏的。为了解决这个问题，我们首先对输入进行线性变换，然后再加上一个softmax，这样得到的是一个非稀疏的门控函数。对应下图标号2。
  
• 但是这样依然不够，因此我们在进行softmax之前，先使用一个topk函数，只保留最大的k个值，其他都设为-∞。这样对于非TopK的部分，由于值是负无穷，这样在经过Softmax之后就会变成 0，不会被选中，就得到了稀疏性。在这个基础上，我们在输入上再加上一个高斯噪声。此处对应下图标号3。
  

随机路由是另一种解决问题的方法，例如，在top-2 设置中的“最佳”专家是使用标准 softmax 函数选择的，而第二个专家则是半随机选择的（每个专家被选中的概率与其连接的权重成正比）。因此，排名第二的专家最有可能被选中，但不再保证其肯定被选中。
另一种解决该问题的方法是根据专家容量的阈值来做调整。这种方法会设置一个阈值，该阈值定义了任何一个专家可以处理的最大token数量。还是以 Top-2的专家选择为例，如果 top-2 中选择的任何一个专家都已达到容量，则选择下一个专家（Top3）来处理后续token。但是，这可能导致token溢出，即超出容量的token无法被指定专家处理。另外，也有研究提出了DSelect-k，这是一种平滑的top-k门控算法，其平滑特性优于传统的top-k方法。
3.2 专家

在 MoE 架构中，专家是指训练好的子网络（神经网络或层），它们专门处理特定的数据或任务。专家和门控机制都通过梯度下降与其他网络参数一起进行联合训练。MoE里的“专家”是一种拟人的形象化的说法，其实，专家在本质上是基于某种人类先验“知识”或“策略”的“跨范畴采样”。
3.2.1 特点

专家具备如下特点：

• 架构。在实际应用中，一般来说，MoE 中的每个专家都是具有相同架构的前馈神经网络。但是，我们也可以使用更复杂的体系结构。我们甚至可以通过将每个 Experts 实现为另一个 MoE 来创建“分层”MoE 模块。在某些情况下，并非所有 FFN 层都被 MoE 取代，例如Jamba模型具有多个 FFN和MoE 层。
  
  
  
• 参数子集：FFN层被分解为多个专家，每个专家实际上是FFN参数的一个子集。专家并不是对FFN的平均切分，实际上我们可以任意指定每个expert的大小，每个expert甚至可以大于原来单个FFN层，这并不会改变MoE的核心思想：对于一个token，部分专家的计算量要小于所有专家的计算量。
  
• 输入分割：不同的专家会专注于不同的主题。用更专业的术语来说，输入空间被“区域化”了（或者说更精细地划分知识空间）。假设某个 LLM 可能收到的请求是一个“完整的知识空间”，而MoE将输入数据根据任务类型分割成多个区域，并将每个区域的数据分配给一个或多个专家模型。
  
• 专注学习。每个专家模型可以专注于处理自己接受到的输入数据，学习数据中的一种特定模式或特征。由于这些专家从一开始就存在，在训练过程中，每个专家都会在某些主题上变得更加专业，而其他专家则会在其他主题上变得更加博学。例如，在图像分类任务中，一个专家可能专门识别纹理，而另一个专家可能识别边缘或形状。
  
• 灵活扩展和组合作战。在MoE范式下，只有相关的专家被激活以处理给定输入，由于只有相关的专家被激活，因此可以减少不必要的计算（帮助我们在扩大模型规模的同时保证计算量是非线性增加的），从而加快模型的推理速度并降低运算成本。而且，MoE可以在减少计算开销、未相应增加计算成本的情况下扩展模型的参数空间，从而受益于大量专业知识。用户不必聘请一位“无所不知”的专家，而是组建一个拥有特定专业领域的团队。这种分工有助于整个模型更高效地处理问题，因为每个专家只处理它最适合的数据类型。另外，MoE这也使得模型能够更加灵活地适应不同的任务，因为不同的任务可能需要不同专家的组合来达到最优的预测效果。
  

我们再用一个示例来看看专家学习到了什么。
从目前的研究成果来看，专家并不专攻“心理学”或“生物学”等特定领域。它最多只是在单词层面学习句法信息：更具体地说，它们擅长于在特定上下文中处理特定的 tokens。专家学习的信息比整个领域的信息更加细粒度。因此，有时将它们称为“专家”可能会产生误导。
Mixtral 8x7B 论文中可以找到一个很好的例子，论文作者测量了所选专家在the Pile验证数据集不同子集上的分布（token分布比例）。下图显示了第0、15和31层的结果（第0层和第31层分别是模型的第一层和最后一层）。论文在根据主题分配专家时没有观察到明显的模式。例如，在所有层次上，ArXiv论文、生物学和哲学文档的专家分配分布都非常相似。真有针对数学的专家的分布才略有不同。这种差异可能是数据集的合成性质及其对自然语言的有限覆盖的结果，在第一层和最后一层尤其明显，其中隐藏状态分别与输入和输出嵌入非常相关。这表明门控网络确实表现出一些结构化的句法行为。

因此，尽管专家似乎没有专业知识，但他们似乎确实被一致地用于某些类型的token。下图显示了来自不同领域（Python代码、数学和英语）的文本示例，其中每个标记都用与其所选专家对应的背景颜色突出显示。该图显示，Python中的“self”和英语中的“Question”等单词经常通过同一个专家传递，即使它们涉及多个标记。同样，在代码中，缩进标记总是分配给相同的专家，特别是在隐藏状态与模型的输入和输出更相关的第一层和最后一层。我们还从图中注意到，连续的token通常分配给相同的专家。事实上，论文作者确实在The Pile数据集中观察到了一定程度的位置局部性。

因此，也有研究认为，专家提升的是记忆效果而不提升推理能力。比如论文“Mixture of Parrots: Experts improve memorization more than reasoning”研究了Mixture-of-Experts（MoE）架构性能和推理上的理论局限性，探讨了与标准密集型Transformer在记忆和推理方面的性能差异。研究发现，随着专家数量的增加，MoE模型在记忆任务上的表现提升，而在推理任务上达到饱和。论文也通过实证证明了MoE在特定记忆密集型任务上的优越性。
另外，MoE模型实际上还提供了一种细粒度的方式来研究和理解模型内部的工作机制。通过观察哪些专家被激活以及它们如何随着时间变化，研究人员可以更深入地洞察模型是如何学习和泛化知识，以及它是如何处理不同的输入特征的。
3.2.2 种类

专家的网络类型通常有如下几种：

• 前馈网络（Feed-Forward Network）：因为FFN层比自注意力层更加稀疏，且展示出更多的领域特定性，所以目前大多数MoE都是FNN的替代品。比如，有研究人员发现，大多数输入仅激活FFN的少量神经元，突显了FFN的内在稀疏性。对于同样输入，FFN层仅激活20%的专家，而自注意力层激活了80%的专家。预训练Transformer中的模块化涌现现象（Emergent Modularity）也揭示了神经元激活与特定任务之间的显著关联，支持了MoE结构反映预训练Transformer模块化特性的观点。另外，从参数量的角度我们也可以看到选择的原因，因为随着模型规模增长，FFN的计算开销呈现急剧上升趋势。例如，在拥有5400亿参数的PaLM模型中，约90%的参数分布在FFN层。
  
• 注意力（Attention）：尽管MoE研究主要集中在Transformer架构的FFN层，也有研究人员提出了多头注意力专家混合（Mixture of Attention Heads, MoA），将多头注意力层与MoE结合，以提升性能并降低计算成本。MoA使用两组专家（查询投影和输出投影），通过共同的门控网络选择相同的专家。为降低计算复杂度，MoA在所有注意力专家间共享\(W_K\)和\(W_V\)投影权重，专家仅在各自的查询（\(q_tW^q_t\)）和输出投影权重（\(o_{i,t}W^O_i\)）上有所区别，从而实现键（\(KW_K\)）和值（\(VW_V\)）序列的预计算共享。
  
• 其他类型。有些研究人员还探索了使用卷积神经网络（CNN）作为专家，也有将参数高效微调（PEFT）技术与MoE结合的努力，例如采用低秩适应（LoRA）作为专家。
  

3.2.3 位置

我们接下来看看专家如何嵌入到Transformer架构中。下图给出了一些实例。

• (a)展示了MoE与注意力机制中的Key和Value模块的集成。
  
• (b)表示MoE在FFN中的应用。
  
• (c)指的是MoE在Transformer块层级的集成，其中应用了两组不同的专家到注意力和FFN层，分别为每个层分配专家，并通过各自的门控机制进行调控。
  
• (d)展示了MoE在每一层的集成，其中每个Transformer层视为一个统一体，门控机制协调专家之间的交互。
  

3.3 分类

之前的学习中，我们大致了解到了MoE有不同种类，此处我们再从最本质的角度（门控函数）来看看如何把MoE分类。门控函数是MoE架构的核心，它负责协调专家网络的参与并整合其输出。根据输入处理方式的不同，门控机制可分为稀疏型、稠密型和软性三类。稀疏门控只激活部分专家，稠密门控激活所有专家，而软性门控包括全微分方法，如输入token合并和专家合并。三种门控函数的特点如下：

• 稀疏门控：仅激活部分专家，包括基于token选择的top-k门控策略，以及使用辅助损失函数来促进专家间token均匀分布。
  
• 稠密门控：激活所有专家，在LoRA-MoE微调中表现出色，因为它可以有效地将多个LoRA整合到各种下游任务中。
  
• 软性（soft）门控：通过token或专家合并的方式实现完全可微性，避免了离散专家选择的问题，例如SMEAR、Lory和Omni-SMoLA。
  

根据门控函数的设计，MoE层可以大致分为以下两类：稠密MoE和稀疏MoE。

• 稠密MoE层在每次迭代中激活所有专家网络\(f_1...f_N\)。稠密MoE能充分利用所有参数，捕获潜在的复杂的模式和关系，因此通常能获得更高的预测精度，但计算开销较大。因此这种方法在早期研究中被广泛采用，近期有研究（EvoMoE、MoLE、LoRAMoE和DSMoE）重新探讨了稠密MoE的应用。因为稀疏激活专家虽然在计算效率上有优势，但当总参数量相同时，往往会导致性能损失。而在LoRA-MoE微调中，由于LoRA专家的计算开销较小，稠密激活表现出色。这种方法能够有效地将多个LoRA整合到各种下游任务中，既保持了原始预训练模型的生成能力，又维持了每个任务特定LoRA的独特性
  
• 稀疏MoE层在每次前向传递中仅激活选定的专家子集。稀疏型MoE不是汇总所有专家的输出，而是通过仅计算前k个专家输出的加权和来实现稀疏性。稀疏激活实际上是计算需求与模型性能之间的一种权衡策略。
  

下图展示了两种MoE的特点。右侧是示意图，左侧是门控函数以及负载函数。

3.3.1 稠密 vs 稀疏

我们通过例子来对稠密和稀疏MoE进行比对。
以人类分工为例。稠密就是类似手工业时代的生产模式。在这种模式下,每个工人（即神经元）都需要参与处理所有类型的任务，就像手工业时代的工匠需要精通产品制作的各个环节，掌握所有的生产技能。这种方法虽然直观且易于实现,但在面对复杂多变的任务时,往往效率低下且难以扩展。稀疏MoE则是工业革命之后的分工模式：每个岗位（专家）只需要完成一部分生产任务。这种分工的方式大大提高了生产效率，推动了工业化的进程，开启了机器大工业的时代。

3.3.2 软性（soft）门控

论文”A Survey on Mixture of Experts“中提出软性（soft）门控这种类型，以突出其通过门控加权合并输入token或专家来缓解计算需求的特点。为了为每个输入token分配适当的专家，稀疏性MoE通常需要启发式辅助损失来确保专家参与的平衡。在涉及分布外数据的场景中，这些问题变得更加突出。与密集MoE类似，软MoE方法通过利用所有专家处理每个输入来保持完全的可微性，从而避免了离散专家选择所固有的问题。

• token合并（Token Merging）：Soft MoE摒弃了传统的稀疏和离散门控机制，采用软分配策略来合并token。该方法计算所有token的加权平均值，权重取决于token和专家的关系，然后用相应专家处理每个加权平均结果。然而，token合并使其难以应用于自回归解码器，因为在推理时无法获取用于加权平均的未来token。
  
• 专家合并（Expert Merging）：SMEAR框架通过加权平均合并所有专家参数来避免离散门控。SMEAR的作者认为，传统稀疏MoE模型难以匹敌参数量相当的稠密模型或使用非学习启发式门控函数的模型，可能是由于非可微、离散门控决策的训练模块中存在梯度估计偏差。SMEAR通过单个合并专家处理输入token，既不显著增加计算成本，又支持标准梯度训练。
  

3.4 比对

论文“A Uniffed View for Attention and MoE”把注意力机制和MoE做了比对。Attention 结构的作用其实就是使用加权和的形式来聚合不同 token 的信息。MoE 其实是通过线性投影来学习一个 Router, 基于 Router 来聚合不同 expert 的信息. 本质上来说, router 的输出经过 top-k+softmax 来学习一个加权求和的权重, 最终聚合不同 token 的信息。
两者的相似之处:

• 都使用了 softmax 来归一化权重，本质上也是希望学习稳定。
  
• 本质都是对新特征的加权和。
  
• 都是动态权重, 根据输入来自适应地聚合信息。
  

二者的区别在于

• Attention 是聚合不同 token 的信息，MoE 是聚合不同地 Expert。
  
• Attention 会用上所有 token (Softmax 结果为正数，故所有 token 都会用上)；MoE 是一种稀疏选择，只会选择部分结果。
  

Attention和MoE都是学习一个权重来聚合信息，可以总结到如下图标号3所示，f是权重, g 是学习到的新特征。具体如下图。

从计算机体系结构的视角来看也有启发。注意力机制（含FFN）类似把所有上下文都放在内存中。MoE则可以理解为页表系统，需要的时候才把专家对应的页表放到内存中。
0x04 计算流程

4.1 算法

下图是top-k门控函数算法的伪代码，具体流程如下。

• 给定输入x，使用门控函数计算出得分score = G(x)。
  
• 选出前k的分数对应的索引。
  
• 遍历索引，使用索引对应的专家来进行计算，得到推理结果。
  
• 使用分数作为权重对专家计算的结果值进行修正。
  
• 综合这k个专家的结果作为最终推理结果。
  

4.2 流程

MoE的整个计算过程如下图所示：
<ul>Routing。Routing也叫“experts selection”或者稀疏性激活，本质上是对门控函数的使用，是MoE的核心理念。Routing是一个对输入进行多分类的鉴别过程，目的是确定最适合处理输入的专家模型。在语言模型的应用中，当输入token通过MoE层时，Token通过和Router的权重矩阵相乘得到一个Expert Indices（决策矩阵）和一个概率张量，即索引和概率：

• Expert indices是expert-to-token映射，用于指示每个token被分配给了哪个expert。即张量中第i个值代表本token应该分配到第i个专家。
  
• Probabilities张量是分配置信度的概率，其中第i个值代表这个专家对于该token最终结果的权重。
  

Permutation（排列/置换）。根据路由决策（expert-to-token映射）将Token分配给对应的专家。其间可能会依据专家的容量对token进行丢弃操作。
Computation。专家依据分配到的经过Permutation重排序的tokens进行计算。通过使每个专家专注于执行特定任务，这一方法实现了计算的高效性。这种方式允许模型对不同类型的输入数据进行个性化处理，提高了整体效率和性能。每个专家网络并行处理其分配到的token，计算输出。这一步涉及到块稀疏矩阵乘法，其中输入矩阵
来源：程序园用户自行投稿发布，如果侵权，请联系站长删除
免责声明：如果侵犯了您的权益，请联系站长，我们会及时删除侵权内容，谢谢合作！