← 返回 2026-06-11

用流形幂迭代重新设计混合专家路由器 Redesign Mixture-of-Experts Routers with Manifold Power Iteration

Songhao Wu, Ang Lv, Ruobing Xie, Yankai Lin 📅 2026-06-10 👍 90 2026-07-13 08:37
优化器无关 大语言模型 幂迭代 混合专家模型 路由器设计
arXiv HuggingFace PDF

通过将路由器行与专家权重矩阵的主奇异方向对齐，提升MoE模型性能

前置知识

Mixture-of-Experts (MoE)

混合专家模型是一种可扩展的深度学习架构，通过将密集的前馈网络替换为一组稀疏激活的专家模块来实现计算效率。MoE层包含一个路由器，它为每个输入token选择top-K个最相关的专家，然后这些专家的输出加权求和。这种设计允许在不增加推理计算量的情况下增加模型容量，是训练万亿参数LLM的关键技术，如GPT-4、DeepSeek等。

本文的核心贡献就是重新设计MoE路由器，理解MoE的基本架构和工作原理是读懂本文的前提。

奇异值分解 (SVD) 与主奇异方向

奇异值分解是将矩阵A分解为A = UΣV^T，其中U和V是正交矩阵，Σ是对角矩阵包含奇异值。主奇异方向是对应最大奇异值的列向量，它保留了矩阵中最密集的信息。从几何角度看，主奇异方向是矩阵在单位球面上的最大伸展方向。对于专家权重矩阵W，其主奇异方向是最能代表该专家特征的压缩表示。

本文的核心思想是将每个路由器行对齐到对应专家权重矩阵的主奇异方向，这是理解方法动机和理论分析的关键。

幂迭代 (Power Iteration)

幂迭代是一种数值算法，用于计算矩阵的主特征值和特征向量。给定矩阵M，从随机向量v0开始，通过vk+1 = Mvk / ||Mvk||迭代，vk会收敛到M的主特征向量。这个方法只需要矩阵-向量乘法，计算复杂度远低于完整的SVD分解。本文用一次幂迭代步骤来近似计算专家权重矩阵的主奇异方向。

本文提出的方法核心就是使用幂迭代作为轻量级替代方案来获取主方向，避免了昂贵的SVD计算。

流形优化

流形优化是指在约束空间（如球面、Stiefel流形）上进行的优化。当参数需要满足特定约束（如L2范数固定）时，常规梯度下降会将参数移出约束空间，需要额外的投影步骤。本文的路由器retraction步骤就是将权重投影回球面流形，确保L2范数保持为常数C。这种方法在数学上可以分析为在球面上的最速上升优化。

本文的理论分析将方法解释为流形上的最速上升优化，理解流形优化概念有助于理解retraction步骤的数学原理。

研究动机

现有的MoE路由器设计存在根本性的设计缺陷。传统路由器使用一个简单的线性权重矩阵R，其中第i行R[i]应该作为第i个专家的代理向量。对于输入token x，通过计算w = Softmax(TopK(xR^T))得到门控权重，选择top-K个专家。然而，这种设计缺乏明确的约束来确保路由器行能够忠实地编码专家的内在特征。理想情况下，R[i]应该最大限度地保留专家权重矩阵Wi*的几何结构，但现有方法中这样的约束完全缺失。这可能导致次优的收敛性，使得训练过程需要更多token才能达到良好的性能。

本文的目标是本文的目标是重新设计MoE路由器，使其每个行向量R[i]能够自动对齐到其对应专家权重矩阵Wi*的主奇异方向。从数学直觉来看，主奇异方向保留了矩阵中最密集的信息，是刻画该矩阵的最优压缩表示。通过这样的对齐，token与路由器的内积能够更准确地反映token-专家的亲和度，从而实现更精确的专家选择和更有效的MoE训练。作者希望通过这种原则性的设计，在保持标准MoE接口的同时，带来内在的性能提升。

