OmniMoE: Scalable MoE Framework

• OmniMoE is a Mixture-of-Experts framework that employs vector-level atomic experts for fine-grained specialization and efficient routing.
• It introduces Cartesian Product Routing and Expert-Centric Scheduling to reduce computational complexity and optimize hardware usage.
• Empirical results demonstrate significant accuracy gains and inference speedups compared to traditional coarse- and fine-grained MoE architectures.

OmniMoE is a system–algorithm co-designed framework for Mixture-of-Experts (MoE) architectures that introduces vector-level atomic experts and scalable routing for efficient and precise expert specialization. Designed to address longstanding trade-offs between MoE granularity and hardware efficiency, OmniMoE enables massive-scale deployment of fine-grained experts without sacrificing inference speed or parameter utilization. The framework includes two primary innovations: the Cartesian Product Router (reducing routing complexity) and Expert-Centric Scheduling (restructuring execution for hardware efficiency), and has influenced subsequent multimodal large model systems through its architectural principles (Shi et al., 5 Feb 2026).

1. Motivations and Context

Traditional MoE architectures face a balance between specialization granularity and computational efficiency. Coarse-grained MoEs, such as DeepSeekMoE, achieve high throughout by routing tokens to large feed-forward modules, maximizing compute unit utilization via dense matrix multiplication but inefficiently activating many irrelevant parameters. In contrast, fine-grained MoEs (e.g., PKM, PEER) achieve exceptional parameter efficiency and specialization by retrieving from millions of small experts, typically vectors, but suffer from high routing costs and hardware inefficiency due to scattered memory access and limited expert expressivity.

OmniMoE pushes expert granularity to its logical extreme by adopting atomic experts, each consisting of an input vector $w_i^{\mathrm{in}}\in\mathbb{R}^d$ and an output vector $w_i^{\mathrm{out}}\in\mathbb{R}^d$. For any input token $x$, an atomic expert computes $$E_i(x) = \sigma(x {w_i^{\mathrm{in}}}^\top)w_i^{\mathrm{out}}$$, where $\sigma$ is the SwiGLU activation. Dynamic Expert Assembly (DEA) then composes the top-$K$ ranked atomic experts per token, enabling the model to construct highly specialized, nonlinear, token-dependent subspaces while retaining maximum parameter efficiency (Shi et al., 5 Feb 2026).

2. OmniMoE Layer Design and Atomic Expert Assembly

The OmniMoE unit is composed of two branches per layer:

• Shared dense MLP branch: A standard MLP, providing broad generalization capacity and stable training, is always applied to every token.
• Sparse routed atomic expert branch: For each input $x\in\mathbb{R}^d$, the model retrieves top-$K$ atomic experts $\{I_1,\ldots,I_K\}$ and associated weights $g_i$, assembling their outputs as:

$w_i^{\mathrm{out}}\in\mathbb{R}^d$0

where $w_i^{\mathrm{out}}\in\mathbb{R}^d$1 are the input/output matrices for the selected atoms, and $w_i^{\mathrm{out}}\in\mathbb{R}^d$2 contains the normalized expert weights.

All atomic experts are stored in two global matrices $w_i^{\mathrm{out}}\in\mathbb{R}^d$3 and $w_i^{\mathrm{out}}\in\mathbb{R}^d$4, encoding the input and output vectors, respectively. The final layer output is $w_i^{\mathrm{out}}\in\mathbb{R}^d$5. This dual-branch architecture balances generalization and specialization by ensuring dense coverage and precise adaptation (Shi et al., 5 Feb 2026).

3. Cartesian Product Routing for Efficient Large-Scale Expert Selection

Routing to $w_i^{\mathrm{out}}\in\mathbb{R}^d$6 atomic experts via a naive projection would incur $w_i^{\mathrm{out}}\in\mathbb{R}^d$7 computational and memory costs, which is prohibitive for very large $w_i^{\mathrm{out}}\in\mathbb{R}^d$8. OmniMoE addresses this with the Cartesian Product Router: it factorizes the $w_i^{\mathrm{out}}\in\mathbb{R}^d$9 expert IDs into a 2D grid ($x$0, where $x$1), and models the joint selection probability as $x$2. Routing is implemented by learning two projections ($x$3) for input $x$4:

$x$5

