💬 NLP

OmniMoE: An Efficient MoE by Orchestrating Atomic Experts at Scale

OmniMoE 是一个系统-算法协同设计的框架，通过利用诸如笛卡尔积路由（Cartesian Product Router）和以专家为中心的调度（Expert-Centric Scheduling）等新颖组件，将混合专家（Mixture-of-Experts）的粒度推向向量级，从而实现了高精度与显著推理加速的统一。

原作者： Jingze Shi, Zhangyang Peng, Yizhang Zhu, Yifan Wu, Guang Liu, Yuyu Luo

发布于 2026-07-02

📖 1 分钟阅读☕ 轻松阅读

CC BY 4.0

原作者： Jingze Shi, Zhangyang Peng, Yizhang Zhu, Yifan Wu, Guang Liu, Yuyu Luo

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

想象一下，你正在经营着一家规模宏大、高速运转的图书馆，数以百万计的书籍（数据）需要被即时处理。在人工智能的世界里，这个图书馆就是一个“模型（Model）”，而书就是它需要理解的信息。

长期以来，这些图书馆都有两种主要的工作方式，且两者都存在巨大的问题：

“大团队”模式（粗粒度）： 想象一下，对于你提出的每一个问题，你都要召集一整支由 256 名图书管理员组成的团队。即使你只需要一个特定的事实，所有 256 名管理员都会开始工作。这种方式很快，因为他们可以成组协作，但它极大地浪费了能量，因为其中大多数人对你的特定问题毫无贡献。
“单人专家”模式（细粒度）： 为了节省能量，你雇佣了数百万名微小、高度专业化的图书管理员。每当你提问时，你只召唤那一位完全了解该事实的完美专家。这种方式在资源利用上极其高效，但也是一场物流噩梦。在图书馆里寻找数百万个分散的小型专家非常耗时，这些管理员花在跑向办公桌上的时间比实际读书的时间还要长。

走进 OmniMoE：终极图书管理员系统

这篇论文介绍了一种名为 OmniMoE 的新系统，它结合了两者的优点。它就像是一个既精准又极速的图书馆。以下是它的工作原理，分为三个简单的部分：

1. “原子级”图书管理员（微型专家）

OmniMoE 并没有雇佣整个团队，而是雇佣了“原子专家（Atomic Experts）”。这些是知识中最小的单位——你可以把它们想象成单个、完美的句子或事实，而不是整个章节。

创新之处： 当你提问时，系统不仅仅是挑选一名图书管理员；它会针对你的特定问题，动态地组建一支定制化的微型专家团队。这就像是为你输入的每一个单词都构建一个定制的拼图碎片。

2. “笛卡尔积”地图（智能地址系统）

拥有数百万名微型专家的难题在于：如何快速找到他们？如果你必须检查一份包含 1000 万个名字的名单，那会太慢了。

解决方案： OmniMoM 使用了“笛卡尔积路由（Cartesian Product Router）”。想象一下，图书馆不再是一张巨大的名单，而是一个巨大的网格（就像一个有行和列的电子表格）。
作用： 系统不再是在数百万个名字中搜索特定的人，而是直接找到该图书管理员所在的“行”和“列”。这就像是说“去第 5 行，第 3 列”，而不是搜索“图书管理员鲍勃”。这把一个庞大、缓慢的搜索过程变成了一个微小的、瞬间完成的计算。

3. “以专家为中心”的班车时刻表（交通调度员）

即便有了智能地图，如果你派一辆巴士去一个接一个地接送 1000 名分散的专家，这辆巴士也会陷入交通拥堵（这被称为“内存瓶颈”）。巴士会一直处于停停走走的状况。

解决方案： OmniMoE 改变了操作顺序。系统不再是为每个问题逐一接送专家，而是先将专家进行分组。
比喻： 想象一下，司机说：“好了，所有要去‘科学区’的人，请先上车。然后我们再去接‘历史区’的人。” 巴士一次性装载一整组专家，并将他们运送到工作区，让他们作为一个高效的大型整体协同工作。这把混乱、断断续续的交通拥堵变成了平滑、高速的高速公路。

结果：快速、廉价且聪明

