Advancing Expert Specialization for Better MoE

2025-05-28

AI Summary

In Mixture-of-Experts (MoE) models, conventional load-balancing losses often induce expert overlap and routing convergence, undermining expert specialization and degrading downstream performance. To address this, we propose a synergistic optimization framework comprising orthogonality loss and variance loss: the former enforces representational orthogonality across experts to enhance differentiated modeling capacity, while the latter increases routing decision variance to mitigate excessive uniformity. Both losses are architecture-agnostic—requiring no structural modifications—and fully compatible with existing load-balancing mechanisms. We further ensure training stability through gradient-level compatibility analysis. Extensive experiments across diverse MoE architectures (e.g., Switch, GLaM) and benchmarks (e.g., WikiText-103, C4) demonstrate that our method improves expert specialization metrics by up to 23.79%, yields consistent gains in downstream task performance, and maintains strong load balancing.

Technology Category

• Machine Learning: Mixture of Experts (MoE)
• Search and Optimization: Learning to Search
• Natural Language Processing: Learning & Optimization for NLP

Application Category

• User Modeling, Personalization and Recommendation: On-Device user modeling, personalization, and recommendation
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs

Abstract

Mixture-of-Experts (MoE) models enable efficient scaling of large language models (LLMs) by activating only a subset of experts per input. However, we observe that the commonly used auxiliary load balancing loss often leads to expert overlap and overly uniform routing, which hinders expert specialization and degrades overall performance during post-training. To address this, we propose a simple yet effective solution that introduces two complementary objectives: (1) an orthogonality loss to encourage experts to process distinct types of tokens, and (2) a variance loss to encourage more discriminative routing decisions. Gradient-level analysis demonstrates that these objectives are compatible with the existing auxiliary loss and contribute to optimizing the training process. Experimental results over various model architectures and across multiple benchmarks show that our method significantly enhances expert specialization. Notably, our method improves classic MoE baselines with auxiliary loss by up to 23.79%, while also maintaining load balancing in downstream tasks, without any architectural modifications or additional components. We will release our code to contribute to the community.

Problem

• Reducing expert overlap in MoE models for better specialization
• Improving routing decisions to enhance model performance
• Maintaining load balancing without architectural changes

Innovation

• Orthogonality loss reduces expert overlap
• Variance loss enhances routing discrimination
• Compatible with existing auxiliary loss