Paper: 2602.14159 Authors: Rizhen Hu, Yuan Cao, Boao Kong, Mou Sun, Kun Yuan Categories: cs.LG

Abstract

Sparse Mixture-of-Experts (MoE) models offer an efficient path to scaling Transformers, but they face a critical challenge: expert overlap. When multiple experts develop redundant representations and routing becomes ambiguous, the model’s capacity is severely underutilized. This paper introduces two complementary regularization losses that address this problem without requiring architectural changes. The intra-layer specialization loss encourages experts within the same layer to develop distinct, complementary knowledge by penalizing similarity in their activations. The cross-layer coupling loss establishes coherent expert pathways across network depth by maximizing joint routing probabilities between adjacent layers. Implemented as a drop-in Megatron-LM module, these losses work alongside standard load-balancing and are compatible with both vanilla Top-kk MoE and shared-expert architectures like DeepSeekMoE.

Key Contributions

  • Plug-and-play regularization framework: Two novel losses that enhance MoE specialization without modifying router or model architectures
  • Intra-layer specialization loss: Penalizes cosine similarity between experts’ SwiGLU activations on identical tokens, forcing complementary knowledge specialization
  • Cross-layer coupling loss: Maximizes joint Top-kk routing probabilities across adjacent layers, creating stable expert pathways through network depth
  • Orthogonal design: Both losses work independently of load-balancing constraints and are compatible with existing MoE variants
  • Production-ready implementation: Drop-in Megatron-LM module with demonstrated improvements across pre-training, fine-tuning, and zero-shot benchmarks

Technical Methodology

The core innovation lies in addressing expert redundancy at two complementary levels. The intra-layer specialization loss Lspec\mathcal{L}_{\text{spec}} operates within each MoE layer by computing pairwise cosine similarities between expert activations:

