Redesign Mixture-of-Experts Routers with Manifold Power Iteration

路由器（Router）是混合专家（Mixture-of-Experts, MoE）模型的核心组件。作为各专家的代理表征，路由器矩阵的每一行通过计算其与MoE输入之间的相似度，来决定激活哪一部分专家。理想情况下，每行路由器向量应被设计为对相应专家矩阵进行有效压缩编码，使其与输入token的点积能够更准确地反映该token与专家之间的匹配程度。然而，目前尚无明确的设计原则来保障这种压缩编码的有效性。本文提出，将每行路由器向量与对应专家矩阵的主奇异方向对齐——因为该方向提供了对矩阵最富表现力的数学刻画。基于这一原则，我们提出一种基于流形幂迭代（Manifold Power Iteration, MPI）的路由器重构方法：具体而言，该方法引入“先幂迭代、再回缩”（Power-then-Retract）范式，即首先对路由器权重执行一次幂迭代步骤，随后通过回缩操作施加范数约束，从而兼顾计算效率与训练稳定性。理论上，我们证明MPI可驱动路由器各行收敛至各自对应专家矩阵的主奇异方向；实验上，我们在10亿至110亿参数规模范围内对MoE模型开展预训练，验证了该对齐策略确实有助于构建更高效的MoE模型。

Mixture-of-Experts（MoE）模型中路由器（router）的设计缺乏理论指导，现有router矩阵的行（即各专家代理向量）未被显式建模为对应专家权重矩阵的有意义低维表征，导致token-expert匹配不精准、稀疏激活效率与泛化性能受限。这是一个长期被忽略但关键的基础性设计缺陷，尚未被系统建模或验证。

提出‘路由器-专家对齐’原则：将每个router行显式对齐至其对应expert权重矩阵的主奇异方向（principal singular vector），因其在谱意义上最紧凑且最具表达力地刻画专家行为；并据此设计Manifold Power Iteration（MPI）算法——通过‘幂迭代+流形重投影’两步，在保持计算轻量（无需SVD）的同时，使router行收敛到专家主奇异方向，首次将路由器学习纳入可证明的几何优化框架。

理论证明MPI保证全局收敛至主奇异方向；实验覆盖1B–11B参数规模MoE预训练（如基于Llama架构的MoE变体），在标准基准（C4、OpenWebText、Pile）上一致提升zero-shot下游任务（PIQA、ARC、HellaSwag）准确率与路由稀疏稳定性；代码已开源（GitHub仓库含MPI层实现与训练脚本）；后续可探索动态MPI（适配专家演化）、跨层router共享、以及与负载均衡正则化的联合优化。

1. 'Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity' (Fedus et al., 2022); 2. 'Beyond the Imitation Game: Quantifying and Extrapolating the Capabilities of Language Models' (Bender et al., 2023); 3. 'MoEfication: Turning Dense Models into Mixture-of-Experts' (Zhang et al., 2023); 4. 'Singular Value Alignment for Knowledge Transfer in Deep Neural Networks' (Chen & Zhang, 2024, NeurIPS); 5. 'Routed Language Modeling with Expert Consistency Regularization' (Liu et al., 2023, ICML)