Technical Summary: Hyperparameter Transfer with Mixture-of-Experts Layers

1. Problem Statement

Mixture-of-Experts (MoE) layers have become a critical mechanism for scaling modern neural networks by decoupling the total number of trainable parameters from the active parameters during the forward pass. However, training sparse MoE models introduces significant complexity regarding hyperparameter (HP) selection. Unlike dense models, MoE architectures introduce new trainable parameters (router weights) and new architectural dimensions (number of experts, expert size) that require careful tuning.

Directly tuning hyperparameters (such as learning rate, initialization scale, and weight decay) at large scales is computationally prohibitive. While HP transfer techniques exist for dense transformers—allowing optimal hyperparameters found on small models to be extrapolated to larger ones—these methods have not been rigorously adapted for the specific scaling dynamics of sparse MoE layers. The core challenge is to determine how to scale hyperparameters when increasing model width, depth, expert count, and expert size simultaneously, without re-tuning at every scale, while ensuring training stability and performance.

2. Methodology

2.1. Proposed Parameterization

The authors propose a new parameterization for transformer models with MoE layers, extending the CompleteP parameterization (previously developed for dense transformers) to include MoE-specific scaling rules. The goal is to define a set of rules that predict how raw hyperparameter values (initialization standard deviation $\sigma$ and learning rate $\eta$) should change as model dimensions scale, such that the training dynamics remain consistent.

The derivation relies on the max-update parameterization ($\mu$P) principle, which requires that network components (pre-activations and residual updates) remain $O(1)$ at initialization and receive $\Theta(1)$ updates per training step. The authors extend this to MoEs by requiring that the max-update conditions hold not just for the layer output, but for individual expert components (mixing coefficients and expert outputs).

Key scaling rules derived for the MoE module include:

• Router Weights: Learning rate scales as $\eta \propto n_{embd}^{-1}$. Initialization scales as $n_{embd}^{-\gamma}$ (with $\gamma \geq 0.5$).
• Expert Biases: Initialized at zero with a constant learning rate $\eta \propto 1$ (independent of expert count), provided sparsity is fixed.
• Expert MLP Weights:
  • Up-projection ($W_{up}$): $\sigma_{init} \propto n_{embd}^{-1/2}$, $\eta \propto n_{embd}^{-1}$.
  • Down-projection ($W_{down}$): $\sigma_{init} \propto \alpha_{ffn}^{-1} n_{embd}^{-1/2}$, $\eta \propto \alpha_{ffn}^{-1} n_{embd}^{-1}$.
  • Here, $\alpha_{ffn}$ represents the expert hidden size multiplier relative to the embedding dimension.

Crucially, the authors fix the sparsity ratio $\kappa = n_{act}/n_{exp}$ (the fraction of active experts) as a constant while scaling the total number of experts ($n_{exp}$) and active experts ($n_{act}$). This contrasts with approaches that fix the number of active experts while increasing the total pool.

2.2. Theoretical Grounding: Dynamical Mean-Field Theory (DMFT)

To justify these heuristic scaling rules, the authors employ Dynamical Mean-Field Theory (DMFT). They analyze the training dynamics of residual networks with MoE layers in the simultaneous limit of infinite width ($n_{embd}$), depth ($L$), expert size ($n_{hid}$), and expert count ($n_{exp}$), while maintaining a constant activation sparsity $\kappa$.

The analysis reveals a novel three-level mean-field hierarchy:

1. Residual Stream Representations: Mean-field over expert outputs.
2. Expert Outputs: Mean-field over individual expert neurons.
3. Individual Neurons: Mean-field within each expert.

The DMFT analysis demonstrates that under the proposed parameterization, the limiting training dynamics are:

• Independent of the FFN ratio ($\alpha_{ffn}$): The dynamics do not depend on the specific size of the experts relative to the embedding dimension, provided the joint scaling limit is taken.
• Dependent only on sparsity $\kappa$: The dynamics are consistent across all scaling parameters as long as the sparsity ratio is fixed.
• Scale-Invariant: The evolution of network summary statistics (e.g., layerwise feature kernels) is consistent across scales, theoretically guaranteeing that hyperparameters transfer reliably.

