Complete-muE

TL;DR

👤 If you're new to MoE, watch the video pairs first, then read why MoE tuning is so painful and continue top-to-bottom.
🛠️ If you train models for a living, jump straight to the recipe and the worked example.
🔬 If you want the math, see how it works (theory) and the "Under the hood" boxes inside it.

See it first: dense vs. MoE samples at identical training cost

Skip the math for a moment. The eight pairs below run on the same text prompts at identical training compute — dense baseline ($\sim$0.62B non-embedding parameters) on the left, MoE 128e8a1s ($\sim$6.3B total, $\sim$0.62B active per token) on the right. The two models are matched in per-token compute — the MoE simply expands total capacity through sparse experts without raising the per-token cost. Both share one hyperparameter setting derived from a single small-dense calibration sweep; there is no per-architecture sweep between them. The MoE side tracks the prompts more faithfully (motion, composition, fidelity) at no extra training budget.

If your reaction after watching is "how would you ever find good MoE hyperparameters without sweeping at this scale?" — that's exactly the question this post answers. The next section names the pain; the one after that shows the numerical proof that the recipe avoids it.

Why MoE tuning is so painful

The MoE samples above look like a free lunch from the outside, but the standard path to them is anything but. A modern Mixture-of-Experts transformer block has many knobs you can turn:

• How wide each expert is ($h$)
• How many experts the router activates per token ($a$, "activated experts")
• How many experts exist in total ($N$, "capacity")
• Granularity (more, smaller experts vs. fewer, larger ones)
• Shared experts that always fire
• Group-balanced routing
• ...all on top of the usual width, depth, batch size, and training duration

The catch: each of these choices changes two things at once.

This double-change — shifting update sizes

• $\mu\text{P}$¹ handles changes to a fixed parameterized model very well — it gives you zero-shot width transfer. Subsequent work extends it to depth, practical training, and diffusion transformers^2,3,10. But plain $\mu\text{P}$ has no slot for the per-expert workload.
• SDE-based batch/duration rules⁴ handle changes in tokens-per-step for a fixed architecture. They don't apply when the architecture itself changes.
• Recent MoE-specific transfer studies^5,6 preserve learning-rate ranges under expert-count sweeps but stop short of one compositional rule for the full MoE design space.

So when you go from dense FFN to MoE, or change the MoE capacity, neither tool alone tells you what learning rate to use. Practitioners end up doing what they always did — another full hyperparameter sweep at the new architecture — and that gets expensive fast. Complete-muE is the compositional rule that closes the gap — and the next section is the numerical proof.

Large-scale evidence: the numbers behind the samples

The samples track what the metrics say. At large scale we trained dense and MoE side by side on both LLM and multi-modal diffusion, all at one hyperparameter setting that came out of a single small-dense calibration sweep. We lead with the LM benchmark numbers — they are the most directly interpretable form of "the recipe worked" — then drop into the underlying training curves and convergence speedups that produced them.

Downstream LLM benchmarks: +4.8 to +6.3 points on average, zero retuning

After 100k steps of LM pretraining at the dense-tuned hyperparameters ($\text{LR} = 5\!\times\!10^{-4}$, $\text{WD} = 0.05$, base init std $= 10^{-2}$), both MoE variants substantially improve the dense baseline's average across 13 standard 5-shot benchmarks:

Average score: 44.3 → 49.1 for the group-balanced variant and 44.3 → 50.6 for the non-grouped variant — a +4.8 to +6.3 point improvement, at the dense-tuned hyperparameters, with no MoE-specific sweep. Same compute, same optimizer state, just a different architecture wired up with Complete-muE's deterministic rule.

Training-loss curves and convergence speedup across modalities

One last observation about these runs: pretraining with Complete-muE is remarkably stable. Across both the diffusion and LLM settings, training-loss and gradient-norm traces stay smooth for the entire run — essentially free of spikes. A well-calibrated hyperparameter setting does more than improve final loss; it keeps the optimizer in a well-behaved regime where per-step updates are neither too aggressive nor too conservative, avoiding the erratic dynamics that otherwise force costly mid-run interventions or restarts at large MoE scale.

Bonus: large MoEs train at near-dense cost

