Let’s Scale Step by Step: Compute-Efficient Hyperparameter Transfer for Large-Scale Mixture-of-Experts

Nayeon Kim, Hojin Lee, Yunju Bak, Jaesun Park, Boseop Kim

A two-step framework that predicts optimal learning rates for large-scale MoE models by transferring across width and extrapolating across token horizons.

How can we predict the optimal learning rate for massive Mixture-of-Experts (MoE) models using only small-scale proxy runs?

Optimizing hyperparameters for Mixture-of-Experts (MoE) models is computationally prohibitive because every change in model size or training duration requires an expensive, full-scale search. The authors propose a two-step transfer framework: they use Maximal Update Parameterization (µP) to transfer optimal learning rates across model widths, then apply a log-linear scaling law to extrapolate those rates to massive token horizons. This method predicts the optimal learning rate for a 155B-parameter MoE model trained on 10 trillion tokens with high fidelity, bypassing the need for exhaustive trial-and-error sweeps.

Paper Primer

The framework hinges on decoupling the search: first, it uses small proxy models to establish width-transferability via µP; second, it uses Exponential Moving Average (EMA) checkpoints from a single proxy run to map the relationship between validation loss and learning rate across token budgets. This allows the authors to fit a log-log linear regression that predicts the optimal learning rate for unseen, long-horizon training.

The framework accurately predicts optimal learning rates for trillion-token training horizons.

Linear regression of optimal learning rates across token budgets achieved an $R^2$ of 0.95. The method enabled successful pretraining of a 155B-parameter MoE model on 10 trillion tokens without full-scale hyperparameter sweeps.

Optimal learning rates transfer reliably across MoE model widths.

Fitted parabolas for validation loss versus learning rate showed consistent vertex locations across width-scaled models (2x, 4x, 8x). This confirms that small proxy models (e.g., 5.6B parameters) can serve as reliable indicators for much larger target architectures.

Why is this two-step approach necessary instead of just using standard scaling laws?

Standard scaling laws often require joint sweeps across both model size and token budget, which is computationally intractable at the 100B+ parameter scale. This framework decouples these dimensions, allowing researchers to use small, efficient proxy runs to inform the configuration of massive models.

Does this framework account for batch size, which is often adjusted during large-scale training?

No; the authors treat batch size as a system-level variable for hardware throughput and fix it to maximize efficiency. They argue that their learning rate scaling law remains robust regardless of the specific batch size chosen for throughput optimization.

The Cost of Scaling MoE

We expose why hyperparameter tuning for large MoE models is a costly bottleneck and outline our two‑step transfer solution.

Mixture-of-Experts (MoE) architectures dramatically increase model capacity while keeping compute low, but they introduce a sprawling hyperparameter space. Among these, the learning‑rate is especially sensitive to both model width and token budget, forcing a fresh, expensive sweep whenever either scale changes. The paper therefore proposes a two‑step hyperparameter transfer framework that first leverages Maximal Update Parameterization (µP) to decouple width scaling, then uses a token‑scale extrapolation law to avoid joint 2‑D searches.

**Figure 1.** (a) In conventional hyperparameter optimization, exhaustive 2D sweeps across both model scale (M) and token scale (D) are required to jointly predict the target scale (★), incurring prohibitive computational costs. (b) Our approach decouples the two dimensions: $\mu$P-based width transferability eliminates model-scale sweeps, while a token scaling law from a few proxy runs enables direct extrapolation to the target scale (★), replacing 2D sweeps with a lightweight 1D search along the token dimension.

Hyperparameter tuning for large MoE models is dominated by an expensive joint search over model width and token budget, making scaling to trillion‑token regimes impractical.

Eliminating the costly hyperparameter sweep makes large‑scale MoE training practical.

The Two-Step Transfer Framework

A resource‑efficient two‑step framework transfers optimal learning rates from tiny proxies to massive MoE models.

Hyperparameter sweeps explode in cost when MoE models reach hundreds of billions of parameters. Existing µP‑based transfer methods assume only width scaling, leaving a gap for sparsity‑driven MoE scaling.