论文测试了该系统与当前最佳方法的对比情况：

速度： 它比之前的最佳“微型专家”系统快了 10.9 倍。它将处理时间从 73 毫秒缩短到了仅 6.7 毫秒。
智慧： 它也更加准确。通过在处理罕见、特定事实的微型专家旁边保留一个“共享稠密 MLP（Shared Dense MLP，即处理常识和语法的通用大脑）”，它并不会丧失推理能力。
效率： 它能更好地利用计算机内存，避免了通常会拖慢此类系统的“交通拥堵”。

总结

OmniMoE 是一种聪明的技巧，它让 AI 模型能够使用数百万个微小且高度专业的专家，而不会降低速度。它通过将专家组织成网格（以便快速定位）并像安排班车时刻表一样对工作进行分组（以便快速移动）来实现这一点。其结果是，AI 既能掌握极其具体的细节，又具备了投入实际应用的极速性能。

技术摘要：OmniMoE

1. 问题陈述

混合专家（Mixture-of-Experts, MoE）架构已成为通过将模型容量与单 Token 计算解耦来扩展语言模型的标准设计。然而，现有设计面临着专家粒度（expert granularity）与硬件执行效率之间的根本权衡：

粗粒度 MoE：（如 DeepSeek-V3, KIMI-K2）利用大型、稠密的前馈网络（FFN）块作为专家。虽然这些设计受益于硬件友好的稠密矩阵乘法和连续内存访问，但它们存在激活不精确的问题。激活一个大型专家块往往涉及计算与特定 Token 无关的参数，从而导致计算浪费和僵化的扩展约束。
细粒度 MoE：（如 PEER, PKM）利用数百万个轻量级专家（例如静态向量）来最大化参数效率和路由精度。然而，将这些规模扩展到巨大量级时会引入严重的系统瓶颈：
1. 表达能力有限： 将专家简化为静态向量会剥夺对复杂语言依赖至关重要的 Token 相关非线性变换（如 MLP 投影）。
2. 路由开销： 路由到数百万个专家会产生极高的计算和内存成本（ $O(N)$ ）。
3. 硬件效率低下： 细粒度路由会触发分散、非连续的内存访问，使执行瓶颈从计算转向内存带宽（受限于内存），导致 GPU 利用率低且延迟高。

核心挑战在于：是否可能将细粒度模型的参数效率与粗粒度架构的硬件效率结合起来。

2. 方法论：OmniMoE

OmniMoE 是一个**系统-算法协同设计（system-algorithm co-designed）**的框架，它在保持硬件效率的同时，将专家粒度推向了逻辑上的极致。它采用了一种混合架构，结合了一个共享的稠密 MLP 和一个海量的路由“原子专家（Atomic Experts）”池。

2.1 核心架构组件

原子专家（Atomic Experts）： 计算的基本单元是“原子专家”，定义为一个由一对向量（ $w_{in}, w_{out}$ ）参数化的最小可路由单元。单个原子专家执行轻量级的非线性变换： $E_i(x) = \sigma(x w_{in}^T) w_{out}$ 。
动态专家组装（Dynamic Expert Assembly, DEA）： OmniMoE 并非存储数百万个独立的专家块，而是将所有 $N$ 个原子专家整合进两个全局参数矩阵（ $W, V$ ）中。对于每个 Token，系统会动态检索这些矩阵中的一个稀疏子集行，并将其组装成一个针对特定 Token 的、融合的计算块。这实现了极高的参数效率，同时通过动态组合保留了高水平的表达能力。
共享稠密 MLP： 一个标准的稠密 MLP 分支被普遍激活以处理所有输入，用于处理通用的语义推理并提供稳定的骨干，以此补充专注于长尾知识检索的路由分支。

2.2 算法创新

为了管理数百万个原子专家的规模，OmniMoE 引入了两种关键的算法机制：

