稀疏概率映射在训练后如何影响专家混合路由？

精确零为何没有如预期保留

专家混合路由器为各专家打分，概率映射把分数转成权重，top-K 算子保留权重最高的 K 个专家，并用重新归一化的权重组合它们的输出。使用 softmax 时，每个 token 恰好使用 K 个专家，未被选中的概率质量直接被截断丢弃。sparsemax、alpha-entmax 和 normmax 等诱导稀疏的映射可以赋予精确的零，因此被选中的专家可能获得零权重并退出混合，即使预算 K 固定，实际贡献的专家数也会随数据变化。

论文追问的是：路由器训练完成后，这种能力是否真的带来随 token 变化的参与度。复杂之处在于路由器分数是学出来的：改变映射也会改变映射接收到的分数，因此训练后的 sparsemax 路由器和 softmax 路由器不仅在如何变换分数上不同，在学到产生哪些分数上也不同。一个能够输出零的映射，训练后可能很少真正输出零。

三个容易混淆的度量

研究区分了三个常被当作一回事的量。保留量是被选中专家保留的概率质量；一个 token 的丢弃质量是未被选中专家的质量，由于重新归一化恰好把这部分质量移到被选中专家上，它也等于路由分布与其选中混合之间的总变差距离。参与度是权重非零的被选中专家数，记作 K+，它至多为 K，在 top-2 选择下只要某个被选中专家权重为零就等于 1。预测质量则是验证损失和下游准确率。

论文还区分了两种专家负载直方图：名义负载统计每一次选择，包括填充项；贡献者负载只统计权重为正的选择，每次分配计 1/K。由于去掉填充项也会改变均值，两种口径可能给出不同的路由器排序，因此只有在说明口径之后，关于负载均衡的说法才有意义。

匹配模型隔离出概率映射

作者在完全相同的架构、数据、随机种子和均衡目标下，训练了匹配的 300M 和 1B top-2 MoE 语言模型，分别使用 softmax、1.5-entmax、sparsemax 和 2-normmax，因此概率映射是唯一变化的组件。训练使用 DCLM-baseline 和 Meta-Llama-3 分词器，名义预算为 10B token、共 19,073 次更新，并使用系数 0.01 的 Switch 式均衡损失，在四个 MoE 层上取平均；没有路由器 z-loss。在固定 alpha 和 gamma 时，四种映射不增加参数。

路由诊断记录每个训练后模型的路由器分数，并测量其离散程度和 top-2 分数差相对于各映射稀疏化阈值的位置。检查点诊断在 FP32 下对 61 篇 Paloma 测试文档运行每个模型，每层得到 285,107 个目标，每个模型得到 1,140,428 个目标-层事件。目标和层是来自每个映射单个检查点的相关观测，而不是训练重复。

每个映射在固定分数差处丢弃专家

使比较可行的机制是：每个映射仅根据分数就有一个固定的稀疏化阈值。当 top-2 分数差达到 2 时，1.5-entmax 会丢弃一个被选中的专家；达到 1 时，sparsemax 和 normmax-2 会丢弃。因此，训练后模型的单专家率不过是其学到的分数差中达到该映射阈值的比例，这也构成对记录分数与记录路由之间一致性的检验。

由于阈值只是对分数的条件，路由器既可以通过重塑标准化后的分数差，也可以通过让分数尺度相对阈值移动，来改变丢弃专家的频率。把分数乘以一个正因子会放大每个分数差，却不改变分数差分布的形状，于是单专家率变成缩放后分数差达到阈值的概率。普通反向传播中没有任何机制固定这个尺度，因此论文把它当作一个学到的量。

训练后的路由器学到不同分数尺度

在 1B 规模下，三个稀疏路由器学到的分数离散程度都明显小于 softmax 路由器，中心化分数 RMS 在 0.32 到 0.38 之间，而 softmax 为 0.86。对 1.5-entmax 而言，后果是决定性的：其第 99 百分位分数差为 1.45，远低于阈值 2，因此训练后的路由器只在 0.05% 的状态下产生单专家，而 softmax 路由器有 7.8% 的分数差达到该阈值。于是 1.5-entmax 路由器丢弃的概率质量比 softmax 少 30%，却几乎从不丢弃被选中的专家。