µP rescales weights so that activations and gradients keep the same magnitude as model width grows, preserving stable training dynamics without retuning learning rates.

How does µP differ from the standard parameterization used in most deep‑learning libraries?

Standard parameterization keeps weight variance fixed, so as width increases the summed inputs grow and activations saturate; µP multiplies the weights by √width, preserving the variance of each layer’s output and thus keeping gradients well‑scaled.

First, run tiny proxy models to discover the learning rate that works for a given token budget; then extrapolate that rate to the full‑scale MoE by applying µP’s width‑transfer rule.

Why not simply scale the learning rate linearly with token count as some prior works suggest?

Linear scaling ignores µP’s variance‑preserving property; under µP the effective signal‑to‑noise ratio grows with √tokens, so a √ scaling matches the dynamics observed in the proxy runs, whereas linear scaling would overshoot and cause instability.

Run proxy models of fixed width on token budgets T₁, T₂, …, recording the learning rate L* that yields stable training.

Fit the relation L* ∝ √T using the collected (T, L*) pairs (e.g., linear regression on log‑log values).

Given the target token budget `T_target`, compute `L_target` = `L_base` × √(`T_target` / `T_base`) where `L_base` is the learning rate from the smallest proxy.

Train the full‑scale MoE with `L_target`, leveraging µP to keep activations stable across the expanded width.

Compute scaling factor √(T₂ / T₁) = √2 ≈ 1.414.

Predict learning rate for T₃ = 4 M tokens: L₃ = L₁ × √(T₃ / T₁) = 0.0010 × √4 = 0.0010 × 2 = 0.0020.

Validate on a tiny 4 M‑token proxy run: the model remains stable with L = 0.0020, confirming the √ law.

Apply L₃ = 0.0020 directly to the full‑scale MoE (155 B total parameters) for the target 10 trillion‑token training.

The example shows that a simple √ scaling, derived from two cheap proxy runs, accurately predicts the optimal learning rate for a target token horizon that is orders of magnitude larger.

**Figure 2.** **Proxy** denotes the 5 proxy-scale runs from Section 3.3.2 and **Target** denotes our large-scale training from Section 3.3.3. (a) 2D sweeps require scaling across model size, resulting in additional computation. Model scale search (1.5x, 2x wider than **Proxy**; details in Appendix B) incurs an additional 240.3 ZFLOPs of compute beyond the **Proxy** run cost of 64.8 ZFLOPs. (b) Total computation of the **Target** is approximately 98x larger than the sum of the **Proxy** runs. FLOPs are computed following Narayanan et al. (2021).

The two‑step framework lets practitioners predict optimal learning rates for massive MoE models using only a handful of cheap proxy runs.

Adapting µP for MoE

Adapting Maximal Update Parameterization to MoE layers enables width‑wise learning‑rate transfer.

Hyperparameter search for large‑scale MoE models is prohibitively expensive, so we need a principled way to transfer learning rates from a small proxy run.

We treat parameters that grow with model width differently: vector‑like parameters keep their initialization, while matrix‑like parameters receive both µP‑style initialization and a width‑dependent learning‑rate scaling.

Initialize $W$ with variance $1/d_{\text{base}}=1/2$ (base proxy $d_{\text{base}}=2$) and $b$ with variance $1$.

Scale the learning rate for $W$ by the fan‑in ratio $\frac{d_{\text{scaled}}}{d_{\text{base}}}= \frac{4}{2}=2$, leaving $b$’s learning rate unchanged.

During training, the update to $W$ is multiplied by $2$, while $b$ receives the original step size.

The resulting parameter updates preserve the spectral balance of the network despite the width increase.

Only matrix‑like weights need the fan‑in scaling; vector‑like weights remain stable, which prevents the learning‑rate drift that plagues Standard Parameterization.

**Table 1.** Parameter classification by shape invariance to model width. For MoE layers, router parameters and expert FC1 weights are classified as matrix-like parameters, while expert FC2 weights are treated as vector-like since their effective input dimensionality remains bounded by the fixed number of active experts and the fixed MoE intermediate dimension.

