PaperlayerView exact revision on arXiv

Source-linked AI summary

Redesign Mixture-of-Experts Routers with Manifold Power Iteration

Songhao Wu, Ang Lv, Ruobing Xie, Yankai Lin

arXiv:2606.12397v1cs.LGcs.AIcs.CL

TL;DR

MoE routers lack an explicit constraint to encode experts’ intrinsic features, limiting their use as expert proxies. MPI aligns router rows with experts’ principal singular directions through power iteration and retraction, and experiments up to 11B parameters report faster convergence, stronger downstream performance, and better load balancing.

  • Problem

    MoE routers lack an explicit constraint ensuring that router rows faithfully encode experts’ intrinsic features for token–expert affinity.

  • Method

    MPI performs online power iteration followed by L2-norm retraction, driving router rows toward the principal singular directions of their associated experts.

  • Results

    Experiments up to 11B parameters show that MPI consistently accelerates convergence, improves downstream performance, and enhances load balancing.

  • Takeaways & Limitations

    MPI provides a scalable, theoretically grounded router design that explicitly connects router representations with expert parameters.

  • Takeaways & Limitations

    The design assumes that routing logits should remain bounded at a constant scale to prevent explosion.

Abstract

from arXiv · show

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.

1 Gaoling School of Artificial Intelligence, Renmin University of China · 1 Introduction

The paper redesigns MoE routers by aligning each router row with the principal singular direction of its corresponding expert weight matrix. Its Manifold Power Iteration method combines one power iteration with norm-preserving retraction, with theory and experiments supporting improved router–expert alignment across models up to 11B parameters.

  • 1 Introduction: MoE replaces Transformer FFNs with expert ensembles and uses sparse, per-token routing to scale model capacity under a constrained computation budget.The architecture selects experts for each token rather than activating all expert modules.
  • 1 Introduction: The router computes token–row similarity scores and dispatches each token to experts associated with the highest-scoring rows.Routers are typically parameterized as linear matrices.
  • 1 Introduction: Conventional routers lack a constraint forcing limited-expressivity rows to encode their experts’ intrinsic features, weakening token–router affinity as a proxy for token–expert assignment.The paper identifies this missing encoding constraint as a design problem.
  • 1 Introduction: The proposed design aligns each router row with its expert’s principal singular direction, which preserves the highest density of matrix information as a compressed representation.This principle applies because expert modules are parameterized by weight matrices and routers encode them into single vectors.
  • 1 Introduction: MPI performs one power iteration on router weights per training step, then retracts them to maintain a constant L2-norm scale and prevent explosion or collapse.This “Power-then-Retract” paradigm defines Routers with Manifold Power-Iteration.
  • 1 Introduction: Theoretical analysis characterizes MPI as steepest ascent maximizing router projection onto expert weights under minimal updates, driving router rows toward principal singular directions.The update rule is presented as producing adaptive convergence during training.
  • 1 Introduction: Pretraining experiments across MoE scales using billions of tokens, including models up to 11B parameters, evaluate MPI as a router redesign while retaining the standard router interface.The paper reports intrinsic improvements and frames MPI as a new perspective on router–expert interplay.

2 Background: Mixture-of-Experts

In MoE-based LLMs, a router projects each input into gating weights that select a sparse subset of experts. However, the standard router lacks an explicit constraint ensuring its rows faithfully preserve experts’ intrinsic features.

  • Router mechanism: The router dispatches each input to a sparse subset of experts using a 2D linear weight matrix R ∈ R^N×D.The input x ∈ R^D is projected by R.
  • Router mechanism: The router produces a gating-weight vector over N experts, and the top-K experts are selected.The selected experts determine the MoE layer’s weighted-sum output.
  • Router limitation: Standard routers impose no explicit constraint ensuring each row faithfully reflects its associated expert’s intrinsic features.Ideally, each router row would preserve the expert-weight geometry and serve as a representative feature vector for affinity computation.

3 Methodology

MPI redesigns MoE routers by aligning each router row with its expert’s principal singular direction through efficient power iteration and norm-constrained retraction. Its optimization interpretation explains this update as adaptive constrained ascent that progressively rotates router rows toward the most informative expert features.

  • Motivation: MoE router rows should encode their associated expert matrices, with principal singular directions capturing the matrices’ most essential traits.This alignment makes router–token inner products better reflect token–expert affinity.
  • Manifold Power-Iteration: MPI replaces costly per-step SVD with power iteration, allowing router rows to track and converge toward expert principal singular directions through matrix-vector products.The method is motivated by the computational cost of exact SVD across all expert matrices.
  • Manifold Power-Iteration: MPI follows a “Power-then-Retract” paradigm: power iteration aligns each router row with expert weights, while L2 retraction constrains norms for containment and numerical stability.Retraction also mitigates expert bias caused by router norm disparities that can inflate gating weights and overload experts.
  • Design Principle: The retraction scale is chosen as C := C′ √N so routing logits maintain an O(1) ceiling while scaling effects from the expert count N are decoupled.This follows the design principle that C scales as Θ(1/√N) for scale-invariant inputs before redefinition.
  • Optimization Interpretation: MPI is an optimization tailored to maximum-projection constraints with adaptive stepsizes, steadily rotating router weights toward expert principal singular directions and informative features.As alignment increases, updates become more careful; the SVD view identifies the limiting direction as the dominant singular vector.

