Manifold Power Iteration (MPI)

• Manifold Power Iteration (MPI) is a power-iteration extension that incorporates manifold retraction to solve constrained optimization problems on the Stiefel manifold and spheres.
• The method alternates between an affine update (via SVD or scaling) and a projection step, ensuring convergence to a local optimum while handling orthogonality constraints.
• MPI is applied in orthogonal regression and Mixture-of-Experts router redesign, improving convergence speed, stability, and load balancing in complex model architectures.

Searching arXiv for the cited papers and closely related uses of “manifold power iteration.” Manifold Power Iteration (MPI) denotes a class of updates that extend the power-iteration principle from unconstrained eigenvector computation to constrained optimization on manifolds. In the literature represented here, MPI has two closely related instantiations. In the 2017 work on the quadratic problem on the Stiefel manifold, the method appears as a generalized power iteration (GPI): an affine transformation is followed by projection back to the Stiefel manifold, yielding a power-iteration-like solver for orthogonally constrained quadratic objectives (Nie et al., 2017). In the 2026 work on Mixture-of-Experts (MoE) routing, MPI is the explicit name of a router redesign based on a “Power-then-Retract” paradigm: a power step is applied to each router row using the associated expert matrix, followed by retraction onto a fixed-radius sphere (Wu et al., 10 Jun 2026). This suggests that MPI is both a specific modern method name and a broader algorithmic pattern in which linear or affine ascent steps are alternated with manifold retractions.

1. Conceptual structure

The unifying idea of MPI is that a power-style update can be generalized from Euclidean space to a constrained set by replacing scalar normalization with a manifold-compatible projection or retraction. In standard power iteration for a symmetric matrix $A$, the dominant eigenvector is computed by repeated updates of the form

$w \leftarrow \frac{Aw}{\|Aw\|_2}.$

For the top $k$ eigenvectors, orthogonal iteration uses

$$M \leftarrow AW,\qquad M = QR,\qquad W \leftarrow Q,$$

or equivalently the compact SVD form

$M = USV^T,\qquad W \leftarrow UV^T.$

For positive semidefinite $A$, this solves

$\max_{W^TW=I_k} \operatorname{Tr}(W^T A W).$

The 2017 formulation generalizes this structure to an affine-quadratic objective on the Stiefel manifold, with update

$$W \leftarrow \operatorname{Proj}_{\mathrm{St}(m,k)}(2\tilde A W + 2B),$$

where the constraint set is

$\{W : W^TW=I_k\}.$

The 2026 formulation applies the same principle row-wise to router weights in an MoE model: one power iteration step is taken with respect to an expert-derived positive semidefinite matrix, and the result is retracted to a sphere of fixed radius.

A common misconception is to identify MPI with ordinary repeated multiplication by a matrix. In both cited uses, the manifold operation is essential. On the Stiefel manifold, the SVD-based update is the exact solution of a linear maximization subproblem; in MoE routing, the retraction is introduced to ensure both efficiency and stability and to control the scale of routing logits (Nie et al., 2017).

2. Stiefel-manifold formulation and generalized power iteration

The 2017 paper studies the quadratic optimization problem on the Stiefel manifold

$$\min_{W^TW=I_k} \operatorname{Tr}(W^T A W - 2W^T B),$$

where $w \leftarrow \frac{Aw}{\|Aw\|_2}.$0 has orthonormal columns, $w \leftarrow \frac{Aw}{\|Aw\|_2}.$1 is symmetric, and $w \leftarrow \frac{Aw}{\|Aw\|_2}.$2. The authors call this problem QPSM. The motivation is that many problems in orthogonal regression and Procrustes analysis can be rewritten in exactly this form.

Because $w \leftarrow \frac{Aw}{\|Aw\|_2}.$3 may not be positive definite, the derivation introduces a scalar $w \leftarrow \frac{Aw}{\|Aw\|_2}.$4 such that

$w \leftarrow \frac{Aw}{\|Aw\|_2}.$5

QPSM is then equivalent, up to the constant shift from $w \leftarrow \frac{Aw}{\|Aw\|_2}.$6, to

$w \leftarrow \frac{Aw}{\|Aw\|_2}.$7

The Lagrangian is

$w \leftarrow \frac{Aw}{\|Aw\|_2}.$8

