一次在数千步后突然发散的训练
设想一个团队在训练多语言稀疏 MoE 模型：2048 块加速器上跑了几天，loss 平稳下降，专家利用率看起来也正常。某个检查点之后，loss 突然变成 NaN，梯度范数飙到无法恢复，重启后同样的位置再次发散。任务数据没有变，学习率没有调，唯一持续变化的量是路由器 logits 的幅度——它在几千步里缓慢爬升，直到 softmax 进入饱和区。
直觉上，路由器只要给出正确的专家排序就够了，logits 大一点小一点似乎无所谓。问题在于 softmax 对幅度极其敏感：logits 从 [2.0, 0.5, 0.1] 变成 [20.0, 5.0, 1.0]，路由决策几乎不变，但指数运算的数值条件已经完全不同。这一节要讲清楚的就是：为什么幅度膨胀会同时破坏梯度、破坏专家分配，以及 z-loss、logit 钳制和路由器精度分别在哪一层拦截这个失败。
路由器 logits 为什么会自己长大
MoE 的每一层里有一个路由器（router，也叫门控网络 gating network），它接收 token 的隐藏表示，为每个专家打一个原始分数，也就是 logit。这些 logit 经过 softmax 变成路由概率，再按 top-k 选出真正参与计算的专家。
训练过程中，让某个专家的 logit 变大，通常能降低该 token 的损失，因为被选中的专家获得了更多梯度、变得更专精。于是梯度下降会持续推高“当前看起来正确”的那个 logit。这个过程没有天然的上界：只要继续训练，logits 的绝对值就会缓慢漂移上升。Switch Transformer 的作者在设计稀疏模型时明确指出，MoE 的广泛采用受复杂度、通信成本和训练不稳定拖累，而他们提出的训练技巧正是为了压制这类不稳定。ST-MoE 则系统记录了这种失败模式：大规模 MoE 训练可以平稳进行数千步，然后突然以 NaN 损失和梯度爆炸崩溃，元凶始终是偏离稳定 softmax 所需范围的 logit 幅度。
这里有一个容易被忽略的区分：路由决策和 logit 幅度是两件事。路由器用 [2.0, 0.5, 0.1] 和用 [20.0, 5.0, 1.0] 可以做出同样的排序，但后者的指数计算会挤压浮点数的有效位。z-loss 正是利用这种解耦——它惩罚幅度，而不惩罚承载路由信息的相对差值。
softmax 饱和如何让梯度消失
softmax 把 logits 转成概率：
p_j = exp(z_j) / Σ_k exp(z_k)其中 z_j 是专家 j 的 logit，p_j 是它被选中的概率，N 是专家总数。工程上要关注的是这个式子的梯度行为。当某个 logit 远大于其他 logit 时，它的概率趋近 1，其余趋近 0，softmax 进入饱和区。在饱和区，概率对 logit 的导数趋近于零：路由器几乎不再从任务损失里收到有效的梯度信号，它“锁死”在当前分配上。
更麻烦的是反向传播经过指数运算时的数值放大。logits 幅度越大，exp 的中间值跨度越大，低精度下的舍入误差越容易主导结果。一旦某个 logit 因为舍入被错误放大，softmax 会进一步极化，形成正反馈。这正是 ST-MoE 描述的“平稳数千步后突然崩溃”的机制：幅度在早期只是缓慢漂移，直到越过某个阈值，数值误差与梯度爆炸同时爆发。
专家塌缩是这条链路的终点。当少数专家的 logit 持续领先，softmax 把绝大多数 token 都分给它们，其余专家长期收不到 token，参数得不到更新，逐渐退化成无用分支。此时模型名义上有 N 个专家，实际只有少数在参与计算，稀疏带来的容量优势被抵消。
z-loss 惩罚的是什么
z-loss 的思路是给训练目标加一个正则项，专门惩罚 logits 的尺度。它针对的量是 log-sum-exp：
L_z = (1/B) · Σ_b ( log Σ_k exp(z_k) )²这里 B 是当前 batch 中的 token 数，z_k 是第 k 个专家的 logit。log-sum-exp 是“最大 logit”的光滑近似，它衡量的是 logits 的整体量级，而不是它们之间的相对差异。把它平方后加进损失，就形成了一个持续把 logits 往小幅度方向拉的力。
关键在于这个惩罚项的性质。第一，它对相对排序不敏感：把所有 logits 同时减去一个常数，softmax 结果不变，但 log-sum-exp 会下降，所以 z-loss 会主动抵消这种“整体平移”式的膨胀。第二，它的梯度平滑，即使 logits 已经很大，log-sum-exp 依然可以稳定计算，不会像直接对 logits 平方那样在极端值上产生剧烈梯度。第三，它和负载均衡损失作用在不同维度：负载均衡管的是“token 是否均匀分给专家”，z-loss 管的是“logits 是否处在数值健康的范围”。两者不能互相替代。
ST-MoE 把 z-loss 作为稳定稀疏专家的核心手段之一，其结论是这类正则项让大规模稀疏模型可以稳定训练。需要说明的是，z-loss 的系数需要调：系数太小压不住漂移，系数太大则强行压低 logits，让路由概率趋于均匀，削弱专家的专精能力。
logit 钳制与路由器精度
除了加惩罚项，还有两种更直接的数值手段。
logit 钳制（clamping）是在 softmax 之前把 logits 截断到一个固定区间，例如超过上限就取上限。它简单、开销低，能立刻阻止单个 logit 无限增长。代价是它破坏了梯度：被截断的 logit 梯度为零，路由器在这些维度上不再学习，而且截断阈值是人为设定的，不同层、不同训练阶段的最优阈值并不一致。工程上通常把钳制当作兜底，而不是主要手段。
路由器精度是另一条线。Switch Transformer 的一个结论是，大规模稀疏模型首次可以用较低精度（bfloat16）格式训练。但路由器本身对数值误差特别敏感：softmax 的指数运算会放大舍入误差，而路由决策又是离散的 top-k，误差一旦改变排序，就会把 token 送到完全不同的专家。因此在混合精度训练中，工程上常见的做法是让路由器的 logits 计算和 softmax 在 float32 下进行，其余专家计算仍用 bfloat16。这样既保留了低精度训练的内存和吞吐收益，又避免了路由决策被舍入误差污染。
这三者的分工可以这样理解：z-loss 从训练目标层面持续压制幅度漂移，是预防；logit 钳制是运行时兜底，防止单点爆炸；float32 路由器精度是保证 softmax 和 top-k 在给定幅度下仍能正确计算。缺任何一环，另外两环的压力都会变大。
贯穿场景：一次路由决策的完整流动
回到多语言训练服务这个场景。一个 batch 里混着不同语言的 token，路由器要为每个 token 选专家。下面这张图展示单个 token 从隐藏表示到专家输出的路径，以及三个稳定手段分别在哪一步介入。
flowchart TD
A[token 隐藏表示 h] --> B[路由器线性层 W_r h]
B --> C[float32 计算 logits]
C --> D{logit 钳制}
D --> E[softmax 得概率]
E --> F[top-k 选专家]
F --> G[专家计算]
G --> H[加权合并输出]
C --> I[z-loss 惩罚 log-sum-exp]
I --> J[加入总损失]
J --> B
E --> K[负载均衡损失]
K --> J
流程里有几个真实的状态转折。第一，logits 在 float32 下计算，保证 softmax 的指数运算不被低精度舍入干扰。第二，钳制发生在 softmax 之前，它改变的是进入 softmax 的数值范围，而不是路由排序本身——只要截断不改变相对顺序，路由决策就不变。第三，z-loss 的梯度回传到路由器权重，形成对幅度漂移的持续抑制；负载均衡损失的梯度则作用在概率分布上，纠正 token 分配的偏斜。两条损失路径最终都汇入总损失，但作用于不同对象。
如果去掉 z-loss，只保留负载均衡损失，会出现一种典型现象：专家利用率在训练早期看起来均匀，因为负载均衡确实在纠正分配偏斜；但 logits 幅度仍在缓慢上升，softmax 逐渐饱和。等到饱和严重时，负载均衡损失能提供的梯度也被压缩，它想纠正分配却推不动已经极化的概率。这时 loss 可能突然发散，或者专家利用率在几个检查点内快速向少数专家集中。
与仅用负载均衡损失的方案对比
负载均衡损失解决的是“专家塌缩”的分配症状，z-loss 解决的是导致塌缩的数值病因。两者对比可以看下面这张表。
| 维度 | 仅用负载均衡损失 | 加入 z-loss 与路由器精度控制 |
|---|---|---|
| 作用对象 | 路由概率的分配均匀度 | logits 的整体量级 |
| 对 logits 膨胀 | 不直接约束，幅度可继续漂移 | 持续施加向小幅度收缩的力 |
| softmax 饱和 | 饱和后梯度被压缩，纠正能力下降 | 在饱和前抑制幅度，维持梯度通道 |
| 专家利用率 | 早期有效，后期可能快速失衡 | 更长时间保持多专家参与 |
| 数值稳定性 | 不解决 NaN 与梯度爆炸 | 直接针对发散诱因 |
| 主要代价 | 引入路由偏差，可能牺牲专精 | 需调系数；过大时路由趋于均匀 |
| 实现复杂度 | 低，只需统计 token 计数 | 中，需改损失与路由器计算精度 |
从工程决策看，负载均衡损失是必需的，但它不是训练稳定性的充分条件。ST-MoE 把 z-loss 作为稳定稀疏专家的关键组件，说明在大规模训练里，数值正则和分配正则是两条独立的防线。
失败模式、可观测信号与适用边界
实践中需要盯住几个信号。路由 logits 的绝对最大值和 log-sum-exp 的均值是最直接的指标，它们应该在训练中保持平稳，如果持续单调上升就是漂移的早期信号。专家利用率分布要看长尾：如果 top-1 专家的 token 占比持续上升，而尾部专家长期接近零，塌缩已经开始。梯度范数在发散前往往先出现异常尖峰，可以配合梯度裁剪一起观察。
z-loss 也有失效边界。如果系数设得过大，logits 被压得过小，softmax 输出接近均匀，top-k 的选择变得接近随机，专家失去专精能力，模型退化成“稀疏但无效”。如果训练数据本身高度同质，少数专家确实应该承担大部分 token，此时强行压低 logits 可能损害质量——z-loss 的目标是数值健康，不是强制均匀。此外，z-loss 不能修复已经发生的塌缩：如果某些专家已经长期无梯度、参数退化，加 z-loss 只能防止进一步恶化，无法让它们重新参与。
还有一个前提容易被忽略：z-loss 和路由器精度控制都假设路由器权重本身没有病态。如果初始化不当或学习率过高，路由器可能在第一步就产生极端 logits，此时正则项需要时间才能起作用。因此工程上通常把路由器初始化、学习率预热、梯度裁剪和 z-loss 一起使用，而不是指望单一手段。
尚未完全解决的问题是系数选择。z-loss 系数与模型规模、专家数量、数据分布都有关，目前主要靠小规模实验外推，缺少能直接迁移到万亿参数规模的理论指导。这也是为什么大规模 MoE 训练仍然需要密切监控路由指标，而不能把稳定性完全交给默认配置。