Load Balancing Mixture of Experts with Similarity Preserving Routers

2025-06-16

AI Summary

Sparse Mixture-of-Experts (MoE) models often suffer from capacity waste and performance degradation due to routing bias toward a small subset of experts. Conventional load-balancing methods enforce uniform expert utilization but risk undermining semantic coherence, leading to knowledge redundancy across experts. To address this, we propose a similarity-preserving load-balancing mechanism: a differentiable routing loss grounded in token embedding similarity, which encourages semantically similar tokens to be consistently routed to the same expert—thereby jointly optimizing load distribution and routing consistency. Our approach requires no additional experts or auxiliary modules and integrates seamlessly into standard MoE training pipelines. Experiments demonstrate a 36% acceleration in convergence on benchmark tasks, substantial reduction in inter-expert knowledge redundancy, and improved model generalization and inference efficiency.

Technology Category

• Machine Learning: Mixture of Experts (MoE)
• Planning, Routing, and Scheduling: Planning with Language Models
• Search and Optimization: Learning to Search

Application Category

• User Modeling, Personalization and Recommendation: Fairness-aware retrieval and ranking
• 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: Graph embeddings and representation learning for Web-related graphs

Abstract

Sparse Mixture of Experts (MoE) models offer a scalable and efficient architecture for training large neural networks by activating only a subset of parameters ("experts") for each input. A learned router computes a distribution over these experts, and assigns input tokens to a small subset. However, without auxiliary balancing mechanisms, routers often converge to using only a few experts, severely limiting model capacity and degrading performance. Most current load balancing mechanisms encourage a distribution over experts that resembles a roughly uniform distribution of experts per token. During training, this can result in inconsistent routing behavior, resulting in the model spending its capacity to learn redundant knowledge. We address this by introducing a novel load balancing loss that preserves token-wise relational structure, encouraging consistent expert choices for similar inputs during training. Our experimental results show that applying our loss to the router results in 36% faster convergence and lower redundancy compared to a popular load balancing loss.

Problem

• Prevents expert underutilization in sparse MoE models
• Reduces redundant knowledge learning in expert routing
• Ensures consistent expert assignment for similar inputs

Innovation

• Similarity preserving router for expert selection
• Novel load balancing loss for consistent routing
• Reduces redundancy and accelerates convergence