4 Experiment

Experiments from 1B to 11B show that MoE with MPI converges faster, improves language-modeling and downstream performance, and maintains practical training and inference efficiency. MPI also improves load distribution relative to vanilla MoE, with the balance improvement tentatively attributed to router retraction.

  • 1B experiments: At 1B scale, MPI achieves accelerated convergence and improved downstream performance across AdamW, Muon, AdamH, and MuonH setups.Models are pretrained on 100B tokens, with downstream performance averaged over 25 benchmarks.
  • Convergence and Performance: MPI delivers consistent performance gains on challenging benchmarks at both 3B and 11B scales.These results further validate the effectiveness of the redesigned router.
  • Load Balancing: MPI’s balance loss drops sharply early in pretraining and remains low, while MaxVio provides a more accurate load-balance assessment because retraction may artifactually reduce balance loss.The reported MaxVio values confirm compatibility with load-balancing loss and more equitable load distribution than vanilla MoE.
  • Efficiency Analysis: 0.2% slowdown is incurred by MPI relative to vanilla MoE’s 34.97 billion tokens per day throughput in 11B pretraining.MPI adds zero communication overhead and avoids conflicts with standard training frameworks.
  • Efficiency Analysis: MPI incurs zero inference overhead because router weights can be pre-computed during model loading, preserving compatibility with standard inference engines.This supports deployment at larger scale without modifying standard inference workflows.

5 Method Analysis

Post-hoc analyses show that MPI aligns router rows more closely with experts’ principal singular directions and that a single power iteration is sufficient. Ablations further establish Router Retraction as critical for stable training, while MPI remains effective across hyperparameter and expert-matrix choices.

  • Alignment Analysis: MPI produces significantly larger λ values than vanilla MoE, indicating tighter alignment between router vectors and the principal singular directions of expert weights.λ is normalized by the spectral norm and averaged across experts per layer.
  • Iteration Count: Increasing power iterations to 10 lowers throughput by 5% without improving convergence or downstream performance, supporting the sufficiency of a single iteration.The additional iterations provide a more precise estimation but no further practical benefit.
  • Power Iteration Ablation: Row-wise normalization without power iteration performs nearly identically to vanilla MoE and underperforms MPI, showing that retraction alone does not explain MPI’s gains.The ablated variant nevertheless exhibits a similar balance-loss distribution to MPI.
  • Router Retraction Ablation: Without Router Retraction, power iteration destabilizes pretraining through loss spikes, abnormal gradients, and suboptimal convergence, especially with AdamW and Muon.Figure 5 reports pretraining collapse without retraction for these optimizers, which lack weight constraints.
  • Expert Matrix Choice: Across Wg, Wp, and Wo candidates, no significant differences emerge in pretraining loss or downstream performance; Wg is selected by its marginal advantage.The comparison uses 1B MoE baselines pretrained on 50B tokens.

6 Compatibility of Manifold Power-Iteration with other Router Designs

MPI modifies only router weights while preserving gating-weight computation, making it compatible with alternative router designs. Experiments show compatibility with auxiliary losses and Sigmoid activation, including downstream improvements.

  • Router-design compatibility: MPI preserves gating-weight computation and changes only router weights, making the refinement conceptually orthogonal to most alternative router designs.Compatibility was explored by pretraining 1B-parameter baselines on 50B tokens.
  • Auxiliary loss compatibility: The router design remains compatible with load-balancing loss and router z-loss, with the z-loss variant improving downstream tasks by 0.68 points.The router z-loss used a coefficient of 0.001, with no observed loss or gradient anomalies in small-scale trials.
  • Activation-function compatibility: Using Sigmoid instead of Softmax narrows the pretraining-loss advantage but improves downstream performance from 41.64 to 42.05.The Sigmoid experiment fixes C = 1 to align with the Frobenius norm of MuonH; broader activation-function exploration is deferred.

7 Related Work