and stationarity gives

$w \leftarrow \frac{Aw}{\|Aw\|_2}.$9

This yields the fixed-point structure

$k$0

The next iterate is defined by the inner subproblem

$k$1

If the compact SVD of $k$2 is

$k$3

then the maximizer is

$k$4

The resulting generalized power iteration therefore has the explicit form: $k$5

In this setting, the final step is identified as a manifold retraction on the Stiefel manifold. The method is therefore “power-iteration-like” not merely by analogy but by construction: multiplication by $k$6 in classical orthogonal iteration is replaced by the affine map $k$7, and Euclidean orthonormalization is replaced by Stiefel projection via SVD (Nie et al., 2017).

3. Optimality of the projection step and convergence properties

The SVD update is exact for the linearized subproblem. For

$k$8

if $k$9, then

$$M \leftarrow AW,\qquad M = QR,\qquad W \leftarrow Q,$$0

with equality achieved at

$$M \leftarrow AW,\qquad M = QR,\qquad W \leftarrow Q,$$1

or equivalently, using the compact SVD, $$M \leftarrow AW,\qquad M = QR,\qquad W \leftarrow Q,$$2. This identifies the retraction as the exact optimizer of the local linear objective rather than a heuristic normalization.

The monotonicity analysis relies on the lemma that if $$M \leftarrow AW,\qquad M = QR,\qquad W \leftarrow Q,$$3, then for any $$M \leftarrow AW,\qquad M = QR,\qquad W \leftarrow Q,$$4,

$$M \leftarrow AW,\qquad M = QR,\qquad W \leftarrow Q,$$5

The proof uses

$$M \leftarrow AW,\qquad M = QR,\qquad W \leftarrow Q,$$6

with $$M \leftarrow AW,\qquad M = QR,\qquad W \leftarrow Q,$$7. From this, Theorem 1 states that the algorithm monotonically decreases the original QPSM objective

$$M \leftarrow AW,\qquad M = QR,\qquad W \leftarrow Q,$$8

If $$M \leftarrow AW,\qquad M = QR,\qquad W \leftarrow Q,$$9 is the updated iterate, then

$M = USV^T,\qquad W \leftarrow UV^T.$0

Theorem 2 proves convergence to a local minimum of QPSM. At a fixed point, the update satisfies the KKT condition for the linearized subproblem,

$M = USV^T,\qquad W \leftarrow UV^T.$1

and since $M = USV^T,\qquad W \leftarrow UV^T.$2, this becomes

$M = USV^T,\qquad W \leftarrow UV^T.$3

which matches the KKT condition of the reformulated problem. Because the reformulation is equivalent to QPSM, the fixed point is a stationary or local optimum for QPSM.

The paper does not prove global optimality. It notes that the subproblem $M = USV^T,\qquad W \leftarrow UV^T.$4 has a unique solution when $M = USV^T,\qquad W \leftarrow UV^T.$5 has full column rank, due to SVD uniqueness, and that empirically many random initializations converge to the same objective value. A plausible implication is that the method often behaves more globally than its formal guarantee establishes, but the formal claim remains local convergence (Nie et al., 2017).

4. Special cases: orthogonal regression and Procrustes analysis

A major significance of the 2017 framework is that several classical orthogonally constrained problems become direct instances of QPSM. For orthogonal least squares regression (OLSR),

$M = USV^T,\qquad W \leftarrow UV^T.$6

with $M = USV^T,\qquad W \leftarrow UV^T.$7, $M = USV^T,\qquad W \leftarrow UV^T.$8, $M = USV^T,\qquad W \leftarrow UV^T.$9, $A$0, and $A$1 the all-ones vector, optimizing over $A$2 yields

$A$3

With the centering matrix

$A$4

OLSR reduces to

$A$5

which expands to

$A$6

with

$A$7

The orthogonal Procrustes family is treated similarly. For

$A$8

with $A$9, $\max_{W^TW=I_k} \operatorname{Tr}(W^T A W).$0, and $\max_{W^TW=I_k} \operatorname{Tr}(W^T A W).$1, the balanced case $\max_{W^TW=I_k} \operatorname{Tr}(W^T A W).$2 is the classical balanced OPP and expands to

$\max_{W^TW=I_k} \operatorname{Tr}(W^T A W).$3

