Joint MoE Scaling Laws: Mixture of Experts Can Be Memory Efficient

2025-02-07

AI Summary

Scaling Mixture-of-Experts (MoE) models under fixed memory and compute budgets remains challenging due to unclear trade-offs among expert count, active parameter count, and dataset size. Method: We propose the first unified scaling law jointly governing dense and MoE architectures, theoretically modeling performance dependence on these key factors. Our analysis is grounded in rigorous theoretical derivation and empirically validated across 280+ large-scale experiments—spanning up to 2.7B active parameters and 5B total parameters. Contribution/Results: We demonstrate that, under identical memory constraints, MoE models consistently outperform dense counterparts—refuting the conventional view that MoE gains stem solely from parameter growth. Furthermore, we introduce a practical, principled framework for MoE configuration selection, enabling efficient large-model training. This work provides both theoretical foundations and actionable guidelines for resource-aware MoE scaling.

Technology Category

• Machine Learning: Mixture of Experts (MoE)
• Computer Vision: Large Vision Models
• Natural Language Processing: (Large) Language Models

Application Category

• User Modeling, Personalization and Recommendation: Practical large-scale studies of user experience
• Security and Privacy: Large-scale security measurements
• Search and Retrieval-Augmented AI: Efficiency and scalability of Web search engines

Abstract

Mixture of Experts (MoE) architectures have significantly increased computational efficiency in both research and real-world applications of large-scale machine learning models. However, their scalability and efficiency under memory constraints remain relatively underexplored. In this work, we present joint scaling laws for dense and MoE models, incorporating key factors such as the number of active parameters, dataset size, and the number of experts. Our findings provide a principled framework for selecting the optimal MoE configuration under fixed memory and compute budgets. Surprisingly, we show that MoE models can be more memory-efficient than dense models, contradicting conventional wisdom. To derive and validate the theoretical predictions of our scaling laws, we conduct over 280 experiments with up to 2.7B active parameters and up to 5B total parameters. These results offer actionable insights for designing and deploying MoE models in practical large-scale training scenarios.

Problem

• Memory-efficient Mixture of Experts
• Scalability under memory constraints
• Optimal MoE configuration selection

Innovation

• Mixture of Experts (MoE) scaling laws
• Memory-efficient MoE configurations
• Empirical validation with 280 experiments