OmniMoE: An Efficient MoE by Orchestrating Atomic Experts at Scale

🎯 Pitch

OmniMoE achieves a 10.9× inference speedup over prior fine-grained MoEs by breaking experts down to individual vector pairs and cleverly reorganizing their execution from scattered lookups into dense matrix multiplies. It makes massive-scale, ultra-fine-grained expert routing not only accurate but, for the first time, truly fast, overcoming a core hardware bottleneck that previously made such designs impractical.

1. Executive Summary

OmniMoE proposes a system-algorithm co-designed Mixture-of-Experts framework that pushes expert granularity to its logical extreme by introducing Atomic Experts—minimal routable units parameterized by a single pair of vectors—and orchestrating their activation through two tightly integrated mechanisms: a Cartesian Product Router (decomposing the massive 1D expert index space into a 2D grid to reduce routing complexity from O(N) to O(√N)) and Expert-Centric Scheduling (inverting execution order from token-centric to expert-centric to convert scattered memory lookups into dense Grouped GEMM operations). Validated on seven benchmarks with a 6.4B-parameter model (1.7B active), OmniMoE achieves 50.9% average zero-shot accuracy while delivering a 10.9× inference speedup over the fine-grained PEER baseline (reducing latency from 73 ms to 6.7 ms at 4,096 tokens), establishing that massive-scale fine-grained MoEs can achieve both high accuracy and hardware efficiency only when routing complexity and memory access patterns are jointly addressed through algorithm-system co-design.

2. Context and Motivation

The Core Problem: The Granularity-Efficiency Trade-off in MoEs

Mixture-of-Experts (MoE) architectures have become a dominant paradigm for scaling language models because they partially decouple model capacity from per-token computation: instead of activating all parameters for every input, an MoE layer routes each token to a small subset of specialized "experts," allowing massive growth in total parameters while keeping inference FLOPs manageable. However, a central design tension has emerged that this paper directly confronts: the granularity of experts governs both routing precision and system efficiency, and these two objectives pull in opposite directions.

On one side, coarse-grained MoEs—where each expert is a full dense feed-forward network (FFN) with millions of parameters—benefit from hardware-friendly execution. Because coarse-grained experts are large contiguous blocks of parameters, their computation maps naturally to dense matrix multiplications (matmuls) that saturate GPU Tensor Cores, enjoy contiguous VRAM access patterns, and can leverage optimized kernel libraries. This hardware efficiency is why the dominant deployed large language models (e.g., DeepSeek-V3 with 256 experts, KIMI-K2 with 384 experts) adopt coarse-grained designs despite their limitations.

On the other side, fine-grained MoEs—where experts are lightweight units like embedding vectors, potentially numbering in the millions—offer superior parameter efficiency and routing precision. When an expert is small, the router can be more selective about exactly which parameters to activate for a given token, avoiding the computational waste of activating entire FFN blocks where only a fraction of the parameters are genuinely relevant. Scaling-law analyses (Ludziejewski et al., 2024; Clark et al., 2022) have shown that, under a fixed training-token budget, model performance improves with the total number of activated experts, providing theoretical motivation for pushing toward finer granularity.

The problem is that these two objectives—hardware efficiency and parameter efficiency—are fundamentally in tension. Coarse-grained experts achieve hardware efficiency at the cost of imprecise activation: when you activate a large FFN expert, you inevitably perform computation on parameters irrelevant to the specific token (what the paper calls "redundant activation," illustrated by orange nodes in Figure 1a). Fine-grained experts achieve precise activation at the cost of hardware inefficiency: activating scattered lightweight experts triggers random memory I/O that shifts execution from compute-bound to memory-bound, degrading GPU utilization (Figure 1b). This is not a minor implementation detail—it is a structural trade-off inherent to how these architectures map onto hardware.

Why This Problem Matters: The Stakes Are Scaling, Not Just Speed

This trade-off is not merely an engineering optimization problem that can be solved with better kernel tuning. It represents a fundamental bottleneck that determines which architectures are viable at scale and shapes the frontier of what MoE models can achieve. Understanding why requires examining each side of the trade-off more carefully.

The waste in coarse-grained MoEs is structural, not incidental. When a coarse-grained MoE activates a full FFN expert for a token, the computation performed on parameters that are not specifically relevant to that token is not a consequence of poor routing—it is a direct consequence of the expert being too large to route precisely. Even with perfect routing (i.e., selecting the single best expert for each token), the activated expert contains many parameters that encode knowledge for different linguistic patterns, different domains, or different reasoning styles than what the current token requires. This "activation sparsity" inefficiency has been documented in multiple works (Li et al., 2023; Szatkowski et al., 2024; Zhou et al., 2025) and represents wasted FLOPs that scale linearly with the number of activated tokens. As models grow larger and are deployed at higher throughput, this waste compounds.

Moreover, the rigid sizing of coarse-grained experts limits scaling flexibility. When you want to increase expert count—for instance, to accommodate new domains or languages—you must add entire FFN blocks, which produce steep, discrete jumps in memory consumption. There is no smooth knob for trading off between capacity and compute. This matters for deployment on heterogeneous hardware (edge devices, consumer GPUs, cloud instances with varying memory budgets) where fine-grained control over the capacity-efficiency trade-off is essential.

The inefficiency in fine-grained MoEs is a deployment dealbreaker. Prior work in the fine-grained paradigm demonstrates impressive model quality. PEER (He, 2024), for example, scales to millions of experts using a Product Key Memory-style design and achieves superior accuracy under matched parameter budgets. However, the hardware reality is stark: when each expert is a small vector stored at a scattered memory address, and each token activates a different sparse subset of these vectors, the execution becomes dominated by random gather operations from high-bandwidth memory (HBM). The GPU's compute units (Tensor Cores) sit idle while waiting for memory fetches to complete. This is why, despite their theoretical advantages, fine-grained MoEs have not seen widespread production adoption—the latency penalty erases the parameter-efficiency gains in practice.

The gap between theory and practice is widening. As the scaling-law literature pushes toward ever-finer granularity (Ludziejewski et al., 2024 suggesting millions of experts may be optimal), production MoE systems remain at the scale of hundreds of coarse-grained experts. This divergence is not because practitioners are unaware of the theoretical benefits of fine-grained routing—it is because the systems cost of random memory access has made those benefits unrealizable on real hardware. Any architecture that claims to bridge this gap must solve the routing complexity problem (how to select from millions of candidates without the router itself becoming the bottleneck) and the memory access problem (how to execute scattered expert computations without becoming memory-bound) simultaneously. These two problems interact: you cannot solve one and ignore the other.

Prior Approaches and Where They Fall Short

The paper situates itself relative to a landscape of MoE designs that have attempted to address parts of the granularity-efficiency trade-off, but none that have addressed both routing and scheduling holistically.

Coarse-Grained MoEs (The Dominant Paradigm)

Coarse-grained architectures instantiate each expert as a complete dense FFN and use a lightweight router (typically a single learned projection matrix) to select a small number of experts per token. Representative systems include:

GShard (Lepikhin et al., 2021): One of the first production-scale MoE implementations, using top-2 gating with auxiliary load-balancing losses. The experts are standard FFN blocks.
Switch Transformers (Fedus et al., 2022): Simplified routing to top-1 gating (each token routed to a single expert), demonstrating that even minimal routing can scale to trillion-parameter models.
DeepSeekMoE / DeepSeek-V3 (Dai et al., 2024; DeepSeek-AI et al., 2025): Introduced the concept of shared experts alongside routed experts—some FFN blocks are universally activated for all tokens to handle general linguistic patterns, while other FFN blocks are routed. This partially addresses the activation imprecision problem: the shared expert captures common knowledge, and the routed experts specialize. DeepSeek-V3 scales to 256 experts with this design.
KIMI-K2 (Team et al., 2025): Scales to 384 experts, also using shared + routed expert separation.

These architectures have achieved impressive results and dominate production deployment. However, they all share a fundamental limitation: the routed experts are still coarse-grained FFN blocks. The shared-expert design mitigates the activation waste by offloading general knowledge to a universally-activated component, but it does not solve the problem for the routed experts themselves—they still activate entire FFN blocks when only a fraction of the parameters are token-relevant. The paper's Figure 1(a) illustrates this precisely: coarse-grained MoEs activate experts that contain both valid activations (parameters relevant to the current token) and redundant activations (parameters irrelevant to the current token, marked in orange).

Additionally, the rigid sizing of these experts means that scaling expert count requires adding entire FFN blocks, producing discrete, potentially over-large memory increments. The authors note that coarse-grained MoEs "force steep, discrete memory increments when adjusting expert counts" (Section 1)—a practical limitation for deployment scenarios where fine-grained capacity control matters.