which has a closed-form SVD solution. In the unbalanced case $\max_{W^TW=I_k} \operatorname{Tr}(W^T A W).$4,

$\max_{W^TW=I_k} \operatorname{Tr}(W^T A W).$5

so with

$\max_{W^TW=I_k} \operatorname{Tr}(W^T A W).$6

it becomes QPSM. The corresponding update is

$\max_{W^TW=I_k} \operatorname{Tr}(W^T A W).$7

with $\max_{W^TW=I_k} \operatorname{Tr}(W^T A W).$8 chosen so that

$\max_{W^TW=I_k} \operatorname{Tr}(W^T A W).$9

Empirically, the paper reports experiments on varying dimensions and compares against EB, RSR, LSR, SP, LR, and PMCT for LSQE. The reported conclusions are that GPI converges to the same objective value as competing methods, is consistently faster, often substantially faster, that the advantage becomes more pronounced for larger data dimensions, that for UOPP and LSQE it is especially efficient, that the relaxation parameter $$W \leftarrow \operatorname{Proj}_{\mathrm{St}(m,k)}(2\tilde A W + 2B),$$0 affects convergence speed but not the final convergence behavior, and that the method is robust to random initialization. The low complexity claim is based on forming $$W \leftarrow \operatorname{Proj}_{\mathrm{St}(m,k)}(2\tilde A W + 2B),$$1 or $$W \leftarrow \operatorname{Proj}_{\mathrm{St}(m,k)}(2\tilde A W + 2B),$$2 and performing one compact SVD of an $$W \leftarrow \operatorname{Proj}_{\mathrm{St}(m,k)}(2\tilde A W + 2B),$$3 matrix (Nie et al., 2017).

5. MPI for Mixture-of-Experts router redesign

The 2026 paper uses the term Manifold Power Iteration explicitly for a router redesign in Mixture-of-Experts models. In a standard MoE layer, the router is typically a linear matrix $$W \leftarrow \operatorname{Proj}_{\mathrm{St}(m,k)}(2\tilde A W + 2B),$$4, and for an input token $$W \leftarrow \operatorname{Proj}_{\mathrm{St}(m,k)}(2\tilde A W + 2B),$$5, router logits are computed by

$$W \leftarrow \operatorname{Proj}_{\mathrm{St}(m,k)}(2\tilde A W + 2B),$$6

The paper’s critique is that this design has no explicit principle telling each router row how to represent its expert, even though token–router similarity is used as a stand-in for token–expert affinity.

The proposed design principle is to align each router row with the principal singular direction of the associated expert matrix. For expert $$W \leftarrow \operatorname{Proj}_{\mathrm{St}(m,k)}(2\tilde A W + 2B),$$7, the router row $$W \leftarrow \operatorname{Proj}_{\mathrm{St}(m,k)}(2\tilde A W + 2B),$$8 is treated as a compressed proxy for the expert matrix $$W \leftarrow \operatorname{Proj}_{\mathrm{St}(m,k)}(2\tilde A W + 2B),$$9, and the motivating objective is

$\{W : W^TW=I_k\}.$0

This is presented as a Rayleigh-quotient-style objective: the desired router row is the one whose projection through the expert matrix is maximized.

MPI implements this through the “Power-then-Retract” paradigm. For each router row, with associated expert matrix denoted $\{W : W^TW=I_k\}.$1, the power step is

$\{W : W^TW=I_k\}.$2

Because repeated power steps could cause norm drift, retraction is applied as

$\{W : W^TW=I_k\}.$3

The full router matrix is reconstructed row-wise, and gating is then computed with the normalized rows: $\{W : W^TW=I_k\}.$4

The paper also recommends scaling $\{W : W^TW=I_k\}.$5 with the number of experts $\{W : W^TW=I_k\}.$6 as

$\{W : W^TW=I_k\}.$7

motivated by keeping routing logits bounded at $\{W : W^TW=I_k\}.$8. The retraction is assigned two roles: preventing norm explosion or collapse and avoiding a router row with large norm producing disproportionately large logits and overloading an expert. In this formulation, MPI is not presented as a new dispatch mechanism; gating still uses the standard top-$\{W : W^TW=I_k\}.$9 softmax interface, so the method is a router redesign compatible with existing MoE routing stacks (Wu et al., 10 Jun 2026).

