Redesign Mixture-of-Experts Routers with Manifold Power Iteration

原始摘要

At the heart of MoE lies the router, which is typically parameterized as a linear matrix. For each input token, the router computes similarity scores against the matrix rows and dispatches the token to the experts corresponding to the top-scoring rows. While this design is straightforward and has long been accepted as a matter of course, we challenge this conventional wisdom. Ideally, each individual row in MoE router matrix should faithfully reflect the expert’s intrinsic features. The router matrix can thus better ground the identity of each expert, allowing token–router affinity to serve as a precise proxy for token–expert assignment. However, there lacks a constraint in MoE router to enforce the encoding of expert features into router rows of limited expressivity. This absence may lead to suboptimal router design, compromising both training convergence and competence of MoE models. We propose to align each router row with the principal singular direction of its corresponding expert’s weight matrix. This choice is grounded in a linear algebraic intuition: the principal singular direction preserves the highest density of information within a matrix (Golub and Van Loan, 1996; Halko et al., 2010), making it the optimal compressed representation to characterize that matrix. Since each expert module is parameterized as weight matrices, encoding it into a single router vector is exactly the task of capturing its most informative direction. To avoid the prohibitive cost of exact singular value decomposition (SVD), we leverage power iteration (Halko et al., 2010) as a lightweight alternative to obtain this principal direction online. Specifically, the power iteration scheme uses only standard matrix-vector products to solve for the principal singular vector, obviating the need for expensive full matrix factorization. In practice, we perform only one single power iteration on the router weights during each training step. After that, a retraction step is introduced to regularize the L2 norm of the router weights, main- 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 “Powerthen-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.