Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Hyperparameter Transfer with Mixture-of-Expert Layers

Tianze Jiang, Blake Bordelon, Cengiz Pehlevan, Boris Hanin

TL;DR

This paper tackles the costly problem of hyperparameter tuning in large Mixture-of-Experts (MoE) transformers by developing a DMFT-grounded parameterization that enables HP transfer across width, depth, and MoE dimensions while preserving sparsity. Building on the max-update ($\mu$P) and CompleteP scaling framework, the authors extend these principles to MoE layers, deriving explicit initialization and learning-rate rules for router weights, expert projections, and biases, and proving scale-invariant dynamics in the dynamical mean-field theory (DMFT) limit. Empirically, HPs identified from small models transfer robustly to larger models and longer training horizons, with uniform expert load, and MoE models achieving competitive performance relative to dense baselines under fixed token budgets. The work provides both theoretical and practical foundations for MoE scaling laws and offers concrete guidance for stable, scalable MoE training, opening avenues for longer horizons and broader HP exploration in sparse transformers.

Abstract

Mixture-of-Experts (MoE) layers have emerged as an important tool in scaling up modern neural networks by decoupling total trainable parameters from activated parameters in the forward pass for each token. However, sparse MoEs add complexity to training due to (i) new trainable parameters (router weights) that, like all other parameter groups, require hyperparameter (HP) tuning; (ii) new architecture scale dimensions (number of and size of experts) that must be chosen and potentially taken large. To make HP selection cheap and reliable, we propose a new parameterization for transformer models with MoE layers when scaling model width, depth, number of experts, and expert (hidden) size. Our parameterization is justified by a novel dynamical mean-field theory (DMFT) analysis. When varying different model dimensions trained at a fixed token budget, we find empirically that our parameterization enables reliable HP transfer across models from 51M to over 2B total parameters. We further take HPs identified from sweeping small models on a short token horizon to train larger models on longer horizons and report performant model behaviors.

Hyperparameter Transfer with Mixture-of-Expert Layers

TL;DR

This paper tackles the costly problem of hyperparameter tuning in large Mixture-of-Experts (MoE) transformers by developing a DMFT-grounded parameterization that enables HP transfer across width, depth, and MoE dimensions while preserving sparsity. Building on the max-update (P) and CompleteP scaling framework, the authors extend these principles to MoE layers, deriving explicit initialization and learning-rate rules for router weights, expert projections, and biases, and proving scale-invariant dynamics in the dynamical mean-field theory (DMFT) limit. Empirically, HPs identified from small models transfer robustly to larger models and longer training horizons, with uniform expert load, and MoE models achieving competitive performance relative to dense baselines under fixed token budgets. The work provides both theoretical and practical foundations for MoE scaling laws and offers concrete guidance for stable, scalable MoE training, opening avenues for longer horizons and broader HP exploration in sparse transformers.

Abstract

Mixture-of-Experts (MoE) layers have emerged as an important tool in scaling up modern neural networks by decoupling total trainable parameters from activated parameters in the forward pass for each token. However, sparse MoEs add complexity to training due to (i) new trainable parameters (router weights) that, like all other parameter groups, require hyperparameter (HP) tuning; (ii) new architecture scale dimensions (number of and size of experts) that must be chosen and potentially taken large. To make HP selection cheap and reliable, we propose a new parameterization for transformer models with MoE layers when scaling model width, depth, number of experts, and expert (hidden) size. Our parameterization is justified by a novel dynamical mean-field theory (DMFT) analysis. When varying different model dimensions trained at a fixed token budget, we find empirically that our parameterization enables reliable HP transfer across models from 51M to over 2B total parameters. We further take HPs identified from sweeping small models on a short token horizon to train larger models on longer horizons and report performant model behaviors.
Paper Structure (28 sections, 52 equations, 18 figures, 1 table)

This paper contains 28 sections, 52 equations, 18 figures, 1 table.

Table of Contents

  1. Introduction
  2. Related backgrounds
  3. Scaling recipe for MoEs models
  4. Model setup
  5. Scaling that preserves sparsity versus topK
  6. MoE Parametrization
  7. Dynamical mean field theory
  8. Empirical results and findings
  9. Experiment findings for fixed token budget
  10. Experiments on larger token-horizon
  11. Tradeoff between expert count versus expert size
  12. Conclusion and future directions
  13. More experiment details
  14. Learning rate for biases
  15. Weight decay
  16. ...and 13 more sections

Figures (18)

  • Figure 1: Matching active architecture to be GPT2-small (124M) and 500K batch size, comparison of MoE training loss on our zero-shot Adam HPs (found from tuning 38M activated base models) on FineWeb versus (dense GPT) baseline karpathy2023nanogpt, speedrun AdamW, and speedrun Muon loss curves jordan2025moddednanogpt towards the 3.28 val. loss for baseline at 10B tokens. See \ref{['sec: experiments']} for details. Dense benchmarks are taken before Oct.14, 24, with advanced mods such as zero down-proj. and QK-norms.
  • Figure 2: Global base learning rate (first row) and global base init (second row) transfer trained on 1B tokens (2000 steps) on the Fineweb dataset, with different model sizes from 20M to 1.8B scaling across width, depth, number of experts (fixing sparsity), and expert MLP hidden multiplier. We fix the base config ($n_{\operatorname{ embd}}\text{ (W) }=512$, $L=8$, $(n_{\operatorname{exp}}, n_{\operatorname{act}})=(4, 1)$ ('E4A1' in the figure), $\alpha=1$) and vary one dimension at a time. Error bars in the last row are (max, min, and median) over four independent seeds. See Section \ref{['sec: experiments']} for details.
  • Figure 3: (Finding 1.2): Loss curve collapse (scale invariance) of model scaling dimension (in earlier steps) when scaling up the base model in different ways. Parts (1) and (2): Fineweb dataset on sparsity $1/4$. Parts (3) and (4): C4 dataset on sparsity $1/12$.
  • Figure 4: Fixed token transfer of global base LR (row 1) and global base init (row 2) on $\kappa=1/12$ and the C4 dataset
  • Figure 5: Parts (1) and (2): At a fixed token budget of 1B, scaling up from the base model (1 activated expert and expert $\alpha_{\operatorname{ffn}}=1$), increasing the number of experts is more parameter-efficient than expert size. Error bars are (min, median, max) out of 4 seeds. Part (3): At a 5B token horizon with GPT2-small activated and cosine LR cool-down to zero, having more activated experts (and thus smaller experts) monotonically improves performance. Error bars for $n_{\operatorname{act}}\leq 4$ are (min, median, max) out of 3 seeds.
  • ...and 13 more figures