正文
由于稀疏MoE是当前主要采用的MoE类型，本文后续内容将主要聚焦于稀疏MoE的技术实现。
专家选择与路由机制
专家选择过程基于token的d维向量表示，从N个专家中选择k个专家进行处理。具体的选择流程包括：首先对token向量进行线性变换，生成包含N个值的新向量，每个值对应一个专家；然后应用softmax函数将这些值转换为概率分布；最后根据概率值选择前k个专家。
这一策略源于首次引入稀疏MoE层的开创性研究。
TopK(k)函数保留向量中前k个最大值，将其余值设置为负无穷大(−∞)。经过softmax操作后，设置为−∞的条目接近零值。参数k是根据具体任务选择的超参数，常用值为k=1或k=2。
在稀疏MoE层的训练过程中，通常引入噪声项（R_noise）作为正则化技术。这种噪声有助于系统探索不同专家的使用模式，提高训练过程的稳定性。
基础路由方法存在专家使用不平衡的问题。随着训练进行，模型可能持续选择相同的少数专家，导致这些专家获得更多训练机会，进而在未来更容易被选中。这种现象被研究人员称为路由崩溃（routing collapse）。
原始论文对此问题的描述为："门控网络倾向于收敛到一种状态，即它总是为相同的少数几个专家产生较大的权重。这种不平衡是自我强化的，因为受青睐的专家训练得更快，因此更容易被门控网络选中。"
辅助损失函数的平衡机制
为解决路由崩溃问题，研究人员引入了辅助损失函数（auxiliary losses function）。该机制通过添加奖励工作负载平均分配的损失项来实现专家使用的平衡。
每个专家根据其被选中的频率（路由机制分配的概率总和）获得重要性得分。在批处理过程中，将每个专家的概率进行累加。如果所有专家的得分相近，则损失较小；如果得分差异显著，则损失增加。
通过最小化方差（以变异系数的平方衡量），训练过程促使模型保持重要性得分的平衡。这种"重要性损失（importance loss）"机制帮助路由器更均匀地利用所有专家。
对于包含T个token和N个专家的批次B = {x₁, x₂, …, x_T}，基于先前研究，该批次的辅助负载均衡损失计算如下：
负载均衡损失包含两个组成部分：每个专家接收的路由概率分数，以及实际分配给每个专家的token分数。系统将这两个组件存储为向量，然后计算其点积。当两个组件都趋于均匀分布时，损失值最小化，从而实现重要性和实际负载的双重平衡。
MoE层中的专家容量管理
在MoE模型中，数据处理方式在不同批次间存在变化，这是由于模型在训练和推理过程中使用路由机制决定数据分配。大多数实际部署的MoE模型为每个专家采用固定的批处理大小，以充分利用硬件资源。
专家容量定义为单个专家在一个批次中能够接收的最大token数量。当路由到某个专家的token数量超过其容量限制时，溢出的token将跳过该专家的处理，通过Transformer的残差连接直接传递到下一层。
大多数实现使用容量因子（capacity factor）参数来控制专家容量。当容量因子设置为1时，token在专家间平均分配。将容量因子设置为大于1的值为模型提供处理路由不均匀情况的缓冲空间，允许专家处理超过平均水平的token数量。但是较高的容量因子会增加内存消耗并可能降低系统效率。
研究表明，MoE模型通常在较低容量因子下表现良好，但需要关注因容量限制而丢弃的token数量，过多的token丢弃会影响模型性能。实际应用中，可以为训练和评估设置不同的容量因子。例如，ST-MoE模型在训练期间使用1.25的容量因子，在评估期间使用2.0的容量因子。
MoE层输出生成机制
MoE层的最终输出计算过程包括以下步骤：路由器将每个token分配给一个或多个专家；各专家处理分配给它们的token；来自各专家的输出通过加权平均进行组合，权重由路由器对每个专家的置信度确定；对于每个token，最终结果是其所分配专家输出的加权和。
MoE研究中的一个新兴概念是同时使用共享专家（shared experts）和路由专家（routed experts）。共享专家接收所有token，而路由专家根据标准路由机制接收token。共享专家的目标是避免在多个专家间重复存储相同信息，通过集中存储通用知识来提高效率。通常，共享专家的数量少于路由专家，因为过多的共享专家会削弱模型的稀疏性和效率优势。
使用共享专家时，共享专家的输出与路由专家的输出进行简单相加，使模型能够结合通用信息（来自共享专家）和专门处理（来自路由专家）。
基于MoE层的仅解码器LLM整体架构
在标准仅解码器Transformer模型中，每个模块包含一个前馈网络。而在MoE架构中，通过在模块内部部署多个专家单元而非单个前馈网络来实现修改。每个专家都是原始前馈网络的独立副本。
MoE模块的所有组件，包括自注意力层、专家网络和路由机制，都通过梯度下降法进行联合训练。路由机制通过对token向量进行线性变换来决定专家分配，token通过选定的专家进行处理，专家输出被组合并传递到模型的下一层。
修改后的MoE模块结构包含以下核心组件：自注意力层、残差连接加归一化、为每个token选择专家的路由机制、包含多个独立前馈网络的专家层，以及应用于专家层输出的最终加法和归一化操作。Transformer的其他部分保持不变。
实际实现中，并非需要将每个模块都替换为MoE层。常见做法是仅修改每P个模块中的一个，其余模块保持标准结构。虽然某些模型在每一层都部署专家，但在实践中，每隔一个、三个或五个模块使用MoE层是更常见的方法。这种策略有助于控制模型的总参数数量。
MoE技术实现
（一下代码为演示，如果想查看完整的MOE代码可以搜索我们以前发布的文章）
以下展示专家网络、路由机制和稀疏MoE架构的具体实现：
classNoisyTopkRouter(nn.Module): def__init__(self, n_embed, num_experts, top_k): super(NoisyTopkRouter, self).__init__() self.top_k=top_k self.topkroute_linear=nn.Linear(n_embed, num_experts) self.noise_linear=nn.Linear(n_embed, num_experts) defforward(self, mh_output): logits=self.topkroute_linear(mh_output) noise_logits=self.noise_linear(mh_output)
