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

cs.LGAug 20, 2026

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

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

Organizations: Kakao Corp. · Upstage AI

Abstract

Mixture-of-Experts (MoE) architectures significantly expand model capacity without a proportional increase in computational cost. However, optimizing their hyperparameters---particularly the learning rate---at extreme scales of both model size and token budget via sweeping remains computationally prohibitive. In this paper, we propose a compute-efficient, two-step hyperparameter transfer framework that estimates optimal learning rates for training large MoE models by transferring them across scaling model widths, and subsequently extrapolating to trillion-token horizons. First, we formulate a Maximal Update Parameterization ( $μ$ P) adaptation for MoE architectures utilizing Multi-head Latent Attention (MLA) and the Muon optimizer, demonstrating that optimal learning rates transfer consistently across width-scaled models. Second, we extend this transferability along the token dimension by establishing a predictive scaling law. By applying linear regression to the optimal values derived from small proxy models on limited budgets, we successfully extrapolate the ideal learning rate to massive training horizons (e.g., 10 trillion tokens) with high fidelity ( $R^2=0.95$ ). Consequently, this indicates that proxy training on small models is sufficient to determine the optimal learning rate for the extensive training of large-scale MoEs. We apply the proposed methodology to pretrain our foundation model (155B total, 17B active parameters) from scratch, and the stable training and evaluation results validate that optimal configurations for full-scale target models can be accurately predicted with minimal ablation costs.

Figures & tables

Panel of Figure 1 — 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 ( $\bigstar$ ), 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 ( $\bigstar$ ), replacing 2D sweeps with a lightweight 1D search along the token dimension.

Panel of Figure 2 — 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.5\times,2\times$ 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 $98\times$ larger than the sum of the Proxy runs. FLOPs are computed following Narayanan et al. (2021) .

Parameter Type	Definition	Examples
Vector-like	Parameters that contain exactly one infinitely expandable dimension.	Embeddings, biases, expert FC2 weights
Matrix-like	Parameters that contain two infinitely expandable dimensions.	FFN, attention, router, expert FC1 weights

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.

	I/O Embeddings & All Biases	Hidden Weights	Hidden Weights
	(Vector-like)	(Vector-like)	(Matrix-like)
Init. Var	0.04	$\mathrm{fan\_in_{\text{base}}/\mathrm{fan\_in}}$	$\mathrm{fan\_in_{\text{base}}/\mathrm{fan\_in}}$
LR Scaling Factor	1	1	$\mathrm{fan\_in_{\text{base}}/\mathrm{fan\_in}}$

Table 2:

\mu

P and learning rate scaling factors for dense and MoE models.

\mathrm{fan\_in}

and

\mathrm{fan\_in_{\text{base}}}

are the fan-in values of the width-scaled model and the base proxy model, respectively. Learning rate scaling factors are multiplied by the learning rate when updating the corresponding parameters. Notably, for MLA, the low-rank projection dimensions for query and key-value are kept fixed during width scaling. Since these serve as the

\mathrm{fan\_in}

of the corresponding up-projection matrices, the learning rate scaling applied to these matrices has no effect, as the scaling factor reduces to 1.

Panel of Figure 3 — 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.

Figure 4: Solid lines (

-

) denote the base proxy model (5.6B total, 1.8B active), while dashed lines (

--

) represent the

2\times

width-scaled held-out model (20.7B total, 3.8B active), each trained on 40B, 60B, 80B, and 100B tokens. The

\times

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.

Panel of Figure 5 — 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 $\bigstar$ 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.

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.

Appendix figures & tables7 assets

Supplementary material from the paper’s appendix.

Appendix

Model Scale	Total Params.	Active Params.	$n_{\text{layers}}$	$d_{\text{model}}$	$d_{\text{ff}}$	$n_{\text{heads}}$	$n_{\text{experts}}$	$\lambda_{\text{routed}}$
Models for Section 3.2 and Appendix C
Dense Base proxy	0.24B	0.24B	48	384	1536	4	-	-
Dense $2\times$ scale	0.7B	0.7B	48	768	3072	8	-	-
Dense $4\times$ scale	2.27B	2.27B	48	1536	6144	16	-	-
Dense $8\times$ scale	8.02B	8.02B	48	3072	12288	32	-	-
MoE Base proxy	0.6B	0.3B	48	256	6144	4	16	2.428

Appendix

Table 3: Model configurations for scaling width.

n_{\text{layers}}

d_{\text{model}}

d_{\text{ff}}

, and

n_{\text{heads}}

represent the number of layers, hidden dimension, intermediate dimension of an MLP (active experts for MoE;

d_{\text{ff}}=k\times d_{\text{expert}}

), and attention heads, respectively.

