Advancing Expert Specialization for Better MoE
This addresses a specific bottleneck in scaling large language models efficiently for AI researchers and practitioners, representing an incremental improvement over existing MoE methods.
The paper tackles the problem of expert overlap and overly uniform routing in Mixture-of-Experts (MoE) models, which hinders expert specialization and degrades performance during post-training, and shows that their method improves classic MoE baselines by up to 23.79% while maintaining load balancing.
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.