Mixture of Diverse Size Experts
This addresses a specific bottleneck in scaling large language models efficiently for AI researchers and practitioners, representing an incremental improvement over current MoE designs.
The paper tackles the limitation of uniform expert sizes in sparsely-activated Mixture-of-Experts (MoE) models for large language models, proposing Mixture of Diverse Size Experts (MoDSE) with experts of different sizes, which outperforms existing MoEs on multiple benchmarks while maintaining the same total parameters and number of experts.
The Sparsely-Activated Mixture-of-Experts (MoE) has gained increasing popularity for scaling up large language models (LLMs) without exploding computational costs. Despite its success, the current design faces a challenge where all experts have the same size, limiting the ability of tokens to choose the experts with the most appropriate size for generating the next token. In this paper, we propose the Mixture of Diverse Size Experts (MoDSE), a new MoE architecture with layers designed to have experts of different sizes. Our analysis of difficult token generation tasks shows that experts of various sizes achieve better predictions, and the routing path of the experts tends to be stable after a training period. However, having experts of diverse sizes can lead to uneven workload distribution. To tackle this limitation, we introduce an expert-pair allocation strategy to evenly distribute the workload across multiple GPUs. Comprehensive evaluations across multiple benchmarks demonstrate the effectiveness of MoDSE, as it outperforms existing MoEs by allocating the parameter budget to experts adaptively while maintaining the same total parameter size and the number of experts.