MoE 为什么不能只追求负载均衡：从路由机制到辅助损失与动态偏置

把四个专家的工作量调成一样，是否就得到了更好的 MoE？这听起来像一个调度问题，却会深入训练目标：谁获得输入，谁就有机会学习；为了改变分配而修改路由，又可能改变原本有利于预测的组合。问题不只是让柱状图变平，而是弄清楚我们为这份均匀付出了什么。

本文以《Auxiliary-Loss-Free Load Balancing Strategy for Mixture-of-Experts》的固定 v1 为研究起点，先把计算与梯度拆开，再运行一个可检查的小实验。[1] 后面的数字来自本地 CPU 合成回归，不是语言模型复现；结论也不会预设动态偏置必须获胜。

1. 跟随一个 token，把路由完整算一遍

普通前馈网络让所有输入经过同一组参数。MoE 则先让路由器给专家打分，每个 token 只选其中 K 个参与结果。这里的专家不是预先命名的“数学专家”或“代码专家”，而是一组独立、可学习的函数。路由器不知道这些称号，只能通过任务训练逐渐调整选择。

设一批输入 X 的形状是 [T,d]，T 是 token 数，d 是特征宽度；路由矩阵 W 是 [d,N]，N 是专家数。为避免依赖网站的数学插件，下面用可直接阅读的公式代码块：

Z = X @ W                         # [T,N]，路由 logits
S = sigmoid(Z)                    # [T,N]，每个分数在 0 与 1 之间
I = topk_indices(S, K)             # [T,K]，离散专家编号
A = gather(S, I)                   # [T,K]，取出被选中的原始分数
w = A / sum(A, dim=专家)           # [T,K]，选中权重之和为 1
Y = sum_k(w_k * Expert_{I_k}(X))   # [T,d]

这里特意选 sigmoid：各专家分别评分，并不要求所有 N 个分数之和等于一。最后一行的权重归一化是本教学实现的选择；原论文公式（1）、（3）直接用选中的原始分数加权，没有写这一步，且包含残差连接。[1] 三种实验策略统一使用上面的归一化，避免把权重尺度差异误当成均衡方法的贡献，但因此不能声称逐式复现论文。

现在取 T=1,d=2,N=4,K=2。输入 x=[1,-1]，W 第一行取 [0.847298,0.405465,-0.405465,-1.386294]，第二行全零。矩阵乘法得到同样的四个 logits，经过 sigmoid 得到 [0.7,0.6,0.4,0.2]。不加路由偏置时选专家一、二，权重分别是 7/13、6/13。若两个专家输出 [2,0]、[0,2]，加权结果是 [14/13,12/13]。

为什么不让最大的分数直接胜出、只选一个专家？在本文“选中权重再归一化”的约定下，若 K 等于一，唯一权重恒为一，任务损失就无法通过这个权重向路由分数提供梯度。选择集合本身又是离散的，于是必须另外考虑训练信号。选两个专家不仅展示混合，也保留了一条清楚可检查的任务梯度路径。这不是说所有 top-1 MoE 都无法训练，而是提醒我们：门控函数、归一化位置和梯度路径必须一起讨论，不能只比较 K 的数值。

这条计算链有两个不同的问题：选择谁由排序决定；选中后贡献多少由权重决定。后面动态偏置只修改前一个问题。批量计算时，被选中的专家输出可以组织成 [T,K,d]，权重扩展为 [T,K,1] 后相乘，沿 K 维相加。实验为了容易检查，会计算所有专家输出再 gather，并不真的节省这部分运算。

图中的 stopgrad 只属于离散选择分支。若把 S 整体 detach，再拿它计算输出权重，任务损失就无法通过权重训练路由器；“top-k 不可微”不意味着“路由器没有梯度”。

2. 负载不等于能力，失衡不等于分工

