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

> Complete-muE enables tuning hyperparameters once on a dense model and transferring them to any sparse MoE variant, achieving up to 5.5x faster convergence without retuning.

- **Source:** [arXiv](https://arxiv.org/abs/2605.23893)
- **Published:** 2026-08-18
- **Permalink:** https://picx.dev/p/4MwF44
- **Whiteboard:** https://picx.dev/p/4MwF44/image

## Summary

## 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

1. **Identification of the bottleneck**: Dense-to-sparse MoE transfer couples architecture and workload transfer—neither μP nor SDE rules alone can handle this.
2. **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).
3. **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.

---

_Markdown view of https://picx.dev/p/4MwF44, served by PicX — AI-generated visual whiteboard summaries of research papers._