The results above are about tuning being cheap with Complete-muE. The other half of the practical win is that the training itself — at least along the capacity axis that delivers the largest loss gain (we'll see this in the small-scale evidence further down) — is also cheap. On a single H100 80GB GPU, scaling from $8$ to $256$ total experts (at a fixed $8$ activated) raises step latency only from $87.8$ to $97.0$ ms/step — just $1.08\text{-}1.20\times$ the $81.0$ ms dense SwiGLU baseline. So you can grow MoE parameter count by an order of magnitude while keeping training cost essentially flat.

For context, reaching comparable parameter count by widening a dense model instead would push step latency from $81.0$ to $667.6$ ms ($8\times$ slower) before larger sizes simply run out of memory. The same comparison goes for granularity: fine-grained partitioning ($k = 2 \to 64$) raises latency from $84.1$ to $135.0$ ms/step — not free, but still far cheaper than dense width scaling.

Bottom line: large MoEs trained via capacity scaling are nearly free relative to dense, and Complete-muE makes them tunable for free, too. Both halves of the win compound.

The recipe in three steps

So how do you get the results above? Tune once on a small dense FFN proxy, then compose two short, deterministic tables — one MoE-architectural, one for batch/duration — to transfer those values to any MoE target.

The tuning menu below fits on one screen. You do not need to re-sweep for any MoE setting once the dense reference is tuned.

Tuning menu at a glance

The actual rule is even simpler than it sounds. There is a single active-width $H_{\mathrm{act}}$ that drives every MoE-specific adjustment, plus one routing choice that keeps things clean:

In short: pick a routing choice once, compute $H_{\mathrm{act}}$ for your block, derive three numbers ($A$, $R$, $\sigma_{\mathrm{down}}$), apply the CompleteP batch/duration multipliers if the schedule changed, and reuse every other hyperparameter from the dense baseline.

Worked example: small dense → large MoE (128e8a1s)

This is the recipe applied end-to-end: take the small-dense LM and diffusion calibration values, plug them through the two tables, and arrive at the hyperparameters actually used in the large-scale runs. The summary table below shows every knob at both scales; the three steps after it explain where the "Large" column comes from.

Setup — every hyperparameter at a glance

Step 1 — global batch/duration scaling (Table 2)

The batch is unchanged; only the duration grows ($T_\star = 25\text{k} \to T = 100\text{k}$). This is the "fixed batch, change duration" row of the tuning menu, with $\rho_B = 1$ and $\rho_D = \rho_B \cdot \rho_T = 4$, so

$$\eta,\lambda \;\;\times\;\; \sqrt{\rho_B/\rho_D} \;=\; \sqrt{1/4} \;=\; 1/2.$$

That halves the base LR and WD between the two columns of the table above — LM: $10^{-3} \to 5\!\times\!10^{-4}$, $0.10 \to 0.05$; diffusion: $4.52\!\times\!10^{-3} \to 2.26\!\times\!10^{-3}$, $0.02 \to 0.01$. Init std is untouched by Step 1 (batch and duration don't enter the initialization rule).

Step 2 — MoE layer-level rule (Table 1)

The 128e8a1s block has 8 routed-active experts plus 1 shared expert, giving an FFN expansion equivalence of $4.5\times$ and an active width $H_\mathrm{act} = 9d$ when shared and routed contributions are combined. The Table 1 lookup for hybrid blocks then gives:

• Output multiplier $A = d/H_\mathrm{act} = 1/9$, applied to the MoE down-projection output.
• Route scale $R = 8$ on the normalized routed sum; the shared branch keeps $R = 1$.
• Down-projection init-std factor $\sqrt{H_\mathrm{act}/d} = \sqrt{9} = 3$, multiplied onto the dense-baseline init std for that one parameter.

No MoE-specific multiplier touches the LR or WD — those changes are fully absorbed by $A$, $R$, and the init-std factor.

Step 3 — standard $\mu\text{P}$ per-layer scaling ($d$: $128 \to 1024$)

Standard $\mu\text{P}$ scales every 2D linear projection inside the transformer block — the FFN/MoE up, gate, and down projections, the router readout, the attention $Q$, $K$, $V$, and $O$ projections, plus anything else that is a learned $d_\text{in} \!\times\! d_\text{out}$ matrix. All of them inherit the same two width factors:

$$\eta_{\text{per-layer}} \;\propto\; \rho_d^{-1} = 1/8, \qquad \sigma_{\text{per-layer}} \;\propto\; \rho_d^{-1/2} = 1/\sqrt{8}.$$

In words: every such linear layer at $d = 1024$ uses $1/8$ the LR and $1/\sqrt{8}$ the init std of the same layer at $d = 128$, before any MoE-specific factors are applied on top.

$\mu\text{P}$ also specifies a separate output multiplier for the final embedding-to-logits (or embedding-to-output) projection — the LM head for language modelling, the noise/velocity prediction head for diffusion. That multiplier is $1/\rho_d$ (the inverse width-expansion ratio). With $\rho_d = 8$ here, the final output projection gets an extra $\times 1/8$ scaling factor, independently of the per-layer LR/init rules above.

Effective per-parameter values at the large MoE

Composing Step 2 (MoE layer rule) with Step 3 ($\mu\text{P}$ width factor) gives the actual per-parameter LR and init std for each parameter group:

The final output projection gets both the per-layer $\mu\text{P}$ scaling (Step 3, same as every other 2D linear) and the additional output multiplier $\times 1/\rho_d$ on its output activations. On top of that, the MoE-specific multipliers from Step 2 apply only to the routed-output branch inside the MoE block: output multiplier $A = 1/9$ on the down-projection output, and route scale $R = 8$ on the normalized routed sum (the shared expert keeps $R = 1$).

There is no per-architecture sweep between the small calibration and the large run. Every number in the "Large" column of the setup table is mechanical — it follows from the small-proxy values by applying the two tables in sequence.

Small-scale evidence

The recipe works because of three layered facts: the dense calibration sweep gives broad, well-separated optima (so reuse is forgiving); fixing those values and scaling only the MoE architecture gives consistent loss reductions (the cleanest possible "tune-once" test); and sweeping any AdamW hyperparameter at fixed MoE architecture keeps the optimum in the same window. We show each one in order.

The small-dense calibration sweep

Step 1 of the recipe is a single sweep on a small dense FFN proxy. For both modalities we scan three hyperparameters — learning rate, weight decay, and initialization std — at proxy backbone width $d_\star = 128$. The optima are broad and well-separated, which is what makes them easy to reuse downstream:

Fixed-LR scaling: the cleanest tune-once test

The most direct check of the recipe is to fix all AdamW hyperparameters at the small-dense calibration values (LM: $\text{LR}=10^{-3}$, $\text{WD}=0.10$, init std $=10^{-2}$; Diffusion: $\text{LR}=4.52\!\times\!10^{-3}$, $\text{WD}=0.02$, init std $=2\!\times\!10^{-2}$) and then scale only the MoE architecture along four axes. If the recipe really is tune-once-and-transfer, loss should drop monotonically and consistently as the architecture grows — no per-setting re-tuning. That is what we see:

One pattern stands out across the four panels: capacity scaling (more total experts at fixed activated count) drives substantially lower loss than granularity scaling at the same fixed hyperparameters. Among the MoE quality knobs, total-expert count is the strongest lever per unit of training compute — which, conveniently, is also the cheapest one to scale, as we'll see in the systems benchmark below.

Activated experts: optima align across LR, WD, init, and the initial loss curve

The fixed-LR test above is the headline. The next two subsections answer the natural follow-up question: why is fixed-LR scaling so clean? Because the LR/WD/init optima are themselves stable across MoE configurations. Start with the most theory-loaded axis: the activated-expert count $a$ (Bridge II in the recipe). We sweep the three AdamW hyperparameters (LR, WD, init std) plus the initial-training-loss trajectory, all at fixed total experts and per-expert width:

LR optima persist across the other MoE architectural axes

Bridge II covered activated experts. The follow-up question is whether the same LR transfer holds for the other MoE architectural knobs — capacity (total experts $N$), fixed-density granularity, shared experts, group-balanced routing, depth, and backbone width. We sweep each axis individually, with Complete-muE applied first and only the base LR scanned. The optima stay aligned across all six axes for both LM and diffusion:

Combined with the activated-experts sweep above, this covers every MoE architectural axis the paper studies. The takeaway is uniform: once Complete-muE plugs the dense-tuned hyperparameters into the new architecture, the LR optimum stays put. There is no per-axis re-tuning step.

What's actually going on under the recipe

If you walked through the worked example, you already used both bridges: Step 2 (the MoE layer rule) is Bridge I — active-width $\mu\text{P}$ — and Step 1 (the global batch/duration rescaling) plus the implicit per-expert workload accounting is what Bridge II handles. The next two sections name the two operations, explain why each works at first order, and then flag the one second-order effect — a bounded $\sigma_0$ drift — that keeps the transfer honest rather than miraculous.

The core insight: active width + expert workload

Complete-muE's main observation is that you can take the messy MoE design space and decompose it into two clean, well-understood steps. Each step changes one thing.

Decomposing the problem this way is the entire trick. Each step is a known problem with a known answer. The hard MoE-specific cases — total expert count, granularity, hybrids with shared experts and group-balanced routing — turn out to be compositions of these two steps. Nothing new needs to be invented for them.

Two bridges from dense to any MoE

The teaser figure at the top of this post is the picture of these two bridges: the two green diagonal arrows of the triangle are Bridge I (left, marked $\mu\text{P}$) and Bridge II (right, marked $\mu\text{P}$ + SDE); the red dashed bar at the bottom is the direct dense-to-sparse jump that doesn't transfer.

Bridge I: Dense FFN ↔ Dense MoE

The first bridge says: a dense FFN of width $H$ and a "Dense MoE"

• Use an output multiplier $A = d/H$ on the down-projection (where $d$ is the residual width). This is the standard $\mu\text{P}$ factor for a width-$H$ FFN.
• Use a route scale equal to the number of activated experts. Normalized routing averages expert outputs, which would otherwise shrink the update by $1/a$; the route scale undoes that.

Bridge II: Dense MoE ↔ Sparse MoE

The second bridge handles the move from a Dense MoE (all $N$ experts active) to a sparse MoE (only $a < N$ active). Once Bridge I is applied at the active width $H_a = ah$, the only remaining change is the per-expert data exposure.

Here's the elegant cancellation: when you switch from $a$ activated experts to $a'$ at fixed global batch size and steps, each expert's batch changes by the same factor as its training duration. Both go up or down together by $a'/a$. The standard SDE correction cancels:

What does change is the expert-side noise-to-signal level $\sigma_0(a) = \eta \cdot \sigma_{\mathrm{exp}}(a)$, where $\sigma_{\mathrm{exp}}(a)$ is the per-expert gradient-noise standard deviation (which scales like $1/\sqrt{B_\mathrm{exp}(a)}$ with the per-expert batch). Making the MoE denser (larger $a$) lowers the per-step gradient noise, which is why denser MoEs can reach lower loss at the same hyperparameters — we return to this caveat in the next section.

Everything else is composition

That's it for new primitives. Every other MoE setting — total capacity, granularity, hybrid blocks with shared experts and grouped routing — is built by composing Bridge I and Bridge II:

• Total capacity ($N \to N'$): first widen the Dense-MoE companion (Bridge I), then re-sparsify back to $a$ active experts (Bridge II). The two factors cancel exactly.
• Granularity at fixed density ($s = a/N$): the active width is preserved by construction, so the same sparse-layer rule applies.
• Hybrid blocks (shared + grouped + routed): use one common active-width multiplier over the total active width $H_{\mathrm{tot}}$ and apply route scale only to routed branches.

Bridge II in practice: a second-order drift caveat

The first-order SDE correction in Bridge II cancels exactly, but there's a second-order effect we want to be upfront about: the per-step expert-side noise $\sigma_0(a)$ shifts when $a$ changes. The change is bounded:

$$\sigma_0(a') = \frac{\sigma_0(a)}{\sqrt{\rho_B^{\mathrm{exp}}}}.$$

This means Bridge II is not a strict SDE invariance — it's a relatively stable transfer with some mild residual drift in the optimal hyperparameters. The same kind of bounded drift shows up in two other places in the paper:

• Capacity scaling (varying $N$ at fixed $a, h$), because the composition that builds it uses Bridge II as a sub-step.
• Batch-size transfer at fixed training iterations, where the gradient noise $\sigma_0$ shifts even though the optimization horizon $H_{\mathrm{SDE}} = T\eta^2$ is preserved.

The empirical question is whether this drift is small enough to ignore in practice. The evidence in "Small-scale evidence" and "Large-scale evidence" — the activated-experts sweep, the architectural-axis sweeps, the fixed-LR scaling, and the large-scale runs at one shared hyperparameter setting — and the paper's calibration both answer yes.

Takeaways

• Small dense sweep, large MoE deployment. A single small-dense-proxy sweep is enough — no architecture-specific re-tuning at large scale.
• Two bridges, everything else by composition. Bridge I (active-width $\mu\text{P}$) + Bridge II (per-expert workload) cover the whole MoE design space without new primitives.
• The math is honest about its limits. The SDE cancellation is first-order, with a bounded second-order $\sigma_0$ drift. Small-scale sweeps and large-scale runs both confirm this drift is small enough in practice.
• Practical wins are real and measurable. $4.5\times$ on $240$P $5$s video and $5.3\text{-}5.5\times$ on LLM, all at the dense baseline's hyperparameters.

Why this matters

For teams training large MoE models, the cost of "what learning rate should I use for this new MoE configuration?" has historically been a hyperparameter sweep at scale. Complete-muE replaces that sweep with one table lookup. The savings compound every time a team tries a new sparsity, granularity, or hybrid routing variant.

For researchers, the contribution is conceptual: capacity, granularity, and hybrid routing aren't new primitive hyperparameter-transfer rules — they're compositions of two well-understood bridge cases. That decomposition is the part most likely to outlive any specific architecture trend.

The paper

Complete-muE: Optimal Hyperparameter Transfer and Scaling for MoE Models — Hongwu Peng, Ohiremen Dibua, Yuanjun Xiong, Yifan Gong, Jianming Zhang, Yan Kang (Adobe Research). Read the preprint on arXiv: arXiv:2605.23893.

Citation

Please cite this work as:

Peng, Hongwu and Dibua, Ohiremen and Xiong, Yuanjun and Gong, Yifan and Zhang, Jianming and Kang, Yan, "Complete-muE: Optimal Hyperparameter Transfer and Scaling for MoE Models", arXiv preprint arXiv:2605.23893, 2026.

Or use the BibTeX citation:

@article{peng2026completemue,
  title={Complete-muE: Optimal Hyperparameter Transfer and Scaling for MoE Models},
  author={Peng, Hongwu and Dibua, Ohiremen and Xiong, Yuanjun and Gong, Yifan and Zhang, Jianming and Kang, Yan},
  journal={arXiv preprint arXiv:2605.23893},
  year={2026}
}

Questions? Email hongwup@adobe.com.

References

Foundational works Complete-muE builds on: $\mu\text{P}$-style transfer, SDE-based optimizer scaling, MoE-specific transfer, and the broader sparse-MoE literature. Citations marked in the body link to the corresponding margin note; the full list is here.

1. Yang, G., Hu, E. J., Babuschkin, I., Sidor, S., Liu, X., Farhi, D., Ryder, N., Pachocki, J., Chen, W., & Gao, J. (2022). Tensor Programs V: Tuning large neural networks via zero-shot hyperparameter transfer. NeurIPS. arXiv:2203.03466.
2. Dey, N. S., Zhang, B. C., Noci, L., Li, M., Bordelon, B., Bergsma, S., Pehlevan, C., Hanin, B., & Hestness, J. (2025). Don't be lazy: CompleteP enables compute-efficient deep transformers. NeurIPS. arXiv:2505.01618.
3. Mlodozeniec, B., Ablin, P., Béthune, L., Busbridge, D., Klein, M., Ramapuram, J., & Cuturi, M. (2025). Completed Hyperparameter Transfer across Modules, Width, Depth, Batch and Duration. arXiv:2512.22382.
4. Malladi, S., Lyu, K., Panigrahi, A., & Arora, S. (2022). On the SDEs and scaling rules for adaptive gradient algorithms. NeurIPS. arXiv:2205.10287.
5. Małaśnicki, J., Ciebiera, K., Boruń, M., Pióro, M., Ludziejewski, J., Stefaniak, M., Krutul, M., Jaszczur, S., Cygan, M., Adamczewski, K., et al. (2025). Mu-Parametrization for Mixture of Experts. ES-FoMo III Workshop. arXiv:2506.16962.
6. Jiang, T., Bordelon, B., Pehlevan, C., & Hanin, B. (2026). Hyperparameter Transfer with Mixture-of-Expert Layers. arXiv preprint.
7. Shazeer, N., Mirhoseini, A., Maziarz, K., Davis, A., Le, Q., Hinton, G., & Dean, J. (2017). Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer. ICLR. arXiv:1701.06538.
8. Fedus, W., Zoph, B., & Shazeer, N. (2022). Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity. JMLR. arXiv:2101.03961.
9. Dai, D., Deng, C., Zhao, C., Xu, R. X., Gao, H., Chen, D., Li, J., Zeng, W., Yu, X., Wu, Y., et al. (2024). DeepSeekMoE: Towards Ultimate Expert Specialization in Mixture-of-Experts Language Models. ACL. arXiv:2401.06066.
10. Zheng, C., Lou, A., Liu, C., Wei, X., Liu, Z., & Ermon, S. (2025). Scaling Diffusion Transformers Efficiently via $\mu\text{P}$. NeurIPS. arXiv:2505.15270.
11. DeepSeek-AI. (2024). DeepSeek-V3 Technical Report. arXiv:2412.19437.
12. Orvieto, A., De, S., Gulcehre, C., Pascanu, R., & Smith, S. L. (2025). In Search of Adam's Secret Sauce. NeurIPS. arXiv:2510.08198.