与已有工作不同的是，本文的独特切入角度是从线性代数的角度重新审视路由器设计，而不是像大多数现有工作那样专注于辅助损失、路由策略或专家结构。作者观察到，理想的路由器行应该编码专家的最本质特征，而主奇异方向正是这样一个理论上最优的表示。然而，直接对每个专家的权重矩阵进行SVD分解在计算上是不可行的。本文的创新点在于使用幂迭代作为轻量级替代方案，并通过retraction步骤确保数值稳定性。这种设计既是原则性的（基于数学理论），又是实用的（计算开销可忽略）。

核心方法

本文提出的方法名为流形幂迭代路由器，采用Power-then-Retract范式。直觉上，每个专家的权重矩阵包含大量信息，而路由器行只是一个压缩的代理向量。最合理的代理应该是专家权重矩阵的主奇异方向，因为这是在约束向量空间中能最大化保留矩阵信息的方向。技术路线包括两个步骤：首先对路由器权重执行一次幂迭代步骤，使其对齐到专家主方向；然后执行L2范数retraction，将权重约束到固定尺度以防止数值爆炸或崩溃。整个过程在每个训练步骤中在线进行，开销可忽略，但理论上等价于在球面流形上的最速上升优化。

核心创新点是将路由器行的更新解释为在球面流形上最大化投影约束的优化问题。具体来说，目标是将R[i]对齐到Wgi Wgi^T的主特征向量，这等价于最大化||R[i]Wgi||2^2 / ||R[i]||2^2，即Rayleigh商。传统的梯度更新可能在球面流形上不稳定，而本文提出的更新规则具有自适应步长：当R[i]已经对齐到主方向时，分母增大，更新变得谨慎；当对齐较差时，更新步长更大。这种结构上的差异使得MPI能够更有效地将路由器推向主奇异方向。

方法步骤详情

方法的完整步骤如下。输入：原始路由器矩阵R，专家权重Wgi，全局超参数C。输出：更新后的门控权重w。步骤1：对路由器的每一行R[i]执行一次幂迭代：R[i]^ = R[i]Wgi Wgi^T。这个步骤通过两次矩阵乘法将R[i]向Wgi Wgi^T的主特征向量方向推进一步。步骤2：执行L2范数retraction：R[i] = C * R[i]^ / ||R[i]^||2，其中C = C/sqrt(N)。这个步骤将更新后的行向量约束到固定尺度，防止数值不稳定，并避免不同专家之间的尺度偏差影响负载均衡。步骤3：重新计算门控权重：w = Softmax(TopK(xR^T))。最终使用更新后的路由器R进行专家选择和输出计算。整个过程在每个训练步骤中在线进行，无需额外的离线预计算。

技术新颖性

技术新颖性体现在多个方面。首先，这是首次从数学原理出发重新设计MoE路由器，而不是基于经验调优。其次，理论分析表明MPI更新等价于在球面流形上的最速上升优化，具有自适应步长特性，这解释了其有效性。第三，方法在保持标准MoE接口的同时带来内在改进，与现有路由器设计正交，可以与辅助损失、z-loss等结合使用。第四，MPI是优化器无关的设计，在AdamW、AdamH、Muon、MuonH等多种优化器上均展现出一致的性能提升。最后，方法计算开销极小（仅相当于N个额外token的计算量），且推理时零开销（可以预计算），具有出色的可扩展性。

实验结果

实验在多个尺度上验证了MPI的有效性。在1B规模上，使用不同优化器（AdamW、AdamH、Muon、MuonH）预训练，MoE with MPI在所有情况下都实现了更快的收敛和更好的下游性能。以MuonH为例，预训练损失降低了0.013，25个基准测试的平均准确率从42.78提升到43.98。在3B和11B规模上，使用MuonH优化器预训练350B token，11B模型在FineWeb-Edu上的预训练损失持续低于基线，下游任务平均准确率从40.92提升到42.76。在具体任务上，如GSM8K数学推理任务，11B MoE with MPI的准确率从17.89大幅提升到27.60。负载均衡方面，MPI在3B模型上的MaxVioBatch从1.133降低到1.024，MaxVioGlobal从0.964降低到0.711，意外地改善了负载分布。效率分析显示，MPI在训练时仅引入0.2%的throughput下降（从34.97B tokens/day），推理时零开销。消融实验证实，单独的行归一化（无幂迭代）无法带来改进，而retraction步骤对于训练稳定性至关重要，尤其是在AdamW和Muon等缺乏范数约束的优化器上。

