OmniMoE: An Efficient MoE by Orchestrating Atomic Experts at Scale

2026-02-05

AI Summary

This work addresses the challenge of balancing expert granularity and hardware efficiency in existing Mixture-of-Experts (MoE) architectures. The authors propose OmniMoE, a novel approach that integrates system-aware and algorithmic co-design to introduce vector-level atomic experts within a single MoE layer while retaining a shared dense MLP branch, thereby significantly enhancing parameter efficiency and inference performance. Key innovations include the first use of vector-level atomic experts, a Cartesian product router that reduces routing complexity to O(√N), and an expert-centric scheduling mechanism that transforms sparse memory accesses into dense matrix operations. Experiments demonstrate that OmniMoE achieves 50.9% zero-shot accuracy across seven benchmarks and reduces inference latency from 73 ms to 6.7 ms compared to PEER, yielding a 10.9× speedup.

Technology Category

• Machine Learning: Mixture of Experts (MoE)
• Search and Optimization: Mixed Discrete/Continuous Search
• Planning, Routing, and Scheduling: Planning with Language Models

Application Category

• Systems and Infrastructure for Web, Mobile and WoT: Applied ML and AI for Web-based mobile applications
• Search and Retrieval-Augmented AI: Efficiency and scalability of Web search engines
• Economics, Online Markets and Human Computation: Incentives in network design for Web infrastructures and ecosystems

Abstract

Mixture-of-Experts (MoE) architectures are evolving towards finer granularity to improve parameter efficiency. However, existing MoE designs face an inherent trade-off between the granularity of expert specialization and hardware execution efficiency. We propose OmniMoE, a system-algorithm co-designed framework that pushes expert granularity to its logical extreme. OmniMoE introduces vector-level Atomic Experts, enabling scalable routing and execution within a single MoE layer, while retaining a shared dense MLP branch for general-purpose processing. Although this atomic design maximizes capacity, it poses severe challenges for routing complexity and memory access. To address these, OmniMoE adopts a system-algorithm co-design: (i) a Cartesian Product Router that decomposes the massive index space to reduce routing complexity from O(N) to O(sqrt(N)); and (ii) Expert-Centric Scheduling that inverts the execution order to turn scattered, memory-bound lookups into efficient dense matrix operations. Validated on seven benchmarks, OmniMoE (with 1.7B active parameters) achieves 50.9% zero-shot accuracy across seven benchmarks, outperforming coarse-grained (e.g., DeepSeekMoE) and fine-grained (e.g., PEER) baselines. Crucially, OmniMoE reduces inference latency from 73ms to 6.7ms (a 10.9-fold speedup) compared to PEER, demonstrating that massive-scale fine-grained MoE can be fast and accurate. Our code is open-sourced at https://github.com/flash-algo/omni-moe.

Problem

• Mixture-of-Experts
• expert granularity
• hardware efficiency
• parameter efficiency
• routing complexity

Innovation

• Atomic Experts
• Cartesian Product Router
• Expert-Centric Scheduling
• Mixture-of-Experts
• System-Algorithm Co-design