MoE 负载均衡的数学
为什么均匀分布是最小值
稀疏 MoE 的路由天然想"抱团"——少数专家被挤爆、多数专家闲置。负载均衡损失就是把这种倾向按住的数学约束。本文逐符号拆开 L_{aux}=\alpha N\sum_i f_i P_i，用三种独立方法严格证明它在专家负载均匀时取到全局最小，再讲清 z-loss、专家容量、Top-k 与 DeepSeek 无辅助损失偏置法，最后给出实际调参。
- 负载均衡损失惩罚的是"分派频率 f_i 与路由概率 P_i 同时集中在少数专家"这件事。
- 它的形式是两个概率向量的内积 \sum_i f_i P_i，乘以专家数 N 做归一化，再乘权重 \alpha。
- 在负载完全均匀（f_i=P_i=1/N）时，这个内积取到全局最小值 1/N——可用凸性、柯西-施瓦茨、拉格朗日三种方法各自证明。
- 它是可微代理：真正想最小化的是离散的负载方差，但那不可导，于是用 f_i（不可导）与 P_i（可导）相乘，让梯度只经 P_i 回流。
- DeepSeek 的 loss-free 方案干脆不加损失、不污染梯度，改用一个每专家偏置 b_i 直接调路由，用反馈控制把负载推平。
01路由为什么会塌缩
先把问题讲成一句大白话：稀疏 MoE 里，一个 token 只被送进 N 个专家中的少数几个，选谁由一个可训练的路由器（router）决定。麻烦在于——路由器和专家是一起训练的。
训练早期，某几个专家碰巧初始化得好、被多分到几个 token，于是它们学得更快、变得更"好用"；路由器发现送给它们收益高，就更爱送给它们；它们又因此学得更快……这是一个正反馈。数学上叫赢者通吃（winner-take-all）：
最终少数专家承担几乎所有 token（过载 + 超出容量被丢弃），多数专家几乎不被激活（dead experts，梯度稀少、参数浪费）。后果有三：
所以我们需要一个数学手段，在训练目标里加一项，让"把 token 均匀摊到所有专家"变成对模型有利的事。这就是负载均衡损失（load balancing / auxiliary loss）。
02Top-k 路由与门控
要写出均衡损失，先把路由的记号钉死。设一层 MoE 有 N 个专家，一个 batch 有 T 个 token。
路由器是一个线性层 W_r\in\mathbb R^{d\times N}，对 token 表示 x_t 打出 N 个 logit，再过 softmax 得到该 token 分给每个专家的概率：
只保留概率最大的 k 个专家（Switch 用 k=1，GShard/Mixtral/DeepSeek 常用 k=2 或更多），其余置零，选中的再归一化做加权：
负载均衡只关心两个"每专家"的量，都在一个 batch 内对 token 求平均：
两者天差地别：f_i 数的是"硬"分派次数——含 argmax/TopK，对参数不可导；P_i 是"软"概率的均值——处处可导。这个"一硬一软"的区别，是理解损失设计的全部关键，第 5 节会用到。
03辅助损失逐符号拆解
这是 Switch Transformer / Mixtral 用的经典负载均衡损失，也是本文要讲透的主角：
加到语言模型主损失上：\mathcal L_{total}=\mathcal L_{LM}+\mathcal L_{aux}。下面把每个符号是什么、为什么在那里，逐个讲清。
逐符号解释
为什么是"频率 × 概率"这个乘法形式
关键在于 f_i 与 P_i 分工不同：
把两者相乘 f_i P_i，等于用不可导的 f_i 当作一个常数权重（stop-gradient）去加权可导的 P_i：哪个专家实际过载（f_i 大），就给它的概率 P_i 施加更大的下压梯度。于是损失"知道"该压谁，梯度又"能够"压下去。这一步的梯度推导见第 5 节。
GShard 的等价写法
GShard（更早）写成 mean-importance 与 mean-load 的乘积和，本质一致——把上面的 P_i 叫 importance、f_i 叫 load：
两篇论文的形式相同，只是命名与个别归一化系数差异。下面的最优性证明对二者都成立。
04均匀分布取最小值：三种证明
这是全文的数学核心，也是你的问题所在。我们要严格证明：在约束下，\sum_i f_i P_i 在负载完全均匀时取到全局最小值。三种方法各自独立成立，从三个角度看清同一件事。
证法一 · 拉格朗日乘子（找驻点）
为了把两个向量的耦合讲透，先看一个更干净的对称情形：路由器若把软概率与硬分派对齐（训练收敛时 f_i\propto P_i），问题退化为最小化 \sum_i P_i^2（下面证法二会严格处理一般情形）。这里先对 \sum_i P_i^2 在约束 \sum_i P_i=1 下求最小。
逐分量求导，二次项导数是 2P_i，约束项导数是 \lambda：
关键：右边 \lambda/2 与下标 i 无关——所以所有 P_i 必须相等。这一步已经逼出"均匀"。
把 P_i=\lambda/2 代入 \sum_i P_i=1：N\cdot\tfrac\lambda2=1，于是 \tfrac\lambda2=\tfrac1N：
目标 \sum_i P_i^2 的 Hessian 是 2I（正定），函数严格凸，约束是线性的——凸函数在凸集上的驻点必是全局最小。所以 P_i=1/N 是全局最小点，最小值 \sum_i(1/N)^2=1/N。
证法二 · 柯西-施瓦茨（一步给下界，最一般）
这个方法不需要假设 f=P，直接对内积给下界，是最干净的证明。对任意非负向量 \mathbf f,\mathbf P：
取全 1 向量 \mathbf 1=(1,\dots,1)，柯西-施瓦茨给出：
左边 \big(\sum P_i\big)^2=1^2=1（用约束），所以：
柯西-施瓦茨取等号当且仅当两向量成比例，即 P_i\propto 1——所有 P_i 相等，结合 \sum P_i=1 得 P_i=1/N。这就从下界的角度证明了：任何偏离均匀的分布，\sum P_i^2 都严格大于 1/N。
证法三 · 与均匀分布的距离（几何直觉）
把 \sum_i P_i^2 按"到均匀分布的偏差"展开，最直观。令 \bar P=1/N 为均匀值，写 P_i=\bar P+\delta_i，其中 \sum_i\delta_i=0（因为两者和都为 1）：
交叉项因 \sum_i\delta_i=0 消失，只剩 \tfrac1N（常数）加上 \sum_i\delta_i^2\ge 0（偏差平方和，恒非负）。所以：
这行公式把话说死了：损失 = 一个固定基线 + 一个"离均匀多远"的平方惩罚。惩罚项当且仅当所有 \delta_i=0（即完全均匀）时为零。均匀分布是唯一的全局最小，任何不均衡都被平方放大地惩罚。这也解释了为何它对付赢者通吃格外有效——越集中，\delta_i 越大，惩罚二次增长。
| 方法 | 核心一步 | 给出什么 |
|---|---|---|
| 拉格朗日 | 2P_i-\lambda=0\Rightarrow P_i 与 i 无关 | 驻点位置 = 均匀；凸性保证是最小 |
| 柯西-施瓦茨 | (\sum P_i)^2\le N\sum P_i^2 | 下界 1/N + 等号条件 = 均匀（最一般） |
| 偏差展开 | \sum P_i^2=\tfrac1N+\sum\delta_i^2 | 损失 = 基线 + 偏离的平方惩罚（最直观） |
05为什么这样设计梯度
证明说清了"最小值在哪"，但没说"优化器怎么走到那"。这一节把梯度算出来，看清损失到底在推路由器做什么。
反向传播时 f_i 是常数（含 argmax，梯度视为 0），梯度只经 P_i：
再往里，P_i=\tfrac1T\sum_t p_{t,i}，而 p_{t,i} 是 softmax，其对 logit z_{t,j}=x_t^\top W_{r,j} 的导数是标准 softmax 雅可比 \partial p_{t,i}/\partial z_{t,j}=p_{t,i}(\delta_{ij}-p_{t,j})。链式合起来，对某个 token 的 logit：
注意梯度前面有 p_{t,j} 因子：路由器对某专家概率越高，这次修正越大。这让"当前正被偏爱的专家"承受最强的下压，正好对症赢者通吃。
06z-loss 与专家容量
光有均衡损失还不够稳。大规模训练里还有两个配套机制，一软一硬。
router z-loss：按住 logit 不爆炸
softmax 的输入 logit 若整体越来越大，会导致数值不稳定（fp16/bf16 溢出）和路由过度自信。Switch-XXL 引入 z-loss，惩罚 log-partition 的平方：
\log\sum_j e^{z_{t,j}} 是 logsumexp，正比于 logit 的整体量级。把它平方后惩罚，等于把路由 logit 往 0 附近拉、防止 softmax 饱和成 one-hot。总损失变成三项：
专家容量 capacity：硬性的溢出闸门
均衡损失是"软"约束——它降低不均衡的概率，但不保证某个 batch 里不会有专家被挤爆。所以工程上还有一道硬闸门：给每个专家设一个容量上限。
超过容量的 token 被丢弃（dropped）——不进该专家的 FFN，只走残差连接直通。容量因子越大越少丢 token，但每个专家要预留更大的缓冲显存、计算浪费越多。这是"丢 token 的信息损失"与"显存/算力浪费"之间的权衡。
07DeepSeek 无辅助损失偏置法
辅助损失有个副作用：它往主损失里塞了一个"和语言建模无关"的梯度，多少会干扰模型质量（论文称 interference gradient）。DeepSeek 的 Loss-Free Balancing（arXiv:2408.15664）提出：不加损失、不动梯度，改用一个每专家的可调偏置直接调路由。
给每个专家一个偏置 b_i。选 Top-k 时用加了偏置的分数，但真正的门控加权权重仍用原始分数 s_{t,i}：
关键：b_i 只影响"哪些专家入选"这个离散决策，不进入前向输出的加权，因此不产生梯度、不污染主损失。它是一个纯粹的路由调节旋钮。
每个 step 后统计各专家的负载 c_i，与平均负载 \bar c 比较，欠载就抬高偏置（下步更容易入选）、过载就压低：
| 辅助损失法（Switch/GShard） | 无损失偏置法（DeepSeek） | |
|---|---|---|
| 均衡手段 | 往损失加 \alpha N\sum f_iP_i | 路由分数加偏置 b_i |
| 是否动梯度 | 是，产生干扰梯度 | 否，梯度纯净 |
| 核心超参 | \alpha（损失权重） | u（偏置更新率） |
| 质量影响 | 轻微负面（需调 \alpha 平衡） | 基本无损 |
08实际参数怎么调
把上面的量落到训练配置里，四个旋钮加两条准则。
| 参数 | 典型值 | 调大 → | 调小 → |
|---|---|---|---|
| \alpha aux 权重 | 10^{-2} | 更均衡，但路由忽略语义、质量降 | 更贴语义，但易塌缩、dead experts |
| \beta z-loss 权重 | 10^{-3} | logit 更小更稳，可能欠自信 | logit 易爆、数值不稳 |
| capacity_factor | 1.0\!\sim\!1.25（训）/ \ge2（推） | 少丢 token，费显存/算力 | 省显存，丢 token 多 |
| k Top-k | 1（Switch）/ 2（主流） | 质量↑、计算↑、更难均衡 | 更省算力、更易塌缩 |
| u 偏置更新率 | 10^{-3} | 均衡收敛快、可能震荡 | 平稳、收敛慢 |
09资料来源
Switch Transformer — Fedus, Zoph, Shazeer. Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity. JMLR 2022. arXiv:2101.03961（负载均衡损失 \alpha N\sum f_iP_i、z-loss、容量因子的出处）
GShard — Lepikhin et al. GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding. ICLR 2021. arXiv:2006.16668（mean-importance / mean-load 均衡项）
DeepSeek Loss-Free — Wang et al. Auxiliary-Loss-Free Load Balancing Strategy for Mixture-of-Experts. 2024. arXiv:2408.15664（偏置反馈法 b_i\leftarrow b_i+u\,\mathrm{sign}(e_i)）
ST-MoE — Zoph et al. Designing Effective Sparse Expert Models. 2022. arXiv:2202.08906（router z-loss 的系统分析）
公式记号做了统一（N 专家数、T token 数、k top-k、f_i/P_i 频率/概率），与各原论文一一对应，未改变数学内容。