Load Balancing Mixture of Experts with Similarity Preserving Routers¶
会议: NeurIPS 2026
arXiv: 2506.14038
代码: https://github.com/lucidrains/st-moe-pytorch
领域: 大模型效率
关键词: 混合专家模型、负载均衡、正交路由器、专家冗余、成对专家相似度
一句话总结¶
针对传统 MoE 负载均衡损失强制专家均匀分布导致专家知识冗余的问题，本文提出相似度保持路由器（SimBal），通过对路由器权重施加软正交性正则化，在防止专家坍缩的同时保持输入语义几何结构，实现训练收敛加速 36% 并显著降低专家冗余。
研究背景与动机¶
随着大语言模型（LLM）参数规模与预训练数据量的指数级攀升，兼具高模型容量与低计算开销的稀疏混合专家（Mixture of Experts, MoE）架构成为前沿基座模型的核心选择。MoE 通过在特定层将密集前馈网络替换为多个平行的子网络（即“专家”），并利用轻量级门控路由器为每个输入词元（token）动态激活极少数专家（例如 32 选 4），从而在保持活跃参数量和推理计算开销基本不变的前提下成倍扩充总参数量。然而，MoE 模型在预训练初期极易陷入专家利用率极度不均衡的“路由坍缩（routing collapse）”陷阱——未经约束的路由器往往倾向于反复选择极少数表现较好的专家，导致其余专家长期闲置，极大削弱了稀疏架构的容量扩展优势。
为了稳定 MoE 训练并防止坍缩，现有主流方案普遍引入序列级辅助负载均衡损失（Load Balancing Loss, LBL，如 Switch Transformer 等），通过惩罚每个专家的选中使用频率与平均路由权重的乘积，强行将词元分配推向各专家之间的离散均匀分布。然而，这种强加的均匀分布机制存在严重的内在矛盾：当两个在语义表征空间中极其相近的 token 出现在同一个批次中时，为了维持全局或序列级别的均匀配额，路由器不得不将它们机械地分流至完全不同的专家网络。在漫长的预训练过程中，不同专家频繁接触到高度同质化的 token 输入，导致模型被迫消耗大量宝贵的参数容量在不同专家之间重复学习完全相同的冗余知识。实证研究表明，专家间知识冗余的上升直接导致了特化能力的退化与整体模型表现的下滑。
针对“强行均匀化导致路由不一致与知识重复”的核心痛点，本文提出了一个截然不同的切入视角：优良的路由器应当在专家分配过程中保持 token 之间的语义几何关系，即相似的输入词元应当稳定路由至相似的专家组合，从源头上遏制冗余知识的扩散。核心 idea：通过对路由器权重矩阵施加软正交正则化（SimBal），在保证度量保持特性的同时利用正交基的方向分散性自然防止专家坍缩，以无数据依赖的极低计算开销实现一致性路由与去冗余特化。
方法详解¶
整体框架¶
在标准 MoE Transformer 架构中，第 \(l\) 层的输入 token 隐状态向量为 \(\mathbf{x} \in \mathbb{R}^{D_M}\)。门控路由器参数由线性变换矩阵 \(R \in \mathbb{R}^{D_M \times E}\) 构成，其中 \(D_M\) 为模型隐藏层维度，\(E\) 为该层专家总数（通常 \(E \ll D_M\)）。路由器首先计算原始路由分数 \(\mathbf{s} = \mathbf{x} R \in \mathbb{R}^E\)，随后通过 Softmax 激活函数生成所有专家上的概率分布 \(\mathbf{r} = \text{Softmax}(\mathbf{s})\)，最后选取概率最大的前 \(A\) 个专家（Top-\(A\) 稀疏激活，例如 \(A=4, E=32\)）并对其输出进行加权聚合。
SimBal 的整体框架并不改变前向推理的数据流，而是在训练期间重塑路由器的几何映射属性。为确保输入空间中两词元之间的夹角与点积在投影至路由器权重空间时得以保真，路由器权重列向量应当构成一组正交基。SimBal 在训练主目标函数之上，直接附加对路由器 Gram 矩阵偏离单位矩阵的 \(L_1\) 范数惩罚，使各专家的表征向量在潜空间中彼此正交。同时，作者构建了非侵入式的成对专家相似度（Pairwise Expert Similarity, PES）评估流水线，对专家输出表征的多样性与冗余度进行全局监控。
%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
A["输入 Token 隐状态 x ∈ ℝ^{D_M}"] --> B["路由器矩阵投影 xR<br/>R ∈ ℝ^{D_M × E}"]
B --> C["Top-A 稀疏门控选择<br/>Top-4 专家动态激活与加权"]
C --> D["MoE 专家并行前馈计算<br/>∑ a=1..A r_a · FFN_{i_a}(x)"]
D --> E["输出聚合与残差连接"]
B -.-> F["Gram 矩阵对齐正则化 (SimBal)<br/>‖R^T R - I_E‖_1 反向传播"]
D -.-> G["成对专家相似度评估 (PES)<br/>余弦多样性与冗余度量"]关键设计¶
1. 基于正交投影的相似度保持路由：以几何保真解决路由不一致性
传统 LBL 强推专家均匀选取的做法割裂了 token 语义相近性与专家分配的映射关系，造成相邻上下文中同类语义概念被随机掷骰子般指派给不相干的专家。SimBal 的几何理论根基在于保距/保角变换：若矩阵 \(R\) 的各列向量相互正交且模长归一（即 \(R^\top R = I_E\)），则对于任意两个输入 token \(\mathbf{x}_1, \mathbf{x}_2 \in \mathbb{R}^{D_M}\)，它们在路由器投影空间中的内积严格等于原始空间内积在子空间上的正交投影，即 \(\langle \mathbf{x}_1 R, \mathbf{x}_2 R \rangle = \mathbf{x}_1 (R R^\top) \mathbf{x}_2^\top\)。此时，路由器对 token 之间角距离的度量保真使得语义邻近的 token 必定获得高度一致的路由打分与专家指派；与此同时，\(E\) 个正交列向量在 \(D_M\) 维空间中指向完全相互垂直的独立方向，天然将整个表示空间均匀划分为 \(E\) 个正交子空间，在不破坏语义关联的前提下自发实现了全局专家负荷的均匀覆盖，彻底避免了单一专家垄断的坍缩现象。
2. 软约束 Gram 矩阵对齐正则化：兼顾低算力开销与数值稳定性
在深度神经网络中直接施加硬性正交参数约束（例如通过 QR 分解、Cayley 变换或 Stiefel 流形优化）在工程上极其昂贵，不仅每次迭代都需要高阶矩阵分解计算，且在现代大语言模型通用的低精度格式（如 bfloat16）下极其容易发生下溢或数值崩塌。为了解决这一系统瓶颈，SimBal 规避了显式正交化算法，转而采用纯净的软正则化损失：
$$
\mathcal{L}_{\text{orth}} = \left| R^\top R - I_E \right|_1
$$
由于专家数目 \(E\)（如 32）远小于模型维度 \(D_M\)（如 768 或 1536），计算 Gram 矩阵 \(R^\top R \in \mathbb{R}^{E \times E}\) 仅涉及极小的矩阵乘法，FLOPs 开销几乎可以忽略不计。更重要的是，该损失函数完全由路由器权重参数自发计算，不依赖于批次中的输入数据，从而彻底摆脱了传统 LBL 和无辅助损失策略（Loss-Free, LF）对超大 Batch Size 及跨卡分布式全归约通信同步的严重依赖。结合正交初始化策略（OrthoInit），路由器在训练启动阶段便能以近乎零的步数收敛到 Gram 偏差极小的正交状态。
3. 成对专家相似度（PES）度量：高拓展性全局专家特化审计
传统评估 MoE 专家特化程度的方法主要依赖于组合式专家剪枝消融（Expert Dropout），即暴力剔除使用频率最高的前若干个专家并在验证集上重新推理评估性能跌幅。然而这种方式面临严重的组合爆炸问题，对于 32 专家的系统难以穷尽评估且极其耗时。本文创新地设计了成对专家相似度指标（Pairwise Expert Similarity, PES），通过一次前向传播即可精准量化所有专家输出的多样性：
$$
\text{PES}{\text{model}} = \frac{1}{|B|} \sum} \in B} \mathcal{C{\text{expert}}(\mathbf{x}), \quad \mathcal{C})\right)
$$
式中 }}(\mathbf{x}) = \frac{2}{E(E-1)} \sum_{i=1}^E \sum_{j=i+1}^E \cos\left(\mathbf{f}_i(\mathbf{x}), \mathbf{f}_j(\mathbf{x\(\mathbf{f}_i(\mathbf{x})\) 表示第 \(i\) 个专家网络对输入 \(\mathbf{x}\) 的输出向量。PES 直接衡量所有专家表征之间的平均余弦相似度。PES 越低，表明专家之间的表征越独立、多样性越高、参数冗余越低。该指标无需多次推理，可直接在批次计算中向量化完成，为 MoE 特化程度提供了鲁棒、连续且可微的全局评测基准。
一个完整示例¶
考虑一个批量中包含两个来自相近科技语境的输入 token：“Transformer”与“Attention”，二者隐向量高度相似（夹角趋近于 0）。 1. 路由打分阶段：输入向量进入维度为 \(1536 \times 32\) 的路由器矩阵 \(R\)。由于 \(R\) 在 SimBal 约束下满足 \(R^\top R \approx I_{32}\)，32 个专家的特征基底彼此正交。“Transformer”与“Attention”由于内积高度接近，在 32 个专家维度上得到完全一致的激活分数峰值（例如专家 3 和专家 17 的分数显著最高）。 2. Top-4 专家分发：两个 token 均被指派至专家 3、17、22 和 30。这使得专家 3 和 17 能够持续吸收并精细建模自注意力机制相关的专精知识，而不是由 LBL 为了维持全局均匀而将“Attention”随机丢给处理生物医学领域的专家 8。 3. 空间正交展开：若此时输入一个截然不同的文学类 token “Shakespeare”，由于其与技术 token 在隐空间中近似正交，它在专家 3 和 17 上的投影极小，反而稳定激活了正交方向上的专家 9 与 14。所有 32 个专家由此在保持语义局部一致性的同时获得了空间维度的全面利用。
损失函数 / 训练策略¶
模型预训练的总优化目标由语言模型自回归交叉熵损失 \(\mathcal{L}_{\text{LM}}\)、SimBal 正交正则化损失 \(\mathcal{L}_{\text{orth}}\) 以及用于防止 logits 数值溢出的 Z-loss \(\mathcal{L}_z\) 线性组合而成： $$ \mathcal{L}{\text{total}} = \mathcal{L}_z $$ 其中 Z-loss 系数设定为 }} + \lambda_{\text{orth}} \cdot \mathcal{L}_{\text{orth}} + \lambda_z \cdot \mathcal{L\(\lambda_z = 10^{-5}\)；\(\mathcal{L}_{\text{orth}}\) 默认系数为 \(\lambda_{\text{orth}} = 0.1\)。消融实验证实，模型表现对 \(\lambda_{\text{orth}} \in [0.01, 1.0]\) 完全不敏感，无需繁琐调参。
在训练优化层面，模型采用 AdamW 优化器（权重衰减 0.01），前 2000 步进行线性预热，随后采用余弦退火策略衰减至峰值学习率的 10%。对于中等规模模型 MoE-M（活跃参数 230M，总参数 627M），峰值学习率设为 \(5 \times 10^{-4}\)，在 8 张 NVIDIA A100 (40GB) 上训练 19.9B tokens；对于大规模模型 MoE-L（活跃参数 761M，总参数 3.14B），峰值学习率设为 \(3 \times 10^{-4}\)，在 8 张 AMD MI300X (192GB) 上训练 78.6B tokens。全流程参数均以 bfloat16 格式运行。
实验关键数据¶
主实验¶
论文在 MoE-M 与 MoE-L 两个规格上全面对比了密集基线（Dense）、传统 LBL 平衡策略以及 SimBal 方案的困惑度与多样性表现，并在 78.6B token 预训练完成的 MoE-L 上评估了 6 项主流下游自然语言基准测试。
| 模型架构 | 负载均衡策略 | 活跃参数量 | 总参数量 | 验证集困惑度 (PPL) ↓ | 最小专家相似度 (Min PES) ↓ | 训练收敛加速比 |
|---|---|---|---|---|---|---|
| Dense-M | – | 230M | 230M | 19.468 | – | – |
| MoE-M | LBL (\(\alpha=0.01\)) | 230M | 627M | 14.086 | 0.0255 | 基准 |
| MoE-M | SimBal | 230M | 627M | 13.685 | 0.0044 | +36.2% |
| Dense-L | – | 761M | 761M | 10.047 | – | – |
| MoE-L | LBL (\(\alpha=0.01\)) | 761M | 3.14B | 8.517 | 0.0241 | 基准 |
| MoE-L | SimBal | 761M | 3.14B | 8.304 | 0.0028 | +36.4% |
在下游基准测试中，SimBal-L 相较于 LBL-L 展现出全面的性能领先：
| 评估基准 (Benchmark) | LBL-L (均值 ± 标准误) | SimBal-L (均值 ± 标准误) | 绝对提升 |
|---|---|---|---|
| ARC Challenge | 22.44% ± 1.22% | 23.21% ± 1.23% | +0.77% |
| ARC Easy | 40.49% ± 1.01% | 41.16% ± 1.01% | +0.67% |
| HellaSwag | 35.45% ± 0.48% | 35.74% ± 0.48% | +0.29% |
| PIQA | 66.49% ± 1.10% | 66.81% ± 1.10% | +0.32% |
| WinoGrande | 49.72% ± 1.41% | 52.49% ± 1.40% | +2.77% |
| GLUE Benchmark | 45.10% ± 1.98% | 51.73% ± 1.97% | +6.63% |
| 全任务平均 (Mean) | 43.28% | 45.19% | +1.91% |
消融实验与分析¶
针对正交化实现方式、负载均衡稳定性指标以及推理期专家剪枝进行了系统消融。
正交化方案与数值误差对比（100 次试验均值 ± 标准差，形状 \(1536 \times 32\)）：
| 方法类别 | 最大偏离度 (Max Dev) ↓ | 平均 L1 距离 (L1 Dist) ↓ | 计算开销与格式兼容性 |
|---|---|---|---|
| 显式参数化 (Param, Lezcano-Casado 2019) | \(2.00 \times 10^{-4} \pm 2.31 \times 10^{-5}\) | \(4.80 \times 10^{-5} \pm 1.60 \times 10^{-6}\) | 昂贵，不支持直接 bfloat16 优化 |
| 正交初始化 (OrthoInit, Saxe 2014) | \(1.93 \times 10^{-4} \pm 1.88 \times 10^{-5}\) | \(4.62 \times 10^{-5} \pm 1.79 \times 10^{-6}\) | 仅初始步有效，随训练漂移 |
| SimBal 优化（Trained，100 优化步） | \(1.03 \times 10^{-5} \pm 2.76 \times 10^{-6}\) | \(8.52 \times 10^{-7} \pm 5.84 \times 10^{-8}\) | 最轻量，直接原生 bfloat16 训练 |
负载均衡与正交化度量（MoE-L）：
| 均衡策略 | 序列级专家利用率 (SEU) | 路由分布熵 (Entropy) | Gram 偏差均方误差 \((R^\top R - I)^2\) |
|---|---|---|---|
| LBL | 1.000 | 1.268 | \(3.11 \times 10^{-2}\) |
| SimBal | 0.991 | 1.168 | \(2.121 \times 10^{-8}\) |
推理期 Dynamic-K 专家剪枝协同效应（MoE-L 完整验证集，丢弃分配概率 \(< P(E)\) 的专家）：
| 剪枝阈值 \(P(E) <\) | SimBal 困惑度 (PPL) | SimBal 推理耗时 (s) | LBL 困惑度 (PPL) | LBL 推理耗时 (s) | 性能与吞吐对比 |
|---|---|---|---|---|---|
| 0 (无剪枝) | 8.317 | 620.927 | 8.536 | 619.657 | SimBal PPL 领先 0.219 |
| 0.10 | 8.364 | 571.147 | 8.542 | 575.121 | SimBal PPL 仍优于无剪枝 LBL |
| 0.15 | 8.598 | 503.065 | 8.621 | 543.027 | SimBal 提速 7.4% 且 PPL 更优 |
| 0.20 | 9.380 | 472.915 | 9.057 | 495.200 | LBL 因冗余表现出更高的鲁棒性 |
关键发现¶
- 训练收敛速度实现大幅跃升：在 MoE-L 规模下，SimBal 仅消耗约 50B tokens 即可达到传统 LBL 在 78.6B tokens 时的验证集损失，收敛速度加快 36.4%；在 MoE-M 上，SimBal 仅需 12.7B tokens 即可匹敌 LBL 训练 19.9B tokens 的效果（提速 36.2%）。
- 专家冗余度数量级下降：在 MoE-L 每一层中，SimBal 的最小成对专家相似度（Min PES）仅为 0.0028，相比 LBL 的 0.0241 降低了一个数量级；层级 Gram 偏差均方误差从 LBL 的 0.0311 骤降至 \(2.12 \times 10^{-8}\)，表明专家表征实现了近乎完全的正交解耦。
- 统计显著性极高：在 MoE-M 上进行的 5 随机种子实验显示，LBL 困惑度均值为 \(14.051 \pm 0.026\)，SimBal 困惑度均值为 \(13.691 \pm 0.039\)。SimBal 的均值落后 LBL 达 13 个标准差以上，具备极其显著的统计优势。
- 与动态剪枝高度协同：由于 SimBal 促使路由器产生更尖锐、更具特化确定性的概率分布，在丢弃权重低于 0.15 的长尾专家时，SimBal 在保持 8.598 优异困惑度的同时将推理延迟从 543 秒压缩至 503 秒（提速 7.4%）。
亮点与洞察¶
- 从机械均匀转向几何保真：深刻揭示了经典 LBL 强推均匀分布背后的“表征污染”隐患，将负载均衡重新定义为“保持输入流形局部度量的正交映射问题”，立意深刻且切中 MoE 特化本质。
- 无数据依赖与零通信开销：SimBal 纯粹基于路由器静态权重进行 Gram 矩阵正则化，不参与 token batch 统计，完美契合任意小批次及高度并行的去中心化分布式训练，规避了传统方法对大批次同步的苛刻依赖。
- 即插即用且极具超参鲁棒性：在 \(0.01 \sim 1.0\) 的三个数量级跨度内，SimBal 的正则化损失权重无需任何精细调优即可稳定发挥最优效益，工程复现门槛极低。
局限与展望¶
- 训练总 Token 规模受限：受制于算力，MoE-L 训练步数停止在 78.6B tokens。虽然相较于基线展现了明显的相对收益，但未在万亿（Trillion-token）级别的工业界规模下考察持续特化能力的长期演变。
- 单一路由拓扑假定：研究主要聚焦于无 Token 丢弃（Dropless）且基于 Top-A 稀疏选择的标准架构，尚未探讨在容量因子截断（Capacity Factor / Token Dropping）以及跨节点 Expert Parallelism 通信受限场景下的交互影响。
- 未来方向：探索动态自适应正交投影与层次化路由（Hierarchical MoE）的结合，以及利用正交路由器先验指导预训练后期的自适应专家合并与量化压缩。
相关工作与启发¶
- vs Switch Transformers (Fedus et al., 2022)：Switch Transformer 采用序列级辅助损失强制实现 token-to-expert 的离散均匀配额，而 SimBal 采用权重空间的正交 Gram 约束，在确保无专家闲置的前提下消除了同质 token 乱分流带来的知识冗余。
- vs Auxiliary-Loss-Free (DeepSeek-V3 / Wang et al., 2024)：Loss-Free (LF) 机制通过在推理分数上动态累加专家偏置项实现平衡，但在小批次及序列级利用率（SEU）上波动剧烈且对超参敏感。SimBal 与 LF 可互补融合，在保持优异 SEU 的同时避免困惑度崩塌。
- vs OMoE / MOORE (Liu et al., 2024; Hendawy et al., 2024)：OMoE 直接对专家网络参数的更新方向施加正交优化约束，计算开销庞大且限制了专家本身的拟合能力。SimBal 将正交性巧妙下放至参数量仅占 0.018% 的轻量级路由器中，以最小代价撬动了整个 MoE 系统的表示解耦。
评分¶
- 新颖性: ⭐⭐⭐⭐⭐ [跳出传统基于数据分布强制拉平负载的固有思路，首次从路由器权重正交性与度量保持几何出发解决专家冗余]
- 实验充分度: ⭐⭐⭐⭐☆ [双模型规模充分对标，包含下游 benchmark、5-seed 显著性检验与动态剪枝分析，若具备万亿 token 工业级实验更佳]
- 写作质量: ⭐⭐⭐⭐⭐ [逻辑严密，理论机理与实证指标（PES）闭环完整，问题剖析透彻直击痛点]
- 价值: ⭐⭐⭐⭐⭐ [通用且计算友好的 MoE 核心组件升级方案，训练提速 36% 且完美契合分布式系统，具有极高的工程落地价值]