2.3. Experimental Setup

The authors validate their parameterization empirically using decoder-only Transformer language models on the FineWeb and C4 datasets.

• Base Models: Tuned on small models with approximately 38M activated parameters.
• Scaling: Scaled up to models with up to 2B total parameters, varying width, depth, expert count, and expert size.
• Token Budget: Experiments were conducted on a fixed token budget of 1B tokens (2000 steps) to isolate early-training dynamics, as well as longer horizons (up to 7.5B tokens).
• Optimizer: Standard Adam optimizer.
• Load Balancing: An auxiliary-loss-free strategy was used, updating expert biases directly to encourage load balancing without adding a regularization term to the loss function.

3. Key Contributions

1. MoE Parameterization: The paper extends the CompleteP parameterization to sparse MoE models, providing explicit scaling rules for router weights, expert biases, and expert MLP weights across width, depth, expert count, and expert size.
2. Theoretical Justification via DMFT: The authors provide a rigorous theoretical foundation for their parameterization using DMFT. They derive an explicit description of training dynamics in the infinite-width/depth limit, proving that the dynamics converge to a stable, scale-invariant system that depends only on sparsity, not on the specific scaling of expert dimensions.
3. Empirical Validation of HP Transfer: The study demonstrates that optimal hyperparameters (learning rate and initialization scale) identified on small base models (38M active params) transfer reliably to much larger models (up to 2B total params) across various architectural dimensions.
4. Architectural Insights: The authors empirically verify that, under their parameterization, increasing the number of experts (while keeping total parameters fixed) yields better performance than increasing the size of individual experts. This finding aligns with recent literature but is achieved here without the need for expensive hyperparameter sweeps at each scale.

4. Results

• Reliable HP Transfer: Under the proposed scaling rules, optimal base learning rates and initialization standard deviations transfer effectively across models ranging from 51M to 2B total parameters. The loss curves of scaled-up models collapse with the base model in early iterations before diverging (with larger models achieving lower loss).
• Stability: The parameterization ensures stable training dynamics, including uniform expert load balancing, even when scaling the number of experts. The authors note that MoE pre-training is particularly sensitive to constant-scale hyperparameters (multipliers treated as $\Theta(1)$), and tuning these is crucial for stability.
• Performance: Models trained using zero-shot hyperparameters (transferred from small models) achieve competitive performance against dense baselines (e.g., GPT-2 small/medium) when matched on active parameter counts.
• Expert Count vs. Size: Experiments confirm that increasing the number of experts is more parameter-efficient than increasing expert size at a fixed parameter count. This benefit persists even in longer training horizons (5B tokens).
• Sparsity Sensitivity: The study highlights that hyperparameter transfer is valid only when the sparsity ratio $\kappa$ is fixed. Scaling the number of experts while fixing the number of active experts (sending $\kappa \to 0$) breaks the transferability of optimal hyperparameters.

5. Significance and Claims

The paper claims to provide a practical yet rigorous framework for scaling MoE models. By combining heuristic parameterization with DMFT analysis, the authors offer a method to:

• Reduce Training Costs: Enable the selection of optimal hyperparameters for large-scale MoE models by tuning only small, cheap base models.
• Ensure Stability: Provide rules that prevent training instability (e.g., expert collapse or divergence) which is common in sparse MoE training.
• Guide Architecture Design: Offer theoretical and empirical evidence that increasing expert count is preferable to increasing expert size under fixed sparsity, aiding in the design of efficient large-scale models.

The authors acknowledge limitations, noting that the current work focuses on learning rate and initialization, leaving other hyperparameters (batch size, weight decay, LR schedules) for future investigation. They also note that while the DMFT analysis supports the infinite-width limit, the behavior of small-width transfers requires further theoretical study. Furthermore, the paper does not claim to solve the "compute-optimal" scaling laws for MoEs (e.g., Chinchilla exponents), as the FLOP-to-performance trade-off in MoEs differs significantly from dense models due to sparsity-induced hardware constraints.

In summary, this work establishes a foundational parameterization that allows for the reliable extrapolation of hyperparameters from small to large MoE models, grounded in a novel three-level mean-field theory of training dynamics.