OmniMoE: Efficient Mixture-of-Experts via Atomic Expert Orchestration

Paper: 2602.05711 Authors: Jingze Shi, Zhangyang Peng, Yizhang Zhu, Yifan Wu, Guang Liu, Yuyu Luo Affiliations: BAAI (Beijing Academy of Artificial Intelligence) Categories: cs.CL, cs.AI

Abstract

OmniMoE addresses a fundamental challenge in Mixture-of-Experts architectures: balancing expert specialization granularity with hardware efficiency. The framework introduces vector-level Atomic Experts that enable fine-grained routing while maintaining computational speed, achieving 50.9% zero-shot accuracy across seven benchmarks with 1.7B active parameters and a 10.9-fold speedup in inference latency.

Key Contributions

• Vector-level Atomic Experts: Fine-grained expert decomposition at the vector level rather than full FFN blocks
• Cartesian Product Router: Reduces routing complexity from O(N) to O(√N)
• Expert-Centric Scheduling: Converts scattered memory operations into dense matrix computations

The Granularity-Efficiency Tradeoff

Traditional MoE architectures face a fundamental tension:

• Coarse-grained experts (full FFN blocks): Hardware-efficient but limited routing flexibility
• Fine-grained experts (neuron-level): Better specialization but scattered memory access patterns that kill GPU throughput

OmniMoE resolves this by decomposing experts at the vector level — finer than FFN blocks but coarser than individual neurons — and introducing hardware-aware scheduling.

Architecture

Atomic Experts

Each expert is decomposed into vector-level atomic units that can be independently routed. This enables combinatorial expert compositions without the memory overhead of maintaining many full-sized experts.

Cartesian Product Router

Instead of routing over N experts directly (O(N) complexity), the router operates over two smaller sets whose Cartesian product defines the expert space. This reduces routing complexity to O(√N) while maintaining the same expressiveness.

Expert-Centric Scheduling

Converts the scattered memory access patterns inherent in fine-grained routing into dense matrix operations that GPUs can execute efficiently. This is the key to achieving the 10.9x speedup.

Results

• 50.9% zero-shot accuracy across seven benchmarks with 1.7B active parameters
• 10.9x speedup in inference latency (73ms → 6.7ms) vs comparable fine-grained approaches
• Open-source code available

Takeaways

1. Granularity doesn’t require inefficiency: Hardware-aware design can make fine-grained routing practical
2. Routing complexity is reducible: Cartesian product decomposition is an elegant solution
3. Memory access patterns matter: Expert-centric scheduling is as important as the routing algorithm itself