（2025|JSAI|东京大学，连续/无限专家，基于 mask 的 FFN 调制）∞-MoE：将混合专家模型推广到无限专家模型

本文提出 ∞-MoE，它将 MoE 从有限的专家集合扩展到连续（实际上无限）的专家空间。它仅为每个 token 激活少数几个采样专家，保持了类似 MoE 的效率，同时提高了准确性。在 GPT-2 Small 和 Medium 上，∞-MoE 优于 Switch Transformer 和标准 MoE。

EDPJ

546人浏览 · 2026-01-28 22:35:29

EDPJ · 2026-01-28 22:35:29 发布

∞-MoE: Generalizing Mixture of Experts to Infinite Experts

论文地址：https://arxiv.org/abs/2601.17680

进 Q 学术交流群：922230617 或加 CV_EDPJ 进 W 交流群

6.1 扩展到 GPT-2 Medium 之外

1. 引言

大型语言模型（LLMs）最近在自然语言处理（NLP）任务（如机器翻译、问答、代码生成）中取得了显著性能。这些进步主要受模型参数量、训练数据和计算资源的规模扩展所驱动。然而，单纯增加模型规模会导致巨大的计算和内存开销，这推动了对更高效扩展策略的研究。

专家混合模型（Mixture of Experts，MoE）因其在保持相对较低的单 token 计算成本的同时扩展参数数量的能力而脱颖而出。通过将每个输入路由到专门的专家子集，基于 MoE 的架构可以高效地稀疏存储大量知识。最近的大规模模型如 DeepSeek、Mistral 和 Phi 已成功采用 MoE 设计，证明稀疏路由能在不产生过高计算开销的情况下显著提高性能。

近期 MoE 研究的一个显著趋势是 积极增加专家数量以实现细粒度专业化。经验证据表明，更大的专家池能提高整体容量，并通常能在相似或更低的计算成本下获得更高的准确性。例如，PEER 能处理数百万个专家，最近的理论工作也证实 MoE 性能可随专家数量进行可预测的扩展。

遵循此趋势，一个自然的问题出现：能否通过进一步将专家数量增加到无限来获得更好的性能？原则上，拥有更多专家应能实现更专业的表示，潜在地提升跨不同任务的泛化能力。

本文引入了 ∞-MoE，它将专家从离散集合扩展到连续域，理论上允许无界的专家容量。在此框架中，每个输入从连续的专家空间中采样，将 “增加专家” 的概念推向了极致。

尽管可能有无穷多的专家，本文提出的 ∞-MoE 在计算上仍然是可处理的，因为 它在任何给定时间仅稀疏激活少数几个采样专家。这种设计保留了稀疏路由的效率，同时显著增强了模型容量。

通过在 GPT-2 Small 和 Medium 上的实验，本文观察到基于 GPT-2 Small 的 ∞-MoE 变体（129M活跃参数，186M 总参数）实现了与具有 350M 参数的密集 GPT-2 Medium 模型相当的性能。

此外，在推理时增加样本数量能带来额外的准确率增益，而减少样本数量仍能保持比标准 MoE 高 2.5% 的准确率，实现了速度与准确性之间的灵活权衡。

3. 方法

3.1 专家混合模型（MoE）

设 Z = {1,2,…,n} 为 n 个专家的离散索引集。令 x ∈ R^{d_in} 表示输入。每个专家是一个函数：

其中，i ∈ Z，是专家索引。路由器（router）在专家上产生概率分布 p(i∣x)。

MoE 输出是期望的专家输出：

与标准 MoE 的联系。标准 MoE 可被视为一个特例，其中通用专家函数 f(x,i) 简单地从一组 n 个预定义专家函数 {e_1(x), …, e_n(x)} 中选择第 i 个专家；即 f(x,i) = e_i(x)。路由器通常使用 softmax 函数计算选择专家 i 的概率：

其中 g(x) ∈ R_n 是路由器网络产生的分数向量。使用 top-k 操作选择专家的子集 K，最终输出为：

这一结果清楚地表明，标准 MoE 是这种离散公式的一个特例。

3.2 ∞-MoE：无限专家

∞-MoE 将离散 MoE 扩展到连续的、可能不可数无限的专家空间 Z ⊆ R^{d_z}。路由器在 Z 上定义概率密度 p(z∣x)。模型输出为：

通过蒙特卡洛采样来近似这个积分：采样 z ∼ p(z∣x) 并使用 f(x,z) 作为 y 的估计量。

路由器设计。对路由器使用高斯密度：

其中一个小型神经网络根据 x 预测 μ(x) 和 Σ(x)（即所有非对角线条目为零）。

在训练时，采样 z^(k) ∼ p(z∣x) K 次 (k = 1,…,K)，使路由器学会将概率质量分配到 Z 的适当区域。

