MoE Router 如何训练：矩阵维度、Top-K 离散选择与梯度路径

回顾 2026-09-17 的问题，整理核查于 2026-09-18。公式和小例子为教学推导，不是用户模型的训练记录。

问题起点与维度纠正

你问 Wr h 为何被称为向量，以及“没选中的专家没有试过，梯度怎么知道下一次该选谁；换专家时 loss 会不会跳变”。

设单 token 隐藏状态 h 是 d 维列向量，路由矩阵 W_r 是 N×d，乘积 z=W_r h 是 N 维向量。若一批有 T 个 token，以行向量堆成 H，H W_r 的转置版本应写为 H W_r^T，结果是 T×N。矩阵是参数，向量是对单 token 的打分结果，两者没有矛盾。

一条可分析的路由公式

教学版本先算 p=softmax(z)，选择集合 S=TopK(p)，再计算 y=Σ(i∈S) g_i E_i(h)。其中 g_i 可能直接使用 p_i，也可能在选中集合内重新归一化；真实模型还可能用 sigmoid、分组、偏置或额外平衡机制。稀疏门控与负载平衡的经典背景见MoE 原论文，工程配置差异见Megatron Router 实现。

梯度怎样穿过“选择”

离散索引通常不按连续变量求导。在本次前向选中的集合固定时，损失仍可以沿专家输出和连续门控权重反传。选中的专家获得这条 token 路径的梯度；未选专家一般没有该 token 对其专家参数的梯度，但它们会被其他 token 选择并训练。

不能笼统说“没选中的 router 分数一定没有梯度”。若使用全体专家上的 softmax，p_i 对未选 z_j 的偏导为 -p_i p_j，分母带来竞争耦合。若先 Top-K 再仅在集合内归一化，则梯度范围不同。

一个关键边界：Top-1 后若把唯一门控权重严格归一化成常数 1，这一权重路径不再提供同样的路由梯度。实际训练必须看是否保留原始门控值、平衡损失或其他学习机制，不能把 Top-2 公式不加修改地套到所有 Top-1 实现。

换专家时是否不平滑

会存在离散边界。即便两位专家分数接近，专家输出也未必相同，切换后输出可能跳变；不能声称 softmax 让整个稀疏网络处处光滑。优化利用当前区域的连续梯度推进参数，token 表示、专家和路由器在大量样本上共同演化，跨过边界后采用新的路径。

Router 没有在每个 token 上穷举所有专家再挑最优。多 token 的覆盖、Top-K、噪声/探索和平衡机制有助于形成分工，但不保证任何初始化都能找到全局最优，也不保证专家天然均匀负载。

可复现的学习实验

用两个小线性专家和一个二维 router，对同一输入分别实现全局 softmax 后选取、Top-K 内归一化。保留 logits、门控权重与梯度，观察未选 logit 是否非零。再连续扫描一个 logit，画输出与损失，比较边界两侧。该实验可解释梯度机制，不能据此推断大型 GLM 的训练细节。

训练中的路由学习与推理系统中的 EPLB 分工不同：前者学习选择，后者调整专家的物理映射。DP2、TP8 与 MoE EP：两节点部署为何不一定是两份独立模型。