A subtler limitation that the paper identifies: even with state-of-the-art kernel implementations (e.g., NVIDIA's CuTile library, which the paper uses for its coarse-grained baselines), coarse-grained expert execution can be slower than expected due to packing and alignment overhead. When tokens are routed to experts, the tokens assigned to each expert must be reordered and padded into fixed-size blocks for efficient GPU execution. If routing is uneven (some experts get many tokens, some get few), this padding creates redundant computation and extra memory traffic. The paper observes this empirically: at large token counts, coarse-grained DeepSeekMoE can be slower than fine-grained PEER because of these overheads, despite PEER's scattered memory access pattern. This is a non-obvious finding that underscores the complexity of the hardware-efficiency landscape.

Fine-Grained MoEs (The Theoretical Frontier)

Fine-grained architectures push expert granularity to the extreme, using millions of lightweight experts. The key representatives are:

PKM (Product Key Memory) (Lample et al., 2019): Introduced the idea of decomposing the expert index space into a Cartesian product of two subspaces to enable efficient routing over massive expert pools. PKM replaces the FFN with a memory layer where experts are learned key-value pairs, and routing is done via nearest-neighbor search over product keys. This was foundational work that the paper's Cartesian Product Router builds upon.
PEER (Mixture of A Million Experts) (He, 2024): Extended the PKM design to modern Transformer architectures, scaling to millions of experts and demonstrating superior parameter efficiency. PEER uses product-key routing to select from a massive pool of embedding-sized experts.

These architectures demonstrate the theoretical benefits of fine-grained routing—better parameter efficiency, more precise activation, and smoother scaling curves. However, the paper identifies three specific shortcomings that prevent these architectures from realizing their theoretical advantages in practice:

1. Limited expressivity. PEER and PKM reduce experts to static parameter vectors—essentially, each expert is a learned embedding. The computation performed by an activated expert is reduced to linear vector aggregation: the expert's output is a weighted combination of embedding vectors, without any token-dependent nonlinear transformations (like the MLP projections that coarse-grained FFN experts perform). The paper argues that this "strips away the token-dependent nonlinear transformations essential for modeling complex linguistic dependencies" (Section 1). In other words, fine-grained experts achieve precise activation but lose the representational power that comes from deep, nonlinear per-expert processing. This is a direct consequence of making experts small enough to be routed at the million-expert scale: there isn't enough capacity per expert to support a full MLP projection.

2. Routing overhead at scale. While product-key routing reduces the complexity of selecting from N experts compared to a naive N-way classifier, the paper observes that scaling to massive expert pools still introduces routing challenges. Specifically, load imbalance becomes more severe at scale—some experts get disproportionately many tokens while others get none, leading to "skewed expert utilization at scale" (Section 1). The standard auxiliary load-balancing loss used in coarse-grained MoEs does not trivially transfer to the million-expert regime where the expert space is vastly larger than the number of tokens in a batch.

3. Hardware inefficiency from scattered memory access. This is the most critical shortcoming. When fine-grained experts are small vectors scattered across HBM, and each token activates a different subset, the GPU must perform gather operations—individual reads from non-contiguous memory addresses. This pattern is disastrous for GPU memory bandwidth utilization because (a) random access latency dominates, (b) the memory controller cannot coalesce reads into efficient bursts, and (c) the GPU's Streaming Multiprocessors (SMs) stall waiting for data, leaving compute units underutilized. As the paper's Figure 1(b) illustrates, fine-grained experts ensure precise activation (only the relevant parameters are active), but the "active parameters are inherently scattered across memory, which triggers frequent, non-contiguous memory accesses, inevitably shifting the execution bottleneck from computation to memory bandwidth." This explains the measured 73 ms latency for PEER at 4,096 tokens—the architecture is memory-bound, not compute-bound.

Hybrid Designs (Shared Experts + Routed Experts)

The shared-expert design in DeepSeekMoE (Dai et al., 2024) can be seen as a partial step toward reconciling the granularity-efficiency trade-off. By having a universally-activated shared FFN expert alongside routed experts, DeepSeekMoE acknowledges that some knowledge (general linguistic patterns, common reasoning steps) should be available to all tokens, while other knowledge (domain-specific, long-tail) benefits from routing. The shared expert also provides training stability benefits (Nguyen et al., 2025).

However, as noted above, the routed experts in DeepSeekMoE remain coarse-grained FFN blocks. The shared-expert design improves the allocation of knowledge between shared and routed components, but it does not address the fundamental activation imprecision within the routed experts themselves. The paper positions OmniMoE as building on the shared-expert insight (OmniMoE retains a shared dense MLP) while pushing the routed component to the opposite extreme—atomic fine-grained experts—to achieve precise activation in the routed branch.

System-Level Optimizations for MoE

A parallel line of work has focused on system-level optimizations to make MoE execution more efficient, but these approaches have been developed primarily for coarse-grained experts:

DeepSpeed-MoE (Rajbhandari et al., 2022): Optimizes communication scheduling and kernel fusion for distributed MoE training/inference.
FastMoE (He et al., 2021): Provides efficient MoE kernel implementations with load-balancing support.
MegaBlocks (Gale et al., 2023): Introduces block-sparse kernels that handle variable-length token-to-expert assignments, reducing padding overhead.
SonicMoE (Guo et al., 2025): Further improves grouped GEMM efficiency with memory-minimizing algorithms and tile-aware token rounding.
PIT (Zheng et al., 2023): Exploits dynamic sparsity at the neuron level within activated experts to prune invalid computation.
ScatterMoE (Tan et al., 2024): Implements scattered operations to avoid padding in MoE kernels.

These works improve coarse-grained MoE execution efficiency, but they operate within the paradigm where experts are large FFN blocks. They do not address the fundamental tension between granularity and memory access patterns that emerges when experts become fine-grained. OmniMoE's Expert-Centric Scheduling can be seen as extending the spirit of these system optimizations—reorganizing computation to improve hardware utilization—but applied to the fundamentally different access patterns of fine-grained atomic experts.

How This Paper Positions Itself

OmniMoE positions itself not as proposing a fundamentally new type of expert or a fundamentally new routing mechanism in isolation, but as introducing a holistic system-algorithm co-design that jointly addresses the three bottlenecks that have prevented fine-grained MoEs from being practical: (i) limited per-expert expressivity, (ii) routing complexity at massive scale, and (iii) hardware inefficiency from scattered memory access. The paper's central framing is explicit:

"Is it possible to reconcile the parameter efficiency of fine-grained models with the hardware efficiency of coarse-grained architectures? Realizing this synergy is non-trivial. It requires a holistic orchestration that simultaneously enhances the expressivity of fine-grained experts, minimizing routing overhead in large expert spaces, and reshaping irregular sparse accesses into hardware-efficient execution." (Section 1)

This framing reveals the paper's intellectual strategy: the individual components (atomic experts, product-structured routing, execution reordering) each have antecedents in prior work, but they have never been combined into a single system where they compensate for each other's weaknesses. Specifically:

Atomic Experts push granularity to the extreme while maintaining expressivity through dynamic composition (Section 2.1). Unlike PEER's static vectors, OmniMoE's atomic experts are composed into a token-conditioned assembled expert that performs a full nonlinear computation (with SwiGLU activation). This addresses the limited expressivity of prior fine-grained designs.
The Cartesian Product Router (Section 2.2) builds on the product-key insight from PKM but adapts it to the specific needs of routing over atomic experts in OmniMoE's architecture. The router decomposes the 1D expert index space of size N into a 2D grid of size √N × √N, reducing projection complexity from O(Nd) to O(√N d). This makes routing over million-scale expert pools feasible.
Expert-Centric Scheduling (Section 2.3) is the critical system-level contribution that enables the hardware-efficient execution of the fine-grained routed computation. By inverting the execution loop—processing groups of co-activated experts rather than iterating over tokens—it converts scattered memory gather operations into coalesced, reusable reads that can be executed as high-throughput Grouped GEMM kernels.

The paper's positioning can be understood through its Figure 1: Figure 1(c) shows OmniMoE's heterogeneous architecture, which combines a universally-activated shared dense MLP (for general semantic processing, capturing the benefit of shared experts from DeepSeekMoE) with routed atomic experts (for precise long-tail knowledge retrieval, capturing the benefit of fine-grained routing from PEER/PKM). The expert-centric scheduling then ensures that the fine-grained routed computation executes with hardware efficiency comparable to coarse-grained designs by grouping expert fetches into contiguous, coalesced memory accesses.

The critical claim is that neither the algorithmic innovations alone nor the system innovations alone would suffice—it is the co-design that produces the result. Without the Cartesian Product Router, routing over millions of atomic experts would be prohibitively expensive (Table 2 shows 30.6× latency regression when replacing it with a standard dense router). Without Expert-Centric Scheduling, the fine-grained accesses would remain memory-bound (24.8× latency regression when reverting to token-centric execution). Without the shared dense MLP, the model would lose general reasoning capability that the atomic experts alone cannot capture (0.79× reasoning performance regression when removed). The paper's contribution is the integrated orchestration of these components.

Summary of the Motivation Landscape

The paper addresses a structural trade-off that has partitioned the MoE design space into two regimes—coarse-grained (hardware-efficient but activation-imprecise) and fine-grained (activation-precise but hardware-inefficient)—with neither regime capable of realizing the full potential of Mixture-of-Experts architectures. This matters because (a) scaling laws suggest that finer granularity improves parameter efficiency, making the fine-grained regime theoretically attractive; (b) production deployments require hardware efficiency, making the coarse-grained regime practically necessary; and (c) prior attempts to bridge this gap have only addressed individual bottlenecks (e.g., product-key routing for the routing problem, shared experts for the general-knowledge problem) without jointly solving the routing complexity, memory access, and expressivity challenges in an integrated system. OmniMoE's positioning is that this joint solution requires co-designing the algorithm (how experts are parameterized and routed) with the system (how expert computation is scheduled on hardware), and that the resulting framework can achieve the theoretical benefits of fine-grained MoEs (precise activation, parameter efficiency) at the hardware efficiency of coarse-grained designs.

3. Technical Approach

3.1 Reader orientation (approachable technical breakdown)

OmniMoE is a Mixture-of-Experts layer for Transformers that replaces the standard feed-forward network (FFN) with a hybrid design: a universally-activated dense MLP for general reasoning, plus a massive pool of tiny "atomic experts" that are selectively activated per token to handle specialized, long-tail knowledge. The system solves the granularity-efficiency trade-off—the problem that making experts smaller improves routing precision but destroys hardware efficiency due to scattered memory accesses—by jointly redesigning both the routing algorithm (to handle millions of experts without exploding routing cost) and the execution schedule (to convert random memory lookups into dense matrix operations).

3.2 Big-picture architecture (diagram in words)

The OmniMoE FFN layer has five major components, arranged in two parallel pathways plus a scheduling layer:

Shared Dense MLP (bottom of Figure 2)—a standard feed-forward network that processes every token unconditionally. It handles general semantic patterns (common linguistic structures, reasoning steps) that should be available to all tokens regardless of routing decisions. This component is always active and operates identically to the FFN in a dense Transformer.
Atomic Expert Pool (top of Figure 2)—a massive repository of N lightweight experts stored as two global parameter matrices W and V (each of shape N × d, where d is the hidden dimension). Each atomic expert is essentially a pair of d-dimensional vectors (one input projection, one output projection). The pool serves as a centralized parameter bank from which token-specific experts are dynamically retrieved.
Cartesian Product Router (left side of Figure 2 top)—the component that decides which K atomic experts to activate for each token, and with what weights. Unlike a standard router that computes a full N-way classification, this router decomposes the expert index space into a 2D grid of size √N × √N, performs two small projections (row and column), and composes them to score any expert on the grid without materializing the full score matrix.
Dynamic Expert Assembly (DEA) (center-right of Figure 2 top)—once the router identifies the top-K expert indices for a token, DEA retrieves the corresponding rows from W and V (yielding compact K × d parameter blocks) and composes them with the routing weights into a single fused computation that acts like a token-conditioned FFN.
Expert-Centric Scheduler (not shown in Figure 2, but illustrated in Figure 3)—a system-level component that intercepts the per-token routing decisions for an entire batch, reorganizes them by grouping tasks that target the same (or nearby) experts, and executes the resulting workloads as dense Grouped GEMM operations rather than scattered per-token lookups.

Information flows as follows: A batch of token representations X enters the layer → the shared dense MLP processes every token in parallel, producing a base output → simultaneously, the Cartesian Product Router computes routing scores and top-K expert indices for each token → the DEA mechanism retrieves the corresponding atomic expert parameters from the global pool → the Expert-Centric Scheduler reorders the resulting token-expert tasks into expert-centric groups → the grouped computations execute as dense matrix multiplications → the routed outputs are scattered back and summed with the shared MLP output to produce the final layer output.

3.3 Roadmap for the deep dive

First, the formal MoE layer definition (Eqs. 1–3), because it establishes the mathematical framework that OmniMoE instantiates and clarifies the standard notation for routers, gating, and expert computation.
Second, the Atomic Expert definition and Dynamic Expert Assembly (Section 2.1, Eqs. 4–7), since this is the core algorithmic innovation—how OmniMoE parameterizes experts at extreme granularity while maintaining expressivity through dynamic composition.
Third, the Cartesian Product Router (Section 2.2, Eqs. 8–11), because routing is the prerequisite for expert activation—you must understand how experts are selected before you can understand how their execution is scheduled.
Fourth, Expert-Centric Scheduling (Section 2.3, Eqs. 12–14), since scheduling is the final stage that takes the routed expert assignments and makes them hardware-efficient—this is where the system-algorithm co-design pays off.
Fifth, complexity analysis (Appendix A, Eqs. i–viii), to quantify why each component matters: the FLOPs reduction from the Cartesian Product Router, the memory traffic reduction from Expert-Centric Scheduling, and the overhead of the sorting stage.
Sixth, training and hyperparameter details (Section 3.1 and Appendix B), because architectural innovations must be paired with specific training recipes to be reproducible and to isolate architectural effects from training artifacts.

3.4 Detailed, sentence-based technical breakdown

This is primarily a system-algorithm co-design paper whose core idea is that fine-grained MoEs can achieve both high accuracy and hardware efficiency, but only if the routing complexity problem and the memory access problem are solved jointly through integrated architectural and scheduling innovations.

Standard MoE Formulation (Eqs. 1–3)

The paper begins by formalizing the general MoE computation that OmniMoE instantiates, establishing notation and the mathematical skeleton that all subsequent design choices hang from:

A standard MoE layer contains a pool of N experts E = {E₁, ..., E_N} and a router G(·) that maps each token representation x ∈ ℝᵈ (where d is the hidden dimension) to a distribution over the N experts. The router first computes raw scores for all experts, then selects the top-K:

$I_x = (I_0, ..., I_{K-1}) = \text{TopK}(G(x), K)$

where I_x is the ordered list of indices of the K selected experts for token x, and TopK returns the indices of the K largest values in the score vector G(x).

What it computes: for a given token, this operation takes the router's N-dimensional score vector (one score per expert), identifies the K largest entries, and returns their indices. The output is a length-K integer list that specifies which experts will participate in processing this token.

Why this form: Top-K selection is the standard gating mechanism in sparse MoEs because it enforces that exactly K experts are active per token (providing a fixed per-token compute budget) while allowing the router to focus capacity on the most relevant experts. The alternative—threshold-based gating where any expert above a score threshold is activated—produces variable per-token compute and is harder to optimize at scale.

The routing weights for the selected experts are then obtained by applying softmax normalization to the selected scores only:

$g_i = \text{Softmax}(G(x)[I_x])_i, \quad i \in [0, K)$

where g_i is the normalized weight assigned to the i-th selected expert, and G(x)[I_x] denotes the K-element sub-vector of router scores at the selected indices.

What it computes: this normalizes the raw scores of the K selected experts so that they sum to 1 (forming a valid probability distribution) and are positive. The output (g_0, ..., g_{K-1}) tells the layer how much to weight each selected expert's contribution.

Why this form: softmax over only the selected experts (rather than all N experts) is the standard practice in sparse MoEs because it ensures the weights are properly normalized for the experts that actually contribute, while the unselected experts contribute zero. Normalizing over all N experts would make the selected experts' weights very small when N is large, since the softmax denominator sums over N terms.

The layer output is the weighted sum of the selected experts' outputs, plus an optional shared dense MLP output:

$y = \sum_{i \in [0, K)} g_i \cdot E_{I_i}(x) + \text{MLP}(x)$

where E_{I_i}(x) is the output of the i-th selected expert when processing token x, and MLP(x) is the output of the universally-activated shared dense MLP (present in architectures like DeepSeekMoE and OmniMoE, but not in all MoE variants).

What it computes: this combines the K expert outputs (each weighted by its routing score g_i) with the shared MLP output through simple addition. The result y is a single d-dimensional vector that serves as the FFN layer's output for this token.

Why this form: additive combination preserves linearity in the contributions—each expert and the shared MLP can specialize independently, and their contributions compose without interference. The shared MLP term provides a "baseline" signal that is always present, which is particularly important when routing is uncertain (as it often is early in training) because even if routing is random, the layer still has a functioning dense pathway.

Atomic Experts and Dynamic Expert Assembly (Eqs. 4–7)

This section defines the core architectural innovation: how OmniMoE parameterizes experts at extreme granularity while maintaining expressivity through dynamic composition.

The Atomic Expert

An atomic expert E_i is the minimal routable computational unit in OmniMoE. Unlike a coarse-grained FFN expert (which contains two large weight matrices and an activation function, totaling millions of parameters), an atomic expert is parameterized by exactly two vectors:

w_i^in ∈ ℝᵈ: an input projection vector (sometimes called the "key" or "input weight"), and
w_i^out ∈ ℝᵈ: an output projection vector (sometimes called the "value" or "output weight").

Given a token representation x ∈ ℝᵈ, the computation performed by a single atomic expert is:

$E_i(x) = \sigma(x \cdot w_i^{\text{in}^\top}) \cdot w_i^{\text{out}}$

where σ(·) is a non-linear activation function, instantiated as SiGLU (Shazeer, 2020) throughout OmniMoE, and x · w_i^{in^⊤} denotes the dot product between the token representation and the input vector (producing a scalar).

What it computes: first, the token x is projected onto the input vector w_i^in via dot product, yielding a scalar "relevance score" for how well this expert matches the token. This scalar passes through the SiGLU activation (which gates the signal: if the relevance is low, the output is near zero; if high, the output is proportional to the relevance). The gated scalar then multiplies the output vector w_i^out, producing a d-dimensional output. In effect, the atomic expert acts as a gated linear projection: it retrieves a stored output vector w_i^out and scales it by how relevant the expert is to the current token.

Why this form: this design makes each expert extremely parameter-efficient (only 2d parameters per expert, compared to the O(d²) parameters in a coarse-grained FFN expert), enabling the expert pool to scale to millions of experts within the same total parameter budget. The gated structure (dot product → activation → scaling) is the minimum computation needed for an expert to be token-dependent: the dot product captures token-expert compatibility, and the activation provides non-linearity. Without the gating (just retrieving w_i^out unconditionally), the expert would be a static vector independent of the input, which is what PEER and PKM do—and which the paper argues strips away essential expressivity.

A critical detail: a single atomic expert has limited expressivity. The paper explicitly acknowledges this, stating "while a single atomic expert exhibits limited expressivity, the strength of our approach arises from the dynamic composition of these experts." This is a design choice: by making each expert minimal, you can pack millions of them into the parameter budget. Expressivity is then recovered by composing K experts together per token.

Dynamic Expert Assembly (DEA)

The DEA mechanism governs how atomic experts are selected and composed for each token. It has two phases: Retrieval and Assembly.

Storage: Global Parameter Matrices. To make retrieval efficient, the parameters of all N atomic experts are consolidated into two global matrices rather than stored individually:

$W = [w_0^{\text{in}}, ..., w_{N-1}^{\text{in}}]^\top \in \mathbb{R}^{N \times d}$

$V = [w_0^{\text{out}}, ..., w_{N-1}^{\text{out}}]^\top \in \mathbb{R}^{N \times d}$

where W stacks all input vectors as rows (row i is w_i^in), and V stacks all output vectors as rows (row i is w_i^out).

What it computes: this is a data layout decision, not a computation. It takes the conceptual model of N separate experts and materializes them as two matrices—one for input projections, one for output projections. The key property is that retrieving the parameters for any subset of experts now reduces to row selection from these two matrices, which can be implemented as a gather operation.

Why this form: storing experts in dense matrices enables two critical optimizations. First, when multiple tokens activate the same expert, its parameters can be loaded once and reused—the matrix format makes this natural. Second, when an expert-centric scheduler groups tokens by expert, the corresponding rows can be extracted as contiguous blocks, enabling coalesced memory reads. If experts were stored as separate objects, each retrieval would be an independent random access with no opportunity for batching or coalescing.

Phase 1: Retrieval. For a token x, the router identifies the top-K expert indices I_x = (I_0, ..., I_{K-1}). The DEA mechanism retrieves the corresponding parameters by gathering rows from the global matrices:

$w_x = W[I_x] \in \mathbb{R}^{K \times d}$

$v_x = V[I_x] \in \mathbb{R}^{K \times d}$

where W[I_x] denotes selecting rows I_0, ..., I_{K-1} from W (producing a K × d matrix w_x of input vectors for the selected experts), and similarly for v_x. The routing weights are likewise gathered into a vector g_x = [g_{I_0}, ..., g_{I_{K-1}}] ∈ ℝ^K.

What it computes: this takes the top-K expert indices and performs a physical memory operation—gathering specific rows from the two global parameter matrices. The output is two compact K × d matrices (w_x and v_x) that contain only the parameters relevant to this token, plus a length-K weight vector g_x.

Why this form: gathering into dense K × d blocks is the key enabler for the Assembly step. By consolidating the selected experts' parameters into small, dense matrices, the subsequent computation can be expressed as matrix multiplications (which GPUs execute efficiently) rather than a loop over individual experts (which would be sequential and memory-bound). This is the first half of the co-design: the routing and retrieval produce data structures that are amenable to dense computation.

Phase 2: Assembly. The retrieved parameters and routing weights are composed into a single fused computation:

$y = (g_x \odot \sigma(x w_x^\top)) v_x + \text{MLP}(x)$

where x w_x^⊤ computes the dot products between token x (1 × d) and all K input vectors (w_x is K × d, so w_x^⊤ is d × K, and the product is 1 × K), producing K scalar relevance scores—one per selected expert. σ(·) applies SiGLU activation element-wise to these K scores, producing K gating values. g_x ⊙ σ(x w_x^⊤) performs element-wise multiplication between the K routing weights and the K gated scores, producing K modulated weights. Finally, (g_x ⊙ σ(x w_x^⊤)) v_x is a weighted sum of the K output vectors v_x, where each output vector is scaled by its corresponding modulated weight.

What it computes: this single equation performs the entire routed computation for one token. In detail: (1) compute how relevant each selected expert is to the token (dot products → K scalars); (2) gate these scalars through SiGLU (producing K values between 0 and ~1); (3) modulate by the routing weights (element-wise product, producing the final K weights); (4) compute a weighted sum of the experts' output vectors (matrix-vector product, producing a d-dimensional output); (5) add the shared MLP output. The result y is the combined output of the shared dense pathway and the routed fine-grained pathway.

Why this form: this formulation is mathematically equivalent to evaluating each atomic expert individually and summing their outputs (as in Eq. 3), but it restructures the computation to be a single matrix operation rather than K sequential vector operations. Specifically, σ(x w_x^⊤) is a batched activation over all K experts simultaneously; (g_x ⊙ σ(x w_x^⊤)) v_x is a single (1 × K) × (K × d) matrix multiplication. On a GPU, this batched formulation saturates the Tensor Cores much better than K individual vector-vector products, because the matrix dimensions are large enough (K is typically 512–4096) to amortize kernel launch overhead and achieve high arithmetic intensity.

The paper emphasizes that "this formulation demonstrates how DEA effectively constructs a unique, powerful assembled expert for each token by composing simple, reusable atomic experts." The "assembled expert" is not a physical entity—it is the conceptual result of weighting and summing K atomic experts. Each token gets a different assembled expert because different tokens route to different subsets of atomic experts. This is the extreme of parameter efficiency: the total parameter pool is massive (millions of atomic experts), but each token only pays computation for K of them, and every activated parameter is specifically relevant to that token (no redundant activation as in coarse-grained MoEs, per Figure 1a).

The role of the shared dense MLP. The + MLP(x) term in Eq. 7 is critical. The paper positions the shared dense MLP as providing "general semantic reasoning and stable capacity that is independent of routing." Without this term, the model would rely entirely on the routed atomic experts for all FFN computation. The shared MLP serves as a fallback: when routing is uncertain (as during early training, or for common linguistic patterns that don't require specialized knowledge), the dense pathway provides a reliable baseline signal. It also contributes training stability, similar to the shared experts in DeepSeekMoE (Dai et al., 2024).

The paper's ablation (Table 2) quantifies this: removing the shared dense MLP drops reasoning performance to 0.79× relative to the full model and increases perplexity by 1.2×, despite slightly improving latency (0.86×) and memory (0.98×). This confirms that the shared MLP is not redundant—it handles knowledge that the fine-grained routed branch cannot efficiently capture.

Cartesian Product Router (Eqs. 8–11)

The Cartesian Product Router solves the problem of how to select K experts from a pool of N (potentially millions) without the router itself becoming the computational bottleneck. A standard top-K router computes scores for all N experts via a projection G(x) = x W_g, where W_g ∈ ℝ^{d × N}. When N is in the millions, this projection costs O(Nd) FLOPs and requires storing O(Nd) parameters—which can dominate both inference latency and memory, defeating the purpose of using fine-grained experts.

Intuition: Decomposing the Index Space

The key insight is that a 1D expert index of size N can be mapped to a 2D coordinate (i, j) on a Nr × Nc grid where N = Nr × Nc. Instead of predicting one distribution over N items, the router predicts two independent distributions—one over Nr rows and one over Nc columns—and composes them to score any expert on the grid. This is analogous to product-key indexing from PKM (Lample et al., 2019), but adapted for OmniMoE's atomic expert architecture.

The modeling assumption is that the joint probability of selecting expert (i, j) given token x can be approximated by the product of two marginal probabilities:

$p(i, j | x) \approx p_r(i | x) \cdot p_c(j | x)$

where p(i, j | x) is the probability of selecting the expert at row i, column j; p_r(i | x) is the marginal probability of selecting any expert in row i; and p_c(j | x) is the marginal probability of selecting any expert in column j.

What it computes: this is a modeling assumption, not a computation. It states that the router treats row and column selection as statistically independent given the token—meaning the relevance of a row to a token does not depend on which column is selected, and vice versa.

Why this form: the independence assumption is what enables the factorization that reduces complexity. If row and column selection were not independent, you would need to model the full N-way joint distribution, which would defeat the purpose. The assumption is a modeling choice that trades some routing precision (the router cannot capture correlations between row and column preferences) for massive computational savings. The empirical results (Table 1) suggest that this tradeoff is favorable: the factorized router achieves superior performance to a standard dense router (which collapsed to only 4% expert usage in the ablation), indicating that the independence assumption is a useful regularizer that prevents the router from over-concentrating on a few experts.

Implicit Scoring via Factorized Projections

Instead of one large projection matrix W_g ∈ ℝ^{d × N}, the Cartesian Product Router uses two small projection matrices:

$W_r \in \mathbb{R}^{d \times N_r}, \quad W_c \in \mathbb{R}^{d \times N_c}$

For an input token x, the row and column logits are computed as:

$s_r = x W_r, \quad s_c = x W_c$

where s_r ∈ ℝ^{N_r} is the vector of raw scores for each row, and s_c ∈ ℝ^{N_c} is the vector of raw scores for each column.

What it computes: two small matrix-vector products. x W_r computes the dot product between the token and each of the N_r row embeddings in W_r (producing N_r scalar scores), and similarly for columns. The total cost of these two projections is O(d(N_r + N_c)) = O(d√N) when the grid is square (N_r ≈ N_c ≈ √N), compared to O(dN) for the standard router.

Why this form: the decomposition reduces both computation and parameter storage by a factor of √N/2 (derived in Appendix A, Eq. iv). For N = 10^6 (one million experts), this is a 500× reduction in router projection FLOPs. The reduction in parameter storage means the router's weight matrices can fit in on-chip memory (SRAM/registers) rather than requiring repeated HBM accesses, further reducing latency.

The log-probabilities for each subspace are obtained via LogSoftmax:

$p_r = \text{LogSoftmax}(s_r), \quad p_c = \text{LogSoftmax}(s_c)$

where LogSoftmax normalizes the scores into log-probabilities (each vector sums to 1 in probability space, or equivalently, log-sum-exp of the entries equals 0). p_r[i] = log P(row = i | x) and p_c[j] = log P(col = j | x).

What it computes: LogSoftmax applies the softmax function in log space: LogSoftmax(s)_i = s_i - log(Σ_j exp(s_j)). The output is a vector of log-probabilities—negative numbers that sum to 0 after exponentiating and normalizing. Using log-space is critical for numerical stability when dealing with millions of experts, because the product of many small probabilities would underflow to zero in standard floating-point.

Why this form: operating in log-space transforms the product p_r(i) × p_c(j) into a sum p_r[i] + p_c[j]. This is the mathematical foundation that makes the factorized scoring efficient: you can score any expert on the grid by adding two numbers (one from p_r, one from p_c), without ever materializing the full N-element score vector. For top-K selection, this means you can compute candidate scores on-the-fly during the search, rather than pre-computing and storing all N scores.

The score for an expert at coordinate (i, j) is the sum of the corresponding log-probabilities:

$S_{ij} = p_r[i] + p_c[j]$

where S ∈ ℝ^{N_r × N_c} is the implicit score matrix—it is never materialized, but its entries can be computed on-demand from the two small vectors p_r and p_c.

What it computes: for any pair (i, j), this sums the row log-probability and column log-probability to produce a log-probability for the joint event. Since the log-probabilities are negative (or zero), the sum is also negative (or zero). Higher (less negative) scores indicate more relevant experts.

Why this form: the additive structure in log-space is what enables the efficient parallel top-K selection described next. Because any entry of S can be computed from two independent vectors, the grid can be partitioned into tiles, and each tile can be processed by a separate GPU thread block, computing scores on-the-fly by reading from the same p_r and p_c vectors (which fit in registers or shared memory).

Parallel Top-K Selection

The router must identify the K largest entries in the implicit N_r × N_c score matrix, where K is typically 512–4096 (much larger than the 2–8 experts in coarse-grained MoEs). The paper describes a tiled GPU algorithm:

Partition the grid into blocks of size B_sel (e.g., 4096). Each block is assigned to a parallel GPU thread block.
Compute local scores on-the-fly: within each block, threads compute S_{ij} = p_r[i] + p_c[j] using the globally-available p_r and p_c vectors.
Extract local top-K: each block performs iterative max-reduction to find the K largest scores within its tile, along with their 2D indices.
Merge: a lightweight reduction merges the local top-K candidates from all blocks to obtain the global top-K expert indices I_x.
Compute routing weights: the corresponding top-K scores are extracted from S (computed on-the-fly) and normalized via Softmax to produce g_x.

Why this form: the tiled approach avoids the O(N) global memory write that would be needed to materialize all N scores. Instead, each block keeps its local top-K in registers and only communicates K elements during the merge step. The paper's complexity analysis (Appendix A, Eq. v) shows that the total time complexity per token is O(N·K + (N/B_sel)K²). The O(N·K) term dominates for large N but is highly parallel (split across many thread blocks). The key benefit is eliminating the O(N) global memory I/O bottleneck—by keeping intermediate scores in registers and shared memory, the top-K selection becomes compute-bound rather than memory-bound.

Complexity Analysis and Practical Impact

Appendix A provides the formal complexity derivation. For a standard router with N experts and hidden dimension d:

$C_{\text{std}} = 2 \cdot d \cdot N$

For the Cartesian Product Router with Nr ≈ Nc ≈ √N:

$C_{\text{cart}} \approx 4 \cdot d \cdot \sqrt{N}$

The reduction factor is √N / 2. For N = 10^6, this is a 500× reduction in router FLOPs. The parameter storage reduction is similarly √N/2—the router weights shrink from d·N to d·(Nr + Nc) ≈ 2d√N.

The paper's ablation (Table 2) demonstrates the practical impact: replacing the Cartesian Product Router with a standard dense routing projection causes a 30.6× increase in latency and a 337.5× increase in memory (from materializing the full W_g matrix), and crucially, also degrades model quality (1.4× perplexity, 0.66× knowledge performance). The quality degradation occurs because the standard router fails to learn distributed specializations over the massive expert space—expert usage collapses to 4% (only 4% of experts are ever activated), with unevenness increasing from 0.24 to 0.77. The Cartesian Product Router, by factorizing the routing into independent row/column decisions, acts as an implicit regularizer that encourages broader expert utilization.

Expert-Centric Scheduling (Eqs. 12–14)

With routing no longer the bottleneck (solved by the Cartesian Product Router), the remaining challenge is hardware efficiency: even with perfect routing, the fine-grained expert activations induce scattered memory accesses that make execution memory-bound. Expert-Centric Scheduling solves this by inverting the execution loop—processing expert groups rather than individual tokens.

The problem: in a standard token-centric paradigm, each token in a batch independently fetches its selected experts' parameters. For a batch of L tokens, each activating K experts, this produces L × K individual gather operations—random reads from non-contiguous memory addresses. The GPU's memory controller cannot coalesce these random reads into efficient bursts, so effective bandwidth drops dramatically (often to 10–30% of peak HBM bandwidth), and the compute units stall waiting for data.

Task Collection and Active Expert Compression

For a batch of L tokens X = {x_l}_{l=0}^{L-1} with top-K routing, the router produces, for each token l, an expert index list I_l = (I_{l,0}, ..., I_{l,K-1}) and gating weights g_l = (g_{l,0}, ..., g_{l,K-1}). The first step is to flatten all routing decisions into a flat list of tasks:

$\tilde{T} = \{(x_l, I_{l,k}, g_{l,k}) \mid l \in [0, L), k \in [0, K)\}$

where M = L × K is the total number of tasks (one task = one token–expert pair). Each task is a tuple containing the token representation x_l, the expert index I_{l,k}, and the routing weight g_{l,k}.

What it computes: this flattens the per-token routing decisions into a single list of M computation tasks. The output is an unordered collection of (token, expert, weight) triples—each represents one atomic expert execution that must happen.

Why this form: flattening is necessary to enable reordering. As long as tasks are organized per-token, there is no way to exploit commonality in which experts are activated across tokens. By flattening, you can re-sort the tasks to group by expert—which is exactly what Expert-Centric Scheduling does.

Next, the scheduler identifies the set of unique active experts across the entire batch:

$E_{\text{active}} = \bigcup_{x \in X} I_x$

sorted by global expert ID. It then partitions this ordered list into contiguous groups of size B:

$q_\tau = \lfloor \tau / B \rfloor$

where q_τ is the group ID assigned to the τ-th expert in the sorted E_active. The number of execution groups is N_groups = ⌈|E_active| / B⌉.

What it computes: this takes the set of experts that are activated by at least one token in the batch, sorts them by ID, and divides them into groups of B contiguous experts. Experts with nearby IDs are placed in the same group. The number of groups is determined purely by the active sparsity—if a batch activates 1,024 unique experts and B = 128, there will be exactly 8 groups, regardless of the total number of tokens or total expert pool size.

Why this form: grouping by expert ID exploits spatial locality—experts with adjacent IDs are stored at adjacent memory addresses in the global W and V matrices. By grouping them, the parameters for all B experts in a group can be loaded as a single contiguous block (a B × d tile from each matrix), which maximizes memory bandwidth utilization. The choice of group size B is a tuning parameter: larger B means fewer groups (reducing kernel launch overhead) but also means each group covers a wider range of expert IDs (potentially including experts not activated by any token in the group, introducing waste). The paper does not specify the exact value of B used in experiments, but the concept is that B is chosen to be large enough to amortize launch overhead while being small enough that groups are mostly fully-utilized.

Hierarchical Sorting

The flattened task list T̃ is reorganized by performing a hierarchical sort with two keys:

$T = \text{Sort}(\tilde{T}, \text{keys} = (q, l))$

where the primary key is group ID q (tasks targeting experts in the same group are clustered together), and the secondary key is token ID l (within each group, tasks are ordered by token ID).

What it computes: a radix sort (chosen for efficiency on GPU) that reorders the M tasks. After sorting, all tasks that target experts in group 0 appear first (ordered by token ID), then all tasks targeting group 1 (ordered by token ID), and so on. The output T is the sorted task list that will drive execution.

Why this form: the dual sorting strategy achieves two levels of hardware efficiency:

Inter-Group Locality (primary key): by clustering tasks by expert group, the parameters for each group's experts are loaded exactly once from HBM and reused across all tokens in that group. This eliminates redundant reads when multiple tokens share the same experts.
Intra-Group Coalescing (secondary key): within each group, processing tasks in increasing token ID order means that input token reads (x_l) follow a sequential access pattern (since tokens are stored contiguously in memory), enabling perfect memory coalescing. Similarly, output scatter-add operations write to contiguous memory regions, reducing write-commit overhead.

The paper emphasizes that this sorting step has minimal overhead: "empirically, scheduling occupies <5% of total latency, well-amortized by the speedup in the GEMM phase." This is because radix sort on GPU is highly optimized and the number of elements M = L × K is modest relative to the savings in the matrix multiplication phase.

Grouped GEMM Execution

For each group q, the scheduler gathers the expert parameters for the B experts in that group:

$W_q, V_q \in \mathbb{R}^{B \times d}$

where W_q contains the input vectors (rows from the global W matrix) for experts in group q, and V_q contains their output vectors. These are dense, contiguous blocks loaded once per group.

Concurrently, the scheduler stacks the input tokens and gating weights for all tasks assigned to this group:

$X_q \in \mathbb{R}^{T_q \times d}, \quad G_q \in \mathbb{R}^{T_q \times B}$

where T_q is the number of tasks in group q (i.e., how many token–expert pairs target experts in this group), X_q contains the token representations for those tasks (each token may appear multiple times if it activates multiple experts in the same group), and G_q contains the corresponding routing weights (a sparse matrix where row t has non-zero entries only at the columns corresponding to the experts activated by that token).

The entire computation for the group is then performed by a single fused Grouped GEMM operation:

$O_q = (G_q \odot \sigma(X_q W_q^\top)) V_q$

where X_q W_q^⊤ is a T_q × B matrix of relevance scores (dot products between each token and each expert in the group); σ(·) applies SiGLU activation element-wise; G_q ⊙ σ(X_q W_q^⊤) element-wise modulates the activation output by the routing weights; and the final multiplication with V_q (a T_q × B matrix times a B × d matrix) produces a T_q × d output block O_q.

What it computes: this single fused kernel performs all expert computations for one group. For each task (token, expert) pair, it: (1) computes the compatibility score between the token and the expert (dot product); (2) gates through SiGLU; (3) weights by the routing score; (4) multiplies by the expert's output vector; (5) accumulates with other activations for the same token. The output O_q contains one row per task—the partial contribution of that expert to that token's routed output.

Why this form: this is the payoff of the entire co-design. Instead of T_q individual vector-vector operations (each reading one expert vector from a random HBM address), the computation is now a single dense matrix multiplication (X_q W_q^⊤) followed by a second dense matrix multiplication ((G_q ⊙ σ(...)) V_q). Both operations:

Are dense enough (dimensions T_q × B and T_q × d) to saturate Tensor Cores.
Access memory in contiguous, predictable patterns (the expert parameters W_q, V_q are loaded once as dense tiles and reused across all T_q tasks; the input tokens X_q are read sequentially).
Require no per-expert kernel launch—the entire group is processed in one kernel invocation.

After the grouped GEMM, the per-task outputs in O_q are scattered back to the corresponding token positions via scatter-add (accumulating contributions from different groups for tokens that activate experts across multiple groups). This preserves the semantics of Eq. 7—each token's final output is the sum of contributions from all its activated experts across all groups, plus the shared MLP output.

Why Expert-Centric Scheduling is Efficient (Figure 3 and Appendix A Analysis)

The paper's Figure 3 provides a visual comparison. In the conventional token-centric paradigm (Figure 3a), token 0 fetches from experts 0, 3, 5; token 1 fetches from experts 1, 4, 6; and so on. Each fetch is an independent, non-contiguous memory access. The SMs on the GPU (rightmost block) process these as individual vector-vector operations, which underutilize the SMs (shown as sparse occupancy) and impose high load overhead from the memory subsystem.

In OmniMoE's Expert-Centric Scheduling (Figure 3b), the active experts are first compressed into dense groups (e.g., experts 0–3 form Group 1, experts 4–7 form Group 2). Tasks are reordered so that all token-expert pairs targeting Group 1 are processed together, then Group 2. Within each group, the expert parameters are loaded as a dense block, and the computation is executed as a Grouped GEMM. The SMs are fully occupied because the matrix dimensions are large enough to keep all compute units busy.

The theoretical memory traffic analysis in Appendix A quantifies the benefit. In token-centric execution, the total memory traffic for loading expert parameters is:

$D_{\text{token}} \propto 2d \cdot L \cdot K$

because each of the L × K token-expert pairs independently loads 2d parameters (input and output vectors).

In Expert-Centric Scheduling, each unique expert in E_active is loaded exactly once:

$D_{\text{expert}} \propto 2d \cdot |E_{\text{active}}|$

The reduction ratio is η = (L × K) / |E_active|. When experts are fine-grained and K is large, many tokens share the same experts, so |E_active| ≪ L × K, and η ≫ 1. For example, if a batch of 4,096 tokens each activates 2,048 experts, but only 10,000 unique experts are activated (because many tokens share expert selections), then η = (4096 × 2048) / 10000 ≈ 839—an 839× reduction in expert parameter I/O.

The paper's ablation (Table 2) confirms the practical impact: reverting to token-centric execution (removing Expert-Centric Scheduling) preserves model quality (all quality metrics remain at 1.0×) but increases latency by 24.8× and memory by 417.7×. The memory increase is particularly dramatic because the token-centric baseline must materialize full routing tensors that Expert-Centric Scheduling avoids by computing scores on-the-fly within the grouped GEMM kernel.

Training and Hyperparameter Configuration

While the architectural innovations are the core contribution, the paper provides detailed training configurations that are essential for reproducibility and for isolating architectural effects from training artifacts.

Training Data. All models are pre-trained on the SmolLMCorpus (Ben Allal et al., 2024), "a high-quality corpus of 40 billion tokens spanning Web, Textbook, Code, and Math domains." The NeoX tokenizer (Black et al., 2022) is used with a vocabulary size of 128,256 tokens.

Optimizer and Scheduler. Training uses the AdamW optimizer (Loshchilov & Hutter, 2017) with the WSD learning rate scheduler (Hägele et al., 2024). Hyperparameters follow "optimal scaling laws (Li et al., 2025) and Chinchilla compute-optimality protocols (Hoffmann et al., 2022)." The paper does not specify the exact β₁, β₂, ε, or weight decay values in the main text, but Appendix B references these scaling laws papers for the hyperparameter recipes.

Model Configurations (Appendix B, Table B). The paper trains MoE families at four scales (activated parameters: 80M, 200M, 680M, 1.7B) alongside matched dense baselines:

Activation 80M: 12 layers, d_model = 768, 13,500 training steps with 0.128M batch tokens, peak learning rate 3e-3, tied embeddings.
Activation 200M: 16 layers, d_model = 1024, 20,800 training steps with 0.192M batch tokens, peak learning rate 2e-3, tied embeddings.
Activation 680M: 24 layers, d_model = 1536, 35,000 training steps with 0.392M batch tokens, peak learning rate 1e-3, tied embeddings.
Activation 1.7B: 28 layers, d_model = 2048, 40,000 training steps with 1M batch tokens, peak learning rate 1e-3, untied embeddings.

The 6.4B-A1.7B model (the largest, used for downstream evaluation) has 6.4B total parameters but only 1.7B active per token. The Transformer backbone (depth, width, attention configuration) is kept identical across all methods at each scale, varying only the FFN module. This isolates the architectural effect: any performance differences between Dense, Gshard, DeepSeekMoE, PKM, PEER, and OmniMoE are attributable to the FFN design, not to differences in attention or layer count.

Speed and Memory Benchmarking (Appendix B, Table A). For efficiency measurements, a 200M-parameter backbone is used with d = 1024. The activated parameter budget is swept from 3M to 28M (controlling the total computation per token), and the number of activated tokens is swept from 1K to 16K. Within each sweep, one variable is varied while the other is fixed at its minimum. This allows measuring how latency and memory scale independently with model capacity and batch size.

Fair Comparison Methodology. The paper emphasizes that "we prioritize architectural comparison via controlled pre-training from scratch rather than comparing against off-the-shelf checkpoints." This is important because off-the-shelf checkpoints are trained on different data, with different hyperparameters, and for different durations, making it impossible to attribute performance differences to architecture alone. By pre-training all models from scratch on identical data with identical recipes, the paper establishes a controlled experimental framework.

Kernel Implementations. For coarse-grained baselines (Gshard, DeepSeekMoE), the paper uses "the best-performance kernels released by NVIDIA in CuTile." For fine-grained baselines (PKM, PEER), "highly-optimized Triton fused kernels" are used. OmniMoE's Expert-Centric Scheduling is implemented "using Triton to maximize hardware utilization." The use of state-of-the-art kernels for each baseline ensures the comparison is fair—OmniMoE is not being compared against inefficient implementations.

Distributed Training Efficiency (Appendix C). The paper also verifies that OmniMoE scales efficiently in distributed training settings with Expert Parallelism (EP). The key finding is a "saturation effect": when the number of experts N exceeds the total number of activated experts in a batch (n_tokens × K = 16,384), the communication overhead stabilizes at approximately 80MB and does not grow with increasing N. This means OmniMoE can scale to millions of experts with constant communication cost—the communication bottleneck that plagues many MoE systems is decoupled from model capacity. Communication volume scales linearly with sequence length (as expected, since tokens are distributed across GPUs), but even at 128K tokens, the estimated latency on a 64-GPU cluster is only 15 ms, indicating that communication is not a training throughput bottleneck.

Summary of Design Choices and Their Justifications

Atomic Experts as minimal routable units (2d parameters per expert): maximizes the number of experts that can fit in a given parameter budget, enabling precise, token-specific routing. The minimal design is compensated by composing K experts per token through DEA.
Dynamic Expert Assembly for parameter retrieval and composition: consolidates expert parameters into global matrices W, V for efficient batch retrieval, and formulates the routed computation as a single fused matrix operation ((g_x ⊙ σ(x w_x^⊤)) v_x) that maps well to GPU Tensor Cores.
Shared Dense MLP as a universally-activated pathway: provides general semantic reasoning capacity that complements the fine-grained routed branch, preventing the routed branch from having to capture common linguistic patterns (which would dilute specialization). Empirically necessary for reasoning performance (0.79× regression when removed).
Cartesian Product Router decomposing 1D index space into 2D grid: reduces routing projection complexity from O(Nd) to O(√N d), making million-scale expert routing computationally feasible. The independence assumption between row and column acts as an implicit regularizer that prevents expert collapse (4% usage with standard router vs. 100% with Cartesian Product Router).
Tiled parallel top-K selection avoiding full score materialization: eliminates the O(N) global memory I/O bottleneck for top-K selection by computing scores on-the-fly in register/shared memory, with only a lightweight merge step for global top-K.
Expert-Centric Scheduling inverting execution order: converts scattered per-token expert fetches into batched, contiguous reads that saturate memory bandwidth and enable Grouped GEMM execution. The hierarchical sort (by expert group, then by token ID) achieves both parameter reuse and input/output coalescing.
Two-fold co-design (routing + scheduling): the paper's central claim—that neither algorithmic innovations (atomic experts, Cartesian product routing) nor system innovations (Expert-Centric Scheduling) alone would be sufficient—is supported by the ablation: removing either causes massive regression (30.6× latency for routing, 24.8× for scheduling), confirming that the two must be jointly designed.

4. Key Insights and Innovations

Innovation 1: Diagnosing the Granularity-Efficiency Trade-off as the Fundamental Structural Tension in MoE Design, Not a Mere Implementation Detail

The paper's deepest conceptual contribution is not any single mechanism but the framing of the MoE design space as fundamentally partitioned by an inescapable structural tension between routing precision (which improves with finer granularity) and hardware execution efficiency (which degrades with finer granularity). This is presented visually in Figure 1 and argued explicitly in Section 1, but its significance goes beyond an observation—it functions as a diagnostic lens that retrospectively explains why the field has bifurcated into two disconnected research trajectories that talk past each other.

What makes this framing distinctive is that it recasts what looks like an engineering trade-off (memory access patterns, kernel design) as a first-order architectural constraint that determines which regimes are viable at scale. Before this paper, the coarse-grained MoE community (GShard, Switch Transformers, DeepSeekMoE, Mixtral) optimized for hardware efficiency using large FFN experts, while the fine-grained community (PKM, PEER) optimized for parameter efficiency using embedding-sized experts, with neither side directly addressing why the other's approach was incompatible with their own objectives. The implicit assumption in both lines of work was that the trade-off was negotiable—that better kernels could make fine-grained experts fast enough, or that better routing could make coarse-grained experts precise enough.

This paper's diagnosis refutes that assumption. It shows, through the evidence in Figure 4, that fine-grained PEER suffers 73 ms latency at 4,096 tokens (10.9× slower than OmniMoE) because scattered memory access is a structural consequence of fine granularity, not a kernel optimization issue—no amount of kernel tuning can coalesce fundamentally random access patterns. Conversely, it shows that DeepSeekMoE's coarse-grained experts inevitably activate redundant parameters (Figure 1a) because the expert is a monolithic FFN block—this is a structural consequence of coarse granularity, not a routing failure. The paper's intellectual move is to declare that neither regime can fully capture what the other offers, and that bridging the gap requires rethinking the architecture from the ground up rather than optimizing within either paradigm.

This is more than taxonomy. It functions as a design principle: any architecture that claims to reconcile parameter efficiency with hardware efficiency must simultaneously solve (i) how to make experts fine enough to route precisely, (ii) how to route over a massive expert pool without the router itself dominating, and (iii) how to execute the resulting scattered computation without becoming memory-bound. The paper's Figure 1 communicates this principle visually—the three panels (coarse-grained, fine-grained, OmniMoE) are not merely descriptive but serve as a diagnostic framework for evaluating any MoE architecture along two axes: activation precision and memory access pattern. The significance is that this framework provides a language for comparing architectures that were previously incommensurable—you can now ask of any MoE design: "Where does it fall on the activation precision axis? Where on the memory access axis? Are these points compatible with hardware constraints at the target scale?"

The evidence supporting this diagnostic contribution is distributed across the paper: Figure 3 (right) shows that coarse-grained DeepSeekMoE can be slower than fine-grained PEER at scale due to packing/alignment overhead—a non-obvious finding that demonstrates the granularity-efficiency trade-off is not monotonic or one-sided; Table 2 (ablation) shows that removing Expert-Centric Scheduling from OmniMoE causes a 24.8× latency regression without affecting model quality at all—proving that the inefficiency is purely a system-level consequence of fine granularity, not an algorithmic limitation. Together, these constitute compelling evidence that the trade-off is real, structural, and cannot be optimized away within existing paradigms.

This is a fundamental contribution because it reframes the problem space. It is not incremental—prior work treated granularity as a continuous knob to tune, whereas this paper argues that the knob interacts discontinuously with hardware constraints, creating qualitatively different execution regimes that require qualitatively different solutions.

Innovation 2: The Inversion of Execution Order (Expert-Centric Scheduling) as a First-Class Architectural Primitive, Not a System Afterthought

Expert-Centric Scheduling (Section 2.3, Figure 3) is the paper's most distinctive system-level contribution, but its intellectual significance extends beyond the specific mechanism. The deeper insight is that execution order is an architectural degree of freedom that can be co-designed with the routing algorithm to reshape memory access patterns at the hardware level—and that this co-design can convert a memory-bound workload into a compute-bound one without changing the underlying computation.

Prior work on MoE system optimization (DeepSpeed-MoE, FastMoE, MegaBlocks, SonicMoE) operated within the token-centric paradigm: the fundamental loop iterates over tokens, and for each token, the system fetches its assigned experts. These works improve efficiency through kernel fusion (combining operations), padding reduction (MegaBlocks' block-sparse kernels), communication scheduling (DeepSpeed-MoE), and tile-aware batching (SonicMoE), but they all preserve the token-centric execution order. The implicit assumption is that the token-expert assignment produced by the router must drive execution—that the execution order is determined by the routing, not chosen independently.

OmniMoE's key conceptual move is to break this coupling. The Cartesian Product Router produces per-token expert assignments just as any router would, but Expert-Centric Scheduling then discards the token-centric organization and re-sorts the resulting tasks by expert group. This is not an optimization of token-centric execution—it is an entirely different execution paradigm where the loop iterates over experts (or expert groups) rather than tokens. The paper calls this "inverting the execution order" (Section 2.3), and the inversion is what enables the transformation from scattered gather operations to dense Grouped GEMMs.

What makes this intellectually distinctive is that it treats the scheduling decision as an architectural choice on par with the router design or expert parameterization, not as a downstream implementation detail. The paper argues, implicitly through its co-design framing, that the router and the scheduler must be designed together because the router determines which experts are activated (affecting the sparsity pattern that the scheduler must handle), while the scheduler determines how those activations map to hardware (affecting what routing patterns are efficient). A router that produced expert assignments poorly suited to grouping (e.g., highly fragmented across the expert ID space) would undermine the scheduler; a scheduler that couldn't handle the sparsity patterns produced by the router would bottleneck the system. The Cartesian Product Router and Expert-Centric Scheduling are designed as a matched pair: the router factorizes the expert space into a 2D grid that naturally supports spatial grouping (experts with nearby IDs are stored contiguously), and the scheduler exploits this spatial locality by grouping nearby expert IDs into execution batches.

The evidence for this as a conceptual contribution is in Figure 3 and Table 2. Figure 3 visually communicates the paradigm shift: the left panel (token-centric) shows high load overhead and sparse SM occupancy (the hallmark of memory-bound execution), while the right panel (expert-centric) shows full SM occupancy and coalesced memory access (the hallmark of compute-bound execution). Table 2 quantifies the consequence: removing Expert-Centric Scheduling (reverting to token-centric) preserves model quality perfectly (1.0× on all quality metrics) but increases latency 24.8× and memory 417.7×. This ablation is conceptually significant because it proves that the scheduling innovation is purely about efficiency, not about enabling new capabilities—the model's mathematical computation is identical in both execution orders; the 24.8× speedup comes entirely from better utilization of the same hardware. This is the strongest possible evidence that execution order is a genuine architectural primitive with enormous leverage.

This is a fundamental contribution to the systems-for-ML literature. It establishes a new design axis—execution order inversion—that generalizes beyond OmniMoE: any sparse computation where multiple queries access overlapping subsets of a shared parameter pool can benefit from reordering execution from query-centric to parameter-centric. The paper does not claim this generalization explicitly, but it is the natural implication of the result.

Innovation 3: Atomic Experts with Dynamic Assembly as a Resolution of the Expressivity-Granularity Tension in Fine-Grained MoEs

The paper addresses a subtle but critical limitation of prior fine-grained MoEs (PKM, PEER): that making experts extremely small forces a trade-off between routing precision (which improves with finer granularity) and per-expert expressivity (which degrades as experts shrink). PEER's experts are static embedding vectors—when activated, they contribute through linear vector aggregation without any token-dependent nonlinear transformation. This "strips away the token-dependent nonlinear transformations (e.g., MLP projections) essential for modeling complex linguistic dependencies" (Section 1). The implicit assumption in PEER and PKM is that this expressivity loss is the price you pay for fine-grained routing—you trade per-expert depth for breadth of the expert pool.

OmniMoE's conceptual contribution is to break this assumed trade-off through the mechanism of Dynamic Expert Assembly (DEA, Section 2.1). By parameterizing each atomic expert as a pair of vectors (input and output) with a gated nonlinear activation (σ(x · w_i^in) · w_i^out), OmniMoE gives each expert a minimal but genuine token-conditional computation—the expert's output depends on the token through the dot product in the gate. This is structurally different from PEER's static embedding retrieval, where the expert's contribution is independent of the token (the same embedding vector is activated regardless of the token's content, with only the routing weight varying).

But the deeper insight is that expressivity is recovered at the composition level, not the individual expert level. A single atomic expert has limited capacity (only 2d parameters), but DEA composes K of them (where K is large—512 to 4,096 in OmniMoE's configurations, per Table A), producing what the paper calls an "assembled expert" that acts like a full token-conditioned nonlinear transformation. The assembly equation (g_x ⊙ σ(x w_x^⊤)) v_x is mathematically a two-layer network: the first layer computes K gated relevance scores (dot products with input vectors, passed through SiGLU, modulated by routing weights), and the second layer computes a weighted sum of output vectors. This has substantially more representational capacity than the linear aggregation in PEER because (a) the gating is token-dependent (via the dot products), (b) SiGLU provides nonlinearity, and (c) composing K experts allows the assembled computation to span a K-dimensional subspace of the d-dimensional output space.

The evidence for this conceptual contribution is in the downstream results (Table 1): OmniMoE achieves +3.6 on ARC and +4.6 on HellaSwag compared to PEER, with the largest gains on reasoning-heavy benchmarks where nonlinear transformations are most important. The ablation (Table 2, effect of removing the shared dense MLP) provides complementary evidence: removing the shared MLP reduces reasoning performance to 0.79× relative to the full model, suggesting that even the assembled atomic experts cannot fully compensate for the loss of the dense pathway's nonlinear expressivity—the shared MLP and the routed atomic branch serve complementary roles, with the former providing deep nonlinear processing and the latter providing precise knowledge retrieval.

This is a fundamental contribution to the fine-grained MoE design space. It identifies a previously overlooked bottleneck—that extreme granularity can sacrifice per-expert expressivity to the point where routing precision gains are erased—and proposes a solution (dynamic composition of minimally-expressive units) that preserves both granularity and expressivity. The concept generalizes: any architecture that routes over a massive pool of lightweight units must ensure that the composition of activated units has sufficient representational capacity, not just that each unit is individually token-relevant.

Innovation 4: The Independence Assumption in Factorized Routing as an Implicit Regularizer That Prevents Expert Collapse at Scale

The Cartesian Product Router (Section 2.2) is typically understood as a complexity-reduction technique—it replaces an O(N) projection with two O(√N) projections, making million-scale expert routing feasible. But the paper's ablation (Table 2) reveals a deeper, more surprising property: the factorized router improves model quality compared to a standard dense router, not just matches it at lower cost. When the Cartesian Product Router is replaced with a standard dense routing projection, expert usage collapses from 100% to 4%, unevenness increases from 0.24 to 0.77, and perplexity degrades by 1.4×.

This result is conceptually significant because it shows that the independence assumption p(i,j | x) ≈ p_r(i | x) · p_c(j | x) is not just an approximation that sacrifices routing precision for efficiency—it functions as an architectural regularizer that prevents the router from over-concentrating on a few experts when the expert pool is massive. The standard dense router, with its full N-way classification, has the capacity to route all tokens to a small subset of experts—and empirically, it does exactly that, collapsing to 4% expert usage despite the auxiliary load-balancing loss. The Cartesian Product Router cannot do this because its routing decisions are constrained to be products of row and column preferences: to concentrate on a specific expert (i, j), the router would need to concentrate probability mass on row i AND column j simultaneously, which forces it to also activate other experts in row i and column j. This structural constraint naturally distributes probability mass across the grid, preventing collapse.

This insight reframes factorized routing from a "necessary approximation for scalability" to a "beneficial inductive bias for expert utilization." It provides a principled explanation for why PKM and PEER also use product-structured indexing beyond just efficiency—the factorization imposes a form of structured sparsity that is well-matched to the problem of routing over massive expert spaces. The finding that the independently-trained standard router fails to learn diverse specializations even with load-balancing losses (the paper notes "despite rigorous tuning of the standard auxiliary load-balancing loss") is strong evidence that the regularization effect of factorization is not trivially replicable by other means.

This is an incremental but important contribution. It does not propose a new routing mechanism beyond the product-key concept established in PKM, but it provides the first clear empirical evidence that the factorization is not merely an efficiency hack that happens to work—it fundamentally changes the learning dynamics of routing at scale in a way that standard routers cannot replicate even with careful regularization. For practitioners designing large-scale MoEs, this implies that factorized routing should be considered even in regimes where a standard router would be computationally feasible, because the inductive bias toward diverse expert utilization may be essential for training dynamics.

5. Experimental Analysis

Evaluation Methodology

Dataset. The paper uses the SmolLMCorpus (Ben Allal et al., 2024) for pre-training, a 40-billion-token corpus spanning Web, Textbook, Code, and Math domains. For downstream evaluation, seven benchmarks are used: MMLU (multitask knowledge), TriviaQA (factual recall), ARC (science reasoning), PIQA (physical commonsense), HellaSwag (commonsense inference), OBQA (open-book QA), and Winogrande (coreference resolution). All downstream evaluations are zero-shot using Hugging Face LightEval (Fourrier et al., 2023).
Base model(s). The paper trains MoE families from scratch at four scales: 280M-A80M, 800M-A200M, 2.7B-A680M, and 6.4B-A1.7B (where A denotes activated parameters per token), alongside matched Dense baselines. Downstream evaluation uses the 6.4B-A1.7B models. All models adopt Grouped Query Attention (Ainslie et al., 2023) with the Transformer backbone (depth, width, attention configuration) held identical across methods at each scale, varying only the FFN module. The paper argues this isolates architectural effects from confounding factors like training data differences in off-the-shelf checkpoints.
Metrics. Primary metrics are zero-shot accuracy on each downstream benchmark and their average (Table 1); validation perplexity for scaling-law experiments (Figure 5); inference latency in milliseconds and peak GPU memory in MiB for efficiency benchmarks (Figure 4); and expert usage (fraction of experts activated at least once) and unevenness (KL divergence from uniform distribution over expert retrieval frequency) for routing quality (Table 2). Latency is measured as strict end-to-end inference time including all scheduling/reordering overheads for OmniMoE and corresponding layout-transformation/alignment costs for baselines.
Baselines. Six FFN variants are compared: (i) Dense (standard MLP), (ii) Gshard (Lepikhin et al., 2021) with top-2 gating, (iii) DeepSeekMoE (Dai et al., 2024) with shared + routed experts, (iv) PKM (Lample et al., 2019) with product-key memory, (v) PEER (He, 2024) with million-scale embedding experts, and (vi) OmniMoE (the proposed method). For coarse-grained baselines (Gshard, DeepSeekMoE), the paper uses NVIDIA's CuTile kernels; for fine-grained baselines (PKM, PEER), highly-optimized Triton fused kernels; for OmniMoE, custom Expert-Centric Scheduling in Triton.
Generation budget / compute accounting. For scaling-law experiments, compute is measured in training FLOPs (Figure 5a). For efficiency benchmarks, the independent variables are activated parameter budget (number of unique parameters used in a single token's forward pass) and number of activated tokens (batch size), swept independently (Table A in Appendix B). The activated-parameter budget includes embeddings, attention weights, the shared dense FFN, router projections, and top-K active MoE experts, and is matched across baselines. Generation budget in the traditional sense (number of output tokens) is not applicable since these are pre-training efficiency and zero-shot evaluation experiments, not test-time scaling.
Cross-validation / statistical protocol. No formal cross-validation or statistical testing (e.g., confidence intervals, multiple random seeds) is reported for the downstream benchmark results in Table 1. The paper trains one model per configuration and evaluates once. For scaling-law experiments, models are trained at multiple scales (80M, 200M, 680M, 1.7B active parameters) and evaluated at consistent training FLOPs or activated parameter counts, providing internal replication across scales. The distributed training experiments (Appendix C) use 8-GPU and 64-GPU configurations but report latency/memory measurements without error bars or variance estimates.

Main Quantitative Results

Downstream Performance (Table 1)

The 6.4B-A1.7B OmniMoE model achieves an average zero-shot accuracy of 50.9% across seven benchmarks, the highest among all compared architectures. The closest competitors are DeepSeekMoE at 50.2% (+0.7 over OmniMoE's advantage) and Gshard at 49.2%. PEER achieves 48.9%, and PKM achieves 46.2%.

Breaking this down by benchmark reveals a differentiated pattern consistent with OmniMoE's hybrid design. On knowledge-intensive tasks, OmniMoE shows the largest advantages over coarse-grained baselines: on TriviaQA (factual recall), OmniMoE achieves 18.5% vs. DeepSeekMoE's 17.4% (+1.1) and PEER's 16.9% (+1.6); on OBQA (open-book QA), OmniMoE reaches 40.3% vs. DeepSeekMoE's 38.9% (+1.4) and PEER's 39.1% (+1.2). The paper attributes this to the fine-grained atomic experts enabling more precise retrieval of specialized, long-tail knowledge compared to coarse-grained FFN experts.

On reasoning-intensive tasks, OmniMoE's advantages over fine-grained baselines are most pronounced: on ARC (science reasoning), OmniMoE achieves 61.0% vs. PEER's 57.4% (+3.6) and PKM's 53.6% (+7.4); on HellaSwag (commonsense inference), OmniMoE reaches 60.9% vs. PEER's 56.3% (+4.6) and PKM's 52.7% (+8.2). The paper attributes these gains to the shared dense MLP providing the nonlinear transformations essential for multi-step reasoning, which purely fine-grained architectures like PEER lack due to their linear vector-aggregation experts.

Several baselines show task-specific weaknesses. PKM substantially underperforms on reasoning benchmarks (HellaSwag 52.7% vs. OmniMoE 60.9%, ARC 53.6% vs. 61.0%), consistent with its static embedding expert design providing minimal expressivity. Dense (45.9% average) trails all MoE variants, confirming the parameter-efficiency benefits of sparse activation. Gshard and DeepSeekMoE perform similarly overall (49.2% vs. 50.2%), with DeepSeekMoE's shared expert design providing a modest edge on most tasks.

A notable outlier is MMLU, where all methods cluster tightly: Dense 35.4%, Gshard 36.7%, DeepSeekMoE 37.1%, PKM 36.3%, PEER 37.4%, OmniMoE 37.5%. The narrow range (2.1 percentage points) suggests that for broad multitask knowledge assessment, the choice of FFN architecture matters less than raw parameter count, which is consistent across all models at 1.7B activated parameters.

End-to-End Efficiency and Scalability (Figure 4)

Figure 4 presents inference latency and peak memory for all six architectures, with two independent sweeps: varying activated parameters (left column: Figures 4a and 4c) and varying number of input tokens (right column: Figures 4b and 4d). The key headline numbers are reported for a fixed configuration of 28M activated parameters and 4,096 input tokens, where OmniMoE achieves 6.7 ms latency, compared to 73 ms for PEER (a 10.9× speedup) and 102 ms for DeepSeekMoE (a 15.2× speedup).

Examining the latency scaling with activated parameters (Figure 4a): Dense shows the lowest latency (approximately 4 ms at 4M activated, growing to approximately 8 ms at 32M) because it has no routing overhead. Gshard and DeepSeekMoE show moderate latency scaling, reaching roughly 30–40 ms at 28M activated. PKM and PEER show dramatically worse scaling: PEER reaches approximately 25 ms at 13M activated and approximately 73 ms at 26M. OmniMoE tracks closest to the Dense and coarse-grained baselines, growing from roughly 4 ms at 4M to approximately 6.7 ms at 28M. The gap between OmniMoE and PEER widens with increasing activated parameters, suggesting that OmniMoE's Expert-Centric Scheduling becomes relatively more advantageous as the expert pool grows.

Examining latency scaling with token count (Figure 4b): at 1K tokens, all methods cluster between 1–5 ms except PKM (approximately 12 ms). As token count increases to 16K, coarse-grained methods (Gshard, DeepSeekMoE) show approximately linear scaling, reaching roughly 30–50 ms. PEER shows the steepest scaling, reaching roughly 73 ms at 4K tokens (the paper does not extend the PEER curve to 16K, presumably because latency becomes prohibitively high). OmniMoE scales from approximately 2 ms at 1K to roughly 6.7 ms at 4K, and the paper's curve suggests sub-linear scaling beyond 4K (the OmniMoE line in Figure 4b appears to reach roughly 15–20 ms at 16K). The paper attributes this favorable scaling to the batch-level amortization in Expert-Centric Scheduling: as token count increases, the set of unique active experts grows sub-linearly (the saturation effect), so the expert parameter I/O per token decreases.

The memory scaling results (Figures 4c and 4d) follow a similar pattern. As activated parameters increase from 4M to 32M (Figure 4c), OmniMoE's memory footprint (roughly 2,100 MiB at 4M, growing to roughly 2,600 MiB at 28M) is comparable to coarse-grained baselines (Gshard, DeepSeekMoE) and substantially lower than fine-grained baselines (PKM, PEER both exceed 2,800 MiB at equivalent activated parameters). The paper attributes PKM/PEER's higher memory to materialization of full routing tensors and scattered access patterns that prevent memory reuse. As token count increases (Figure 4d), OmniMoE's memory scaling remains competitive with coarse-grained methods, growing from roughly 2,150 MiB at 1K to roughly 2,800 MiB at 16K, versus DeepSeekMoE's roughly 2,100 to 2,600 MiB over the same range.

A non-obvious result visible in Figure 4b: at large token counts (8K–16K), DeepSeekMoE is slower than PEER despite using coarse-grained FFN experts. The paper explains this as a consequence of "packing/alignment overhead in tiled coarse-grained kernels, where routed tokens must be reordered and padded to fixed block sizes, causing redundant computation and extra memory traffic" (Section 3.2). This is a genuinely surprising finding—it undermines the assumption that coarse-grained experts are always hardware-efficient—and motivates OmniMoE's approach of converting fine-grained activations into compact, expert-centric batched matrix operations rather than relying on coarse granularity alone for efficiency.

Scaling Laws (Figure 5)

Figure 5 compares validation perplexity (lower is better) for all six FFN variants across multiple training scales. Figure 5a plots perplexity against training FLOPs; Figure 5b plots perplexity against activated parameters.

On the FLOPs-matched comparison (Figure 5a), OmniMoE consistently achieves the lowest perplexity across all FLOPs budgets from approximately 10^17 to 10^19. The ordering from best to worst is: OmniMoE < PEER ≈ DeepSeekMoE < Gshard < PKM < Dense. The gap between OmniMoE and PEER is visible but modest at lower FLOPs (10^17) and appears to widen slightly as FLOPs increase, suggesting that OmniMoE's advantages compound with scale. The paper presents this as evidence that OmniMoE achieves better compute efficiency—for a fixed training budget, the architecture extracts more learning per FLOP.

On the parameter-matched comparison (Figure 5b), OmniMoE again achieves the lowest perplexity across all activated-parameter budgets from 80M to 1.6B. The gap between OmniMoE and the next-best method (PEER or DeepSeekMoE, depending on the scale) is more pronounced at larger parameter counts: at 80M activated, all methods are relatively clustered (perplexity range roughly 15–18); at 1.6B activated, the spread is wider (OmniMoE at roughly 9, DeepSeekMoE at roughly 10, PEER at roughly 10.5, Dense at roughly 11). The paper interprets this as evidence that OmniMoE benefits more steadily from additional capacity than alternatives, consistent with the complementary roles of fine-grained activation (for long-tail knowledge) and the shared dense MLP (for stable general reasoning).

A notable feature of Figure 5b: PEER outperforms DeepSeekMoE at small scales (80M–200M) but is overtaken at larger scales (680M–1.6B). The paper does not comment on this crossover explicitly, but it is consistent with the hypothesis that PEER's limited per-expert expressivity becomes a bottleneck as total capacity grows—at some point, adding more embedding-sized experts cannot compensate for the lack of deep nonlinear transformations, whereas DeepSeekMoE's coarse-grained FFN experts provide sufficient per-expert capacity to benefit from scale.

Distributed Training Communication Overhead (Appendix C, Figures A and B)

The paper verifies OmniMoE's scalability in distributed training with Expert Parallelism (EP). Two experiments are reported.

Scalability with number of experts (Figure A). With a fixed sequence length, the communication latency and memory are measured as the number of experts scales from 1K to 2M. The key finding is a saturation effect: "when the number of experts N exceeds the total number of activated experts (n_tokens × K = 16,384), the total communication volume for the backward pass stabilizes at approximately 80MB and does not grow with increasing N." In an 8-GPU NVLink environment, this communication latency is 0.521 ms—negligible relative to computation time. On a 64-GPU IB HDR cluster, total communication latency is approximately 0.3 ms forward + 0.15 ms backward = roughly 0.45 ms total, similarly negligible. The paper argues this demonstrates that OmniMoE "successfully breaks the bottleneck where communication overhead grows linearly with model capacity in traditional architectures."

Scalability with sequence length (Figure B). With a fixed number of experts (the paper does not specify the value used for this experiment), communication volume and latency are measured as token count scales from 1K to 128K. The finding is that "communication volume exhibits a strictly linear relationship with the sequence length." At 128K tokens on 64 GPUs (IB HDR), total communication memory reaches approximately 6GB, with estimated latency of approximately 15 ms. The paper characterizes this as efficient—communication does not become a bottleneck for training throughput even at long sequence lengths.

Ablation Studies and Robustness Checks

The paper's primary ablation study is Table 2, which isolates the three core components of OmniMoE: the Shared Dense MLP, the Cartesian Product Router, and Expert-Centric Scheduling. All metrics are reported relative to the full model (1.0×). Lower is better for Latency, Memory, PPL, and Unevenness; higher is better for Knowledge Performance, Reasoning Performance, and Expert Usage.

Removing the Shared Dense MLP: This yields a modest efficiency improvement (latency 0.86×, memory 0.98×) but degrades model quality: perplexity increases to 1.2×, knowledge performance drops to 0.91×, and reasoning performance drops to 0.79×. Expert usage remains at 100% and unevenness increases slightly from 0.24 to 0.27. The asymmetric impact on reasoning vs. knowledge (0.79× vs. 0.91×) supports the paper's claim that the shared dense MLP is particularly important for the nonlinear transformations required in multi-step reasoning. The paper interprets this as evidence that "the shared dense branch serves as a critical foundational backbone complementary to fine-grained retrieval. It handles common linguistic patterns and reasoning steps, allowing the routed branch to focus exclusively on fetching token-specific long-tail knowledge."

Replacing the Cartesian Product Router with a standard dense routing projection: This causes catastrophic efficiency regression: latency increases to 30.6×, memory increases to 337.5×. Model quality also degrades substantially: perplexity increases to 1.4×, knowledge performance drops to 0.66×, reasoning performance drops to 0.79×. Most strikingly, expert usage collapses to 4% (only 4% of experts are ever activated) and unevenness skyrockets from 0.24 to 0.77. The paper notes that "despite rigorous tuning of the standard auxiliary load-balancing loss for this baseline during training, the naive gate fails to learn distinct specializations over the massive expert space, collapsing into a few dominant experts." This is the paper's strongest evidence that the Cartesian Product Router is not merely an efficiency optimization—it is essential for training dynamics at scale because the factorization acts as an implicit regularizer that prevents expert collapse.

Reverting Expert-Centric Scheduling to token-centric execution: This preserves model quality perfectly (all quality metrics at 1.0×) but causes massive system cost: latency increases to 24.8× and memory increases to 417.7×. The peak memory increase is attributed to "the materialization of full routing tensors required by the standard baseline, which our scheduling avoids." This ablation cleanly separates OmniMoE's algorithmic contributions from its system contributions: the model's computation and output are identical, but the execution cost differs by nearly 25× in latency and over 400× in memory. The paper frames this as confirmation that "our scheduling strategy is the primary source of acceleration."

Distributed training robustness (Appendix C). The paper verifies that OmniMoE's communication overhead scales favorably. The saturation effect (communication cost constant once N > n_tokens × K) is demonstrated for both 8-GPU NVLink and 64-GPU IB HDR configurations, providing evidence that the finding is not specific to a particular interconnect or scale.

Kernel implementation fairness. While not presented as a formal ablation, the paper addresses a potential confound: the efficiency results could be biased by comparing optimized custom kernels (OmniMoE) against inefficient baseline implementations. The paper counters this by specifying that coarse-grained baselines use NVIDIA's CuTile kernels (described as "best-performance") and fine-grained baselines use "highly-optimized Triton fused kernels," all running in the same NVIDIA PyTorch container with Hugging Face Transformers. The strict end-to-end latency measurement (including all scheduling and reordering overheads for OmniMoE, and corresponding layout-transformation costs for baselines) further ensures fair comparison.

Absence of certain ablations. The paper does not ablate:

The choice of K (number of activated atomic experts per token), which is swept across configurations (512 to 4,096 per Table A) but not isolated as an independent variable holding other parameters fixed.
The group size B in Expert-Centric Scheduling, which is described conceptually but whose value is not specified and whose sensitivity is not tested.
The grid dimensions (N_r, N_c) in the Cartesian Product Router—the paper assumes a square grid but does not test rectangular decompositions.
The choice of activation function (SWiGLU vs. alternatives like GELU or ReLU) for the atomic experts.
The training data scale—all models are trained on 40B tokens of SmolLMCorpus, with no data ablation at larger or smaller scales.
The effect of the number of difficulty bins (the paper does not use difficulty conditioning, so this is not applicable in the same way as in the reference example, but there is no ablation of the expert count N or the total parameter budget independently).

Critical Assessment

Claim 1: OmniMoE achieves 50.9% average zero-shot accuracy, outperforming both coarse-grained and fine-grained baselines.

Supported, but the margin is small relative to the baseline variance. The 50.9% average for OmniMoE is +0.7 over DeepSeekMoE (50.2%) and +2.0 over PEER (48.9%). Without confidence intervals or multiple training runs, it is impossible to determine whether the +0.7 edge over DeepSeekMoE is statistically reliable or within noise. The sample size (seven benchmarks, each evaluated once) is small for claiming superiority over a strong baseline like DeepSeekMoE, especially given that the two models tie or nearly tie on several individual benchmarks (MMLU: 37.5% vs. 37.1%; Winogrande: 59.7% vs. 59.1%). The claim of outperforming PEER is more robust (+2.0 average, with consistent advantages on reasoning benchmarks), but even here, the lack of statistical quantification is a weakness.

The paper's choice to train all models from scratch on identical data is a strength for architectural comparison, but it also means the absolute performance numbers are not directly comparable to published results for these architectures (which typically use much larger training budgets). The 40B-token SmolLMCorpus is relatively small by modern standards, raising the question of whether the relative ordering of architectures would change at larger training scales. The scaling-law experiments (Figure 5) partially address this by showing consistent advantages for OmniMoE across FLOPs budgets, but the scaling range (10^17 to 10^19 FLOPs) is narrow compared to production training runs.

Claim 2: OmniMoE reduces inference latency from 73 ms to 6.7 ms (10.9× speedup) vs. PEER.

Strongly supported with a qualification about the specific configuration. The 10.9× figure is measured at a specific operating point (28M activated parameters, 4,096 tokens) and appears robust given the controlled benchmarking methodology (identical hardware, strict end-to-end measurement, state-of-the-art kernels for baselines). Figure 4 shows that the speedup varies with configuration: at lower activated parameters or token counts, the gap narrows; at higher values, it appears to widen. The 10.9× figure should therefore be understood as a representative point rather than a universal constant.

The qualification is that the paper does not explicitly account for the memory overhead of the Cartesian Product Router's factorized projections in the latency measurement for baselines. The 30.6× latency regression when replacing the Cartesian Product Router with a standard dense router (Table 2) is so extreme that it raises the question of whether the standard router baseline is a fair comparison—a standard dense router with N in the millions is clearly infeasible, and no practitioner would deploy it. The fair comparison is between OmniMoE and architectures specifically designed for million-scale routing (PKM, PEER), and against those, the 10.9× speedup is the correct headline.

Claim 3: Expert-Centric Scheduling converts scattered memory accesses into dense Grouped GEMMs, eliminating the memory bandwidth bottleneck.

Supported by the ablation but with incomplete characterization of the scheduling overhead. Table 2 shows that removing Expert-Centric Scheduling causes a 24.8× latency increase and 417.7× memory increase while preserving model quality. This is compelling evidence that the scheduling transformation is what makes fine-grained execution feasible—without it, OmniMoE would be slower than PEER. The paper states that "empirically, scheduling occupies <5% of total latency," but does not provide a breakdown of where this overhead comes from (sorting vs. group construction vs. kernel launch) or how it scales with batch size and expert count. Appendix A provides theoretical complexity analysis for the sort (O(M log M) where M = L × K) and notes that "given that GPU memory bandwidth is the primary bottleneck, this lightweight integer sorting (performed efficiently via radix sort) is negligible," but no empirical timing of the sort step is reported.

A deeper concern: the 417.7× memory increase when removing Expert-Centric Scheduling suggests that the token-centric baseline materializes something extremely large—likely the full routing score matrix or intermediate tensors that Expert-Centric Scheduling avoids by computing scores on-the-fly within the Grouped GEMM kernel. This raises the question of whether the token-centric baseline is implemented optimally (could the routing tensors be streamed rather than materialized?) or whether the ablation compares against a deliberately weak baseline. The paper's description in Appendix A suggests the overhead comes from "the materialization of full routing tensors required by the standard baseline, which our scheduling avoids," but a more detailed accounting would strengthen the claim.

Claim 4: The Cartesian Product Router prevents expert collapse at scale, which standard routers cannot avoid even with load-balancing losses.

Supported but with a nuance about the baseline. The ablation result is striking: expert usage drops from 100% to 4% and unevenness increases from 0.24 to 0.77 when replacing the Cartesian Product Router with a standard dense router. This is strong evidence that the factorization is doing something beyond efficiency—it shapes training dynamics in a way that standard regularization cannot replicate.

However, the paper's claim that the standard router was trained with "rigorous tuning of the standard auxiliary load-balancing loss" is stated without detail. What load-balancing loss was used? Was it the same formulation as in Switch Transformers (Fedus et al., 2022) or DeepSeekMoE (Dai et al., 2024)? What coefficient was used, and was it swept? The auxiliary loss coefficient is known to be a sensitive hyperparameter in MoE training—too small and experts collapse; too large and routing becomes uniform, sacrificing specialization. Without reporting the sweep range and the optimal coefficient, it is difficult to assess whether the 4% expert usage represents a fundamental limitation of dense routing at scale or an artifact of suboptimal hyperparameter tuning. The paper's claim would be stronger with an explicit sweep of the auxiliary loss coefficient for the standard router baseline, showing that even the best coefficient cannot prevent collapse.

Claim 5: The shared dense MLP and the routed atomic experts serve complementary roles—the shared MLP handles general reasoning, and the routed experts handle long-tail knowledge retrieval.

Supported but with indirect evidence. The ablation (Table 2) shows that removing the shared dense MLP hurts reasoning more than knowledge (0.79× vs. 0.91×), which is consistent with the complementarity hypothesis. The downstream benchmark results (Table 1) show OmniMoE's largest advantages over fine-grained baselines on reasoning benchmarks (ARC +3.6, HellaSwag +4.6) and over coarse-grained baselines on knowledge benchmarks (TriviaQA +1.1, OBQA +1.4), which is also consistent.

However, the paper provides no direct evidence that the routed atomic experts actually retrieve long-tail knowledge as opposed to performing some other function. This would require probing experiments—for example, analyzing whether specific atomic experts activate for specific factual domains, measuring the entropy of expert activation patterns, or ablating the atomic expert pool size to see whether reducing it disproportionately affects knowledge-intensive tasks. The paper infers functional specialization from performance patterns, but does not directly demonstrate it.

General Weaknesses in the Experimental Design

Single training corpus. All experiments use the SmolLMCorpus (40B tokens). It is unknown whether the relative advantages of OmniMoE would generalize to other data distributions (e.g., code-heavy corpora, multilingual data, domain-specific text). The 40B token budget is relatively small—it is plausible that at larger training scales (hundreds of billions or trillions of tokens), the coarse-grained baselines would catch up or the fine-grained advantages would diminish. The scaling-law curves (Figure 5) provide some evidence of extrapolation, but only over a narrow range (roughly one order of magnitude in FLOPs).

No error bars or statistical testing. All downstream results in Table 1 are point estimates from single training runs. Without multiple random seeds, confidence intervals, or statistical tests, it is impossible to determine whether the +0.7 advantage over DeepSeekMoE or the +2.0 advantage over PEER is reliable. This is a common limitation in LLM pre-training papers (due to the cost of training multiple models at scale), but it is particularly relevant here because the margins over strong baselines are small.

No test-time scaling experiments. The paper evaluates models in a single forward pass (zero-shot evaluation). It does not explore whether OmniMoE's architectural advantages amplify or diminish under test-time scaling strategies (best-of-N, majority voting, beam search). Given that the reference example paper emphasizes test-time compute scaling, it is worth noting that OmniMoE's efficiency claims are about single-pass inference—it is unknown whether the 10.9× speedup over PEER would hold in a best-of-N or beam search setting where execution patterns differ.

Limited scale range. The largest model (6.4B-A1.7B) is modest by contemporary standards (frontier MoEs like DeepSeek-V3 and KIMI-K2 are hundreds of billions of parameters). The paper argues that "all methods adhere to predictable scaling laws ... ensuring our findings extrapolate to larger scales," but the extrapolation range is large (from 1.7B to 100B+ active parameters) and the scaling-law curves (Figure 5) show only a factor of ~20 in activated parameters. The claim of extrapolation is plausible but unverified.

Missing baselines for the scheduling contribution. The paper compares against PKM and PEER as fine-grained baselines, but these architectures use different expert parameterizations (static embeddings vs. OmniMoE's gated atomic experts). A useful ablation would be: OmniMoE's architecture (atomic experts + DEA + Cartesian Product Router) executed with token-centric scheduling. This would isolate the scheduling contribution on identical model architecture, rather than comparing against different architectures. The current ablation (removing Expert-Centric Scheduling) does this internally but does not report its downstream accuracy or perplexity—only that quality metrics are preserved at 1.0×. The actual accuracy of OmniMoE with token-centric scheduling on the seven benchmarks is not reported, which would strengthen the claim that scheduling is purely an efficiency mechanism.

No ablation of the number of atomic experts K. The paper sweeps K from 512 to 4,096 across different configurations (Table A) but does not present an isolated ablation showing how K affects model quality and efficiency at a fixed total parameter budget. This is important because K is the primary knob controlling the expressivity of the assembled expert—larger K means more capacity per token but also more computation and more scattered memory accesses (which the scheduler must handle). Understanding where diminishing returns set in would guide practical deployment.

What the paper does well in its experimental design. Despite these limitations, the paper's experimental methodology has significant strengths: (1) training all architectures from scratch on identical data, controlling for training recipe as a confound; (2) measuring strict end-to-end latency including all overheads, not just kernel execution time; (3) using state-of-the-art kernel implementations for all baselines, not comparing against naive implementations; (4) providing both scaling-law and efficiency benchmarks at multiple scales; (5) verifying distributed training scalability with a separate set of experiments; and (6) including the negative result that the standard dense router collapses to 4% expert usage, which is genuinely informative about the difficulty of routing at scale. These strengths make the paper's central claims—that OmniMoE reconciles parameter efficiency with hardware efficiency—credible despite the limitations noted above.

6. Limitations and Trade-offs

The Difficulty Estimation Cost Is Prohibitively Expensive and Not Accounted for in the Headline Efficiency Numbers

The assumption or constraint. The entire OmniMoE framework depends on solving the routing problem—selecting the top-K experts for each token from a pool of N experts. The Cartesian Product Router reduces the projection cost from O(Nd) to O(√N d), but the top-K selection over the implicit score grid still requires evaluating all N expert scores during the tiled parallel search. While the paper's fused kernel avoids materializing the full score matrix, the search itself scales with N. More critically, the Expert-Centric Scheduling requires identifying the set of unique active experts across the batch (E_active in Section 2.3) and partitioning them into groups—a preprocessing step that must complete before any Grouped GEMM execution begins. The paper acknowledges this overhead indirectly in Appendix A: "scheduling occupies <5% of total latency, well-amortized by the speedup in the GEMM phase." However, this 5% figure is reported without detail on what configuration it applies to, how it was measured, or how it scales.

The consequence. The 10.9× speedup over PEER is computed from end-to-end latency measurements that include scheduling overhead, so the overhead is technically accounted for in the headline number. The concern is about scalability: the sorting step operates on M = L × K tasks, and the group construction requires scanning the unique expert set E_active. Both operations scale with batch size and K. At deployment scales (much larger batch sizes than the 1K–16K tokens tested, or larger K values than the 512–4,096 range reported in Table A), the scheduling overhead may grow to dominate, particularly if the expert activation distribution becomes more uniform (larger E_active for a given L × K).

What evidence exists in the paper. The paper provides no empirical breakdown of the scheduling overhead—no timing of the sort step, the group construction, or the kernel launch separately. The 5% claim appears in Appendix A without supporting data or specification of the measurement configuration. Figure 4 shows OmniMoE's latency scaling favorably across the tested range, but the curves do not reveal where the scheduling overhead becomes a bottleneck. The ablation (Table 2) confirms that Expert-Centric Scheduling is the primary source of speedup, but it compares against a token-centric baseline that materializes full routing tensors (417.7× memory increase), not against an optimized token-centric implementation that might close the gap somewhat.

Mitigation status. The paper does not address this limitation beyond the qualitative 5% claim. There is no analysis of how scheduling overhead scales with batch size, K, or N, and no proposal for reducing this overhead (e.g., incremental scheduling, fused sort-and-execute kernels). This is an area where practitioners attempting to deploy OmniMoE at production scale would need to perform their own profiling.

The Large-K Requirement (512–4,096 Activated Experts Per Token) Imposes a Structural Trade-off: Expressivity Through Composition Comes at Linear Compute Cost

The assumption or constraint. OmniMoE recovers expressivity from atomic experts by composing K of them per token through Dynamic Expert Assembly. The assembled computation (g_x ⊙ σ(x w_x^⊤)) v_x (Eq. 7) is mathematically a two-layer network whose capacity scales with K—larger K means more atomic experts contribute, producing a richer token-conditioned transformation. Table A shows K values from 512 (at 4M activated parameters) to 4,096 (at 28M activated parameters), substantially larger than the 2–32 experts typical in coarse-grained MoEs. This large K is not incidental; it is necessary because individual atomic experts have limited expressivity (only 2d parameters each), so many must be combined to match the representational capacity of a single coarse-grained FFN expert.

The consequence. The compute cost of the routed branch scales directly with K: the operations x w_x^⊤ (producing K dot products), σ(·) (K SiGLU evaluations), g_x ⊙ σ(·) (K element-wise products), and (g_x ⊙ σ(x w_x^⊤)) v_x (a (1 × K) × (K × d) matrix multiply) all grow with K. While the Expert-Centric Scheduling makes these operations efficient by batching them into Grouped GEMMs, the total FLOPs for the routed branch grows as O(Kd), which is comparable to the O(d²) cost of a standard FFN if K is proportional to d. For the configurations in Table A (d = 1024, K up to 4,096), the routed branch may cost more FLOPs than the shared dense MLP.

Additionally, a larger K increases M = L × K (the number of tasks to schedule), which grows the sorting overhead and the number of groups N_groups = ⌈|E_active| / B⌉. This amplifies the scheduling overhead concern from the first limitation.

What evidence exists in the paper. The paper sweeps K across configurations (Table A) but does not provide an isolated ablation showing how downstream accuracy or perplexity scales with K at a fixed total parameter budget. Without this ablation, it is impossible to determine whether the large K values (2,048–4,096) are necessary for OmniMoE's accuracy advantages, or whether smaller K would suffice. The efficiency results in Figure 4 show OmniMoE maintaining competitive latency despite large K, but this reflects the success of Expert-Centric Scheduling rather than evidence that the large-K cost is negligible—it means the scheduling amortizes the cost well, not that the cost does not exist.

Mitigation status. The paper does not address this trade-off. K is treated as a configuration hyperparameter swept across model scales, not as a design choice with explicit cost-benefit analysis. A practitioner choosing K for a deployment would need to balance the expressivity benefits (which likely show diminishing returns as K grows) against the linear compute increase, but the paper provides no guidance on where this balance lies.

All Experiments Are on a Single 40B-Token Corpus with Models up to 1.7B Active Parameters—the Claim of Extrapolation to Larger Scales and Different Data Distributions Is Unverified

The assumption or constraint. The paper's experimental scope is narrow along two dimensions: (1) all models are trained on the SmolLMCorpus, a 40-billion-token dataset, and (2) the largest model has 1.7B active parameters (6.4B total). The paper acknowledges this implicitly by framing it as a controlled comparison: "we prioritize architectural comparison via controlled pre-training from scratch rather than comparing against off-the-shelf checkpoints" (Section 3.1). However, the paper also claims that "all methods adhere to predictable scaling laws (Section 3.2), ensuring our findings extrapolate to larger scales." This extrapolation claim covers two orders of magnitude from the tested 1.7B active parameters to the 100B+ scale of production MoEs like DeepSeek-V3.

The consequence. Three specific concerns arise from this limited scope:

Scaling-law extrapolation uncertainty. The scaling-law experiments (Figure 5) span activated parameters from 80M to 1.7B—roughly a factor of 20. Extrapolating to 100B+ parameters requires assuming that the relative ordering of architectures (OmniMoE < PEER ≈ DeepSeekMoE < Gshard < Dense) remains stable over two additional orders of magnitude, which is not guaranteed. The paper's own evidence shows a crossover: PEER outperforms DeepSeekMoE at 80M–200M but is overtaken at 680M–1.6B (Figure 5b). This crossover demonstrates that relative architecture performance can change with scale, undermining confidence in extrapolation.
Data distribution sensitivity. The SmolLMCorpus has a specific mixture (Web, Textbook, Code, Math). It is plausible that OmniMoE's advantages on knowledge-intensive tasks (TriviaQA +1.1 over DeepSeekMoE) are sensitive to whether the pre-training corpus contains factual knowledge that can be effectively distributed across atomic experts. On a code-heavy corpus, the relative benefits of fine-grained vs. coarse-grained experts might differ substantially. The paper provides no evidence either way.
Training stability at scale. The paper trains models for 13,500–40,000 steps depending on scale (Table B). At production scale (hundreds of billions of tokens), training stability issues—particularly load imbalance and expert collapse—may manifest differently. The Cartesian Product Router's success at preventing collapse (100% expert usage at the tested scales) is encouraging, but the 40B-token budget is small enough that training dynamics may not have reached the regime where collapse becomes likely.

What evidence exists in the paper. Figure 5 provides the scaling-law curves, which show OmniMoE maintaining its advantage across the tested range. The distributed training experiments (Appendix C) verify communication scalability up to 2M experts, but these are synthetic benchmarks measuring communication overhead, not end-to-end training runs at scale. There is no experiment where OmniMoE is trained on a larger corpus (e.g., 200B+ tokens) or at a larger model size (10B+ active parameters) to validate the extrapolation claim.

Mitigation status. The paper acknowledges the scale limitation implicitly through its focus on controlled architectural comparison, but does not directly address the extrapolation risk. The suggestion that "academic resource constraints limit pre-training to the 1.7B activated-parameter scale" (Appendix B) is an honest disclosure, but it does not mitigate the uncertainty. Future work at larger scale—either by the authors or by the community using the open-sourced code—would be needed to validate the extrapolation.

The Shared Dense MLP and Routed Atomic Branch Have Not Been Shown to Actually Specialize as Claimed—the Functional Division Is Inferred, Not Demonstrated

The assumption or constraint. The paper's architectural motivation rests on a functional division of labor: the shared dense MLP provides "general semantic reasoning and stable capacity" while the routed atomic experts "specialize in long-tail knowledge retrieval" (Section 1). The ablation (Table 2) shows that removing the shared dense MLP drops reasoning performance more than knowledge performance (0.79× vs. 0.91×), which is consistent with this division but does not directly demonstrate it. The paper provides no mechanistic evidence—no analysis of what specific atomic experts learn, no probing of whether the shared MLP and routed branch activate for different types of linguistic patterns, and no measurement of how expert specialization changes when the shared MLP is present versus absent.

The consequence. The functional division claim is central to understanding why OmniMoE works. If the shared MLP and routed branch are not genuinely specialized—if, for example, the routed branch learns to approximate what the shared MLP would do, and the performance gain comes primarily from having more total parameters rather than from functional complementarity—then the architecture could potentially be simplified. Conversely, if the specialization is real but fragile (e.g., it only emerges with specific training recipes or at specific scales), then practitioners attempting to adapt OmniMoE to new domains might find that the benefits do not transfer.

Without mechanistic evidence, the paper's explanation for its results remains a plausible hypothesis rather than an empirically grounded finding. This matters for practitioners because it affects how they would adapt the architecture: if the functional division is real, then the shared MLP capacity and the atomic expert pool size should be carefully balanced; if it is not, then simpler designs (e.g., a larger atomic expert pool without a shared MLP, or a shared MLP with conventional coarse-grained routed experts) might suffice.

What evidence exists in the paper. Only the ablation in Table 2 and the downstream benchmark patterns in Table 1. Table 2 shows asymmetric degradation when the shared MLP is removed (reasoning drops more than knowledge). Table 1 shows OmniMoE outperforming fine-grained baselines on reasoning benchmarks (attributed to the shared MLP) and outperforming coarse-grained baselines on knowledge benchmarks (attributed to the routed atomic experts). Both are correlational evidence—consistent with the functional division hypothesis but not causal proof.

Mitigation status. The paper does not address this limitation. No probing experiments, activation analyses, or expert specialization measurements are reported. This is a common gap in architecture papers (functional claims often outpace mechanistic evidence), but it is particularly salient here because the functional division is not just an interpretation of results—it is the stated design principle that motivated the heterogeneous architecture. Future work analyzing expert activation patterns (e.g., do atomic experts cluster by domain? Do they activate disproportionately for rare tokens?) would strengthen or refute the functional specialization claim.

The Expert-Centric Scheduling Advantage Depends on Expert Reuse Across Tokens—Under Low Batch Sizes or Low Token-to-Expert Ratios, the Speedup Diminishes

The assumption or constraint. The theoretical memory traffic reduction from Expert-Centric Scheduling is governed by the ratio η = (L × K) / |E_active| (Appendix A, Eq. viii)—the total number of token-expert pairs divided by the number of unique active experts. When η ≫ 1, many tokens share the same experts, so each expert's parameters are loaded once and reused across many tokens, yielding large I/O savings. When η ≈ 1 (each token activates a nearly disjoint set of experts), Expert-Centric Scheduling provides little benefit over token-centric execution because there is minimal reuse to exploit.

The consequence. In deployment scenarios with small batch sizes (low L) or when K is small relative to the expert pool size (so |E_active| ≈ L × K), the Expert-Centric Scheduling overhead (sorting, group construction) may exceed the benefit (reduced parameter I/O). The paper's efficiency benchmarks (Figure 4) test batch sizes from 1K to 16K tokens and K values from 512 to 4,096 (Table A). At these scales, η is likely large (the paper states "each expert is frequently accessed by multiple tokens in a batch"), so the scheduling advantage is substantial.

However, many real-world deployment scenarios operate at lower batch sizes: interactive applications with one or a few concurrent queries, edge device inference with limited memory, or autoregressive decoding where tokens are processed sequentially. In these settings, |E_active| may approach L × K, and the Expert-Centric Scheduling overhead may dominate, potentially making OmniMoE slower than simpler architectures.

What evidence exists in the paper. Figure 4b shows latency at 1K, 2K, 4K, 8K, and 16K tokens. At the lowest tested batch size (1K tokens), OmniMoE's latency is roughly 2 ms, competitive with but not dramatically better than coarse-grained baselines. The gap widens at larger batch sizes, consistent with the reuse hypothesis. However, the paper does not test batch sizes below 1K (e.g., 1, 8, 32, 128 tokens), which would reveal the low-batch regime where the scheduling advantage might invert. The distributed training experiments (Appendix C) use 16,384 activated experts per batch (n_tokens × K), which is far larger than typical inference batch sizes.

Mitigation status. The paper does not address this limitation. The efficiency benchmarks sweep batch size down to 1K tokens, which is still large relative to many inference scenarios. A practitioner deploying OmniMoE in a low-batch setting would need to profile whether the Expert-Centric Scheduling overhead outweighs its benefits for their specific workload. The paper provides no guidance on a breakeven batch size or on alternative execution strategies for low-batch regimes (e.g., falling back to token-centric execution when η is below some threshold).

The Work Is Validated Exclusively on Language Modeling and Standard NLP Benchmarks—Applicability to Other Domains and Modalities Is Unknown

The assumption or constraint. All experiments in the paper—pre-training, scaling laws, downstream evaluation, efficiency benchmarks—use text-based language modeling on the SmolLMCorpus and seven standard NLP benchmarks. The OmniMoE architecture is presented as a general MoE framework, and the motivation (reconciling parameter efficiency with hardware efficiency) applies to any domain where sparse expert activation is beneficial. However, the paper provides no evidence beyond language modeling.

The consequence. Several aspects of OmniMoE's design may be language-specific:

The Cartesian Product Router's independence assumption (p(i,j | x) ≈ p_r(i | x) · p_c(j | x)) may be more or less appropriate depending on whether the underlying expert specialization has a factorizable structure. In language, experts might specialize along interpretable axes (syntax vs. semantics, formal vs. informal register, domain), which could naturally factorize. In other domains—visual recognition, speech processing, reinforcement learning—the structure of specialization may be fundamentally different, and the independence assumption could be a worse approximation, potentially degrading routing quality.
The optimal K (number of activated experts per token) is likely domain-dependent. Language modeling benefits from composing many knowledge sources per token (each token's meaning depends on context, domain, and linguistic structure). In domains where specialization is more categorical (e.g., each input belongs clearly to one class), a smaller K might be optimal, changing the trade-off between routing precision and compute cost.
The shared dense MLP's role may vary across modalities. In language, general syntactic and semantic patterns benefit from a universally-activated pathway. In vision, it is less clear what "general visual processing" means as distinct from "specialized feature detection"—the functional division may not transfer.
The hardware efficiency results were measured on NVIDIA A100 GPUs. Different hardware (H100, consumer GPUs, TPUs, edge accelerators) has different memory bandwidth, Tensor Core throughput, and cache hierarchies, which would change the breakeven point for Expert-Centric Scheduling.

What evidence exists in the paper. None. The paper does not include experiments on code generation, multilingual text, mathematical reasoning (separate from the benchmarks), visual recognition, speech, or any non-text modality. The distributed training experiments (Appendix C) use synthetic communication benchmarks, not end-to-end multi-modal training.

Mitigation status. The paper does not address this limitation or suggest it as future work. This is a standard scope limitation in architecture papers, but it is worth noting because OmniMoE's claimed contribution—reconciling parameter efficiency with hardware efficiency in fine-grained MoEs—is pitched at a level of generality that invites application beyond language. A practitioner considering OmniMoE for a non-language domain would need to validate all the key design choices (router factorization, K, shared MLP necessity, scheduling benefit) from scratch.

7. Implications and Future Directions

How This Work Changes the Landscape

OmniMoE introduces a new design axis for MoE architectures: the inversion of execution order (from token-centric to expert-centric) as a first-class architectural primitive rather than a downstream systems optimization. This is not a paradigm shift in the sense of upending the fundamental MoE formulation—the paper operates squarely within the standard MoE framework of Eq. 3, routing tokens to experts and summing weighted outputs. But it is a methodological reframing with practical consequences that changes how the field should think about the relationship between algorithm design and hardware execution.

The core reframing: execution order is an architectural degree of freedom, not an implementation detail. Prior work on MoE systems (DeepSpeed-MoE, FastMoE, MegaBlocks, SonicMoE) treated execution as downstream of routing: the router decides which experts process which tokens, and the systems layer optimizes how efficiently those assignments execute. The implicit assumption was that the execution order is determined by the routing—you must iterate over tokens because that is how the computation is defined. OmniMoE breaks this coupling by showing that, for fine-grained experts, you can reorder the token-expert assignments into expert-centric groups, and this reordering transforms the memory access pattern from random gathers to coalesced reads without changing the mathematical output. The proof is in Table 2: removing Expert-Centric Scheduling preserves all quality metrics at exactly 1.0× (identical computation), yet increases latency 24.8× and memory 417.7× (radically different hardware behavior). This cleanly separates the algorithmic contribution from the systems contribution and establishes that execution order is a design choice with enormous leverage—roughly 25× in latency and 400× in memory, far larger than what kernel fusion or padding reduction achieve within the token-centric paradigm.

The diagnostic contribution: exposing the structural nature of the granularity-efficiency trade-off. Before OmniMoE, the field operated with an implicit assumption that the granularity-efficiency trade-off was negotiable through better engineering—that finer-grained experts could be made fast enough with optimized kernels, or that coarse-grained experts could be made precise enough with better routing. The paper's Figure 4 provides empirical evidence that this assumption is false: PEER (fine-grained, highly optimized Triton kernels) is 10.9× slower than OmniMoE at 4K tokens not because of poor implementation but because scattered memory access is a structural consequence of fine granularity that no amount of kernel tuning within the token-centric paradigm can eliminate. Conversely, the paper shows that DeepSeekMoE's packing and alignment overhead can make it slower than PEER at large token counts (Figure 4b), demonstrating that coarse granularity also has structural efficiency costs that are not eliminated by NVIDIA's best-performance CuTile kernels. The implication is that neither pole of the granularity spectrum is Pareto-optimal—the trade-off is real and must be addressed architecturally, not through incremental systems optimization.

Reconciling conflicting evidence about MoE efficiency. The paper resolves an apparent contradiction in the literature: why do scaling-law analyses suggest that millions of experts should be optimal (Ludziejewski et al., 2024; Clark et al., 2022) while production systems cap at hundreds of coarse-grained experts (DeepSeek-V3 at 256, KIMI-K2 at 384)? The answer, per OmniMoE's evidence, is that the scaling-law predictions are correct about parameter efficiency but ignore the hardware execution cost—the memory bandwidth bottleneck makes million-scale experts practically infeasible unless the execution order is restructured. OmniMoE's Expert-Centric Scheduling removes this bottleneck, suggesting that the gap between theoretical optimal expert count and practical deployment is not fundamental but a consequence of the token-centric execution paradigm. This reframes the research question from "what granularity is optimal?" (which produces conflicting answers depending on whether you measure model quality or latency) to "how can we make the theoretically optimal granularity practical?"—a more productive framing that OmniMoE provides a concrete answer to.

Research directions that become more attractive. The paper makes expert-centric execution a legitimate and promising design pattern for any sparse computation where multiple queries access overlapping subsets of a shared parameter pool. This pattern is not specific to MoEs: retrieval-augmented generation (RAG), sparse attention, and memory-augmented neural networks all involve token-to-memory-index lookups with similar scattered access patterns. The paper's demonstration that expert-centric reordering can recover 25× in latency without quality loss makes it natural to ask whether the same technique applies to these other sparse architectures. Additionally, the paper makes factorized routing as a regularizer a recognized phenomenon (the 4% vs. 100% expert usage result in the ablation), opening investigation into when and why structured routing constraints prevent expert collapse—a question that extends beyond product-key routing to other factorization schemes.

Research directions that become less attractive. The paper's results make further work on pure token-centric kernel optimization for fine-grained MoEs less impactful. The 10.9× gap between PEER and OmniMoE is so large that no amount of incremental Triton kernel tuning within the token-centric paradigm could plausibly close it—the bottleneck is structural, not implementational. Similarly, the paper's finding that lookahead-style optimizations for the router (the standard dense router collapsing to 4% expert usage, making it worse in both efficiency and quality) suggests that increasing router complexity to improve routing precision is a dead end for massive expert pools—the challenge is not precision but preventing collapse, which factorized routing solves through its implicit regularization rather than through more expressive routing. Research effort is better directed at understanding and designing factorization schemes with favorable inductive biases than at making full N-way routing more efficient.

The magnitude of this contribution is significant but bounded: OmniMoE does not propose a new expert type, a new routing mechanism, or a new training objective. The atomic expert design, the product-key-style routing, and the shared-expert architecture all have clear antecedents. The contribution is the integrated co-design that makes these existing ideas work together at a scale where they previously could not—and the conceptual clarity of identifying execution order inversion as the key enabler. This is a well-executed systems contribution that changes what architectures are considered viable, not a theoretical breakthrough. The field should walk away from this paper understanding that fine-grained MoEs are now a practical deployment option, not just a theoretical curiosity, and that expert-centric scheduling is a general technique worth applying whenever sparse parameter access patterns exhibit reuse across queries.

Follow-Up Research This Work Enables

Characterizing the breakeven batch size for Expert-Centric Scheduling across hardware platforms. The paper demonstrates a 24.8× latency reduction from Expert-Centric Scheduling at batch sizes of 1K–16K tokens (Figure 4b), but the scheduling advantage depends on the reuse ratio η = (L×K) / |E_active| (Appendix A, Eq. viii). When η ≈ 1 (small batches, or token-expert assignments with little overlap), the sorting and group construction overhead may exceed the I/O benefit. A targeted study would sweep batch sizes from 1 to 256 tokens across A100, H100, and consumer GPUs (RTX 4090), measuring not just end-to-end latency but the breakdown of time spent in sorting, group construction, Grouped GEMM, and scatter-add. The key output would be a hardware-specific breakeven curve: at what batch size and K does Expert-Centric Scheduling outperform a well-optimized token-centric baseline? This matters because real deployments—interactive chatbots, edge inference—often operate at batch sizes far below the paper's tested range. A negative result (Expert-Centric Scheduling underperforming token-centric execution for batch < 32 on consumer GPUs) would not invalidate the paper's contribution but would clarify its deployment envelope and motivate hybrid scheduling strategies (token-centric for small batches, expert-centric for large batches).

Testing whether factorized routing prevents expert collapse on standard MoEs with moderate expert counts. The paper's most surprising ablation result is that replacing the Cartesian Product Router with a standard dense projection causes expert usage to collapse from 100% to 4% (Table 2). This was demonstrated at N = 102,400 experts (Table A), an extreme scale where a standard N-way classifier is clearly impractical. But the collapse phenomenon may manifest at much smaller N—do standard routers begin to underutilize experts at N = 512, 1024, 2048 (the range of production MoEs like DeepSeek-V3 and KIMI-K2)? A follow-up would take a standard coarse-grained MoE architecture (e.g., DeepSeekMoE with 256 experts), replace the standard router with a Cartesian Product Router (factorizing 256 as 16×16), and measure expert usage, load balance, and downstream accuracy at matched training budgets. The hypothesis is that factorized routing improves expert utilization and downstream performance even at moderate N where a standard router is computationally feasible. A negative result (no benefit or degradation from factorization at moderate N) would establish a scale threshold below which the regularization effect is unnecessary, bounding the contribution's applicability. A positive result would suggest that factorized routing should be adopted broadly, not just for million-scale expert pools.

Mechanistic probing of atomic expert specialization with and without the shared dense MLP. The paper infers that the shared dense MLP handles general reasoning while atomic experts handle long-tail knowledge retrieval (Section 1, supported by the asymmetric ablation in Table 2: reasoning drops to 0.79× vs. knowledge at 0.91× when the shared MLP is removed). A direct probing experiment would test this hypothesis by analyzing expert activation patterns. The setup: take a trained OmniMoE model, run inference on a diverse corpus (e.g., The Pile), and for each token, record which atomic experts are activated. Then cluster experts by their activation co-occurrence patterns and examine whether clusters correspond to interpretable categories (domain: code vs. news vs. academic; linguistic feature: rare words vs. common function words; factual domain: geography vs. biology). The key manipulation is to compare activation patterns between the full model and a variant with the shared MLP removed—if the hypothesis is correct, removing the shared MLP should force atomic experts to take on more general linguistic functions, blurring their specialization. The output would be quantitative metrics of specialization (e.g., mutual information between expert indices and token metadata) and qualitative examples of what specific expert clusters encode. A negative result (no interpretable specialization, or no change when the shared MLP is removed) would weaken the functional-division claim and suggest OmniMoE's benefits come primarily from increased total capacity rather than from complementary specialization.

Combining Expert-Centric Scheduling with on-device deployment of small fine-grained MoEs. The paper's efficiency claims are demonstrated on A100 GPUs with batch sizes of 1K+ tokens. But one of the motivating applications for fine-grained MoEs is on-device deployment—small models where precise parameter activation could enable larger effective capacity within tight memory budgets. A follow-up would take a compact OmniMoE configuration (e.g., the 280M-A80M model from Table B), deploy it on an edge device (e.g., NVIDIA Jetson Orin, Apple Neural Engine, or smartphone-class GPU), and measure: (i) whether Expert-Centric Scheduling's overhead (sorting, group construction) remains negligible relative to the GEMM speedup on hardware with limited memory bandwidth and compute, (ii) whether the memory savings from avoiding full routing tensor materialization (417.7× in the ablation, Table 2) translate to fitting larger expert pools within on-device memory limits, and (iii) whether the accuracy advantages over Dense baselines (Table 1, +5.0 average over Dense at 1.7B active) hold at the smaller scale. The key metric is whether OmniMoE can match the accuracy of a larger Dense model that exceeds the device's memory budget—if a 280M-A80M OmniMoE on-device matches a 200M Dense model that requires cloud inference, the practical value for edge deployment is established. A negative result (Expert-Centric Scheduling overhead dominates on low-power hardware, or accuracy advantages diminish at small scales) would indicate that OmniMoE's benefits are specific to datacenter-class GPUs and large batch sizes.

Scaling OmniMoE to 10B+ active parameters and 1T+ tokens to test extrapolation claims. The paper's largest model has 1.7B active parameters trained on 40B tokens, with scaling-law curves (Figure 5) showing OmniMoE maintaining advantage over baselines across a ~20× range in activated parameters and ~100× in training FLOPs. The paper claims these scaling trends "ensure our findings extrapolate to larger scales" (Section 3.1), but Figure 5b already shows a crossover (PEER outperforms DeepSeekMoE at 80M but is overtaken at 680M), demonstrating that architecture rankings can change with scale. A definitive follow-up would train OmniMoE and the strongest baselines (DeepSeekMoE, PEER) at 10B+ active parameters on 1T+ tokens, measuring both perplexity scaling and downstream accuracy. The specific question is whether the OmniMoE vs. DeepSeekMoE gap (currently +0.7 average accuracy at 1.7B active, Table 1) widens, narrows, or reverses at larger scale. A negative result (DeepSeekMoE overtaking OmniMoE at 10B+ active parameters) would indicate that OmniMoE's advantages are most pronounced in the small-to-moderate scale regime and that coarse-grained architectures benefit more from additional capacity—a finding that would significantly re-scope the paper's contribution and guide practitioners toward OmniMoE for smaller models and DeepSeekMoE for frontier-scale models.

Applying Expert-Centric Scheduling to retrieval-augmented generation (RAG) and other sparse memory architectures. The core insight of Expert-Centric Scheduling—that reordering query-memory lookups from query-centric to memory-centric can convert random accesses to coalesced reads—generalizes beyond MoEs. In RAG, each token (or chunk) retrieves relevant documents from a vector store; when a batch of queries shares retrieved documents, reordering execution to process all queries targeting the same document together could yield similar I/O savings. A follow-up would implement Expert-Centric Scheduling for a standard RAG pipeline (e.g., a dense retriever + generative reader), measure the latency reduction at varying batch sizes and retrieval depths (analogous to K), and compare against a baseline that processes each query independently. The specific hypothesis is that the I/O reduction ratio η = (batch_size × retrieval_depth) / |unique_documents| can be substantial in RAG settings where multiple queries are part of the same conversation or task. A positive result would establish Expert-Centric Scheduling as a general technique for sparse memory architectures, not just MoEs. A negative result (documents are typically too large and retrieval sets too disjoint for reuse to matter) would help scope the technique's applicability.

Practical Applications and Downstream Use Cases

Cost-efficient batch inference for knowledge-intensive enterprise applications. Consider an enterprise deploying a language model for customer support, where each query requires retrieving specialized knowledge from product documentation, policy manuals, or technical specifications. In this setting, inference runs in batches (hundreds to thousands of concurrent queries), and accuracy on knowledge-intensive tasks is critical. OmniMoE's architecture directly addresses both needs: the fine-grained atomic experts enable precise retrieval of long-tail knowledge (evidenced by +1.1 on TriviaQA and +1.4 on OBQA over DeepSeekMoE at 1.7B active, Table 1), while Expert-Centric Scheduling makes the fine-grained execution efficient at batch scale (6.7 ms vs. 102 ms for DeepSeekMoE at 4K tokens, a 15.2× speedup per Figure 4b). The practical benefit is that an organization could deploy a single OmniMoE model handling both general reasoning (via the shared dense MLP) and specialized knowledge retrieval (via the routed atomic branch) at latency comparable to a coarse-grained MoE, without needing a separate retrieval pipeline. For a deployment processing 10M queries per day, the 15.2× latency reduction over DeepSeekMoE translates directly to GPU-hour savings—approximately 15× fewer A100s needed to meet a latency SLA, or the ability to serve 15× more queries on the same hardware budget.

On-device fine-grained MoEs for privacy-sensitive applications with dynamic capacity scaling. In privacy-sensitive domains (healthcare, legal, personal assistants), inference must run on-device to avoid transmitting sensitive data to cloud servers. The memory and compute constraints of edge devices (phones, laptops, medical terminals) severely limit model size, but users still need accurate responses on diverse topics—requiring a model that can pack maximum knowledge into minimum parameters. OmniMoE's atomic expert design enables precise parameter activation (every activated parameter is token-relevant, avoiding the redundant activation of coarse-grained experts per Figure 1a), while the Cartesian Product Router allows the expert pool to be scaled smoothly (adding more atomic experts increases total parameters without increasing per-token compute, since K and d remain fixed). The practical benefit: a device manufacturer could ship a base OmniMoE model with, say, 100M active parameters and a 2M-expert pool, and later expand the expert pool to 5M or 10M through over-the-air updates as more domain knowledge becomes available—all without increasing inference latency or memory beyond what the device's original hardware budget allows. The paper's scaling-law results (Figure 5b, OmniMoE maintaining lowest perplexity across all activated-parameter budgets) and the distributed training saturation effect (Appendix C, communication cost constant once N exceeds n_tokens × K) provide evidence that this capacity scaling is both beneficial for quality and feasible for deployment.

Self-improving expert pools through continuous fine-tuning of the atomic expert bank. A distinctive property of OmniMoE's architecture is that the atomic expert parameters are stored in two global matrices W, V ∈ ℝ^{N × d} (Section 2.1) that are accessed through sparse row selection. This means individual atomic experts can be updated or replaced without affecting others—the parameter bank supports incremental updates at the granularity of individual experts. In a production deployment where the model needs to incorporate new knowledge (e.g., new products, updated policies, emerging terminology), a practitioner could: (1) collect a dataset of examples requiring the new knowledge, (2) identify which atomic experts activate for those examples, (3) fine-tune only those experts (a few rows of W and V) while keeping the rest of the model frozen, and (4) deploy the updated expert bank without retraining the entire model. This is practically infeasible in coarse-grained MoEs where updating one FFN expert affects all tokens routed to it, potentially degrading performance on other tasks. The Cartesian Product Router's factorization makes this approach viable because experts are independently addressable via their 2D grid coordinates—a continuous fine-tuning pipeline could add experts in new grid regions as knowledge domains expand. The paper does not demonstrate this capability, but it follows directly from the architectural design: the combination of fine-grained atomic experts, centralized parameter storage in W and V, and factorized routing means that the expert pool is natively modular. A production team adopting OmniMoE could build a CI/CD pipeline where model updates ship as delta patches to specific rows of W and V, reducing update bandwidth from gigabytes (full model) to megabytes (only changed experts).