笛卡尔积路由器（Cartesian Product Router）：
- 问题： 向 $N$ 个专家进行标准路由需要一个大小为 $d \times N$ 的投影矩阵，导致 $O(Nd)$ 的复杂度。
- 解决方案： 路由器将一维专家索引空间分解为二维网格（ $N_r \times N_c$ ）。它预测两个独立的低维分布（行和列的 Logits），并通过组合它们来为专家评分。
- 影响： 这种因子分解将路由复杂度和参数存储从 $O(Nd) $降低到$ O(\sqrt{N}d)$，使得大规模路由变得可行。
以专家为中心的调度（Expert-Centric Scheduling）：
- 问题： 标准的“以 Token 为中心”的执行方式会独立获取每个 Token 的参数，导致分散的内存 I/O 并阻碍向量化。
- 解决方案： 执行范式被反转为“以专家为中心”。
  - 任务收集： 将 Batch 中所有的路由决策展平为一个任务列表。
  - 分组与排序： 将任务按专家 ID 的连续范围进行分组，并在每个组内按 Token ID 进行排序。
  - 分组 GEMM（Grouped GEMM）： 这种重组允许系统一次性加载一个专家权重块，并在一个 Batch 的多个 Token 之间重复使用，从而将分散的内存查找转换为连续、高吞吐量的 Grouped GEMM 操作。
- 影响： 这将执行过程从受限于内存（memory-bound）转变为受限于计算（compute-bound），最大限度地提高了 Tensor Core 的利用率和内存带宽。

3. 主要贡献

原子专家公式化： 一种新的最小可路由单元，通过动态组装实现了在扩展至数百万个专家的同时保持非线性表达能力。
笛卡尔积路由器： 一种路由机制，将复杂度相对于专家数量从线性降低到平方根级别，解决了大规模专家池的索引瓶颈。
以专家为中心的调度： 一种系统级调度策略，通过重新排列计算顺序来消除内存带宽瓶颈，使细粒度 MoE 能够实现与粗粒度设计相当的硬件效率。
系统-算法协同设计： 一个整体框架，整合了上述组件，以实现高模型表达能力和大规模下的硬件效率。

4. 实验结果

OmniMoE 在七个常识基准测试和五个扩展推理/编程基准测试上进行了评估，并进行了广泛的效率测试。

准确率： 总参数量为 6.4B / 激活参数量为 1.7B 的 OmniMoE 模型在七个常识基准测试中实现了 50.9% 的平均零样本准确率。这超过了粗粒度基线（如 DeepSeekMoE）0.7 个百分点，以及细粒度基线（如 PEER）2.0 个百分点。
推理与吞吐量： 在扩展基准测试（GSM8K, MATH 等）中，OmniMoE 达到了 42.5 的平均得分，显著优于 PEER（+15.5）和 DeepSeekMoE（+2.0）。在吞吐量方面，OmniMoE 达到了 ~14.2k tokens/second，与稠密基线相当，并且明显快于细粒度的 PEER（~~11.6k）和 PKM（~~7.9k）。
效率： 与 PEER 相比，OmniMoE 展示了 10.9 倍的推理延迟加速（将延迟从 73ms 降低到 6.7ms），同时保持了与粗粒度 MoE 相当的内存占用。
缩放法则（Scaling Laws）： 在匹配的训练 FLOPs 和激活参数预算下，OmniMoE 一致地实现了所有基线模型中最低的困惑度（Perplexity），表明其具有卓越的计算和参数效率。

5. 意义与主张

论文声称，OmniMoE 证明了大规模细粒度 MoE 既可以具备高准确性，也可以具备高效率，前提是必须有整体的系统-算法协同设计支持。

调和权衡： 它成功弥合了细粒度模型的参数效率与粗粒度模型的硬件效率之间的鸿沟。
功能专业化： 消融研究和分支主导地位分析揭示了明确的分工：共享稠密 MLP 处理通用的语言模式和推理连接词，而路由原子专家则专注于事实实体和长尾知识。
系统影响： 该工作强调，路由复杂度和内存访问模式是细粒度 MoE 的主要瓶颈，因此，只有将算法创新（笛卡尔积路由器）与系统级优化（以专家为中心的调度）相结合，才能释放百万级专家池的潜力。

作者指出，目前的实现依赖于针对 NVIDIA GPU 的自定义 Triton 内核，且将规模扩展到 6.4B 以上的总参数量仍是未来验证的领域。