先数清楚什么叫负载。本文统计一个 batch 内，每个专家被 top-k 选中的次数 c。没有容量限制和丢弃时，总分配数一定是 T*K；理想均匀负载是 T*K/N。负载份额定义为 q_i=c_i/(T*K)，四个专家理想上各占四分之一。一个 token 被分配两次，所以这里的分母不能只写 T。

用论文定义的 MaxVio 衡量最忙专家超过平均负载多少：[1]

MaxVio = max(c) / mean(c) - 1

例如八次分配的计数是 [4,2,2,0]，平均值为二，MaxVio 为一。这只表示最忙专家的工作量是平均值的两倍，不表示丢弃了一半 token。是否丢弃还取决于有没有容量上限、溢出如何处理。本文不实现 capacity、token dropping 或分布式通信，因而不会把这个数字翻译成吞吐提升。

还要区分三件事。失衡是一段观察窗口里的计数偏斜；塌缩是路由长期集中到少数专家，使其余专家持续缺乏学习机会；分工则是不同专家对不同输入形成互补能力。一次不均匀的 batch 不足以证明塌缩，一张均匀柱状图更不能证明分工。随机轮流分配也能非常均匀，但可能让每个专家都学习相似内容。

失衡为什么会自我加强？在硬选择下，未选中的专家通常没有来自该 token 的任务梯度；经常被选中的专家得到更多更新，短期内可能更好用，路由于是继续偏爱它们。但这只是可能的反馈，不是每种失衡的唯一原因：输入本身有结构、路由初始化有偏好、各专家学习速度不同，都能产生相似计数。

统计窗口同样重要。某两个 batch 可以分别全压向不同专家，合起来却很均匀。全验证集 MaxVio 描述长期利用情况，逐 batch MaxVio 才更接近短期分配尖峰。本文把两者分别保存，不用汇总均匀掩盖瞬时波动。若观察到路由改变，也只能说选择发生变化，不能直接给专家贴上语义标签。

3. 辅助损失怎样推动均衡，又怎样改变主目标

硬计数无法提供通常意义上的连续梯度。辅助均衡损失因此把“已经发生的分配”和“可微的平均评分”相乘。按照论文公式（2），去掉系数后定义：[1]

f_i = N * c_i / (K*T)                  # detach；sum_i f_i = N
P_i = mean_t(S[t,i])                   # 保留梯度
L_balance_raw = sum_i(f_i * P_i)
L_total = L_task + alpha * L_balance_raw

注意 f 不是和为一的份额，而是相对理想负载：均匀时每个 f 都等于一。它比上一节的 q 多一个 N。若照抄不同资料里的记号却漏掉这个因子，相同 alpha 就不再意味着相同强度。

在当前 top-k 集合不变的局部区域，f 被当作常量。对 sigmoid 路由分数的 logit 求导，得到：

d L_balance_raw / d Z[t,i] = (f_i/T) * S[t,i] * (1-S[t,i])

负载越重，f 越大，降低该专家分数的压力通常越强；但 sigmoid 饱和也会让梯度变小，所以不是只看计数就能判断更新幅度。这个梯度经过路由器，还可能继续传到产生输入表示的网络。本文输入是固定特征，没有可学习的上游编码器；辅助损失不直接依赖专家输出，因此不会直接更新专家参数。

主任务走的是另一条路：被选中的专家通过输出获得梯度，路由器通过原始分数归一化后的权重获得梯度。若两个专家的预测贡献不同，任务可能希望提高其中一个权重；辅助项却可能因为它太忙而压低分数。总梯度是二者相加，方向可以相助也可以冲突，不能把“辅助梯度存在”直接等同于“必然有害”。alpha 决定的是这种附加优化压力，不是一个无代价的调度旋钮。

也可以直接推导主任务为何偏爱某个专家。固定所选集合，记选中评分之和为 A、混合输出为 Y，某个专家输出为 E_i，那么输出对它的分数的导数是 (E_i-Y)/A。接上任务误差梯度后，是否提高这个分数，取决于该专家相对于当前混合输出能否沿降低误差的方向移动。辅助项看的是负载与评分，任务项看的是预测贡献；它们衡量的不是同一件事，冲突由此产生，而不是因为辅助项的名字里有“辅助”二字。