This section situates the paper within optimizer research focused on convergence acceleration, scalability, and norm-constrained optimization. It highlights Muon and recent weight-and-update norm constraints as relevant foundations for the paper’s empirical studies.

  • Muon: Muon orthogonalizes momentum with Newton-Schulz iterations to update parameters and has been validated in pretraining models with up to trillions of parameters.Further analysis interprets Muon as steepest descent under the spectral norm, motivating other norm-constrained optimizer designs.
  • Norm-constrained optimizers: Recent work imposes norm constraints on both weights and updates to support stable, scalable optimization and accelerate convergence across scales.These constraints also allow hyperparameter transfer without further tuning.
  • This paper’s scope: The paper provides preliminary practice of these norm-constrained optimizers and empirically validates their effectiveness.Its empirical insights span models up to 11B parameters trained on 350B tokens, with the goal of extrapolating efficacy to larger scales.

8 Conclusion · A Supplementary Derivations for Approximation in Equation 10 · B Details for Pretraining Experiments

The paper concludes that Manifold Power Iteration (MPI) provides an efficient, theoretically grounded MoE router design connecting router representations with expert parameters. The supplementary derivation explains that power iteration drives the orthogonal component toward negligibility, yielding the approximation used in Equation 10.

  • 8 Conclusion: MPI reframes MoE router design through row-wise expert-proxy representations and connects router representations with expert parameters.The conclusion presents MPI as a principled router redesign motivated by representation geometry in MoEs.
  • 8 Conclusion: MPI uses lightweight iterative updates while maintaining scalability, offering an efficient alternative to conventional router designs.The conclusion characterizes MPI as both efficient and scalable.
  • 8 Conclusion: The paper reports extensive experiments validating MPI across scales.The supplied conclusion passage truncates the specific scales and experimental results.
  • A Supplementary Derivations for Approximation in Equation 10: The supplementary derivation decomposes the router-weight update Δr into components in the subspace spanned by R′[i] and its orthogonal complement.The derivation introduces the update and its projection onto the relevant subspace and complement.
  • A Supplementary Derivations for Approximation in Equation 10: As power iteration proceeds, R′[i] approaches the dominant subspace, making the orthogonal component negligible.This asymptotic behavior motivates the approximation that follows in the derivation.
  • A Supplementary Derivations for Approximation in Equation 10: A rearrangement of terms produces the approximation used in Equation 10 and completes its derivation.The supplied passages state that the approximation follows by rearranging terms and completes the derivation of Eq. 10.

B.1 Implementation Details · B.2 Optimizer Setup Details

The experiments use a TorchTitan-based, fully sharded MoE training pipeline and scale a 3B model to 11B total parameters by expanding its expert count. Optimizer configurations are calibrated against AdamW through update-RMS alignment, with specific learning-rate scaling rules for Muon and Hyperball Optimization.

  • B.1 Implementation Details: Scaling experts from 64 to 256 expands the 3B model into an 11B-parameter sparse MoE with 470M activated parameters.The resulting model increases expert count while keeping only a subset of parameters active.
  • B.1 Implementation Details: The training pipeline uses TorchTitan, PyTorch SDPA for attention, MegaBlocks MLP for MoE implementation, and Fully Sharded Data Parallel across pretraining experiments.These components provide the framework, attention implementation, efficient MoE implementation, and model-parallel training strategy.
  • B.1 Implementation Details: Table 7 summarizes the hyperparameters used across the model-architecture experiments.The table provides the architecture-level configuration accompanying the reported scaling experiments.
  • B.2 Optimizer Setup Details: AdamW hyperparameters for 1B MoE pretraining were selected through a search over key configurations and used as the reference for other optimizers.The setup is intended to support fair convergence comparisons without repeating extensive searches.
  • B.2 Optimizer Setup Details: For Muon, update-RMS alignment scales the learning rate of Muon-optimized parameters by 0.2 × p max(din, dout).The scaling is applied relative to the AdamW reference configuration.
  • B.2 Optimizer Setup Details: Hyperball Optimization fixes the weight-matrix Frobenius norm at √dout and uses a learning-rate scaler of 0.2 × √din for update-RMS alignment.The din value from a 1B model is substituted into the formula, and the resulting scale is held constant across model scales.

C Evaluation Setup

The evaluation uses OLMES for downstream assessment, with nine core pretraining tasks to quickly measure checkpoint capabilities. Unless otherwise specified, evaluation follows Olmo et al. (2025) on a 25-task multiple-choice benchmark, with detailed task scores omitted.

  • All downstream task evaluations use OLMES.
  • Nine core tasks assess fundamental checkpoint capabilities during pretraining: ARC-Easy, ARC-Challenge, MMLU, CommonsenseQA, SocialIQA, HellaSwag, WinoGrande, PIQA, and SciQ.
  • Unless otherwise specified, evaluation follows Olmo et al. (2025) on a benchmark comprising 25 multiple-choice tasks.The complete task list appears in Table 9.
  • Detailed scores for the 25 benchmark tasks are omitted unless explicitly specified.
Loading 2606.12397v1…