Synergistic Intra- and Cross-Layer Regularization Losses for MoE Expert Specialization
Rizhen Hu, Yuan Cao, Boao Kong, Mou Sun, Kun Yuan
Abstract
Sparse Mixture-of-Experts (MoE) models scale Transformers efficiently but suffer from expert overlap-redundant representations across experts and routing ambiguity, resulting in severely underutilized model capacity. While architectural solutions like DeepSeekMoE promote specialization, they require substantial structural modifications and rely solely on intra-layer signals. In this paper, we propose two plug-and-play regularization losses that enhance MoE specialization and routing efficiency without modifying router or model architectures. First, an intralayer specialization loss penalizes cosine similarity between experts' SwiGLU activations on identical tokens, encouraging experts to specialize in complementary knowledge. Second, a crosslayer coupling loss maximizes joint Top-k routing probabilities across adjacent layers, establishing coherent expert pathways through network depth while reinforcing intra-layer expert specialization. Both losses are orthogonal to the standard load-balancing loss and compatible with both the shared-expert architecture in DeepSeekMoE and vanilla top-k MoE architectures. We implement both losses as a drop-in Megatron-LM module. Extensive experiments across pre-training, finetuning, and zero-shot benchmarks demonstrate consistent task gains, higher expert specialization, and lower-entropy routing; together, these improvements translate into faster inference via more stable expert pathways.
Ask about this paper
Your agent reads all of it.
Lune indexed this paper to the last equation, along with the top-tier papers that cite it. Ask a question and the answer quotes them.
Your agent calls
Luneget_paper_fulltext
Free to start. No credit card required.
Terminal
Install the CLIlune papers fulltext 5a2a5511-33cc-4c48-b79c-2c7a4cb56f33Cited by top-tier papers2
- CR-Net: Scaling Parameter-Efficient Training with Cross-Layer Low-Rank StructureBoao Kong, Junzhu Liang, Yuxi Liu, Renjia Deng et al.ICLR 2026 · 7 citations
- Grouter: Decoupling Routing from Representation for Accelerated MoE TrainingYuqi Xu, Rizhen Hu, zihan liu, Mou Sun et al.ICML 2026 · 7 citations
Builds on20
- Measuring Massive Multitask Language UnderstandingDan Hendrycks, Collin Burns, Steven Basart, Andy Zou et al.ICLR 2021 · 7,905 citations
- Let's Verify Step by StepHunter Lightman, Vineet Kosaraju, Yuri Burda, Harrison Edwards et al.ICLR 2024 · 3,045 citations
- PIQA: Reasoning about Physical Commonsense in Natural LanguageYonatan Bisk, Rowan Zellers, Ronan Le Bras, Jianfeng Gao et al.AAAI 2020 · 2,916 citations
- GShard: Scaling Giant Models with Conditional Computation and Automatic ShardingDmitry Lepikhin, HyoukJoong Lee, Yuanzhong Xu, Dehao Chen et al.ICLR 2021 · 1,954 citations
- GLaM: Efficient Scaling of Language Models with Mixture-of-ExpertsNan Du, Yanping Huang, Andrew M. Dai, Simon Tong et al.ICML 2022 · 1,173 citations
Related papers
- ReMoE: Fully Differentiable Mixture-of-Experts with ReLU RoutingZiteng Wang, Jun Zhu, Jianfei ChenICLR 2025
- Advancing Expert Specialization for Better MoEHongcan Guo, Haolang Lu, Guoshun Nan, Bolun Chu et al.NeurIPS 2025 · 48 citations
- Coupling Experts and Routers in Mixture-of-Experts via an Auxiliary LossAng Lv, Jin Ma, Yiyuan Ma, Siyuan QiaoICLR 2026 · 14 citations
- Dirichlet-Prior Shaping: Guiding Expert Specialization in Upcycled MoEsLeyla Mirvakhabova, Babak Ehteshami Bejnordi, Gaurav Kumar, Hanxue Liang et al.ICML 2026 · 1 citation
- Soft Modality-Guided Expert Specialization in MoE-VLMsZi-Hao Bo, Yaqian Li, Anzhou Hou, Rinyoichi Takezoe et al.CVPR 2026