专家设计。

将 z 视为从路由器采样的连续专家索引。直观上，每个不同的 z 值对应于无限专家空间中的一个不同专家。
前馈网络（feed-forward network，FFN）通过用于掩蔽（mask）中间层中某些神经元的掩码进行调制，允许模型动态选择哪些参数子集是活跃的。

形式上，令 W_z ∈ R^{d_ff × d_z}。给定从上式采样的 z，在中间神经元

上应用 top-N% 操作符，仅保留最大的 N% 节点并将其余设为 0：

因为保留的条目保持了它们的原始值，所以生成的掩码对于选定的位置是部分 "软" 的，而所有其他位置则严格为零。给定此掩码，专家输出 f(x,z) 计算为：

其中 Act(⋅) 是非线性激活函数，⊙ 是逐元素乘法，W_1 和 W_2 是可学习的权重矩阵。

通过此机制，每个采样的 z 有效激活 FFN 神经元的一个独特子集，反映了传统 MoE 模型中的稀疏性，但推广到了无限专家。

4. 实验

本文使用 GPT-2 Small（约 124M 参数）和 GPT-2 Medium（约 350M 参数）在广泛的自然语言理解任务上评估了 ∞-MoE 的有效性。

4.1 设置

数据。本文在名为 FineWeb 的大规模网络语料库上对模型进行预训练，从中提取了 10B 个token。对于微调或直接评估，本文在标准 NLP 基准测试上使用了零样本设置。

对比方法。本文比较了四种架构：

Dense (FFN)：所有输入共享单个 FFN 层的标准 Transformer。
Switch Transformer (Top-1)：将每个 token 路由到恰好一个专家。
MoE (Top-2)：经典的稀疏 MoE 设置，为每个 token 激活 top-2 专家。在此配置中，专家总数固定为 4。
∞-MoE：本文提出的具有无限专家空间的方法。在训练和测试期间，抽取两个样本（即 K=2）；仅使用一个样本时，只有 25% 的整体专家空间是活跃的。

4.2 主要结果

上表展示了每种模型在 GPT-2 Small 和 GPT-2 Medium 上的零样本（Zero-shot）性能。

在所有任务中，∞-MoE 始终优于 Dense 基线、Switch Transformer 和标准 MoE。
值得注意的是，对于 GPT-2 Small，∞-MoE 获得了最高的平均分 39.8，在 GPT-2 Medium 变体上也观察到了类似的改进，∞-MoE 获得了最佳平均分 41.3。

5. 消融研究

5.1 随采样数 K 的扩展

上图通过改变 K，在多个任务上比较了 ∞-MoE 与标准 MoE。

在传统设置中，增加 K 可以提高准确性，但在高值时也可能引入不稳定性。
相比之下，∞-MoE 随 KK 的变化更为平滑，能产生稳健的增益，并在较低的 K 下（相较于标准 MoE 实现了 2.5% 的改进）保持强劲性能。此外，将专家视为连续空间实现了灵活的推理，允许用户根据硬件限制或延迟需求调整 K。

这些结果表明，∞-MoE 结合了无界专家集合的表达能力和稀疏 MoE 的效率，使其非常适合各种运行时条件。

5.2 随数据集大小的扩展

为了评估本文提出的 ∞-MoE 方法在增加数据集大小下的有效性，使用 GPT-2 Small 架构作为基础模型进行了实验。

本文在 HellaSwag 数据集上测量准确率，将训练数据规模逐步增加 10B token，直至 100B。
结果如上图所示，∞-MoE 表现出一致性地性能提示。

5.3 连续索引 z 的维度

本文改变了连续专家索引 z 的维度，并在上表中报告了困惑度。

增加维度将性能提升到 d_z = 128，但更大的维度没有带来进一步的增益。
这一结果表明，中等大小的索引维度足以表示多样化的专家。

5.4 路由稳定性

本文使用 归一化路由熵（normalized routing entropy）来测量路由稳定性。

对于每个小 batch，本文采样

并从 top-N% 掩码中获得活跃神经元集合

定义 FFN 神经元上的经验性选择分布（empirical selection distribution）为：

并计算：

更高的 ~H(q) 表示更均匀的利用率。如下表所示，~H(q) 在训练期间迅速稳定在一个较高的值。

5.5 延迟剖析：自定义内核

本文测量了端到端 FFN 前向延迟（毫秒），同时改变了激活率（由掩码保留的 FFN 隐单元的百分比）。所有运行在 warm-up 后使用相同的 batch 大小和序列长度在单个 GPU 上进行。

本文将 PyTorch 参考实现（Ours）与一个融合的 CUDA 实现（Ours w/ custom kernel）进行了比较，后者在一次传递中应用掩码并进行约简。