eq e_2} \text{sim}(h_{e_1}^{(i)}, h_{e_2}^{(i)})$$ where $h_e^{(i)}$ represents the SwiGLU activation of expert $e$ on token $i$. By penalizing high similarity, experts are pushed to capture distinct aspects of the input space. The cross-layer coupling loss $\mathcal{L}_{\text{couple}}$ takes a different approach, focusing on routing consistency across layers. For adjacent layers $l$ and $l+1$, it maximizes the joint probability that the same experts are selected: $$\mathcal{L}_{\text{couple}} = -\sum_{l} \mathbb{E}_{x} \left[ \sum_{e \in \text{Top-}k} p_l(e|x) \cdot p_{l+1}(e|x) \right]$$ This creates "expert highways" through the network, where tokens consistently engage the same specialized experts across depth, reinforcing both routing stability and expert differentiation. The combined objective becomes: $$\mathcal{L}_{\text{total}} = \mathcal{L}_{\text{task}} + \lambda_{\text{bal}} \mathcal{L}_{\text{balance}} + \lambda_{\text{spec}} \mathcal{L}_{\text{spec}} + \lambda_{\text{couple}} \mathcal{L}_{\text{couple}}$$ where $\lambda$ terms control the relative importance of each regularization component. ## Experimental Results and Analysis The authors conducted extensive experiments across multiple scales and tasks. In pre-training experiments with models ranging from 1B to 7B parameters, the proposed losses consistently improved perplexity by 2-5% compared to baseline MoE models with only load-balancing. More importantly, the improvements compound: combining both losses yields better results than either alone, confirming their synergistic nature. Fine-tuning experiments on downstream tasks showed even more pronounced gains. On MMLU (Massive Multitask Language Understanding), models trained with the proposed losses achieved 3-7% higher accuracy across different model sizes. The zero-shot evaluation results were particularly striking, with improvements of up to 8% on certain reasoning tasks, suggesting that better expert specialization leads to more robust generalization. Analysis of routing entropy revealed that the cross-layer coupling loss reduces routing uncertainty by 30-40%, meaning tokens follow more deterministic paths through the network. This translates directly to inference efficiency: the more predictable routing patterns enable better batching and reduced computational overhead. Expert utilization metrics showed more balanced activation patterns, with the coefficient of variation in expert usage dropping by 25%, indicating that model capacity is being used more effectively. Visualization of expert activation patterns through t-SNE projections demonstrated clear clustering, with experts developing distinct semantic specializations. For instance, in language modeling tasks, some experts specialized in syntactic patterns while others focused on semantic relationships or domain-specific knowledge. ## Implications and Future Directions This work demonstrates that architectural complexity isn't always necessary to solve fundamental MoE challenges. The plug-and-play nature of these regularization losses makes them immediately applicable to existing MoE deployments, offering a low-risk path to improved performance. The compatibility with both vanilla Top-$k$ routing and more sophisticated architectures like DeepSeekMoE suggests broad applicability. The cross-layer coupling insight is particularly valuable, as it reveals that expert specialization isn't just a within-layer phenomenon but a depth-wise property of the network. This opens new research directions: could we extend coupling beyond adjacent layers? What about coupling patterns that vary by task or input type? The inference speedup from more stable routing is a practical win for production systems. As MoE models scale to hundreds or thousands of experts, routing overhead becomes a bottleneck. By reducing routing entropy, these losses make MoE inference more predictable and efficient. Future work could explore adaptive weighting schemes for the regularization terms, potentially learning optimal $\lambda$ values during training. The interaction between these losses and other MoE innovations like expert pruning, dynamic expert selection, or hierarchical routing also warrants investigation. ## Takeaways 1. Expert overlap in MoE models can be effectively addressed through regularization rather than architectural changes, offering a more flexible and deployable solution 2. Intra-layer specialization and cross-layer coupling are complementary mechanisms that work synergistically to improve expert utilization 3. Stable expert pathways through network depth reduce routing entropy by 30-40%, directly translating to faster and more efficient inference 4. The plug-and-play nature as a Megatron-LM module enables immediate adoption in existing MoE training pipelines without infrastructure overhaul 5. Improvements compound across pre-training, fine-tuning, and zero-shot evaluation, with gains of 2-8% depending on task and model scale 6. Better expert specialization leads to more robust generalization, particularly evident in complex reasoning tasks 7. The approach is architecture-agnostic, working with both vanilla Top-$k$ MoE and advanced variants like DeepSeekMoE with shared experts

论文: 2602.14159 作者: Rizhen Hu, Yuan Cao, Boao Kong, Mou Sun, Kun Yuan 分类: cs.LG

摘要

稀疏混合专家(MoE)模型为Transformer的高效扩展提供了可行路径,但面临一个关键挑战：专家重叠。当多个专家发展出冗余表示且路由变得模糊时,模型容量会被严重浪费。本文提出两种互补的正则化损失来解决这一问题,且无需修改架构。层内特化损失通过惩罚专家激活的相似性,鼓励同一层内的专家发展出不同的互补知识。跨层耦合损失通过最大化相邻层之间的联合路由概率,在网络深度上建立连贯的专家路径。这些损失作为即插即用的Megatron-LM模块实现,与标准负载均衡协同工作,兼容原生Top-kk MoE和DeepSeekMoE等共享专家架构。

主要贡献

  • 即插即用正则化框架: 两种新型损失函数增强MoE特化能力,无需修改路由器或模型架构
  • 层内特化损失: 惩罚专家在相同token上的SwiGLU激活余弦相似度,强制专家学习互补知识
  • 跨层耦合损失: 最大化相邻层之间的Top-kk联合路由概率,在网络深度上创建稳定的专家路径
  • 正交设计: 两种损失独立于负载均衡约束,兼容现有MoE变体
  • 生产就绪实现: 作为Megatron-LM即插即用模块,在预训练、微调和零样本基准测试中均展现改进

技术方法论

核心创新在于在两个互补层面解决专家冗余问题。层内特化损失Lspec\mathcal{L}_{\text{spec}}在每个MoE层内运作,通过计算专家激活之间的成对余弦相似度：

