Scaling Mixture-of-Experts with Maximal Scale Stability

Abstract Summary

This paper addresses the challenge of scaling Mixture-of-Experts (MoE) architectures while maintaining stable training dynamics. The authors identify that Maximal Update Parametrization (µP)—the standard framework for scaling dense neural networks—fails when applied to MoE models due to a breakdown in the balance of feature updates across expert parameters.

Key Contributions

1. Diagnosis of µPs Failure in MoEs

The paper demonstrates that under µP:

• The effective update to expert weights (computed from gradients) scales as Θ(1) in the forward pass
• However, the propagating update (the influence of updates on downstream activations) scales as Θ(1/√M) in Regime II and Θ(1) in Regime III
• This imbalance causes feature learning to vanish early in training at finite scales

The root cause: when aggregating across M experts, the sum of M independent random updates creates a Central Limit Theorem effect that suppresses the signal by M^-1/2.

2. Maximal Scale Stability (MSSP)

The authors propose a new criterion requiring that every component of the feature update remains Θ(1), not just the aggregate sum. This is formalized as Maximal Scale Stability (MSS).

Key insight: The failure modes of µP are characterized by a sum under the CLT that grows as √M—when the per-expert update is Θ(1), the aggregate becomes too large (Θ(√M)). The fix requires the per-expert pre-aggregation updates to scale as Θ(1/√M).

---

The µ·-Desiderata Framework

The paper recapitulates three core desiderata for well-behaved scaling:

1. µP-1 (Activations): All pre-activations inherit Θ(1) scale from the data, with weights initialized so RMS norms remain Θ(1) throughout.

2. µP-2 (Effective Updates): Training induces feature updates of order Θ(1) relative to initialization.

3. µP-3 (Stability): The combination of initialization and update scales yields a stable fixed point.

---

Mathematical Setting

Architecture

Given a residual stream $\mathbf{h}_t^\ell$ at layer ℓ and time t, the MoE block computes:

$$\bar{h}_t^\ell = h_t^\ell + \frac{1}{M}\sum_{i=1}^{M} \phi_i^\ell(h_t^\ell) \cdot f_i^\ell(h_t^\ell)$$

where:

• $f_i^\ell$ is the i-th expert (typically a 2-layer MLP with hidden width $N_e$)
• $\phi_i^\ell$ is the router/gating function (softmax over top-K experts)

Scaling Regimes Considered

---

The Central Problem: Update Balance in MoEs

Why µP breaks for MoEs

In a MoE layer, the output is computed as:

$$\mathbf{h}_{t}^{\ell+1} = \mathbf{h}_t^\ell + \underbrace{\sum_{i=1}^{M} \phi^\ell_t(i) \cdot \bar{W}^{\ell,\text{out},i}_t \, \sigma\big( \bar{W}^{\ell,\text{in},i}_t \, \bar{h}^{\ell,\text{in},i}_t \big)}_{\text{MoE sublayer}}$$

where the router output (\phi(i)) determines the weighting of each expert's contribution.

Why Does µP Fail?

The critical issue occurs in the update to the input weights of each expert. When computing the propagated update through the router gating:

$$\Delta h^{l+1}_{t} = \sum_{i=1}^{M} \phi_t(i) \cdot \Delta W^{l,\text{out},i}_t \cdot \sigma(h^{l,\text{in},i}_t)$$

The per-expert updates contain a term:

$$\underbrace{W^{l,\text{out},i}_0(W^{l,\text{out},i}_0)^\top}_{\text{Gram matrix}} \delta^{l+1}_t$$

This Gram matrix product creates unbalanced contributions:

• In Regime II (M, K → ∞ with Nₑ constant): Gram matrices across experts become increasingly incoherent, causing a Θ(1/√M) suppression
• In Regime III (M, K, Nₑ → ∞ together): Gram matrices converge to well-conditioned random matrices, restoring Θ(1) scaling but creating cross-expert correlations

2. Maximal Scale Stability (MSS) Criteria

The authors propose three operational desiderata for scaling MoEs:

1. Stability: The RMS norm of every hidden vector, feature map, and backward signal must be Θ(1) throughout training
2. Stable feature learning: At every layer, the RMS norm of the effective update to pre-activations under one gradient step must be Θ(1) at initialization
3. Stable gradient flow: Gradient updates to parameters must induce Θ(1) feature updates in downstream layers (both backward and forward)

Scaling Rules Summary