结果如下图所示，自定义内核始终具有更低的延迟，并且大约有 1.25 倍的加速；对于自定义内核，延迟在激活率范围内几乎保持平坦。

6. 局限性与未来工作

尽管 ∞-MoE 为将 MoE 模型扩展到无限专家空间提供了一个有前景的框架，但仍有几个开放挑战：

6.1 扩展到 GPT-2 Medium 之外

尽管本文的实验集中于 GPT-2 Small 和 Medium，但 ∞-MoE 扩展到更大模型（例如 GPT-3 及更高）时的行为尚未完全明了。特别是，以下情况下性能和效率将如何变化尚不清楚：

在保持活跃（每个 token）参数数量固定的情况下增加总参数数量。
同时扩展活跃参数和总参数数量。

这些场景引出了关于在极端规模下训练和推理中潜在瓶颈与权衡的问题。

6.2 路由器分布

本文当前的实现为了简便采用了单峰高斯路由器。然而，更丰富的分布，如高斯混合或非参数密度估计器，可能提供更具表达力的专家分配，尤其是在高维专家空间中。

尽管这种方法可能改善对多样化输入模式的覆盖，但设计高效的采样和稀疏推理机制变得更加复杂，并且训练中的方差减少仍是一个开放挑战。

6.3 在其他领域的适用性

尽管本文的研究强调了 ∞-MoE 在语言建模中的实用性，但尚不清楚此框架如何容易地推广到其他领域，如视觉（例如 ViT）或多模态视觉-语言模型（VLM）。

实际关注点包括：使连续专家索引适应处理不同的输入模态、确保高分辨率数据的稀疏高效路由，以及在 NLP 之外的任务中保持有竞争力的准确性。

6.4 训练 FLOPs 与可扩展性

从计算角度看，∞-MoE 保持了 MoE 的特性，即每个 token 的 FFN FLOPs 随激活的参数比例缩放。理论上，对于 K 个样本和 FFN 激活比 r（top-N% 掩码），主要的 FFN 计算量约为密集 FFN 的 O(K⋅r)，路由开销很小。然而，较低的 FLOPs 在实践中并不总是转化为更快的训练。即使使用本文融合的 CUDA 内核，端到端吞吐量仍可能比标准 MoE 系统慢，因为掩码取决于输入token。由于活跃的隐藏单元集合在 token 和样本之间变化，很难将计算批处理成固定形状的 GEMM（通用矩阵乘法），也很难重复使用单个高度优化的内核。因此，简单的实现通常依赖于额外的索引和累加步骤，这使得内核融合和向量化执行更加困难，并降低了实际的硬件利用率。

在本文的实验中，本文优先考虑吞吐量而非严格最小化 FLOPs，并使用融合了掩码操作的矩阵乘法管道来实现 FFN 计算。这种设计允许使用跨样本共享权重的固定形状 GEMM，从而能够使用高度优化的内核（如 Tensor Core）并提高实际时钟效率，代价是与理想的稀疏执行相比，执行了额外的算术运算。要缩小理论计算节省与实际速度之间的剩余差距，可能需要进一步的系统级优化，以及对动态稀疏性更好的编译器/运行时或硬件支持。

7. 注：一些文中未提到的问题

在阅读论文的过程中，遇到一些比较困惑的地方，在此做出整理。

7.1 训练

论文中虽然给出了核心方法，但是对于完整的训练过程描述的并不是很清晰，在此做出整理。

∞-MoE 的训练过程其实可以拆成 路由器怎么学 + 专家（FFN）怎么学 两条线，可以把它理解成：用连续随机变量 z 做 “软专家索引”，再用可微的稀疏掩码去学一个无限 MoE。

参数总览（哪些是可学习的）：

1）Router（路由器）参数

预测均值的网络参数：θ_μ（输出 μ(x)）
预测方差的网络参数：θ_Σ（输出 Σ(x)，通常是对角、log-variance）

2）Expert / FFN 相关参数

FFN 第一层权重：W₁ ∈ ℝ^{d_ff × d_in}
FFN 第二层权重：W₂ ∈ ℝ^{d_out × d_ff}

3）Mask 生成参数

掩码投影矩阵：W_z ∈ ℝ^{d_ff × d_z}

所以需要训练的参数集合是：θ = { θ_μ, θ_Σ, W_z, W₁, W₂ }

z 本身不是参数，是每个样本动态采样的随机变量。

下面是 一次训练 step 中真实发生的事情。

Step 1：Router 前向 —— 学一个 “连续专家分布”

给定输入 x（比如一个 token 的 hidden state）：

Router 小网络前向：

μ(x) = Router_μ(x)
Σ(x) = Router_Σ(x)（通常输出 log σ² 再 exponentiate）

定义连续专家分布：p(z | x) = 𝒩(z | μ(x), Σ(x))

