omni-moe-atomic-experts

name	omni-moe-atomic-experts
title	OmniMoE: Efficient MoE by Orchestrating Atomic Experts
version	0.0.2
engine	skillxiv-v0.0.2-claude-opus-4.6
license	MIT
url	https://arxiv.org/abs/2602.05711
keywords	["Mixture-of-Experts","Atomic Experts","Routing","System-Algorithm Codesign","Efficiency"]
description	Scale mixture-of-experts models efficiently by decomposing experts into atomic vector pairs with Cartesian product routing and expert-centric scheduling. Achieves 10.9× speedup and 50% fewer parameters versus fine-grained baselines through system-algorithm codesign that converts scattered memory access into contiguous batched operations.

OmniMoE: Atomic Expert Orchestration for Efficient Scaling

Scaling mixture-of-experts (MoE) systems to millions of experts requires solving two orthogonal challenges: routing complexity and hardware efficiency. Standard MoE scales linearly with expert count, requiring O(N) per-token routing computations for N experts. Additionally, dynamic expert selection creates memory-bound random access patterns that underutilize modern accelerators designed for dense matrix operations.

OmniMoE solves both problems through system-algorithm codesign. Instead of storing full expert networks, experts reduce to minimal vector pairs (atomic form). Routing decomposes the 1D expert space into 2D grids, reducing complexity from O(N) to O(√N). Critically, execution reorders computation from token-centric (fetch different experts per token) to expert-centric (group requests targeting the same experts), converting scattered lookups into contiguous memory access and high-throughput matrix operations.

Core Concept

Traditional MoE for each token computes routing scores for all N experts, selects top-k, then fetches their parameters. This creates bottlenecks:

1. Routing bottleneck: O(N) score computations per token
2. Memory bottleneck: Each token fetches scattered expert parameters, underutilizing hardware

OmniMoE addresses both:

Atomic Experts: Experts reduce to pairs of vectors (U, V) with dimensions [d_hidden, r] and [r, d_out]. A token selects experts and combines their contributions: output = Σ(α_i · U_i · V_i^T · input)

This slashes parameter count—1M atomic experts require less storage than 10K dense experts.

Cartesian Product Routing: The expert index space decomposes: i ∈ [0, N) becomes (i₁, i₂) where i₁ ∈ [0, √N), i₂ ∈ [0, √N). Routing requires computing √N dimensional projections separately, reducing scoring to O(√N) per token.

Expert-Centric Scheduling: Instead of processing tokens independently, the system groups all tokens requesting expert (i₁, i₂) and computes them in a single batched GEMM. This inverts execution order and converts random access into sequential, cache-friendly operations.

Architecture Overview

• Atomic Expert Decomposition: Store experts as vector pairs (U ∈ ℝ^[d×r], V ∈ ℝ^[r×d]) rather than full dense matrices
• Factorized Routing: Decompose N experts into √N × √N grid; compute two separate √N-dimensional projections
• Two-stage Routing: Stage 1 selects top-k₁ experts from first dimension, Stage 2 selects top-k₂ from second, yielding up to k₁·k₂ experts per token
• Expert-Centric Execution: Batch tokens by target expert and execute as fused GEMM operations