Complete-muE: Optimal Hyperparameter Transfer and Scaling for MoE Models

2026-05-22

AI Summary

This work addresses the challenge of hyperparameter adaptation when transferring from dense models to Mixture-of-Experts (MoE) architectures, where structural differences and varying token-per-expert assignments hinder direct transfer. To this end, the authors propose the Complete-muE framework, which enables cross-architecture hyperparameter transfer via a dual-bridge mechanism. Bridge I connects dense models to dense MoE by integrating μP-effective width with normalized routing scaling, while Bridge II links dense MoE to sparse MoE through active-expert scaling, jointly calibrating learning rate and weight decay. This is the first approach to uniformly support hyperparameter transfer across varying numbers of experts, total capacity, granularity, and shared or grouped MoE configurations, while remaining compatible with general Transformer scaling dimensions. Experiments demonstrate that the method significantly enhances hyperparameter stability in pretraining both language and diffusion models, substantially reducing search costs and accelerating convergence.

Abstract

We propose Complete-muE, a framework which targets hyperparameter transfer across dense FFN and any Mixture-of-Experts (MoE) setups in transformer blocks. Existing tools such as $μ$P (requires fixed architectue) or SDE (requires fixed per-step token count) cannot directly solve the hyperparameter transfer problem in MoE setups because Dense to MoE transfer or MoE total experts scaling changes both architecture and tokens per expert. Complete-muE solves this challenge with a two-bridge system: Bridge~I maps between dense FFN and Dense MoE by active-width $μ$P with a normalized router scale. Bridge~II maps between Dense MoE and sparse MoE by activated-expert scaling, where the first-order SDE LR/WD correction cancels while a bounded residual $σ_0$ shift remains. The resulting transfer rule, which we term as Complete muE, covers changes in activated experts, total capacity, granularity, and shared/group-balanced hybrids for MoE models as well as network width/depth, batch size, and duration changes for general Transformer models. Extensive language model and diffusion model pretraining experiments confirm that complete-muE yields relatively stable hyperparameter optima across model architectures and parameter counts -- with only minor drift consistent with the non-strict SDE behavior of Bridge~II. In practice this drift is small enough that hyperparameters tuned on a single dense reference transfer near-optimally to all MoE configurations -- \emph{tune dense once, transfer to all} is the practical recipe at the core of Complete-muE. This enables MoE models to achieve accelerated convergence speedup over dense models when scaling model capacity without costly hyperparameter search.

Problem

• hyperparameter transfer
• Mixture-of-Experts
• dense-to-MoE scaling
• transformer architecture
• optimal hyperparameters

Innovation

• Mixture-of-Experts
• hyperparameter transfer
• Complete-muE
• scaling laws
• transformer optimization