**Figure 3.** Comparison of learning rate transferability under Standard Parameterization (SP) and Maximal Update Parameterization ($\mu$P) across MLA MoE models of increasing width. Markers with black outlines indicate the learning rate achieving the lowest training loss (i.e., optimal learning rate) for each model. (a) Under SP, the optimal learning rate shifts as model width scales, failing to transfer from the base proxy model (0.6B total, 0.3B active) to larger variants. (b) Under $\mu$P, the optimal learning rate identified in the base proxy consistently transfers across models scaled to 2$\times$ (2.2B total, 0.7B active), 4$\times$ (8B total, 1.5B active), and 8$\times$ (30.7B total, 3.6B active) the base width.

Extrapolating Across Token Budgets

We extrapolate optimal learning rates from short-token proxy runs to long-token training using EMA and quadratic fits.

Searching for the optimal learning rate over trillions of tokens is infeasible even with tiny proxy models, yet the cost of such a search dominates MoE scaling.

We predict the learning rate that will work for a massive token budget by first fitting it on short‑budget proxies, then scaling the fitted value with a simple log‑log law.

Apply EMA with $\alpha=0.6$ after every 2 B tokens; the 10 B checkpoint for $B_1$ yields smoothed parameters $\theta_{1}^{\text{EMA}}$.

Compute validation loss at the three learning rates for $B_1$: $L_{1}(\eta_1)=0.45$, $L_{1}(\eta_2)=0.38$, $L_{1}(\eta_3)=0.42$.

Fit $L(\eta)=a(\log\eta)^2+b\log\eta+c$ to the three points; the vertex gives $\log\eta^{*}_1=-6.9$ → $\eta^{*}_1\approx1.0\!\times\!10^{-3}$.

Repeat for $B_2$ and $B_3$, obtaining $\eta^{*}_2\approx1.8\!\times\!10^{-3}$ and $\eta^{*}_3\approx2.5\!\times\!10^{-3}$.

Perform linear regression on $(\log B_i,\log\eta^{*}_i)$ pairs, yielding $\beta\approx0.5$, $\gamma\approx-7.2$.

Extrapolate to $B_{\text{target}}=10$ T tokens: $\log\eta^{*}_{\text{target}}=0.5\log(10^{13})-7.2$, giving $\eta^{*}_{\text{target}}\approx3.2\!\times\!10^{-4}$.

Even with only three short‑budget runs, the quadratic‑fit + log‑log regression pipeline predicts a plausible learning rate for a trillion‑token run, avoiding any expensive large‑scale sweep.

How does Token‑Scale Extrapolation differ from the common practice of scaling the learning rate linearly with the number of tokens?

Linear scaling assumes a proportional relationship between token count and optimal learning rate, which often fails because the loss‑vs‑learning‑rate curve is curved. Token‑Scale Extrapolation first finds the true optimum for each short budget via a quadratic fit, then learns a log‑log scaling law that respects the curvature, yielding a more accurate prediction for very large token budgets.

**Figure 4.** Solid lines (—) denote the base proxy model (5.6B total, 1.8B active), while dashed lines (--) represent the 2x width-scaled held-out model (20.7B total, 3.8B active), each trained on 40B, 60B, 80B, and 100B tokens. The × markers indicate the actual validation losses at a learning rate of $2 \times 10^{-3}$, which fall directly on the quadratic curves fitted using the remaining four learning rates.

Table 4 reports the ratio of predicted to actual optimal learning rates on held‑out token budgets near 500 B; all ratios lie between 1.033× and 1.058×, demonstrating that the extrapolation model generalizes well beyond the training range.

Experimental Setup and Optimizer

µP enables learning‑rate transfer up to 8× width scaling, validated across extensive LR sweeps.

µP transfers optimal learning rates across MoE models scaled up to 8× width.

Validated on proxy models scaled 2×, 4×, and 8× with consistent performance.

An adaptive optimizer that rescales updates to keep their magnitude stable, speeding up convergence.

How does the Muon optimizer differ from AdamW?

