MoE 的路由与加权汇总

Router 为每个 token 选择 4 个专家中的 2 个,并确定其输出的混合权重。

路由结果

图示可左右滑动,保留矩阵比例与文字大小。

01

路由概率

路由分数 xW_r 经 softmax 转为专家概率,图中以柱高表示。

专家 10.10logit -2.30专家 20.50logit -0.69专家 30.30logit -1.20专家 40.10logit -2.30
02

专家选择

取概率最大的两项,保留专家编号,并将对应概率重新归一化。

选中编号(整数)23选中概率(0.50, 0.30)选中概率归一化(0.625, 0.375)
03

分发与汇总

输入 x 分发给选中的专家,各专家输出按归一化权重求和。

输入 x专家 1专家 2专家 3专家 4y = 0.625 f₂(x) + 0.375 f₃(x)
维度
x 的形状为 1 × d,Router 权重为 d × E,取 E = 4。各专家输出均为 1 × d;图中节点表示专家身份。
对象
Top-k 返回按概率降序排列的专家编号,取值为 1 到 4 的整数。选中概率归一化后,两项混合权重之和为 1。
计算
示例采用 Mixtral 的选中概率归一化方式,路由分数取 log(p)。本 token 仅在选中的专家中进行前向计算。