这个推导也解释了为什么不能把整个路由过程当作一条平滑函数。只要排名不变，上述导数就是局部有效的；当第三名超过第二名，混合中出现的是另一套专家输出，梯度表达式随之改变。优化器每次看到的是当前分支的导数，并没有预先知道跨过边界后替补专家一定更好。负载控制正是在这种局部学习之外，尝试修正机会分配。

sigmoid 还带来容易忽略的细节：P 不在所有专家之间归一化。即使 f 全为一，辅助项仍是所有平均评分之和，降低所有分数也能降低它；它并不是纯粹的负载方差。在本实现里，如果选中分数按共同倍数缩小，归一化输出权重完全不变。这为“辅助损失下降，但负载没有更均匀”提供了可能路径。真实训练还同时改变排名和专家参数，不能仅凭这个推导断定最终现象的全部原因。

配套测试检查了 f 与计数断梯度、辅助项不更新专家、任务仍训练路由器，并把解析导数与 autograd 对照。有限差分只在远离 top-k 换位边界的固定选择邻域执行；跨过换位点以后函数发生分支变化，不能拿那里的差分失败指责自动微分。

4. 动态偏置省去的是附加梯度，不是训练影响

论文的方法把负载控制移到选择环节。每个专家有一个状态 b，选择时比较 S+b，加权时仍取 S；b 不作为优化器学习的参数。[1] 对刚才的单 token，给专家三加 0.31，选择分数变成 [0.7,0.6,0.71,0.2]，于是选中三、一，而非二、一。

但是专家三的输出权重不能使用 0.71。原始分数仍是 0.4，所以一、三的权重分别为 7/11、4/11。假设专家三输出 [-1,3]，加权结果是 [10/11,12/11]≈[0.9091,1.0909]；若再接残差 x，就得到 [1.9091,0.0909]。本实验直接输出回归值，不接这个示意残差。数值脚本保存了矩阵、分数、专家输出和最终结果，可逐项复算。

偏置怎样更新？论文算法 1 用本 batch 的分配结果，为下一个 batch 调整偏置：[1]

用 b_t 完成本 batch 的前向、任务反传与参数更新
c = 本 batch 已产生的专家选择计数
b_{t+1} = b_t + u * sign(mean(c) - c)
下一 batch 才使用 b_{t+1}

忙于平均水平的专家减少偏置，闲于平均水平的专家增加偏置，刚好相等时不变。它是带延迟的负反馈：负载是反馈信号，b 是控制状态，u 是每步修正速度。我们不重新路由当前 batch，也不在验证集或测试集上更新 b。论文强调使用历史负载，在自回归任务里避免让当前序列后面的信息反过来参与前面 token 的选择。[1] 本例没有序列，但仍严格保留这条时序。

还可以把这次更新算成数字。单 token 选择一、三之后，计数为 [1,0,1,0]，平均计数为 0.5；速度取 0.01，就把偏置从 [0,0,0.31,0] 改成 [-0.01,0.01,0.30,0.01]。当前输出已经算完，不会因为这个新状态被追溯修改。若把“先看当前计数再选专家”写进同一次前向，就形成了循环依赖，必须另行定义重排算法；那已不是这里实现的历史反馈。

此外，负载统计与参数更新并非天然拥有相同频率。这里没有梯度累积，一个 batch 就更新一次模型和一次偏置；若以后加入多个 microbatch 的累积，应明确究竟每个 microbatch 还是每个优化器步更新偏置。相同的 u 乘上不同的更新次数，产生的总修正量不同，不能仅凭配置里同一个数字宣称比较公平。

