处理丢弃令牌

稀疏专家混合模型的架构，特别是那些采用 Top-K 路由的模型，通常为每个专家分配一个固定的处理容量，记作 $C$。这个容量代表了专家在给定计算步骤（通常是分布式训练中的一个微批次）内可以处理的最大令牌数量。这种固定容量对于在设备之间维持可预测的计算负载和内存使用是必不可少的，尤其是在专家跨不同处理器并行化的大规模分布式设置中。

然而，由门控网络驱动的路由机制并不固有地保证分配给任何单个专家的令牌数量会严格遵守这个容量 $C$。门控网络根据习得的亲和性选择每个令牌的“最佳”专家。由于输入数据的变化以及路由器学习机制的动态性，某些专家被分配的令牌数量超过其指定容量是很常见的，特别是当负载均衡未能完美实现时。

分配给已达到其容量 $C$ 的专家的令牌被称为“丢弃令牌”。

丢弃令牌的后果

当令牌因容量限制被其分配的专家丢弃时，它有效地绕过了 MoE 层旨在为其提供的专业计算。处理这些丢弃令牌的最常用方法很简单：它们的输入表示未经改变地通过 MoE 层。

考虑一个执行残差更新的 MoE 层：

$$y = x + \text{MoE}(x)$$

对于成功路由并由专家 $j$ 处理的令牌 $x_i$，其输出为：

$$y_i = x_i + g_i \cdot \text{专家}_j(x_i)$$

其中 $g_i$ 是门控权重 (weight)。

然而，如果令牌 $x_k$ 被分配给专家 $j$，但专家 $j$ 已达到其容量 $C$，则令牌 $x_k$ 被丢弃。在这种标准的直通方法中，其输出计算实际上变为：

$$y_k = x_k + 0$$

MoE 的贡献为零，因为该令牌从未进入该层的专家计算流程。

这会带来几个负面影响：

1. 专业化能力受损： 令牌错过了由专业专家处理的潜在益处，可能降低模型质量。对于经常导致丢弃令牌的输入，模型学习效果不佳。
2. 学习信号减弱： 由于丢弃的令牌没有通过专家，它不会提供梯度信号来更新分配给它的专家权重。这会阻碍专家的专业化过程。它还向路由器提供了关于其对该令牌路由决策质量的较少信息反馈。
3. 负载不均衡的迹象： 丢弃令牌的百分比很高是负载均衡不佳的直接迹象。这表明前面讨论过的辅助损失项 ($L_{aux}$) 未能充分惩罚不均衡分配，或者路由器难以有效地分配令牌。

容量管理与丢弃处理

处理丢弃令牌主要涉及预防过度丢弃的策略，而不是事后处理它们的复杂机制。

设置专家容量 ($C$)

专家容量 $C$ 是一个重要的超参数 (parameter) (hyperparameter)。它需要足够大以容纳路由过程中固有的合理程度的不均衡性，但又不能太大，以至于在专家未充分利用时导致过度计算浪费（填充）。

容量通常根据每个专家的理想均匀负载来设置，并乘以一个容量因子。如果 $N$ 是正在处理的令牌数量（例如，在专家并行化中每个设备处理的令牌数量），而 $E$ 是本地可用的专家数量，则理想负载为 $N/E$。容量随后设置为：

$$C = \lceil \text{容量因子} \times \frac{N}{E} \rceil$$

容量因子的常见取值范围是 1.0 到 2.0。

• 接近 1.0 的因子（例如 1.1、1.25）将计算开销降至最低，但如果负载稍有不均衡，则会增加丢弃令牌的风险。
• 较大的因子（例如 1.5、2.0）降低了丢弃令牌的可能性，但可能导致分配的计算量超出严格所需，因为专家通常处理的令牌数可能少于 $C$。

选择涉及计算效率与最小化因丢弃令牌造成的信息损失之间的权衡。它通常需要根据训练初期观察到的丢弃率进行调整。

辅助负载均衡损失 ($L_{aux}$) 的作用

如前一部分所述，辅助损失项 $L_{aux}$ 是鼓励路由器均匀分配令牌的主要机制，从而减少导致令牌丢弃的情况。虽然 $L_{aux}$ 在令牌一旦被丢弃后不会直接处理它，但有效调整其在 $L_{total} = L_{task} + \alpha L_{aux}$ 中的系数 $\alpha$ 对于保持令牌丢弃率较低是必要的。较高的 $\alpha$ 通常会促使路由器进行更均衡的分配，从而减少丢弃，但如果设置过高，则可能干扰最佳专业化。

丢弃机制的可视化

下图说明了令牌被路由、容量检查以及某些令牌被丢弃的过程。

令牌通过专家容量有限的 MoE 层的流程。令牌被路由到专家。如果一个专家的容量（此处 C=2）已满，则后续分配给它的令牌会被丢弃，并且通常通过直通机制绕过专家计算。

监控丢弃令牌

训练期间，有必要监控每个 MoE 层或整个模型平均的丢弃令牌百分比。这项指标可作为一项重要的健康检查：

• 持续的高丢弃率（例如，> 1-5%） 通常表明专家容量不足（容量因子过低）、负载均衡效果不佳（$L_{aux}$ 系数 $\alpha$ 需要调整或路由器学习效果不佳），或者数据分布可能存在问题。
• 丢弃率突然飙升 可能预示着训练不稳定或数据特征发生变化，路由器难以适应。

记录此指标有助于诊断训练问题并调整相关超参数 (parameter) (hyperparameter)（容量因子，$\alpha$）。

替代策略（不常见）

虽然直通是标准方法，但研究已考量了其他处理容量溢出的方法，例如：

• 路由到次选专家： 如果使用 Top-1 路由，丢弃的令牌如果次选专家有可用容量，理论上可以发送给它。但这增加了路由的复杂度。
• 溢出专家： 指定特定专家来处理溢出令牌。这需要细致的设计来保证这些溢出专家得到有效使用。

由于增加了路由逻辑和通信模式的复杂度，这些方法在标准大规模 MoE 实现中不被广泛采用。重点仍在于通过容量管理和负载均衡损失来防止丢弃。

总之，在 MoE 训练中处理丢弃令牌，与其说是主动处理它们，不如说是更多地实施预防措施。设置适当的专家容量并有效地调整负载均衡辅助损失是减少令牌丢弃并确保大多数令牌从专业专家计算中获益的主要策略。监控丢弃率对于保持训练稳定性和模型性能很重要。