sparsemax 和 normmax-2 共享阈值 1，因此在任何给定分数向量上，它们对某个 token 保留一个还是两个贡献者总是判断一致，但训练后的模型每个 token 相差 0.14 个贡献者。这一差异只能来自两个路由器学到的分数：normmax-2 路由器的中位分数差为 0.61，而 sparsemax 为 0.40，且其分数差达到 1 的频率是后者的三倍。normmax-2 把 21% 的 token 路由到单个专家，而 sparsemax 在三个稀疏映射中保留的质量最多。

参与度随深度变化而非均匀

层平均掩盖了强烈的深度依赖。normmax-2 在第一个 MoE 层把 54% 的 token 路由到单个专家，但在最后一层只有 7%；sparsemax 呈现同样的模式，幅度更小，从 19% 降到 4%；1.5-entmax 在任何位置都基本没有单专家。早期层学到最稀疏的路由器，这说明有用的参与程度本身依赖深度，也提供了另一个不应随意设定 K 的理由。

论文还检验了均衡目标是否能解释这些差异。它训练了 300M 的 1.5-entmax 和 sparsemax 路由器，仅在辅助损失中使用 softmax，而路由和专家组合仍使用各自的稀疏映射。两次运行的负载 CV 都上升，意味着流量在专家间的分布更不均匀，同时单专家比例下降。因此，逐 token 参与度和负载均衡朝相反方向移动，这并不支持那种简单解释，即 sparsemax 的不均衡完全来自零概率专家得不到直接的均衡梯度。

稀疏映射以质量换取推理稳健性

没有任何稀疏映射在验证损失上优于 softmax，论文也指出，在每种配置只有一个训练种子的情况下，映射之间很小的质量差异（验证损失至多 0.012 nats）不足以确立排序。稀疏映射真正改变的是对推理时预算的敏感度。当以 K=2 训练的模型在冻结权重下以 K=8 运行时，softmax 在 1B 规模下损失 0.58 nats，因为它被迫把权重分散到六个它从未学会组合的专家上；而 sparsemax 只损失 0.02 nats，因为它继续把大多数额外槽位的权重设为零，在 K=8 时平均使用 4.3 个专家。normmax-2 和 1.5-entmax 介于两者之间，分别为 0.064 和 0.250 nats。

这种稳健性是单向的。把预算降到 K=1 会让每个映射损失约 0.10 nats，而且每个映射的损失都在训练预算处最低，因此选择更多专家以保留更多质量也没有帮助。保留量、参与度和质量是三个不同的量，论文自身的表述是：自适应 MoE 路由必须围绕概率映射与所学分数的联合行为来设计，而不是只围绕映射本身。

稀疏路由器中真正该控制什么

这些结果引出的实际优先事项是分数尺度，而不是映射的选择。由于路由器会让尺度与映射共同适应，这个尺度目前被交给优化器，论文认为更好的做法是显式地学习它，正如稀疏注意力中已经做的那样。一个尺度可控的路由器随后可以被信任去自行选择预算：当 K 提高时，稀疏映射已经把大多数额外槽位的权重设为零，因此选择支持集中的每个专家，而不是固定的 top-K，并用至多 k 算子加以限制，距离训练后模型已有的行为只有一小步。

值得持续关注的比较是 sparsemax 与 normmax-2：它们共享阈值 1，因此在任何单个分数向量上做出完全相同的判断，最终却每个 token 相差 0.14 个贡献者，纯粹因为两个路由器学到了不同的分数差分布。这个差距，而不是映射名义上输出零的能力，才是设计者真正能够影响的量。论文的限制条件界定了这一结论的适用范围：贡献者计数衡量的是专家参与度而非计算量，因此零权重专家是否被跳过、能节省多少，取决于论文未测量的调度实现；而分数尺度分析是描述性的，并未识别出产生更小尺度的训练机制。

本文引用的文献