摘要MoE（混合专家模型）依靠稀疏激活实现大参数量、可控计算量的效果，路由器 Router 是整个 MoE 架构的灵魂。本文以 Kimi 的 MoE 实现作为案例，通俗讲解路由器工作原理、负载崩塌问题，给出核心数学公式、PyTorch 极简示例代码，对比各类路由算法的适用边界，剖析工程落地踩坑点。很多人以为 MoE 只是简单堆砌多个专家，实际上路由调度、负载均衡才是决定 MoE 成败的关键。0. 前言现在很多大模型都采用 MoE 混合专家架构，例如 Mixtral、Kimi。MoE 拥有大量专家，但并不是每一个 Token 都会喂给全部专家，而是依靠路由器挑选一部分专家参与计算。不知道大家有没有疑惑过：路由器如何判断一个 Token 交给哪个专家处理？训练过程中为什么会出现部分专家负载爆满，部分专家全程摸鱼的现象？也就是负载崩塌。负载均衡惩罚到底是怎么约束路由器行为？Top‑2、哈希路由、层级路由分别适合什么场景？本文结合 Kimi MoE 的技术特点，从原理、公式、示例代码、工程边界几个维度讲清楚 MoE 路由器。1. 路由器是什么？MoE 的导诊台MoE 模型内部存在众多专家，每个专家擅长处理一类特征的输入 Token。如果每一个 Token 都送入全部专家计算，那就退化成普通稠密模型，完全丧失 MoE 稀疏计算的优势。路由器 (Router/Gate) 就相当于医院的导诊台。病人（Token）进来，导诊护士（路由器）查看症状（Token 特征向量），分配到对应科室（专家），不需要跑遍全部科室。1.1 Top‑k 路由工作流程工业界最普遍使用Top‑k 路由，绝大多数 MoE 模型选用 Top‑2。接收 Token 特征向量，路由器可学习权重计算每个专家得分；Softmax 归一化得到每个专家被选择的概率；选取概率最高前 k 个专家，将 Token 分发过去；将 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）熵正则，鼓励路由器分配更加分散：
摘要MoE（混合专家模型）依靠稀疏激活实现大参数量、可控计算量的效果，路由器 Router 是整个 MoE 架构的灵魂。本文以 Kimi 的 MoE 实现作为案例，通俗讲解路由器工作原理、负载崩塌问题，给出核心数学公式、PyTorch 极简示例代码，对比各类路由算法的适用边界，剖析工程落地踩坑点。很多人以为 MoE 只是简单堆砌多个专家，实际上路由调度、负载均衡才是决定 MoE 成败的关键。0. 前言现在很多大模型都采用 MoE 混合专家架构，例如 Mixtral、Kimi。MoE 拥有大量专家，但并不是每一个 Token 都会喂给全部专家，而是依靠路由器挑选一部分专家参与计算。不知道大家有没有疑惑过：路由器如何判断一个 Token 交给哪个专家处理？训练过程中为什么会出现部分专家负载爆满，部分专家全程摸鱼的现象？也就是负载崩塌。负载均衡惩罚到底是怎么约束路由器行为？Top‑2、哈希路由、层级路由分别适合什么场景？本文结合 Kimi MoE 的技术特点，从原理、公式、示例代码、工程边界几个维度讲清楚 MoE 路由器。1. 路由器是什么？MoE 的导诊台MoE 模型内部存在众多专家，每个专家擅长处理一类特征的输入 Token。如果每一个 Token 都送入全部专家计算，那就退化成普通稠密模型，完全丧失 MoE 稀疏计算的优势。路由器 (Router/Gate) 就相当于医院的导诊台。病人（Token）进来，导诊护士（路由器）查看症状（Token 特征向量），分配到对应科室（专家），不需要跑遍全部科室。1.1 Top‑k 路由工作流程工业界最普遍使用Top‑k 路由，绝大多数 MoE 模型选用 Top‑2。接收 Token 特征向量，路由器可学习权重计算每个专家得分；Softmax 归一化得到每个专家被选择的概率；选取概率最高前 k 个专家，将 Token 分发过去；将 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）熵正则，鼓励路由器分配更加分散：