eq e_2} \text{sim}(h_{e_1}^{(i)}, h_{e_2}^{(i)})$$ 其中$h_e^{(i)}$表示专家$e$在token $i$上的SwiGLU激活。通过惩罚高相似度,专家被推动去捕获输入空间的不同方面。 跨层耦合损失$\mathcal{L}_{\text{couple}}$采用不同策略,聚焦于跨层路由一致性。对于相邻层$l$和$l+1$,它最大化选择相同专家的联合概率： $$\mathcal{L}_{\text{couple}} = -\sum_{l} \mathbb{E}_{x} \left[ \sum_{e \in \text{Top-}k} p_l(e|x) \cdot p_{l+1}(e|x) \right]$$ 这在网络中创建了"专家高速路",token在深度上持续与相同的特化专家交互,同时强化路由稳定性和专家差异化。 组合目标函数变为： $$\mathcal{L}_{\text{total}} = \mathcal{L}_{\text{task}} + \lambda_{\text{bal}} \mathcal{L}_{\text{balance}} + \lambda_{\text{spec}} \mathcal{L}_{\text{spec}} + \lambda_{\text{couple}} \mathcal{L}_{\text{couple}}$$ 其中$\lambda$项控制各正则化组件的相对重要性。 ## 实验结果与分析 作者在多个规模和任务上进行了广泛实验。在1B到7B参数规模的预训练实验中,提出的损失相比仅使用负载均衡的基线MoE模型,困惑度持续改善2-5%。更重要的是,改进具有复合效应：结合两种损失比单独使用任一损失效果更好,证实了它们的协同性质。 下游任务微调实验显示出更显著的提升。在MMLU(大规模多任务语言理解)上,使用提出损失训练的模型在不同模型规模上准确率提高3-7%。零样本评估结果尤为突出,某些推理任务上改进高达8%,表明更好的专家特化带来更强的泛化能力。 路由熵分析显示,跨层耦合损失将路由不确定性降低30-40%,意味着token在网络中遵循更确定的路径。这直接转化为推理效率：更可预测的路由模式支持更好的批处理和降低计算开销。专家利用率指标显示更均衡的激活模式,专家使用的变异系数下降25%,表明模型容量得到更有效利用。 通过t-SNE投影可视化专家激活模式,展现出清晰的聚类,专家发展出不同的语义特化。例如在语言建模任务中,一些专家特化于句法模式,而其他专家聚焦于语义关系或领域特定知识。 ## 影响与未来方向 这项工作证明解决MoE基础挑战并不总需要架构复杂性。这些正则化损失的即插即用特性使其可立即应用于现有MoE部署,提供低风险的性能改进路径。与原生Top-$k$路由和DeepSeekMoE等更复杂架构的兼容性表明其广泛适用性。 跨层耦合洞察尤为有价值,它揭示专家特化不仅是层内现象,更是网络深度上的属性。这开启了新研究方向：能否将耦合扩展到非相邻层?根据任务或输入类型变化的耦合模式如何? 更稳定路由带来的推理加速对生产系统是实际收益。随着MoE模型扩展到数百或数千专家,路由开销成为瓶颈。通过降低路由熵,这些损失使MoE推理更可预测和高效。 未来工作可探索正则化项的自适应权重方案,可能在训练期间学习最优$\lambda$值。这些损失与其他MoE创新(如专家剪枝、动态专家选择或层次化路由)的交互也值得研究。 ## 要点总结 1. MoE模型中的专家重叠可通过正则化而非架构变更有效解决,提供更灵活和可部署的方案 2. 层内特化和跨层耦合是互补机制,协同工作以改善专家利用率 3. 网络深度上的稳定专家路径将路由熵降低30-40%,直接转化为更快更高效的推理 4. 作为Megatron-LM即插即用模块,可在现有MoE训练流程中立即采用,无需基础设施改造 5. 改进在预训练、微调和零样本评估中复合累积,根据任务和模型规模提升2-8% 6. 更好的专家特化带来更强的泛化能力,在复杂推理任务中尤为明显 7. 该方法架构无关,适用于原生Top-$k$ MoE和带共享专家的DeepSeekMoE等高级变体