u 太小，初始偏好需要很多步才能被抵消；u 太大，刚把专家推成热门，下一批又把它压下去，可能来回换位。速度的意义还依赖评分尺度、batch 噪声和模型本身的变化速度，不能把论文中的数值当成跨架构常数。原文第 4.3 节比较了不同速度、按误差幅度而非符号更新、以及乘法偏置；更好的负载指标也没有保证更好的困惑度。[1]

“Loss-Free”准确地说是没有这项均衡辅助损失的反传梯度。偏置改变选择集合，就改变哪些专家接收任务梯度，进而改变下一步参数；所以它仍然改变优化轨迹。把它说成“完全不干扰训练”，会把直接梯度与间接训练影响混成一件事。

5. 三种策略怎样公平地放进一个最小实验

实验用八维输入预测一个连续值。前两维围绕四个中心采样，提供可辨认的输入情境；不同情境对应不同的非线性目标分量，另加共享线性项与噪声。我们不把情境编号交给路由器当标签，也不规定“第几个情境必须找第几个专家”。四种情境只是让输入有结构，专家完全可以共享、混合或另找分工。

模型有四个 8→16→1 的 tanh 专家，路由器为 8→4，top-2，总计 680 个可学习参数。三种策略都采用 sigmoid 评分、选中权重归一化、Adam 和相同学习率。初始路由器的四个截距设置为 [1,0.4,-0.4,-1]，有意制造可控偏好；它们是可学习参数，和动态路由状态 b 不是同一个东西。这个设计使均衡恢复过程更容易观察，也限制了结论对无偏初始化的适用性。

预先固定七个条件：无均衡；辅助系数 0.01、0.1、1；偏置速度 0.001、0.01、0.1。每个条件跑三个种子，每次 1200 步、batch 128；学习率为 0.003，不早停、不挑最佳 checkpoint。输入情境比例分别取均匀和 70/10/10/10，共 42 次完整训练。先用 100 步做计时小试，训练循环耗时约 0.071 秒，再确定上述预算；完整矩阵的脚本墙钟约 22.27 秒。

每个种子保存 4096 个训练样本、1024 个验证样本、1024 个测试样本。不同策略共享该种子的初始权重、数据和 minibatch 序列；不同种子同时改变数据采样与初始化，因此误差棒包含这两种波动。各次运行记录初始权重哈希与数据张量哈希，并分别保存主任务损失、原始辅助项、加权辅助项及每步偏置更新前后的值。

验证集每 50 步用于诊断曲线，不做选参。测试集只在每次训练结束后评估，七组结果全部报告，没有根据测试集挑出“最佳速度”再隐藏其余结果。我们还在固定验证输入上统计相邻检查点之间有多少样本更换了专家集合，并记录部分训练步上的两种路由梯度夹角。这些是可观察行为，不是专家能力的标签。

实现用单线程 CPU、float32，逐专家计算完整 batch 后再选取输出；未实现稀疏执行、共享专家、Transformer、容量约束和丢弃。因而这里的相同预算指相同训练步数与样本数，不是严格相同硬件 FLOPs，更不是分布式系统效率比较。

6. 真正的曲线，没有给出一个包打天下的赢家

下面是三个种子的测试 MSE，写作“均值 ± 样本标准差”；MaxVio 是各次运行分别计算后再平均，不是先把三次负载混合再算。标准差不是置信区间，三个种子也不足以支持很强的显著性判断。

曲线每 50 步取真实验证值，不平滑；阴影表示三个种子的一个样本标准差。首先看均匀输入：小速度偏置把 MaxVio 明显压低，但测试 MSE 均值没有胜过无均衡。更均匀并不自动带来更好的预测。偏斜输入下，速度 0.001 的偏置却同时得到更低的负载偏差和更低的 MSE 均值。两种结果并不矛盾：负载控制是改变学习路径的手段，它的收益依赖具体任务与超参数。

