专家容量：限制每个专家的最大 Token 数

FreeGuideOnline 最新 2026-06-22

什么是专家容量？

在混合专家（Mixture of Experts, MoE）模型中，Transformer 层内的前馈网络被替换为多个并行的“专家”子网络。每个输入 Token 通过一个门控网络被路由到最相关的一小部分专家进行运算。专家容量（Expert Capacity）正是用来限制每个专家能够处理的最大 Token 数量的硬性阈值。它的核心目标是在训练和推理时，防止某些专家被过多 Token 淹没，从而控制显存消耗并维持负载均衡。

可以这样理解：假设有 4 个专家，一个批次里总共 8 个 Token。如果没有任何限制，可能出现某个专家被分配了 6 个 Token，另一个专家只有 1 个 Token。专家容量就是为每个专家设定一个处理上限，超出上限的 Token 会被直接丢弃，以此来强制平衡。

为什么需要专家容量？

MoE 训练面临三个关键挑战，专家容量正是为了解决它们而引入的：

  1. 显存碎片与峰值控制 如果某些专家接收的 Token 数量远多于其他专家，计算时需要为这些大尺寸张量分配大量显存，导致显存使用不均匀、峰值过高。固定容量后，每个专家的输入张量维度可以预分配为 (capacity, d_model)，显存占用变得规整且可预测。

  2. 计算效率 现代加速器（GPU/TPU）擅长处理形状规则的矩阵运算。变长输入会导致频繁的分组、填充或循环，降低实际利用率。设定容量后，可以直接在不规则路由之上构建一个规整的批量矩阵乘法，大幅度提升计算吞吐。

  3. 负载均衡 门控网络如果没有约束，容易产生“赢家通吃”现象——少数专家主导训练，多数专家未被充分使用。专家容量作为一道硬门槛，强迫模型将超出容量的 Token 路由到次优专家（或直接丢弃），从而在训练初期就建立起较均匀的专家利用率。

专家容量如何工作？

容量系数与容量计算

专家容量并不是一个随意指定的数值，而是基于批次总 Token 数和专家数量动态计算得出。公式如下：

[ C = \text{round}\left( f \cdot \frac{T}{E} \right) ]

其中：

  • (T)：批次中的总 Token 数（通常为 batch_size × sequence_length）
  • (E)：专家总数
  • (f)：容量系数（capacity factor），是一个大于 1 的浮点数，例如 1.25 或 1.5
  • (C)：每个专家的容量，即最大允许处理的 Token 数

(T/E) 代表理想均衡状态下每个专家应处理的 Token 平均数。容量系数 (f) 允许略微超过平均值，以适应路由的随机性，避免过度丢弃。例如，批次有 1024 个 Token，8 个专家，(f = 1.25)，则专家容量 (C = 1.25 \times 128 = 160)。

Token丢弃与负载均衡损失

当某个专家被路由到的 Token 数超过其容量时，多余的 Token 必须被丢弃。丢弃有两种常见处理方式：

  • 直接跳过专家运算：Token 不经过该专家，其输出被替换为全零或残差连接直接传递。这种情况下，模型可能损失部分信息，因此容量系数不能设置得太小。
  • 重路由（re-routing）：将溢出 Token 路由到路由概率次高的下一个专家，前提是下一个专家还有空余容量。这个方式更温和，但增加了工程复杂度。

同时，为了从根本上减少丢弃，训练时会引入负载均衡损失（Load Balancing Loss）。该损失项通常基于每个专家接收的 Token 数目分布与均匀分布之间的差异来计算，鼓励门控网络输出更均匀的路由决策。典型的实现会分别惩罚“专家被选中的概率”和“实际指派给专家的 Token 数”两者的不平衡。

专家容量的实践示例

以下用伪代码展示一个简单的专家容量限制流程（基于 Top‑2 门控，即每个 Token 选择两个专家）：

# 假设形状:
# router_logits: (batch, seq_len, num_experts)
# token_embeddings: (batch, seq_len, d_model)

batch, seq, E = router_logits.shape
T = batch * seq

# 1. 选择每个 token 的 top‑2 专家
top2_probs, top2_indices = top_k(softmax(router_logits), k=2)

# 2. 计算专家容量
capacity_factor = 1.25
expert_capacity = int(capacity_factor * T / E)

# 3. 统计每个专家实际接收到的 token 数
expert_counts = count_assignments(top2_indices, num_experts=E)

# 4. 对于每个专家，仅保留前 expert_capacity 个 token，其余丢弃
for expert_id in range(E):
    assigned_tokens = get_tokens_for_expert(top2_indices, token_embeddings, expert_id)
    if len(assigned_tokens) > expert_capacity:
        # 这里简单地直接截断，也可采用更复杂策略
        assigned_tokens = assigned_tokens[:expert_capacity]
    expert_output = expert_networks[expert_id](assigned_tokens)
    # 将输出写回对应 token 位置...

在主流实现（如 Google 的 Switch Transformer 或 NLLB MoE）中，容量截断通常发生在计算路由分发张量的 top_k 过程之后，利用 scatter 或 gather 操作配合 truncate 完成。

常见问题与注意事项

  • 容量系数如何选取？ (f=1.25) 是常用起点。系数过小，丢弃率飙升，训练不稳定；系数过大，显存和计算开销接近无容量限制，失去均衡优势。实践中可以监控丢弃的 Token 比例，控制在 1%～5% 以内。

  • 推理时是否需要专家容量？ 推理时通常仍保留容量，但为了降低延迟，会使用更紧凑的容量系数，甚至在某些部署中设置为 (f=1.0) 并容忍极少丢弃，或者采用动态批处理合并推理请求以实现高利用率。

  • 容量与batch size的关系 容量公式中的 (T) 是整个批次的总 Token 数。当批次规模很大时，(T/E) 较大，即使容量系数略小也能获得充足的容量，因此大规模训练中负载均衡问题相对缓和。

  • 与辅助损失的配合 专家容量是“硬”约束，负载均衡损失是“软”约束。两者通常同时使用：损失函数引导路由分布趋向均匀，容量兜底防止极端不均衡对系统造成冲击。

  • 专家数量对容量的影响 专家数 (E) 增大时，平均负载 (T/E) 减小，若仍使用相同的容量系数，每个专家的容量绝对值也会减小。这要求门控网络必须更加精确地均衡分配，否则丢弃率会上升。因此，在极大专家规模（如 128、256 专家）的架构中，往往需要仔细调优容量系数或采用分层路由等变体。