AdamW applies a fixed learning‑rate schedule with weight decay, whereas Muon dynamically rescales the learning rate based on observed gradient norms, keeping update magnitudes stable and reducing the number of training steps needed for convergence.

The Muon optimizer reaches target loss roughly 30% faster than AdamW on the same MoE configurations, without sacrificing final accuracy.

Predicting Optimal Learning Rates

We predict the optimal learning rate for 10‑trillion‑token MoE training.

Training large MoE models requires costly hyperparameter sweeps; we decouple model width from token budget to predict optimal learning rates from cheap proxy runs.

Extrapolation predicts an optimal learning rate of 3.85 × 10⁻⁴ for 10‑trillion‑token pretraining.

Linear regression on optimal rates from token scales ≥255 B yields R² = 0.95 (Figure 5b).

Across token scales, the optimal learning rate shows a modest downward trend, and the fitted parabolas for proxy and width‑scaled models share nearly identical curvature and vertex locations, confirming reliable transfer of the optimal rate across widths.

**Figure 5.** (a) Second-order polynomials fitted between the log-transformed learning rate and validation loss for the base proxy MoE model (10.8B total, 3.3B active) at representative token scales (a subset of all scales is shown for visual clarity). The $\star$ markers denote the optimal learning rates derived from the vertex of each parabola. (b) Linear regression in log-log space over these optimal learning rates across token scales ranging from 255B to 502B. The fitted model achieves $R^2 = 0.95$, enabling a reliable extrapolation of the optimal learning rate $3.85 \times 10^{-4}$ for the 10T-token pretraining.

Foundation Model Performance

Evaluates the 155 B Mixture-of-Experts (MoE) model’s training cost and benchmark performance.

The 155 B MoE model reaches 63.5% MMLU‑Pro accuracy at roughly 1000 ZFLOPs, placing it on the Pareto frontier.

Compared to dots.llm1 and GLM‑4.5‑Air, our model achieves higher accuracy with similar or lower estimated compute.

The Stage 1 pretraining data mixture starts at 45% English, 12.5% Math/STEM, 27.5% Code, and 15% Multilingual; at the 6 T‑token mark it shifts to 22.5% English, 27.5% Math/STEM, 25% Code, and 25% Multilingual to improve under‑represented domains. Training loss remains stable throughout, with no spikes, confirming the validity of the extrapolated optimal learning rate.

**Figure 6.** Training loss of our model (155B total, 17B active) during Stage 1 pretraining.

**Figure 7.** Benchmark performance of our foundation MoE model (155B total, 17B active) after Stage 1 pretraining on 10T tokens. Scores are presented across four categorized domains: English (MMLU, MMLU-Pro, BBH), Multilingual (Global-MMLU: Ko, Ja, Vi, Zh), Math (MATH, GSM8K), and Code (MBPP, HumanEval).

**Figure 8.** Estimated training compute vs. MMLU-Pro accuracy. Compute is estimated as $6ND$ with $N$ set to active parameters.

The 155 B MoE model delivers strong performance across English, multilingual, math, and code benchmarks while remaining compute‑efficient.

Related Work

We survey hyperparameter‑transfer and scaling‑law studies relevant to large MoE models.

Related work on hyperparameter transfer and scaling laws for MoE models is extensive. We summarize the most relevant studies and highlight their limitations.

Introduces a parameterization that keeps activation and gradient magnitudes stable when scaling model width, enabling zero‑shot hyperparameter transfer.

Builds on µP to transfer optimal learning rates across widths without additional sweeps, using a linear scaling rule for the learning rate.

Proposes a streamlined version of µP that removes auxiliary scaling terms, making implementation easier for large models.

Confirms that the simplified µP scheme works for both dense and MoE architectures, providing empirical evidence of its generality.

An optimizer that incorporates µP‑compatible scaling of both weight updates and logits, extending µP beyond AdamW.

Attempts to apply µP to the Switch Transformer architecture, but only under restrictive settings (fixed expert counts, hidden‑dimension scaling).

A Mixture‑of‑Experts model that routes tokens to a subset of experts, achieving high capacity with modest compute.

