μ-Parametrization for Mixture of Experts

2025/08/13 by Małaśnicki, Jan, Ciebiera, Kamil, Boruń, Mateusz +8
#FOS: Computer and information sciences #Machine Learning (cs.LG)

paper · doi:10.48550/arxiv.2508.09752

Abstract

Recent years have seen a growing interest and adoption of LLMs, with Mixture-of-Experts (MoE) emerging as a leading architecture in extremely large models. Currently, the largest open-source models reach over 1T parameters. At such scales, hyperparameter tuning becomes prohibitively expensive. Precisely for this reason, the μTransfer is becoming a key technique. It allows for seamless transfer of optimal hyperparameters across model scales, resulting in a huge reduction in tuning costs. However, existing work has primarily focused on dense LLMs, leaving MoE architectures unexplored. In this work, we derive a μ-Parameterization for MoE, providing theoretical guarantees for feature learning across model widths. Our experiments demonstrate that the optimal learning rate reliably transfers across model sizes, establishing a foundation for efficient hyperparameter tuning in large-scale MoE models.

Citations

• Don't be lazy: CompleteP enables compute-efficient deep transformers
• Joint MoE Scaling Laws: Mixture of Experts Can Be Memory Efficient
• Scaling Laws for Fine-Grained Mixture of Experts
• A Spectral Condition for Feature Learning
• Tensor Programs VI: Feature Learning in Infinite-Depth Neural Networks
• Mixture-of-Experts with Expert Choice Routing
• ST-MoE: Designing Stable and Transferable Sparse Expert Models
• Unified Scaling Laws for Routed Language Models
• Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity
• Feature Learning in Infinite-Width Neural Networks
• GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding
• Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer
• Attention Is All You Need
• Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer
• Adaptive Mixtures of Local Experts

Related

• Scaling Laws for Fine-Grained Mixture of Experts
• Joint MoE Scaling Laws: Mixture of Experts Can Be Memory Efficient
• Can Mixture-of-Experts Surpass Dense LLMs Under Strictly Equal Resources?
• Toward Inference-optimal Mixture-of-Expert Large Language Models
• Towards a Comprehensive Scaling Law of Mixture-of-Experts