Downstream performance (average accuracy across 25 benchmarks). MPI consistently improves downstream performance across different optimizers.
Table 1: Downstream performance (average accuracy across 25 benchmarks). MPI consistently improves downstream performance across different optimizers.
Perplexity in bits per byte for MoE with MPI.
Table 2: Perplexity in bits per byte for MoE with MPI.
Performance of MoE with Manifold Power Iteration on challenging benchmarks at both 3B and 11B scales.
Table 3: Performance of MoE with Manifold Power Iteration on challenging benchmarks at both 3B and 11B scales.
MaxVio comparisons for 3B MoE with MPI.
Table 4: MaxVio comparisons for 3B MoE with MPI.
Comparison of lambda distributions. Router with Manifold Power Iteration exhibits an enhanced alignment between R[i] and the principal singular direction of expert weights, manifested by significantly larger lambda values.
Table 5: Comparison of lambda distributions. Router with Manifold Power Iteration exhibits an enhanced alignment between R[i] and the principal singular direction of expert weights, manifested by significantly larger lambda values.
Validation perplexity across choices of C.
Table 6: Validation perplexity across choices of C.
Hyperparameters of model architectures.
Table 7: Hyperparameters of model architectures.
Pretraining hyperparameters (1B AdamW).
Table 8: Pretraining hyperparameters (1B AdamW).
Task-specific performance comparisons for 1B MoE with different optimizers.
Table 9: Task-specific performance comparisons for 1B MoE with different optimizers.
Convergence comparisons for MoE with MPI, exemplified by MuonH-1B. Our router design achieves a 0.013 reduction in pretraining loss.
Figure 2: Convergence comparisons for MoE with MPI, exemplified by MuonH-1B. Our router design achieves a 0.013 reduction in pretraining loss.
Convergence and Downstream Performance Comparison. Manifold Power Iteration facilitates faster convergence and superior downstream task performance throughout the entire course of 11B MoE pretraining.
Figure 3: Convergence and Downstream Performance Comparison. Manifold Power Iteration facilitates faster convergence and superior downstream task performance throughout the entire course of 11B MoE pretraining.
Load balancing loss for 3B MoE with MPI.
Figure 4: Load balancing loss for 3B MoE with MPI.
Ablation studies for the key design choices: (1) Power Iteration and (2) Router Retraction. We observe pretraining collapses without Router Retraction when using AdamW and Muon, showcasing that Router Retraction is critical for maintaining training stability.
Figure 5: Ablation studies for the key design choices: (1) Power Iteration and (2) Router Retraction. We observe pretraining collapses without Router Retraction when using AdamW and Muon, showcasing that Router Retraction is critical for maintaining training stability.
Pre-training loss comparison for a 1B MoE model across optimizers (AdamW, AdamH, Muon). MoE with MPI achieves a convergence advantages over all alternative setups.
Figure 6: Pre-training loss comparison for a 1B MoE model across optimizers (AdamW, AdamH, Muon). MoE with MPI achieves a convergence advantages over all alternative setups.
查看结构化数据
任务指标本文基线提升
25个基准测试平均准确率 Average Accuracy (%) 43.98 (MuonH 1B), 42.76 (MuonH 11B) 42.78 (MuonH 1B), 40.92 (MuonH 11B) +1.2 (1B), +1.84 (11B)
GSM8K数学推理 Accuracy (%) 8-shot CoT 27.60 (11B) 17.89 (11B) +9.71
预训练困惑度 Validation PPL (bits per byte) 0.723 (11B) 0.728 (11B) -0.005
ARC-Challenge Accuracy (%) 5-shot 62.24 (11B) 61.54 (11B) +0.7
MMLU Accuracy (%) 5-shot 50.93 (11B) 50.00 (11B) +0.93

局限与改进

