MoE 混合专家路由器深度解析：专家为啥有的累死有的闲死？

摘要

MoE（混合专家模型）依靠稀疏激活实现大参数量、可控计算量的效果，路由器 Router 是整个 MoE 架构的灵魂。本文以 Kimi 的 MoE 实现作为案例，通俗讲解路由器工作原理、负载崩塌问题，给出核心数学公式、PyTorch 极简示例代码，对比各类路由算法的适用边界，剖析工程落地踩坑点。很多人以为 MoE 只是简单堆砌多个专家，实际上路由调度、负载均衡才是决定 MoE 成败的关键。

0. 前言

现在很多大模型都采用 MoE 混合专家架构，例如 Mixtral、Kimi。MoE 拥有大量专家，但并不是每一个 Token 都会喂给全部专家，而是依靠路由器挑选一部分专家参与计算。

不知道大家有没有疑惑过：

1. 路由器如何判断一个 Token 交给哪个专家处理？
2. 训练过程中为什么会出现部分专家负载爆满，部分专家全程摸鱼的现象？也就是负载崩塌。
3. 负载均衡惩罚到底是怎么约束路由器行为？
4. Top‑2、哈希路由、层级路由分别适合什么场景？

本文结合 Kimi MoE 的技术特点，从原理、公式、示例代码、工程边界几个维度讲清楚 MoE 路由器。

1. 路由器是什么？MoE 的导诊台

MoE 模型内部存在众多专家，每个专家擅长处理一类特征的输入 Token。如果每一个 Token 都送入全部专家计算，那就退化成普通稠密模型，完全丧失 MoE 稀疏计算的优势。

路由器 (Router/Gate) 就相当于医院的导诊台。病人（Token）进来，导诊护士（路由器）查看症状（Token 特征向量），分配到对应科室（专家），不需要跑遍全部科室。

1.1 Top‑k 路由工作流程

工业界最普遍使用Top‑k 路由，绝大多数 MoE 模型选用 Top‑2。

1. 接收 Token 特征向量，路由器可学习权重计算每个专家得分；
2. Softmax 归一化得到每个专家被选择的概率；
3. 选取概率最高前 k 个专家，将 Token 分发过去；
4. 将 k 个专家输出结果，依据路由概率加权融合，得到最终输出。

为什么是 Top‑2，不用 Top‑1 或者 Top‑4？

• Top‑1：只有单个专家处理，缺少冗余，一旦路由出错直接影响结果，负载均衡难度很高；
• Top‑4：虽然冗余充足，但是激活专家过多，稀疏带来的算力优势直接消失；
• Top‑2 是工业界反复试出来的折中方案，既有冗余兜底，又可以保留稀疏计算优势。

1.2 核心数学公式

① 路由器计算专家概率分布

(p_i = \text{softmax}(W_g \cdot x)_i)- x：当前 token 的特征向量

• (W_g)：路由器 Gate 层可训练权重矩阵
• (p_i)：第 i 个专家的路由概率

② Top‑K 专家输出加权合并

(y = \sum_{i \in \text{TopK}} p_i \cdot E_i(x))- (E_i(x))：第 i 个专家对于输入x的输出；

• (p_i)作为权重，概率越高，该专家输出话语权越大，类似多医生会诊综合意见。

③ 加入负载惩罚后的总损失函数

(\mathcal{L}{total} = \mathcal{L}{task} + \lambda \cdot \mathcal{L}{balance})- (\mathcal{L}{task})：主任务损失，语言模型一般为交叉熵损失；

• (\lambda)：惩罚系数，用来控制负载均衡惩罚强度；
• (\mathcal{L}_{balance})：负载均衡损失，负载分配越不均衡，该项数值越大。

常见三类负载均衡损失：

1）重要性损失，基于专家平均路由概率：

(\mathcal{L}_{importance} = N \cdot \sum_i p_i^2)
N为专家总数量，衡量概率层面分配均匀程度。

2）负载损失，基于专家真实接收 Token 占比：

(\mathcal{L}_{load} = N \cdot \sum_i l_i^2)
(l_i)代表第 i 个专家实际承载 token 比例，反映真实计算负载。

3）熵正则，鼓励路由器分配更加分散：