计算机基础·cs336·MoE

MoE 混合专家

核心思想

混合专家模型（Mixture of Experts，MoE）是一种先进的神经网络架构，旨在通过整合多个模型或“专家”的预测来提升整体模型性能。MoE模型的核心思想是将输入数据分配给不同的专家子模型，然后将所有子模型的输出进行合并，以生成最终结果。这种分配可以根据输入数据的特征进行动态调整，确保每个专家处理其最擅长的数据类型或任务方面，从而实现更高效、准确的预测。

优势

• 模型索然很大，但是推理速度快。例如同样的FLOPs，由于只激活少数专家，实际激活参数量比较低，推理速度比同样参数的密集模型快，而且可以得到更好的结果。
• 训练起来不像密集架构那样全部参数都要训练，只有少数专家激活并参与训练
• 分布式架构：专家可以分配到不同设备上去。

劣势

• 理论上分布式，在缺乏分布式节点和计算设备时比较鸡肋。
• 训练不稳定：路由算法往往不可微分，路由算法的学习不稳定。

MoE的基本原理

MoE的架构

• 多个不同的MLPvs 多个不同的 Transformer头
• 注意：为了确保参数量不会爆炸，MLP的d_ff通常会大幅度减小，甚至比d_model还小
  

路由算法

• 基本分类：为每一个token选专家vs 为每一个专家选token
  
• 通常是为每一个token选专家：具体策略可以分类为 随机选择，RL算法，TOPK，哈希法。

TOPK算法

• 基本思想：将路由器认为是一个简单的MLP，给定token x，映射为专家选择的概率s，选择前K个专家。
• 将专家计算的结果和选择概率加权求和与原输入残差链接。

路由选择函数的训练目标(损失函数)

• 核心思想：避免只选择特定专家和设备
  
• fif_ifi​：统计所有batch/设备中实际选择专家i的比例
• PiP_iPi​：统计所有batch/设备中路由器选择专家i的概率(想选择专家i的比例)
• 惩罚路由器想选择专家i且实际选择专家i的情况。

• 直接对于softmax分数进行显示调整的方法：