作者承认的局限性包括：C超参数虽然相对不敏感，但仍需要小规模搜索；本文默认使用Softmax作为激活函数，对Sigmoid等其他激活函数的探索有限；幂迭代次数固定为1，虽然实验表明更多迭代不会带来性能提升，但这个选择的理论解释还不充分；负载均衡损失的改善可能是retraction的副作用，具体机制留待未来工作。此外，本文只在GLU风格的专家上验证了方法，对其他专家架构（如MLP）的适用性未明确讨论；实验主要关注预训练和下游任务评估，对于指令微调和对齐阶段的性能缺乏评估；方法的理论分析虽然展示了与最速上升优化的等价性，但收敛速度的理论界尚未建立。

独立分析的弱点

独立分析的弱点包括：第一，方法依赖于GLU专家的特定结构，默认使用Wg进行幂迭代，虽然作者尝试了Wp和Wo但未发现显著差异，但这是否是所有情况下的最优选择尚不清楚。第二，retraction步骤中的常数C需要调整，虽然作者展示了C在{1,2,4,8}范围内的不敏感性，但这个范围是否适用于其他模型架构和训练设置需要验证。第三，方法主要在类GPT架构上验证，对其他架构（如T5、Bert-style MoE）的适用性未知。第四，虽然MPI声称与辅助损失兼容，但与复杂辅助损失策略（如 Expert Specialization Loss）的交互效应未深入研究。第五，理论分析中的近似假设（如高斯输入）在实际数据上的偏差程度未量化。改进方向包括：探索专家权重矩阵的组合（如WgWp^T）作为幂迭代目标；自适应调整C值；在更多架构上验证方法；建立更严格的收敛性理论。

未来方向

作者提出的未来方向包括：深入探索retraction步骤改善负载均衡的机制；评估MPI在更大规模（如100B+参数）上的有效性；探索除Wg外的其他专家矩阵组合作为幂迭代目标。基于本文成果，可以延伸的研究方向包括：将MPI思想应用到其他需要代理向量的场景（如注意力头的query/key投影）；研究MPI与其他MoE改进（如层间路由共享、专家层级结构）的协同效应；探索幂迭代在流形优化中的更广泛应用；将MPI框架扩展到其他模态（如视觉MoE、多模态MoE）；研究MPI在持续学习和模型编辑中的潜在应用，因为更好的专家表示可能有助于知识注入和修改。

复现评估

复现评估方面，论文提供了伪代码和详细的超参数设置。模型架构细节包括：1B模型有8层、1024维度、64个专家；3B模型12层、1536维度、64个专家；11B模型12层、1536维度、256个专家。训练使用FineWeb-Edu数据集，350B tokens预训练，100B tokens中训练。优化器设置详细，包括学习率调度、warmup steps、权重衰减等。然而，论文未明确说明是否开源代码，也未提供权重下载链接。训练在TorchTitan框架上进行，使用PyTorch SDPA和MegaBlocks MLP，需要多GPU设置（使用Fully Sharded Data Parallel）。复现的主要挑战是需要大量计算资源（11B模型预训练350B tokens），以及可能需要处理多GPU并行训练的复杂配置。小规模复现（1B模型在较少tokens上）相对可行，但可能无法完全展现MPI的优势。

论文图表

Figure 1: Pseudo code for Manifold Power-Iteration.

图1展示了MPI的PyTorch风格伪代码实现。代码展示了前向传播中如何执行幂迭代和retraction步骤：首先计算R_hat = (R.unsqueeze(1) @ wg.transpose(1, 2) @ wg).squeeze()，然后执行归一化R_prime = normalize(R_hat, p=2, dim=-1)，接着计算缩放后的logits，最后使用TopK和Softmax得到门控权重。这个伪代码清楚地展示了Power-then-Retract范式在实际代码中的实现方式。

这张图对理解论文非常重要，因为它将抽象的数学公式转化为具体的代码实现，让读者能够直接看到MPI方法如何在实践中应用。伪代码展示了关键的计算步骤和数据形状，有助于理解幂迭代和retraction在深度学习框架中的具体实现细节。