误差棒同样是三个种子的样本标准差。不能把偏斜输入中的 0.2613 对比均匀输入中的 0.3275，然后断言数据偏斜有助于学习：两套测试分布不同，被赋予高权重的输入区域也不同，平均 MSE 不是同一道考题。

强辅助系数给出另一条值得停下来的曲线。均匀输入中，系数一既没有改善均衡，任务误差也更大；其验证平均 sigmoid 分数约为 0.079，而无均衡约为 0.402。这与前面“降低评分尺度也能降低原始辅助项”的分析相容，但还不足以证明全部因果。尤其本实现把选中分数重新归一化，改变了原论文的输出尺度耦合，因此不能拿这组结果宣布辅助损失普遍无效。

不能只看总损失是否下降来评价训练是否变好。辅助系数不同，总损失包含的目标就不同；把一组的“任务加辅助”与另一组的“仅任务”画在同一条质量轴上，会制造虚假的优劣。这里主图只放验证任务 MSE，辅助原值与乘系数后的值留在逐步记录里。偏置策略虽然也计算辅助原值用于诊断，却把其训练权重设为零，不让这个诊断量偷偷参与反传。

梯度记录提供了一个较具体、但仍有限的观察：均匀输入、辅助系数一的三次训练，在共 75 个记录位置中，任务梯度与原始辅助梯度有 32 次夹角余弦为负。也就是说，至少在这些位置存在局部方向冲突，并不是所有步都冲突。余弦只反映方向，不体现附加梯度有多大；所以日志同时保存两种梯度范数和乘上系数后的范数。即便两者方向相反，也不能凭一次夹角预测最终泛化误差，因为后续专家和路由都会继续改变。

堆叠柱显示三个种子的平均分配份额，只用来观察同编号专家的利用情况。无均衡在均匀输入上约为 [32.9%,28.8%,25.6%,12.6%]；偏置 0.001 则接近各占四分之一。然而不同种子中同编号专家未必学到同一种函数，这种平均不是跨运行的语义对齐，也不能消除个别运行的波动。

速度过大时，“路由变活跃”更不能被误读成能力增强。最后两个验证检查点之间，均匀输入下速度 0.1 平均有约 26.6% 的样本更换专家集合，速度 0.001 约为 2.2%；偏斜输入对应约 30.5% 和 3.6%。这是相隔 50 步的集合变化比例，不是单步换路由概率。它与大速度的负载振荡相互印证，但我们没有测专家专长，不能把换位解释成“形成了更丰富的分工”。

7. 边界在哪里，下一步应该追问什么

这个实验说明了一件比“哪种策略最好”更基础的事：任务质量、负载均衡和路由稳定性是不同轴，不能互相代替。 同样的负载分布可以由不同参数、不同输入分组产生；同样的任务误差也可能隐藏不同的专家利用程度。评价路由时，至少要同时保留这几条证据。

下一步最有价值的不是继续把速度网格加密，而是先拆开当前实现的耦合：保持数据与种子不变，对比是否归一化选中权重、sigmoid 与 softmax 评分、以及原始或归一化平均评分构造的辅助项。这能检查强辅助项的现象究竟有多少来自评分尺度自由度。它们是待做的消融，不是本文已经完成的证据。

然后才适合加入容量上限、真正的稀疏 dispatch、跨设备通信和计时，检验负载变化是否转化成真实系统收益；若要讨论分工，还需要按输入区域拆分误差、交换或屏蔽专家，检验专家贡献，而不只是看计数。更大的模型、更长训练和更多随机种子，也可能改写这里的排序。

回到最初的问题：均衡很重要，因为它关系到学习机会与系统约束；但均衡不是预测目标本身。辅助损失通过附加梯度引导路由，动态偏置通过历史负载改变下一次选择。理解两者的区别，不是给一种方法贴上“有干扰”、另一种贴上“零代价”，而是明确它们各自在何处、以什么时序改变了模型。

参考资料

1. Auxiliary-Loss-Free Load Balancing Strategy for Mixture-of-Experts（固定 v1）