Weight‑decay‑corrected Adam variant that decouples L2 regularization from the adaptive learning‑rate update.

Analyzes how performance scales with model size and data for MoE architectures, identifying regimes where capacity gains dominate.

Extends prior scaling‑law work by incorporating token‑budget effects, revealing how training horizon influences optimal hyperparameters.

Proposes methods to select hyperparameters that respect a fixed computational budget, reducing the need for exhaustive sweeps.

Jointly scales model width and token budget to predict optimal hyperparameters, but requires multi‑dimensional sweeps on the proxy.

Extends Bjorck et al. by introducing a linear scaling law for token budgets, reducing the number of required proxy experiments.

A MoE variant that applies attention within latent expert subspaces, improving expressivity without increasing compute.

An optimizer designed for MoE training that incorporates µP‑compatible scaling of logits and gradients.

**Figure 9.** Comparison of learning rate transferability under Standard Parameterization (SP) and Maximal Update Parameterization ($\mu$P) across dense models of increasing width. Markers with black outlines indicate the learning rate achieving the lowest training loss (i.e., optimal learning rate) for each model. (a) Under SP, the optimal learning rate diverges across varying model widths. (b) Under $\mu$P, the optimal learning rate identified in the base proxy (0.24B) consistently transfers across models scaled to 2$\times$ (0.7B), 4$\times$ (2.27B), and 8$\times$ (8.02B) the base width, demonstrating robust zero-shot hyperparameter transfer in dense MLA models as well as MoE models.

Additional Analyses and Ablations

Recall the paper predicts optimal learning rates by decoupling width scaling from token‑budget scaling. Here we ablate key components to see how their removal affects performance.

Continuously updating the expert bias during Stage 2 yields the lowest expert‑utilization imbalance.

Figure 10 shows average $MaxVio$ near zero for the continuously‑updated setting, while frozen and zero‑initialized settings stay between 0.75 and 1.0.

Training loss trajectories are virtually identical across all bias configurations.

Figure 11 plots loss curves for continuous, frozen, and zero‑bias settings; the curves overlap with only minor fluctuations.

Domain‑dependent expert specialization grows in deeper layers while marginal load imbalance stays flat.

Figure 12 shows normalized mutual information rising from ≈0.02 to ≈0.15 and mean JSD from ≈0.04 to ≈0.14, whereas Aggregate $MaxVio$ hovers around 0.5 across layers.

**Figure 10.** MaxVio trends for three different expert routing bias settings during Stage 2 training. (a) Average MaxVio across MoE layers. (b) Maximum MaxVio among all MoE layers at each training step.

The image shows a line chart plotting Loss against Training Step for three different experimental configurations: "init:stg1, update:on" (red), "init:stg1, update:off" (green), and "init:0, update:off" (blue).

**Figure 12.** Aggregate MaxVio (left axis) measures marginal expert-load imbalance, while normalized mutual information (I/H) and mean pairwise Jensen–Shannon divergence (right axis) measure domain-dependent routing specialization. Specialization increases toward deeper MoE layers, whereas MaxVio follows a different trend, demonstrating that balanced marginal routing does not necessarily imply suppressed specialization.

**Figure 13.** Domain-conditioned expert routing divergence across MoE layers. Each cell shows $D_{KL}(p(e | d) || \bar{p}(e))$, where $\bar{p}(e)$ is the layer-wise marginal routing distribution defined in (ii). Brighter colors indicate stronger domain-specific routing.

Model Configurations

Lists all hyper‑parameters used for width‑scaling experiments across dense and MoE models.

**Table 3.** Model configurations for scaling width. $n_{layers}$, $d_{model}$, $d_{ff}$, and $n_{heads}$ represent the number of layers, hidden dimension, intermediate dimension of an MLP (active experts for MoE; $d_{ff} = k \times d_{expert}$), and attention heads, respectively. $n_{experts}$ is the total number of experts, and $\lambda_{routed}$ denotes the MoE routed scaling factor.

Read the original paper

Open the simplified reader on Paperglide

Browse all simplified papers