Redesign Mixture-of-Experts Routers with Manifold Power Iteration
2026-06-10 • Machine Learning
Machine LearningArtificial IntelligenceComputation and Language
AI summaryⓘ
The authors focus on improving the router in Mixture-of-Experts (MoE) models, which decides which experts process each input. They propose aligning each router row with the main direction (principal singular direction) of its expert to better represent expert characteristics. To achieve this, they introduce a method called Manifold Power Iteration (MPI) that adjusts the router weights efficiently and stably. Their tests on large MoE models show that this alignment helps the models perform better.
Mixture-of-Experts (MoE)RouterPrincipal Singular DirectionSingular Value DecompositionManifold Power IterationPower IterationMatrix NormModel PretrainingToken-Expert AffinityRetract Operation
Authors
Songhao Wu, Ang Lv, Ruobing Xie, Yankai Lin
Abstract
Router is the cornerstone component to the Mixture-of-Experts models. Serving as expert proxies, the rows of the router matrix compute their similarity to the MoE inputs to determine which subset of experts is activated. Ideally, each router row is designed to encode the expert matrix into this representative vector, such that its dot-product with token can better reflect token-expert affinity. However, there exists no design principles to enforce this condensation. In this paper, we propose to align each router row with the principal singular direction of the associated expert, as this direction provides the most expressive mathematical description of a matrix. Based on this principle, we propose a router redesign with Manifold Power Iteration (MPI). Specifically, it introduces a "Power-then-Retract" paradigm, where a power iteration step is performed on the router weights, followed by a retraction to impose a norm constraint to ensure both efficiency and stability. Theoretically, we show that MPI drives router rows to converge toward the principal singular directions of associated experts. Empirically, we pretrain MoE model across scales from 1B to 11B parameters to confirm that this alignment facilitates more effective MoE models.