可以理解为：Router 不是选 top-k 专家，而是在 ℝ^{d_z} 的连续空间里圈一片区域，告诉模型 这些 z 对当前 x 比较合适。

Step 2：从无限专家空间中采样

训练时做 Monte Carlo 近似：

对同一个 x，采样 K 次：z^(k) ~ p(z | x), k = 1,…,K

这里一般会用 reparameterization trick：z = μ(x) + Σ(x)^{1/2} ⊙ ε, ε ~ 𝒩(0, I)

这样梯度才能从 loss 反传到 μ、Σ。

Step 3：z → mask —— 连续专家索引变成稀疏激活模式

这是 ∞-MoE 最核心的一步。

线性投影：

Top-N% 操作：

选出 ^m 中最大的 N% 维度
其余全部置 0

mask 是：

对被选中的维度：软值（可导）
对未选中的维度：硬 0（产生稀疏性）

关键理解：

z 决定激活 FFN 的哪一组神经元

不同 z → 不同 mask → 不同“专家”

专家是无限多的，因为 z 是连续的

Step 4：Expert（FFN）前向计算

对每个采样到的 z^(k)，计算一个专家输出：

可以理解为：FFN 的中间层被 z 动态“剪裁”了一次。

Step 5：Monte Carlo 近似积分

理论输出是：

训练时用 MC 估计：

然后拿 ^y 去算任务 loss（例如语言模型的 cross-entropy）。

Step 6：反向传播（谁学到了什么）

1️⃣ W₁、W₂（FFN 权重）

和普通 Transformer FFN 一样
只不过梯度只流经 被 mask 激活的神经元
自然形成稀疏、专家化的参数更新

2️⃣ W_z（mask 投影矩阵）

决定 z 映射到哪一组 FFN 神经元
梯度来源：loss → f(x,z) → mask(z) → W_z

W_z 学到的是：“不同 z 应该激活哪些神经元组合”

3️⃣ Router 参数 θ_μ、θ_Σ

梯度通过 z（reparameterization）反传
如果某些 z 产生的 mask 让 loss 更小：
- μ(x) 会向这些 z 移动
- Σ(x) 会缩小或扩大，控制探索范围

Router 学到的是：“给定 x，应该把概率质量放在哪片专家空间区域”

Step 7：训练动态的整体直觉

训练初期：

Router 分布很宽（Σ 大）
z 随机性强，像在随机激活子网络

训练中期：

W_z + FFN 学到“哪些神经元组合有用”
Router 开始把 μ(x) 推向这些有用区域

训练后期：

对不同 x，μ(x) 落在不同区域
相当于形成了 输入条件化的专家选择
但没有离散专家 ID，也没有负载均衡 loss

【至少 ∞-MoE 论文中未给出是否使用了负载均衡。

由于文章没有给出源码链接，所以不知道有没有使用负载均衡。

但是，在 5.4 节的实验结果显示该方法有较高的归一化路由熵，即绝大部分神经元被选择性调用。

如果 ∞-MoE 未使用负载均衡，那么一种可能性就是该方法隐含了负载均衡。

详细讨论见下一节。】

7.2 负载均衡

一种传统的 MoE 如上图所示，

它使用 MoE-FFN 替换普通 FFN，Attention 层通常不动。
它接收 token 的 hidden state，输出一个分布，然后选概率最高的几个专家。
但该方法训练中会出现一个问题：路由器很容易 “偏爱” 某几个专家，从而导致其他专家变成僵尸参数，即 路由器塌缩（router collapse）
因此，必须加辅助损失，逼路由器把 token 均匀分配，即 负载均衡

∞-MoE 的方法是，

“通过使用掩码调制，为每个采样 z 有效地激活 FFN 神经元的一个独特子集”，有点动态剪枝的意味。
直觉上，该方法会 “偏爱” 某几个 FFN 神经元的独特子集，但 5.4 节的实验结果显示并未出现该问题。
一种可能的原因就是该方法包含了隐式的负载均衡，避免了 神经元专家塌缩。

“负载均衡” 在 MoE 里想解决的是同一件事：

不要让计算和参数更新长期集中在少数子结构上。

传统 MoE 是用一个 显式约束 去做这件事；

∞-MoE 则是靠 机制本身的几何与随机性，在 “平均意义” 上自然发生。

1）先说结论

连续路由变量 z 与基于投影排序的 top-N% 掩码共同形成了一种几何上平滑且具有随机性的专家选择机制。尽管模型中不存在显式的负载均衡约束，但在跨样本与跨训练步的统计意义下，不同方向的 z 会激活不同的神经元子集，从而自然诱导出接近均匀的神经元利用率。这种均匀性并非硬性保证，而是一种由连续性与高维几何共同产生的隐式效果。

也就是说：