Scores are aggregated as $x$6. Top-$x$7 routing is performed by tiled GPU searches over this implicit grid, yielding a complexity reduction from $x$8 to $x$9. For million-expert regimes, this provides substantial practical gains in scalability (Shi et al., 5 Feb 2026).

4. Expert-Centric Scheduling and System–Algorithm Co-Design

Conventionally, MoE implementations loop over tokens and fetch their assigned experts, resulting in scattered small VRAM accesses and compute-bound inefficiency. OmniMoE introduces an expert-centric scheduling approach:

• All token–expert pairs in a batch are flattened.
• Active experts are grouped, sorted, and partitioned into blocks.
• For each block, a fused General Matrix Multiply (GEMM) operation processes all relevant inputs in one batch, enabling contiguous VRAM access and maximal use of hardware accelerators such as Tensor Cores.
• Outputs are scatter-added back to the token positions.

This approach causes VRAM traffic to scale with the unique number of active experts, not the number of token–expert pairs, and transitions the system from memory-bound to compute-bound execution. The two design pillars—Cartesian Product Routing and Expert-Centric Scheduling—are carefully co-developed to eliminate quadratic scaling bottlenecks and optimize for dense high-throughput execution (Shi et al., 5 Feb 2026).

5. Empirical Evaluation and Comparative Performance

OmniMoE, evaluated with 1.7B active parameters (6.4B total), demonstrates leading zero-shot accuracy across seven major NLP benchmarks:

• Average accuracy: 50.9% (MMLU 37.5, TriviaQA 18.5, ARC 61.0, PIQA 78.7, HellaSwag 60.9, OBQA 40.3, WinoGrande 59.7)
• Outperforms both coarse-grained (DeepSeekMoE, 50.2%) and fine-grained (PEER, 48.9%) baselines.
• Delivers a 10.9-fold speedup in inference latency (73 ms → 6.7 ms) over PEER, and 15.2-fold over DeepSeekMoE, for 4 096-token batches at similar peak memory usage.

Ablations confirm the necessity of both core innovations. Removing the shared MLP reduces latency by 14% but degrades perplexity by 20% and average zero-shot accuracy by 9%. Using a full $$E_i(x) = \sigma(x {w_i^{\mathrm{in}}}^\top)w_i^{\mathrm{out}}$$0 gate inflates latency $$E_i(x) = \sigma(x {w_i^{\mathrm{in}}}^\top)w_i^{\mathrm{out}}$$1, memory $$E_i(x) = \sigma(x {w_i^{\mathrm{in}}}^\top)w_i^{\mathrm{out}}$$2, and dramatically worsens expert usage efficiency. Disabling expert-centric scheduling increases latency $$E_i(x) = \sigma(x {w_i^{\mathrm{in}}}^\top)w_i^{\mathrm{out}}$$3 and memory $$E_i(x) = \sigma(x {w_i^{\mathrm{in}}}^\top)w_i^{\mathrm{out}}$$4 but leaves model quality unchanged (Shi et al., 5 Feb 2026).

6. Influence on Multimodal and Unified Architectures

The OmniMoE principle of combining ultra-sparse, high-capacity atomic experts with efficient hardware-aware routing forms the core of subsequent architectures such as Ming-Flash-Omni (AI et al., 28 Oct 2025) and Uni-MoE-2.0-Omni (Li et al., 16 Nov 2025). These successors extend the approach to unified, multimodal large models for vision, speech, and language, integrating modality-specific encoders and expert selection schemes (including dynamic-capacity routing and shared/null experts). Empirical results across benchmarks in ASR, image generation, segmentation, and reasoning tasks validate the scalability and effectiveness of the OmniMoE design, setting new state-of-the-art performance within a single, efficient model (AI et al., 28 Oct 2025, Li et al., 16 Nov 2025).

7. Limitations, Insights, and Future Directions

Current limitations include minor ($$E_i(x) = \sigma(x {w_i^{\mathrm{in}}}^\top)w_i^{\mathrm{out}}$$55%) overhead from the two-stage sort in expert scheduling and the exclusive use of vector-level atomic experts. Extending atomic experts to lightweight linear layers or higher-order structures could increase expressivity. More advanced index decompositions, such as multi-dimensional Cartesian products, may enable further reduction in routing complexity. A plausible implication is that continued holistic system–algorithm co-design will be essential in scaling MoE models to billion-scale experts while maintaining hardware efficiency (Shi et al., 5 Feb 2026).