Redesign Mixture-of-Experts Routers with Manifold Power Iteration

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.

• The paper presents MPI, which aligns each router row with the leading singular vector of its expert weight matrix using a single-step power iteration and norm retraction.
• It shows accelerated convergence and improved load balancing in MoE systems, validated on models ranging from 1B to 11B parameters.
• Empirical evaluations across multiple optimizers demonstrate MPI’s consistent enhancements in downstream performance with minimal computational overhead.

Redesigning Mixture-of-Experts Routers with Manifold Power Iteration

Motivation for Router Redesign in MoEs

The Mixture-of-Experts (MoE) architecture is central to scaling the capacity of LLMs under fixed compute budgets. A critical bottleneck is the router, parameterized as a matrix, which assigns input tokens to expert modules based on affinity scores. The conventional implementation—using unconstrained learned rows for each expert—lacks any alignment between the router vectors and the intrinsic features of the associated expert's weight matrices. This can result in suboptimal routing, hampering both convergence and performance.

This paper introduces a principled approach: each router row is explicitly aligned with the principal singular vector of its corresponding expert's weight matrix. The intuition is that singular vectors encode the most information-dense direction in the weight matrix, thus providing a natural, information-theoretic coupling between experts and routers. To make this alignment computationally tractable, the authors propose a lightweight approximation using single-step power iteration and a norm-constraining retraction mechanism, forming the "Manifold Power Iteration" (MPI) paradigm.

The Manifold Power Iteration Paradigm

The MPI router design couples each router row to the leading singular direction of its associated expert. This is implemented via the following two-stage procedure at each training step:

1. Power Iteration Step: For each router row, a single power iteration is performed with the associated expert's weight matrix, efficiently tracking the dominant singular direction.
2. Retract (Norm Constraint) Step: The updated router row is $L_2$-normalized to a fixed constant to ensure numerical stability and prevent norm explosion or collapse.

This "Power-then-Retract" scheme guarantees that the router vectors remain well-behaved and improves the expressivity of token-expert assignment. Theoretical analysis demonstrates that MPI can be understood as an adaptive steepest ascent procedure under a maximum projection constraint on the spherical manifold—driving each router row to align with the principal component of the expert weights as training progresses.

Convergence and Optimization Analysis

MPI imposes an explicit geometric prior, ensuring that the dot product between an input and its routed expert's vector genuinely reflects the structure of the expert. The adaptive step size property—whereby updates diminish as alignment improves—prevents oscillatory or divergent behavior even as models scale in both model and expert count.

Extensive pretraining experiments, from 1B to 11B parameter models, show that MPI is compatible with a broad array of optimizers, including AdamW, Muon, and Hyperball approaches. Across all tested optimizers, the MPI-enabled router delivers both accelerated convergence and consistent gains in downstream task performance.

Figure 2: Pre-training loss comparison for a 1B MoE model across optimizers (AdamW, AdamH, Muon). MoE with MPI achieves a convergence advantage over all alternative setups.

Empirical Evaluation: Convergence, Load Balancing, and Performance

Accelerated Convergence and Superior Downstream Metrics

MPI yields a substantial improvement in pretraining loss and downstream metrics over conventional routers. For example, in 11B-parameter models, the convergence speed is consistently higher, and accuracy on downstream evaluation suites is uniformly better across all stages of pretraining and mid-training.

Figure 4: Convergence and Downstream Performance Comparison. Manifold Power Iteration facilitates faster convergence and superior downstream task performance throughout the entire course of 11B MoE pretraining.

Improved Load Balancing

Another practical benefit is the enhancement of expert load balancing. The retraction step regularizes router norms, which, in turn, helps prevent expert overload or starvation emerging from pathological router norm inflation. Empirically, the load balancing loss decreases rapidly and remains low during training.

Figure 3: Load balancing loss for 3B MoE with MPI.

This improved balance is not achieved at the cost of significant computational overhead. The per-step costs of MPI are negligible (<0.2% reduction in throughput in 11B-scale experiments), and no additional inference latency is incurred since routers can be precomputed at load time.

Ablation Studies and Sensitivity Analysis

Importance of Power Iteration and Router Retraction

A systematic ablation demonstrates the necessity of both components: omitting router retraction results in unstable training and loss spikes (especially for AdamW and Muon optimizers), while using norm-retraction alone produces no meaningful advantage over the vanilla router. The empirical evidence highlights that router retraction is critical for stability, and power iteration is essential for improved convergence and performance.

Figure 5: Ablation studies for the key design choices: (1) Power Iteration and (2) Router Retraction. Pretraining collapses without Router Retraction when using AdamW and Muon.

Optimizer Robustness and Hyperparameter Stability

MPI was validated across four optimizer configurations (AdamW, AdamH, Muon, and MuonH), with all benefiting from the redesigned router. Notably, the choice of retraction constant $C$ is shown to be robust; the design principle $C \sim 1/\sqrt{N}$, where $N$ is the number of experts, ensures scale invariance and obviates costly hyperparameter tuning even at large scale.

Theoretical and Practical Implications

The direct alignment between router rows and expert principal components introduced by MPI provides a theoretical foundation for more principled router design in MoE architectures. This connection likely increases the interpretability of token assignments and supports the development of more robust auxiliary losses for load balancing and specialization.

From a practical perspective, the methodology has immediate implications for the scalability and stability of very large sparse LLMs, as well as for transfer learning and model deployment. By decoupling routing dynamics from optimizer hyperparameters and model scale, MPI advances the feasibility of robust routing in extremely large MoE deployments.

Prospects for Future Research

Possible future work includes extending the MPI framework to hybrid or multi-matrix expert scenarios, integrating additional spectral or geometric regularization schemes, or exploring its impact in non-language-modality MoEs. An open line of investigation is the further coupling of routers with learned representations for auto-specialization under dynamic expert groupings, given the tight theoretical connection exposed here.

Conclusion

The "Redesign Mixture-of-Experts Routers with Manifold Power Iteration" paper establishes a mathematically principled, computationally efficient, and empirically validated paradigm for router design in MoEs. By utilizing online power iteration with norm retraction to achieve principal singular vector alignment, the authors demonstrate optimizer-agnostic, scalable gains in both convergence and downstream task performance, coupled with improved load balancing and training stability. These contributions set a new standard for router design and indicate promising directions for further advancements in sparse expert-based LLMs.