6. Optimization view, empirical evidence, and limitations in MoE

The 2026 paper also gives an optimization interpretation of MPI as a constrained maximum-projection problem: $$\min_{W^TW=I_k} \operatorname{Tr}(W^T A W - 2W^T B),$$0 A first-order approximation yields a gradient-ascent-like update on the sphere. The tangent-space gradient step is written as

$$\min_{W^TW=I_k} \operatorname{Tr}(W^T A W - 2W^T B),$$1

and the approximate MPI update as

$$\min_{W^TW=I_k} \operatorname{Tr}(W^T A W - 2W^T B),$$2

The intended conclusion is that MPI behaves like a steepest-ascent step under a norm constraint, with an adaptive step size determined by current alignment. The theoretical intuition is classical power-iteration theory: because the update uses $$\min_{W^TW=I_k} \operatorname{Tr}(W^T A W - 2W^T B),$$3, the router row is driven toward the dominant eigendirection of that matrix, that is, the principal singular direction of $$\min_{W^TW=I_k} \operatorname{Tr}(W^T A W - 2W^T B),$$4.

The post-hoc alignment metric is

$$\min_{W^TW=I_k} \operatorname{Tr}(W^T A W - 2W^T B),$$5

where larger $$\min_{W^TW=I_k} \operatorname{Tr}(W^T A W - 2W^T B),$$6 means stronger alignment. Reported values are around $$\min_{W^TW=I_k} \operatorname{Tr}(W^T A W - 2W^T B),$$7–$$\min_{W^TW=I_k} \operatorname{Tr}(W^T A W - 2W^T B),$$8 for vanilla MoE layers and around $$\min_{W^TW=I_k} \operatorname{Tr}(W^T A W - 2W^T B),$$9–$w \leftarrow \frac{Aw}{\|Aw\|_2}.$00 for MoE + MPI layers. The empirical evaluation spans pretraining MoE models from 1B to 11B parameters. At 1B scale, MPI improves both pretraining convergence and downstream average accuracy across AdamW, AdamH, Muon, and MuonH. In the 25-benchmark average, the reported changes are 42.26 to 43.56 for AdamW MoE, 42.59 to 43.93 for AdamH MoE, 43.01 to 43.55 for Muon MoE, and 42.78 to 43.98 for MuonH MoE. One representative 1B convergence plot shows a roughly 0.013 reduction in pretraining loss.

For larger-scale models pretrained on 350B tokens and then midtrained on 100B tokens, the reported perplexity changes are 0.764 to 0.754 on 3B validation, 1.688 to 1.581 on 3B math, 1.376 to 1.296 on 3B code, 0.728 to 0.723 on 11B validation, 1.852 to 1.581 on 11B math, and 1.263 to 1.259 on 11B code. Downstream benchmark averages after midtraining change from 36.37 to 38.70 for 3B and from 40.92 to 42.76 for 11B. The paper also reports lower load-balance loss and lower max violation: $w \leftarrow \frac{Aw}{\|Aw\|_2}.$01

Several ablations clarify what MPI is not. A baseline that only performs row-wise normalization,

$w \leftarrow \frac{Aw}{\|Aw\|_2}.$02

performs close to vanilla MoE and significantly worse than full MPI, so the gain is not just from normalization; the power step is essential. Without retraction, power iteration can destabilize training, with reported loss spikes and abnormal gradients, especially for AdamW and Muon. Increasing the number of power iterations from one to ten yields about 5% lower throughput, no convergence or downstream gain, and even a slight worsening of performance in the report. The method is stated to be compatible with load balancing auxiliary losses, router z-loss, and different router activations, including sigmoid in place of softmax. The reported efficiency figures are about 0.2% training slowdown at 11B scale and zero inference overhead because router weights can be precomputed during model loading (Wu et al., 10 Jun 2026).

Taken together, these two lines of work show that MPI is best understood as a manifold-constrained extension of power iteration. On the Stiefel manifold, it solves quadratic orthogonally constrained problems through affine updates and exact SVD projection. On the sphere, it redesigns MoE routers so that each row acts as a norm-controlled proxy for the associated expert. The shared principle is the same: apply a power-like transformation induced by the local objective, then retract to the geometry dictated by the constraint set.