MOE架构解读

MOE 基本原理

MOE由两个主要的部件构成：第一个是专家expert，每个专家代表一个就诊网络；第二个是路由router，门控网络，通过开关将输入的token路由到不同的专家。
MOE结构里面每一层都有很多专家和路由，每一层的专家和路由数量都可以调整。

回顾transformer架构，每个输入都是一个token，每个token输入到position embedding，然后通过多层decoder堆叠，最后输出token。

其中每个decoder内部如下所示：

首先需要数据归一化，随后进入多头注意力机制，然后进入残差连接，最后进入稠密的前馈神经网络。
MOE就是把稠密的前馈神经网络变成了稀疏的前馈神经网络，通过门控网络来选择哪些专家来处理输入的token。

每次通过路由选择下一层经过哪一个专家

MOE的核心是把FFNN替换成多个系数的FFN，每一个FFN都是一个专家。

路由原理

路由是用来选择专家的，由FFNN和softmax组成，如下图所示。

softmax会输出每一个专家的概率，最后用一个或多个专家来权和输入的token。

稀疏架构和稠密架构

Transformer是一个稠密的架构，每个输入都需要经过所有的专家。MOE是一个分为稠密MOE和稀疏MOE的架构。

稠密MOE是每个输入都需要经过所有的专家。稀疏MOE是每个输入只需要经过一部分的专家。输出时候只激活部分专家
MOE的一般输入是具体的矩阵，包括包括(b,s,h)，其中b是batch size，s是sequence length，h是hidden size。
输入首先经过FFN路由层，输出到H(x)层在经过softmax，计算得到专家的概率分布。随后通过门控网络选择一部分专家来处理输入的token。

负载均衡 load balancing

在deepseek里面专家为256个，在训练时候每个专家的训练速度不一样，pathway会跟倾向于选择算的快的专家，导致算的快的专家学的多，算的慢的专家学的较少，因此引入load balance让每个专家的训练速度一样。

• Keep Top-k
  这里通过 keep Top-k 策略来实现。引入噪声矩阵，来抑制某个专家选择过多时候的得分。
  
  通过在每一层的时候引入噪声矩阵，实现整个网络的负载均衡。

• 辅助损失
  为了进一步提升负载均衡能力，引入auxiliary loss。通过将每个token的专家选择分布求和，来计算每个专家的平均选择次数，然后计算每个专家的平均选择次数和平均选择次数的差值，计算辅助损失，并将其加入到总的损失函数中。

• 专家容量
  专家容量同样是为了解决训练负载均衡问题。
  每个专家在输进去token之前加入每个专家最大出力的token数量，通过增加容量来提升训练稳定性，让专家之间训练更加寻哼