Paper 2606.12397
Redesign Mixture-of-Experts Routers with Manifold Power Iteration
- Published
- Jun 2026
- Research lab
- Independent
- Citations
- 0
- GitHub
- Not linked
01 In brief
Summary
This paper introduces Manifold Power Iteration (MPI), a redesign of Mixture-of-Experts (MoE) routers to align each router row with the principal singular direction of its associated expert's weight matrix.
The method uses a 'Power-then-Retract' paradigm: a single power iteration step on router weights followed by L2 norm retraction to ensure stability.
Theoretically, MPI is shown to be equivalent to a steepest ascent optimization that maximizes the router's projection onto expert weights.
Empirically, MPI is evaluated on MoE models from 1B to 11B parameters, showing faster convergence, lower pretraining loss, and improved downstream performance across multiple optimizers (AdamW, AdamH, Muon, MuonH).
For example, at 11B scale, MPI reduces validation perplexity from 0.728 to 0.723 and improves average accuracy on 7 challenging benchmarks from 40.92 to 42.76.
MPI also improves load balancing and incurs negligible training overhead (0.2% slowdown) with zero inference overhead.
Ablations confirm that both power iteration and retraction are essential, and the method is compatible with auxiliary losses and alternative activation functions.
02 From the paper
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.