MoE 深入：Router 如何挑中两个 Expert

MoE 解决什么矛盾

Transformer 的前馈层占用大量参数和计算。把它做得更宽，容量上升，但每个 token 都运行整个网络。Mixture of Experts（MoE）把一个前馈层换成多个候选专家，再让 Router 只选少数：保存更大总容量，同时让单 token 计算保持稀疏。

INTERACTIVE FIGUREGrok-1：8 位专家，Top-2 上场

100%

当前 token 语义扰动 37

6,144 维 token 表示Router
softmaxTop-2

拖动 token 的语义倾向，router 分数随之变化。最高两名被激活并按权重合并。专家名称只是教学标签；公开源码没有声明每位专家固定专长。

交互图给专家贴了教学标签以方便观察，但真实源码中的 E1–E8 是结构对称的网络。训练可能让它们形成不同分工，官方没有为 Grok-1 发布“代码专家”“数学专家”之类固定语义。

对 token 表示 $h\in\mathbb{R}^{6144}$ ，Router 有一张 $W_r\in\mathbb{R}^{6144\times8}$ 的权重：

s=hW_r, \qquad p=\operatorname{softmax}(s), \qquad I=\operatorname{TopK}(p,2)

$s$ 是 8 个未归一化分数； $p$ 是路由概率； $I$ 是概率最大的两个专家索引。Grok-1 源码把输入转 fp32 计算路由，padding 位置的概率会被 mask。

选中专家 $E_i$ 后，输出是：

y=\sum_{i\in I}p_iE_i(h)

源码直接使用 Top-2 原始概率做 gate 并求和，没有在选中两项之间显式重新归一化。公式表达“谁被选择、各占多少”；真实实现还需要把 token 映射到专家参数，并在设备间安排计算。

Grok-1 的 DenseBlock 是门控前馈网络。简化表达为：

E(h)=W_o\left((W_vh)\odot\operatorname{GELU}(W_gh)\right)

$W_v$ 产生内容分支， $W_g$ 经过 GELU 形成门控， $\odot$ 是逐元素乘法， $W_o$ 投影回 6,144 维。widening_factor=8 决定中间宽度，仓库用 ffn_size() 将结果对齐到 256 的倍数。

这比普通“两层 MLP”多一条门控分支，让网络可以按输入调制通过的特征；代价是更多参数和矩阵乘法。

如果 8 个 Expert 全运行并平均，容量与计算会一起扩大，失去稀疏意义。Top-1 更省，但把 token 完全交给单一专家，路由错误和训练不稳定影响更大；Top-2 让两个专家共同贡献，是容量、稳健和计算的折中。

如果 Router 总偏爱少数 Expert，这些设备会过载，其余专家学不到东西。现代 MoE 通常需要负载均衡损失、容量限制、溢出策略与专家并行。Grok-1 公开仓库是推理示例，不包含完整训练目标，不能从返回的占位值 0 推断生产训练没有均衡机制。

moe_slow_matmul1/2 先计算各专家输出，再用 one-hot expert index 选取需要的行，配合 shard_map 和 psum。这避免依赖自定义 sparse kernel，便于核对 checkpoint，但会做比优化 MoE kernel 更多的工作。

生产系统通常希望只 dispatch 被选 token、把专家放到不同设备、聚合后再送回原 token 顺序。优化目标包括负载、网络拓扑、batch 形状与 kernel 融合，而不只是数学结果。

MoE 的核心不是“多个模型投票”，而是 Transformer 每一层的前馈位置有多个参数分支；Router 对每个 token、每一层选择少数专家。Attention 仍负责 token 之间交流。