The paper provides explicit parameter scaling for each regime:

Regime I: Fixed Experts (M = Θ(1))

This regime behaves essentially like a dense network with a slightly wider effective hidden layer. Standard µP suffices, with router parameters scaled small initially to ensure controlled routing.

Regime II: Fixed Expert Width (Nₑ = Θ(1))

Here M grows while Nₑ stays constant. The authors identify that a finite-width bottleneck forms in the expert computation—the hidden layer of each expert remains Θ(1)-dimensional, causing rank collapse in the per-expert Gram matrices.

Regime III: All Features Scale Together (N, Nₑ, M → ∞)

When embedding width, expert width, and number of experts all grow, the Gram matrices become proportional to identity, restoring a "feature learning" regime.

---

Main Contributions

A. Precise Scaling Prescriptions (Table 1 in paper)

---

Critical Innovations

1. Maximal Scale Stability (MSS)

For Regime II (M ≍ n, Nₑ = Θ(1)):

• Router (gating) initialization: scale to Θ(N⁻¹/²) instead of Θ(N⁻¹)
• Router learning rate: η_router ~ M/N (SGD) or 1/N (Adam)
• This ensures the router operates in a "feature learning" regime where updates remain Θ(1)

For Regime III (Nₑ ≍ n):

• Shared initialization across experts: W₀^(l,out,i) = W₀^(l,out) for all experts i
• This prevents the Gram-matrix incoherence across experts that breaks feature learning

2. Analysis of Update Dynamics

The key mathematical insight involves decomposing the update to the residual stream:

$$\Delta h^{l+1}_t = \sum_{i=1}^{M} \phi_i(x_t) \sum_{j \in \mathcal{S}_t} \Delta W^{l,\text{out},i}_{t} \sigma(W^{l,\text{in},i} h^{l}_{t} + b^{l,i})$$

The authors show this decomposes into three terms with different scaling behaviors:

• Init term: scales as Θ(1/√M) in both Regimes II and III
• Effective update term: scales as Θ(1) in both regimes
• Propagating update term: scales as Θ(1/√M) in Regime II but Θ(1) in Regime III

This imbalance explains why µP fails: the effective update dominates, causing feature learning to be driven primarily by the current batch of data rather than by stable accumulation across training.

3. Maximal Scale Stability (MSS) Parametrization

The paper introduces a new scaling framework that ensures all contributions to the feature update remain Θ(1) with respect to model scale:

Regime II (M, K → ∞; Nₑ = Θ(1))

Regime III (M, K → ∞ with Nₑ → ∞)

Maximal Scale Stability (MSS) Conditions

The core issue: µP only ensures the update stays Θ(1), but not the initialization. The authors introduce the concept of "Maximal Scale Stability" which requires both:

1. Feature updates remain Θ(1) with scale
2. Initialization and propagated signals maintain Θ(1) norms throughout training

The key mathematical condition: For each expert output transformation, the rank-one gradient must interact with the same input norm scale that the forward pass uses. In standard µP, the forward pass of experts uses Nₑ⁻¹/² scaling, but the backward pass produces updates that scale differently.

---

Parametrization Table

---

Key Contributions

1. Problem identification: Show that µP fails in MoE regimes where M scales to infinity.

2. MSS Principle: Introduce "Maximal Scale Stability" — a stronger requirement than maximal update that ensures all contributions in the residual stream remain Θ(1) throughout training.

3. Prescription: Provide hyperparameter scaling rules for both Regime II (small experts, many experts) and Regime III (wide experts).

4. Mean-Field Theory: Derive exact DMFT limits for both µP and MSSP in all regimes.

5. Validation: Extensive experiments on Transformers with MoE layers.

---

Key Results and Intuition

The fundamental problem identified is that in MoE layers, the router and expert parameters interact in a way that breaks the maximal-update property:

• Under µP, the per-expert updates become incoherent (Θ(1/√M)) in certain regimes, causing the effective learning rate to shrink with M
• The solution requires maximal scale stability (MSS): scaling rules that ensure coherent aggregation of updates across experts

For Regime II (expert width bounded):

$$\text{Router init std: } \sigma_{\text{router}} = N^{-1/2}, \quad \eta_{\text{router, SGD}} = \frac{M}{N}$$

For Regime III:

$$\text{Router init std: } \sigma_{\text{router}} = N^{-1/2}, \quad \eta_{\text{router, SGD}} = \Theta(1)$$