Load Balancing Mixture of Experts with Similarity Preserving Routers

TL;DR

This paper introduces a novel load balancing loss for sparse Mixture of Experts models that preserves input similarity, leading to faster training convergence and reduced redundancy in expert utilization.

Contribution

The paper proposes a new load balancing loss that maintains relational structure among inputs, improving expert utilization and training efficiency in MoE models.

Findings

36% faster convergence with the new loss

Lower redundancy in expert usage

Improved load balancing compared to existing methods

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.