Optimal Expert-Attention Allocation in Mixture-of-Experts: A Scalable Law for Dynamic Model Design
Organizations: The Hong Kong University of Science and Technology (Guangzhou) · The Hong Kong University of Science and Technology · Ant Group
Abstract
This paper presents a novel extension of neural scaling laws to Mixture-of-Experts (MoE) models, focusing on the optimal allocation of compute between expert and attention sub-layers. As MoE architectures have emerged as an efficient method for scaling model capacity without proportionally increasing computation, determining the optimal expert-attention compute ratio becomes critical. We define the ratio as the fraction of total FLOPs per token dedicated to the expert layers versus the attention layers, and explore how this ratio interacts with the overall compute budget and model sparsity. Through extensive experiments with GPT-style MoE Transformers, we empirically find that the optimal ratio follows a power-law relationship with total compute and varies with sparsity. Our analysis leads to an explicit formula for , enabling precise control over the expert-attention compute allocation. We generalize the Chinchilla scaling law by incorporating this architectural parameter, providing a new framework for tuning MoE models beyond size and data. Our findings offer practical guidelines for designing efficient MoE models, optimizing performance while respecting fixed compute budgets.
Figures & tables
| Coefficient | a | b | c | d | |||||
|---|---|---|---|---|---|---|---|---|---|
| Estimate | 0.6288 | 0.0453 | 0.4228 | 0.0431 | 13.7354 | 15.12 | 18.62 | 39.55 | 0.0499 |
Appendix figures & tables4 assets
Supplementary material from the paper’s appendix.
Appendix
| Label | Layers ( ) | Heads ( ) | Batch Size | Learning Rate |
|---|---|---|---|---|
| 20M | 8 | 8 | 96 | 0.0015 |
| 30M | 8 | 8 | 160 | 0.0013 |
| 55M | 10 | 10 | 224 | 0.0011 |
| 100M | 14 | 12 | 320 | 0.0009 |
| 200M | 16 | 16 | 512 | 0.0008 |
| Coefficients | Initial Values |
|---|---|
| [0, 10, 20] | |
| [0, 0.25, 0.5, 0.75, 1, 1.25] | |
| 1.5 |
Mixture-of-Experts Models
Mixture-of-experts models use input-dependent routing to activate specialized subnetworks, increasing model capacity without running every parameter for each input.
See all 454 papers