Complete-μE: Tune Dense Once, Transfer to All MoE Configurations

Paper: 2605.23893 Authors: Hongwu Peng, Ohiremen Dibua, Yuanjun Xiong, Yifan Gong, Jianming Zhang, Yan Kang Categories: cs.LG

The Gap

Mixture-of-Experts (MoE) models promise better scaling than dense models by activating only a subset of parameters per token. But there’s a catch: every time you change the MoE configuration (number of experts, routing strategy, capacity), you need to retune hyperparameters from scratch. Existing transfer methods fail here. μP requires fixed architecture (but dense-to-MoE changes architecture). SDE requires fixed tokens-per-step (but MoE changes tokens-per-expert as you scale). The result: practitioners either stick with dense models or burn compute on hyperparameter sweeps for each MoE variant.

Problem: Dense → MoE hyperparameter transfer
   |
   v
Existing tools fail:
   μP: architecture must stay fixed
   SDE: per-step token count must stay fixed
   |
   v
Complete-μE: Two-bridge system
   Bridge I:  Dense FFN ←→ Dense MoE (active-width μP)
   Bridge II: Dense MoE ←→ Sparse MoE (activated-expert scaling)
   |
   v
Evidence: Language + diffusion model experiments
   - Hyperparameters stable across configs
   - Minor drift within acceptable bounds
   |
   v
Conclusion: Tune dense once → transfer to all MoE setups

The Increment

One sentence: Before this paper, scaling MoE models required hyperparameter retuning for each configuration; after, you tune once on a dense model and transfer to any MoE setup with near-optimal performance.

Core Mechanism

Complete-μE splits the transfer problem into two bridges. Bridge I handles the dense-to-MoE transition by treating the MoE as a “wide” dense layer where only some neurons activate. It applies μP’s width-scaling rules but normalizes them by the number of active experts, not total experts. This keeps learning rates stable when you convert a dense FFN into an MoE with the same active capacity.

Bridge II handles scaling within MoE configurations (changing total experts while keeping active experts constant). Here, the key insight is that first-order learning rate and weight decay corrections cancel out when you scale experts proportionally. What remains is a bounded residual shift in initialization scale (σ₀), which empirically causes only minor drift. The framework composes these bridges to cover arbitrary changes: width, depth, batch size, number of experts, routing granularity, and hybrid architectures.

Input: Dense model hyperparameters (LR, WD, init scale)
   |
   v
[Bridge I: Dense → Dense MoE]
   Active-width μP scaling
   Router scale normalization
   |
   v
Dense MoE hyperparameters
   |
   v
[Bridge II: Dense MoE → Sparse MoE]
   Activated-expert scaling
   First-order corrections cancel
   Residual σ₀ shift (bounded)
   |
   v
Output: Sparse MoE hyperparameters (any config)

Think of Complete-μE as a currency exchange system. Bridge I is the exchange booth that converts dollars (dense model) to euros (dense MoE) using a fixed rate based on active purchasing power. Bridge II is the ATM network that lets you withdraw different denominations (sparse MoE configs) — the exchange rate drifts slightly due to fees (σ₀ shift), but it’s predictable and small enough that you don’t need to re-negotiate the base rate. The two-stage conversion preserves value (hyperparameter optimality) across arbitrary denominations (MoE configurations).

Key Concepts

• Active-width μP: Standard μP scales learning rates by 1/width to keep feature learning stable as networks widen. In MoE, “width” is ambiguous — do you count all experts or just active ones? Active-width μP says: scale by active experts only, because those are the neurons actually processing each token. If you have 8 experts but activate 2, treat it like a dense layer of width 2×expert_dim, not 8×expert_dim. This prevents learning rates from collapsing when you add dormant experts.

• Activated-expert scaling: When you scale from 8 to 16 total experts while keeping 2 active, tokens-per-expert halves. SDE theory says this should require learning rate adjustments. But Complete-μE shows the first-order corrections cancel: the LR boost from fewer tokens-per-expert exactly offsets the LR reduction from more total capacity. What’s left is a second-order initialization shift that’s small enough to ignore in practice.

• Router scale normalization: MoE routers produce logits that determine which experts activate. If you don’t scale these logits properly, adding more experts makes routing probabilities collapse toward uniform (entropy increases). Complete-μE normalizes router output scale by √(active_experts), keeping routing sharpness constant as you change the expert count. Without this, transferred hyperparameters would produce degenerate routing.

Framework Shift

Before (mainstream approach):        After (this paper):

Dense model                          Dense model
    |                                    |
    | tune HP                            | tune HP once
    v                                    v
MoE config A                         [Complete-μE transfer]
    |                                    |
    | retune HP                          +---> MoE config A
    v                                    |
MoE config B                             +---> MoE config B
    |                                    |
    | retune HP                          +---> MoE config C
    v                                    |
MoE config C                             +---> (any MoE config)
    
Each config = new HP search            One dense tune → all configs

From per-configuration tuning to universal transfer, the core shift is treating MoE scaling as a composition of two well-behaved transformations rather than a monolithic architecture change.

Expert Assessment

Problem choice: Real gap. MoE adoption is bottlenecked by hyperparameter tuning costs, not model quality. The paper targets a practical pain point that existing theory (μP, SDE) genuinely doesn’t solve. This sits at the intersection of scaling laws and systems efficiency — timely given the industry’s MoE push.

Method maturity: Clever decomposition. The two-bridge structure is elegant because it reuses existing theory (μP for Bridge I, SDE for Bridge II) rather than inventing new scaling laws. The router normalization is a small but critical detail that prior work missed. However, the “first-order cancellation” claim in Bridge II deserves more theoretical rigor — the paper leans on empirical validation where a tighter proof would strengthen the contribution.

Experimental integrity: Solid baselines. The paper tests on both language models (GPT-style) and diffusion models (DiT), showing the method isn’t domain-specific. The experiments sweep across expert counts, routing strategies, and model sizes. One weakness: all experiments use relatively small models (up to a few billion parameters). The “minor drift” in Bridge II might amplify at GPT-4 scale, but we don’t have evidence either way. The comparison to μP and SDE is fair — those methods genuinely can’t handle the MoE case.

Writing quality: The abstract and intro are crisp, but Section 3 (method) buries the key insight under notation. The paper would benefit from leading with the two-bridge intuition before diving into scaling equations. The experimental section is thorough but repetitive — consolidating similar results would improve readability. Figure quality is high, especially the transfer stability plots.

Verdict: weak accept — Solves a real problem with a principled method and solid experiments, but the theoretical justification for Bridge II could be tighter, and large-scale validation is missing.

Takeaways

Decompose hard transfer problems: When direct transfer fails, look for intermediate representations where existing theory applies. Complete-μE’s two-bridge structure is a template: find a “dense MoE” middle ground where one set of rules works, then apply different rules for the second hop.

Normalize by active capacity, not total capacity: In any sparse architecture (MoE, sparse attention, conditional computation), scale hyperparameters by what’s actually active per forward pass, not the total parameter count. This principle likely transfers to other sparse models.

First-order cancellations are real: When scaling two dimensions simultaneously (here: experts and tokens-per-expert), their first-order effects can cancel, leaving only bounded residuals. If you’re doing hyperparameter transfer, check whether your scaling dimensions have this property — it might save you from overcomplicating the transfer rule.

Router entropy matters: The router normalization trick (scale by √active_experts) is a concrete detail worth stealing. If you’re building any gating mechanism, keep the output logit scale constant as you change the number of choices, or the gate will degrade into uniform noise.