MoE 详解：稀疏激活与路由负载均衡

MoE（Mixture of Experts，混合专家）是近几年大模型"参数量暴涨但推理成本没等比上涨"的关键架构。它的核心承诺是：让总参数量很大，但每个 token 只激活一小部分。听起来像免费午餐，但魔鬼藏在路由和负载均衡里。这篇把机制、数学和工程坑都讲透。

直觉：把"每个 token 都过整个网络"改成"按需调用专家"

标准 Transformer 里，每个 token 都要完整经过每一层的前馈网络（FFN）。FFN 通常占了模型大部分参数。MoE 的想法是：把一个大 FFN 拆成 $N$ 个小 FFN（专家），再加一个路由器（router/gate）决定每个 token 该送给哪几个专家。

这样就解耦了两个量：

总参数量（model capacity）= 所有专家参数之和，决定模型能记多少知识。
激活参数量（active params）= 每个 token 实际经过的专家参数，决定单 token 的计算量。

一个典型配置：64 个专家，每个 token 只选 top-2。于是总参数可能是稠密模型的十几倍，但单 token FLOPs 只略增。用更多内存换更强的能力，同时不等比增加计算——这就是 MoE 的经济学。

机制：路由器如何工作

路由是一个小型可学习的门控。对输入 token 表示 $x$ ，路由器算出每个专家的得分，取 top- $k$ ：

def moe_layer(x, experts, W_gate, k=2):
    # x: [d_model], 一个 token 的表示
    logits = x @ W_gate              # [N_experts]，每个专家一个打分
    topk_val, topk_idx = top_k(logits, k)
    weights = softmax(topk_val)      # 只在选中的 k 个上归一化

    out = 0
    for w, idx in zip(weights, topk_idx):
        out += w * experts[idx](x)   # 只前向选中的专家
    return out

注意几个要点：

softmax 只在被选中的 $k$ 个专家上做归一化，是稀疏的。
top- $k$ 是离散选择，不可导。梯度通过被选中专家的门控权重 $w$ 回流（这部分可导），未选中的专家这一步拿不到梯度。这使路由训练天然不稳定。

数学上，第 $i$ 个专家的门控值常写作：

$g_i(x) = \frac{\exp(h_i)}{\sum_{j \in \text{TopK}} \exp(h_j)}, \quad h = x W_g$

层输出为 $y = \sum_{i \in \text{TopK}} g_i(x)\, E_i(x)$ 。

核心难题：负载均衡

MoE 最大的工程难题不是"怎么路由"，而是"怎么让路由别偏心"。如果不加约束，训练会陷入一个恶性循环：

1 2	某几个专家先被多选 → 它们训练得更好 → 路由器更倾向选它们 → 它们被更多选 → 其余专家"饿死"（几乎拿不到 token，永远学不好）→ 专家坍塌

结果是名义上 64 个专家，实际只有少数几个干活，巨大的参数容量被浪费。更糟的是工程后果：MoE 实现通常给每个专家设容量上限（capacity factor）——一个专家在一个 batch 里最多处理多少 token。超出容量的 token 会被**丢弃（dropped）**或绕过该层。负载不均会导致热门专家爆容量、token 被丢，冷门专家空转，吞吐和质量双输。

解决办法是引入辅助负载均衡损失（auxiliary loss）。一种经典形式是：设 $f_i$ 为路由到专家 $i$ 的 token 比例， $P_i$ 为该专家门控概率的平均值，则

$L_{\text{aux}} = \alpha \cdot N \sum_{i=1}^{N} f_i \cdot P_i$

它在 $f_i$ 和 $P_i$ 都均匀（即每个专家约 $1/N$ ）时最小。把它加进总损失，就是在"鼓励路由器把负载摊平"。 $\alpha$ 是个需要调的小系数：太小压不住坍塌，太大会损害模型质量（强迫把不合适的 token 塞给某些专家）。近年也有无辅助损失的做法，比如给每个专家一个可学习/动态调整的偏置项来均衡，避免辅助损失对主任务的干扰。

另外两个工程武器：

noisy top-k：路由打分上加噪声，增加探索，避免过早锁定到少数专家。
专家并行（Expert Parallelism）：不同专家放在不同设备上。token 按路由结果跨设备发送，这引入了 all-to-all 通信——这是 MoE 分布式训练/推理的主要开销和复杂度来源。

工程权衡：MoE 不是纯赚

天下没有免费午餐，MoE 的代价很具体：

显存翻倍：虽然激活参数少，但所有专家权重都得常驻显存。64 专家的 MoE，推理时显存占用接近其总参数量，而非激活量。这直接限制了它在端侧的可行性。
通信成本：专家并行的 all-to-all 在低带宽互联下可能成为瓶颈，吃掉稀疏带来的计算节省。
负载不均的吞吐损失：动态路由意味着各专家负载随 batch 波动，难以像稠密模型那样规整地打满硬件。容量因子设小了丢 token，设大了浪费算力。
训练不稳定：路由的离散性 + 负载均衡的拉扯，使 MoE 往往比同规模稠密模型更难训稳，对学习率、初始化更敏感。

一个常见误区：“MoE 模型 = N 个独立专家各管一个领域（数学专家、代码专家……）”。实际并非如此。可解释性研究普遍发现，专家的分工更多与 token 级、语法级特征相关，而非人类可理解的高层主题。路由是模型自己学出来的统计分工，不要把它拟人化。

小结

MoE 用一个简单的解耦实现了"大容量、低激活成本"：总参数管知识，激活参数管计算。但它把复杂度从"算得多"转移到了"路由准 + 负载匀 + 通信省 + 显存够"。真正决定一个 MoE 好不好用的，不是它有多少专家，而是它的路由是否均衡、丢 token 率是否可控、all-to-all 是否被打满。理解了负载均衡这条主线，你就抓住了 MoE 工程的命门。