MoRE: Scaling mixture of experts with hardware-aware low-rank routing
Honam Wong ⋅ Surbhi Goel ⋅ Enric B Adserà
Abstract
Mixture-of-Experts (MoE) layers are central to frontier language models, with recent architectures pushing toward larger total expert counts and finer-grained experts. In this regime, the standard linear router becomes a bottleneck: with $M$ experts and hidden dimension $h$, its per-token cost $\Theta(Mh)$ grows to dominate the MoE layer once $M$ is large. We introduce MoRE (Mixture of Rank-reduced-routed Experts), which factorizes the $h \times M$ routing matrix through a rank-$r$ bottleneck, reducing routing cost to $O((h + M)r)$. We prove that rank logarithmic in $M$ is sufficient to preserve the memorization capabilities. At matched active FLOPs, the factorization unlocks a $\Theta(h/r)$ factor more experts. However, realizing this gain in wall-clock requires more than the FLOP reduction, so we design a fused Triton kernel that avoids expensive memory operations on HBM. Empirically, MoRE's additional capacity translates into improved memorization on a synthetic phonebook task and stronger performance on knowledge-intensive Q\&A after pretraining, while preserving reasoning ability.
Chat is not available.
Successful Page Load