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

Summary (Overview)

Complete-muE is a novel framework for transferring hyperparameters (learning rate, weight decay, initialization) across dense FFN and Mixture-of-Experts (MoE) transformer architectures, enabling a "tune dense once, transfer to all" recipe.
The framework uses a two-bridge system: Bridge I maps dense FFN ↔ Dense MoE via active-width μP with a normalized router scale, and Bridge II maps Dense MoE ↔ sparse MoE via activated-expert scaling where the first-order SDE LR/WD correction cancels.
Complete-muE covers all MoE scaling axes: activated experts, total capacity, granularity, shared experts, group-balanced routing, as well as general width/depth/batch/duration changes.
Extensive experiments on language models (LM) and diffusion models (DF) confirm stable hyperparameter optima across all MoE configurations, with only minor drift consistent with the non-strict SDE behavior of Bridge II.
Large-scale runs achieve ~4.5× convergence speedup on 240P 5s video diffusion models and 5.3×–5.5× on LLM training with 100k iterations, without per-setting retuning.

Introduction and Theoretical Foundation

Background and Motivation

The paper addresses a critical gap in hyperparameter transfer for MoE models. MoE architectures expose scaling axes absent from dense FFNs—per-expert width h, activated experts a, total experts N, shared experts, and group-balanced routing. Changing these axes also changes expert data exposure: with balanced routing, each expert processes roughly $Ba/N$ tokens per step and $TBa/N$ tokens over training.

The Gap in Existing Tools

Tool	Limitation
μP (Tensor Programs)	Connects dense FFNs to Dense MoE through active width but cannot handle changes in per-expert token batch size
SDE rules	Transfer across token batch size for fixed architecture but cannot handle architecture changes

Neither tool can simultaneously handle both architecture changes and token batch size changes that occur in dense-to-sparse MoE transfer.

Theoretical Foundation

The unified FFN/MoE formulation represents all cases as:

y(x) = A(H_{\mathrm{act}}) \sum_{i=1}^{M} g_i(x) o_i(x), \qquad o_i(x) = W_{\mathrm{down}}^{(i)} u_i(x) \tag{1}

where $H_{act}$ is the hidden width active on token x. Table 1 shows how this specializes:

Case	Choice of $g_i(x)$	Active width $H_{\text{act}}$	Route scale
Dense FFN	One always-active dense block ( $M=1$ ), $g_1(x) = 1$	$H$	1
Dense MoE	All experts active ( $a=N$ ), $g_i(x) = r_N \pi_i(x)$	$Nh$	$r_N = N$
Sparse MoE	Top- $a$ routed experts, $g_i(x) = r_a \pi_i(x)\mathbf{1}\{i \in \mathcal{A}(x)\}$	$ah$	$r_a = a$

Methodology

Bridge I: Dense FFN ↔ Dense MoE via Active-Width μP

Bridge I matches a dense FFN with active width H to a Dense MoE with the same active width. The transfer rules are:

A(H) = \frac{d}{H}, \qquad \sigma_{\text{down}}(d, H) = \left(\frac{H}{d}\right)^{1/2} \sigma_{\text{down}}^{(1)}(d), \qquad \eta_{\text{down}}(d, H) = \eta_{\text{down}}^{(1)}(d) \tag{2}

The normalized route scale is critical: since normalized routing averages expert outputs (shrinking the update by 1/a), the correction is:

r_a = a \tag{3}

Bridge II: Dense MoE ↔ Sparse MoE via Activated-Expert Scaling

The layer-level transfer across activated experts follows:

A_a = \frac{d}{ah}, \qquad r_a = a, \qquad \sigma_{\text{down}}(d, a) = \left(\frac{ah}{d}\right)^{1/2} \sigma_{\text{down}}^{(1)}(d), \qquad \eta_{\text{down}}(d, a) = \eta_{\text{down}}^{(1)}(d) \tag{4}

Expert-side SDE cancellation: Under approximate load balancing, one expert receives:

B_{\mathrm{exp}}(a) \approx B\frac{a}{N}, \qquad D_{\mathrm{exp}}(a) \approx TB\frac{a}{N} \tag{5}

For a change $a \to a'$ at fixed global batch B and fixed optimizer steps T, the expert-side batch and duration ratios are identical ( $\rho_B^{\mathrm{exp}} = \rho_D^{\mathrm{exp}} = a'/a$ ), so the SDE correction cancels:

\eta' \approx \eta\sqrt{\frac{\rho_B^{\mathrm{exp}}}{\rho_D^{\mathrm{exp}}}} = \eta, \qquad \lambda' \approx \lambda\sqrt{\frac{\rho_B^{\mathrm{exp}}}{\rho_D^{\mathrm{exp}}}} = \lambda \tag{6}

