# How to Scale Mixture-of-Experts: From muP to the Maximally Scale-Stable Parameterization

> Maximal Update Parametrization fails for Mixture-of-Experts models, so the authors propose Maximal Scale Stability, new scaling rules ensuring stable feature learning across all experts.

- **Source:** [arXiv](https://arxiv.org/abs/2605.14200)
- **Published:** 2026-08-17
- **Permalink:** https://picx.dev/p/iX3crF
- **Whiteboard:** https://picx.dev/p/iX3crF/image

## Summary

# 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<sup>-1/2</sup>.

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

| Regime | N (width) | Nₑ (expert width) | M (experts) | K (top experts) |
|--------|-----------|-------------------|--------------|------------------|
| I | → ∞ | → ∞ | Θ(1) | Θ(1) |
| II | n → ∞ | Θ(1) | → ∞ scaling with N | Θ(1) |
| III | n → ∞ | Θ(1) | → ∞ | → ∞ |
| III' | n → ∞ | Θ(1) | → ∞ | = 1 |

---

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

| Parameter | Regime I (M=Θ(1)) | Regime II (Nₑ=Θ(1)) | Regime III (Nₑ∝N) |
|---|---|---|---|
| Router init std | Θ(1/N) | Θ(1/√(N)) | Θ(1/√N) |
| Router LR (SGD) | Θ(1/N) | Θ(M/N) | Θ(1) |
| Expert hidden init | Θ(1/√N) | Θ(1/√N) | Θ(1/√N) |

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

| **Component** | **µP** | **MSSP (Proposed)** |
|---|---|---|
| **Router init std** | N⁻¹ for Regime I | N⁻¹/² for Regimes II & III |
| **Router LR (SGD)** | N⁻¹ | MN⁻¹ (Regime II) |
| **Router LR (Adam)** | N⁻¹ | N⁻¹ (same) |
| **Router ϵ (Adam)** | 1 | M⁻¹ |
| **Expert init std** | N⁻¹/² (each layer) | N⁻¹/² (Regime I), shared init in Regime III |

---

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

| Parameter | Init Std | SGD LR | Adam LR | Adam ε |
|---|---|---|---|---|
| Router (gating) | N^{-1/2} | M/N | 1/N | 1/M |
| Expert Layer 1 (in) | N⁻¹/² | M/N | 1/N | Θ(1) |
| Expert Layer 2 (out) | N⁻¹/² | M/N | 1/N | 1/M |

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

| Parameter | Init Std | SGD LR | Adam LR | Adam ε |
|---|---|---|---|---|
| Router | N^{-1/2} | 1 | 1/N | 1/M |
| Expert Layer 1 | N^{-1/2} | 1/N | 1/N | 1/M |
| Expert Layer 2 | N^{-1/2} | 1/N | 1/N | 1/M |

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

| **Hyperparameter** | **Regime I** (M,K=Θ(1)) | **Regime II** (Nₑ=Θ(1)) | **Regime III** (Nₑ→∞) |
|---|---|---|---|
| Word embedding init | Θ(1) | Θ(1) | Θ(1) |
| Word embedding LR (SGD/Adam) | N^{-1/2} | N^{-1/2} | N^{-1/2} |
| Residual init | Θ(1) | Θ(1) | Θ(1) |
| Residual LR (SGD) | 1/N | 1/N | 1/N |
| Residual LR (Adam) | 1 | 1 | 1 |
| **Router (gating) init std** | Θ(N^{-1}) | Θ(N^{-1/2}) | Θ(N^{-1/2}) |
| **Router SGD LR** | N^{-2} | MN^{-1} | 1/N |
| **Router Adam LR** | N^{-2} | N⁻¹ | N⁻¹ |
| Router Adam ε | 0 | L/M | L/M |
| Expert Layer 1 init std | 1 | 1/√N | 1/√N |
| Expert Layer 2 init std | 1 | 1 | 1 |
| Expert SGD LR | 1/N | 1/N | 1/N |
| Expert Adam LR | 1/N | 1/N | 1/N |

---

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

---

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