n_{\text{experts}}

is the total number of experts, and

\lambda_{\text{routed}}

denotes the MoE routed scaling factor.

Panel of Figure 9 — 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.

Token Budget	Predicted Optimal LR	Actual Optimal LR	Ratio
462.7B	$9.57\times 10^{-4}$	$9.26\times 10^{-4}$	$1.033\times$
472.6B	$9.53\times 10^{-4}$	$9.22\times 10^{-4}$	$1.034\times$
482.5B	$9.49\times 10^{-4}$	$8.97\times 10^{-4}$	$1.058\times$
492.4B	$9.45\times 10^{-4}$	$8.99\times 10^{-4}$	$1.052\times$
502.3B	$9.41\times 10^{-4}$	$9.01\times 10^{-4}$	$1.045\times$

Appendix

Table 4: Held-out validation of learning rate extrapolation. The extrapolation model is fitted using only token budgets in 255B–350B and evaluated on unseen budgets near 500B. Ratios refer to Predicted Optimal LR divided by Actual Optimal LR.

Panel of Figure 10 — 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.

Figure 11: All three expert routing bias settings exhibit nearly identical training loss trends.

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_{\mathrm{KL}}\!\left(p(e\mid d)\,\|\,\bar{p}(e)\right)$ , where $\bar{p}(e)$ is the layer-wise marginal routing distribution defined in (ii). Brighter colors indicate stronger domain-specific routing.

Hyperparameter Transfer

Hyperparameter transfer studies how to reuse training settings across model sizes, architectures, tasks, and training horizons.

See all 41 papers

Hyperparameter Scaling Laws Across MoE Sparsity
Sep 8, 2026Changxin Tian, Kunlong Chen, Jia Liu +3Language Model Scaling Laws Sparse Mixture-of-Experts
Ling Team, Ant Group
How to Scale Mixture-of-Experts: From muP to the Maximally Scale-Stable Parameterization
May 13, 2026Leena Chennuru Vankadara, Moritz Haas, Luke Hayward +2Mixture-of-Experts Models Dynamical Mean-Field Theory
1Gatsby Computational Neuroscience Unit, University College London · 2Amazon · University of Tübingen, Tübingen AI Center
Holistic Scaling Laws for Optimal Mixture-of-Experts Architecture Optimization
Mar 23, 2026Weilin Wan, Jingtao Han, Debing Zhang +2Mixture-of-Experts Models AI Scaling Laws
College of Computer Science and Artificial Intelligence, Fudan University · Hi Lab, Xiaohongshu Inc. · Project Numina +1

Sep 8, 2026J/K move · Enter open · S save

Hyperparameter Scaling Laws Across MoE Sparsity

Changxin Tian, Kunlong Chen, Jia Liu, Ziqi Liu+2

Ling Team, Ant Group

Language Model Scaling Laws Sparse Mixture-of-Experts Hyperparameter Transfer Scaling Laws Mixture-of-Experts Models

Open paper AskPDF arXiv

Mixture-of-Experts (MoE) models expand model capacity without a proportional increase in training compute, but increasing sparsity makes reliable hyperparameter transfer challenging. In this work, we show that conventional hyperparameter scaling laws are insufficient for ultra-sparse MoEs: the optimal learning rate and batch size vary with activation ratio, and these shifts cannot be explained by either total or activated parameter count alone. To characterize this dependence, we conduct 1,800 pre-training runs spanning six activated-parameter scales and models with up to 6B total non-embedding parameters, processing approximately 20 trillion tokens at a cost of 200,000 equivalent H800 GPU-hours. Our results reconcile conflicting findings in prior work by revealing two scaling regimes. At fixed sparsity, the optimal batch size follows a power-law relationship with training tokens

D

, whereas the optimal learning rate scales with training compute

C

and remains robust to the allocation between model size and data. Across sparsity levels, the activation ratio

A

enters both relationships as an additional multiplicative power-law factor. These observations lead to unified hyperparameter scaling laws that transfer across MoE sparsity levels. Large-scale evaluation shows that the scaling form outperforms alternative functional forms. On a held-out ultra-sparse MoE with 12B total parameters and only 1/64 of its experts activated, the predicted hyperparameters remain close to the observed optima, supporting joint extrapolation across model scale and sparsity. Further experiments demonstrate transfer across expert granularities and isolate the effect of activation ratio from that of total expert count.

Abstract

Hyperparameter Transfer

Explore similar work

Hyperparameter Scaling Laws Across MoE Sparsity

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

Holistic Scaling Laws for Optimal Mixture-of-Experts Architecture Optimization