Residual σ₀ shift: What does shift is the expert-side signal-to-noise parameter: $\sigma_0(a') = \sigma_0(a)/\sqrt{\rho_B^{\mathrm{exp}}}$ . This produces bounded, minor hyperparameter drift—not exact invariance.

Composition Rules for All MoE Variants

The key insight is that total-expert and granularity changes are not new primitive rules—they derive from composing Bridges I and II:

Total expert count (capacity): Compose Dense-MoE companion width transfer (Bridge I) with reverse-sparsity (Bridge II). The dense-width and sparsification factors cancel, leaving the same sparse-layer rule.
Fixed-density granularity: Since $s = a/N$ is fixed, the sparse active width has the same ratio as the Dense-MoE companion width; no extra SDE correction is needed.
Shared/group-balanced/hybrid blocks: One common FFN-output ABC parametrization applied to total active width $H_{tot} = \sum_{m\in\mathcal{D}} H_m + a\sum_{g\in\mathcal{G}} h_g$ :

y(x) = \frac{d}{H_{\mathrm{tot}}}\left[\sum_{m \in \mathcal{D}} o_m(x) + a\sum_{g \in \mathcal{G}}\sum_{e=1}^{N_g} \pi_{g,e}(x) o_g^{(e)}(x)\right] \tag{7}

Empirical Validation / Results

Controlled Small-Scale Sweeps

Activated experts (Bridge II): Complete-muE keeps LR, WD, and init optima broad and aligned across 64e2a through 64e16a:

LR range: $4 \times 10^{-4}$ – $4 \times 10^{-3}$
WD: 0.01–0.2 (LM), $10^{-4}$ – $3 \times 10^{-2}$ (diffusion)
Init std: $4 \times 10^{-4}$ – $4 \times 10^{-2}$ (LM), $10^{-4}$ – $3 \times 10^{-2}$ (diffusion)

All MoE axes: Capacity, granularity, shared experts, group-balanced routing, depth, and width all preserve similar LR optima. Larger a consistently reaches lower attainable loss (improved expert-side SNR).

Batch/duration transfer: For batch sizes 128–1024, optima stay close at both fixed steps and fixed samples. At fixed samples, $\eta \propto \sqrt{\rho_B}$ is exact; at fixed steps, transfer is approximate with minor drift consistent with the σ₀ shift.

Fixed-Hyperparameter Loss Scaling

Fixing AdamW hyperparameters at dense-tuned values (LM: LR= $10^{-3}$ , init std= $10^{-2}$ , WD=0.1; DF: LR= $1.6 \times 10^{-3}$ , init std= $2 \times 10^{-2}$ , WD= $10^{-2}$ ) and varying only the MoE architecture yields consistent loss reduction across activated experts, total capacity, granularity, and depth—direct evidence for the tune-dense-once recipe.

Systems Benchmark (Single H100)

Scaling Axis	Configuration	Latency Range	vs Dense Baseline (81.0 ms)
Capacity	E=8–256, k=8, h=2048	87.8–97.0 ms/step	1.08×–1.20× slower
Granularity	k=2–64, active width 16,384	84.1–135.0 ms/step	up to 1.67× slower
Dense width	width scaling	81.0–667.6 ms/step	up to 8.2× slower (OOM beyond)

Key finding: Capacity scaling is the cheaper axis when token batch size per device is large enough.

Large-Scale Results

Convergence speedups (dense steps / MoE steps to reach same loss):

256P images: ~2.5×
240P 5s videos: ~4.5×
LLM (100k iterations): 5.3×–5.5×

LLM benchmark evaluation (13 tasks, 100k steps):

Model	Average Score
Dense	44.3
MoE 128e8a4g1s	49.1
MoE 128e8a1s	50.6

Both MoE variants improve over dense; the non-grouped variant is best on average, while the group-balanced variant is strongest on ARC-Challenge.

Theoretical and Practical Implications

Theoretical Contributions

Identification of the bottleneck: Dense-to-sparse MoE transfer couples architecture and workload transfer—neither μP nor SDE rules alone can handle this.
Two-bridge compositional framework: Rather than treating dense-to-sparse as one direct jump, the authors decompose it into two well-grounded bridges (active-width μP + expert-side SDE bookkeeping).
Non-strict SDE invariance: Bridge II reveals that activated-expert transfer is not exact SDE invariance—a residual $\sigma_0$ shift remains, producing bounded, minor drift. This explains why expert-count sweeps show mild optimizer drift in practice.

Practical Implications

"Tune dense once, transfer to all": A single dense-FFN calibration suffices for near-optimal hyperparameters across the entire MoE design space.
Capacity over granularity: H100 benchmarks show capacity scaling is much cheaper than fine granularity at fixed active width.
Shared experts caveat: Increasing shared experts consistently leads to higher loss unless used for computation/communication overlap in expert-parallel training.
Cross-modal transfer: One hyperparameter setting transfers across images (256P/512P), key frames, videos, and language modeling.

Conclusion

Complete-muE provides a compositional AdamW transfer rule for FFN/MoE families, bridging the gap between dense and sparse architectures through two well-grounded bridges. The framework covers activated experts, total capacity, fixed-density granularity, shared/group-balanced hybrids, and width/depth/batch/duration changes.

Key takeaways:

Hyperparameters tuned on a single dense reference transfer near-optimally to all MoE configurations
Mild drift is consistent with Bridge II's non-strict SDE behavior and small enough for practical use
Large-scale MoE training achieves significant convergence speedups (up to 4.5× video, 5.3×–5.5× LLM) with minimum tuning overhead

Future directions likely include: extending to expert-choice routing, deeper analysis of the σ₀ shift for exact invariance, and application to additional modalities and model families beyond the tested LM and diffusion transformer settings.