One-sentence Summary

Kakao Corp. and Upstage AI propose a compute-efficient hyperparameter transfer framework for large-scale MoE models, adapting µP with Multi-head Latent Attention and the Muon optimizer to transfer optimal learning rates across widths and then using a token scaling law $( R ^ { 2 } = 0 . 9 5 )$(R2=0.95) to extrapolate from small proxy models to trillion-token training, successfully pretraining a 155B-parameter model.

Key Contributions

• A Maximal Update Parameterization (µP) adaptation is formulated for Mixture-of-Experts architectures with Multi-head Latent Attention and the Muon optimizer, demonstrating consistent transfer of optimal learning rates across width-scaled models.
• A predictive scaling law extends this transferability along the token dimension, enabling extrapolation of the optimal learning rate to trillion-token training from small proxy models via linear regression with high fidelity (R²=0.95).
• The two-step framework is validated by pretraining a 155B-parameter MoE foundation model (17B active) from scratch, where stable loss and competitive benchmarks confirm that optimal hyperparameters can be accurately predicted with minimal ablation cost.

Introduction

Mixture-of-Experts (MoE) architectures let large language models scale capacity and specialization while keeping per-token active computation low, but tuning their learning rates becomes expensive because optimal values depend on model size, token budget, and expert routing. Existing hyperparameter transfer methods such as Maximal Update Parameterization mainly address dense width scaling, so their reliability for MoE sparsity and large expert counts remains uncertain. The authors propose a two-step, compute-efficient framework that transfers hyperparameters across model width and expert count while extrapolating optimal learning rates from small proxy runs to long token horizons. They validate this approach by pretraining a 155B-total-parameter MoE model with 17B active parameters over 10 trillion tokens, avoiding costly multidimensional sweeps.

Method

The authors propose a two-step hyperparameter transfer framework designed to efficiently determine the optimal learning rate for large-scale Mixture-of-Experts (MoE) models. The first step adapts Maximal Update Parameterization (µP) to MoE architectures, enabling zero-shot transfer of the optimal learning rate across model widths. The second step establishes a scaling law that predicts how the optimal learning rate evolves with the token budget, allowing extrapolation from short proxy runs to full-scale training horizons.

The µP adaptation for MoE begins by classifying model parameters according to their behavior under width scaling. Since none of the parameters have shapes tied to fixed constants such as vocabulary size, scalar-like parameters are excluded. Vector-like parameters receive µP initialization only, while matrix-like parameters are subject to both µP initialization and learning rate scaling. Following prior findings that µP properties can be preserved by scaling learning rates solely for linear layers, the authors apply learning rate scaling exclusively to matrix-like hidden weights. Depth scaling is avoided due to known instabilities; only width scaling is considered, with the number of layers kept fixed. For attention, the head dimension remains constant while the number of heads scales proportionally with the hidden dimension.

When scaling MoE models, the number of active experts per token and the MoE intermediate dimension are held fixed, while the total number of experts and the hidden dimension increase. This strategy raises sparsity as a key scaling axis alongside width, coupling total scale, active scale, and sparsity. The coupling remains compatible with µP because changing the active ratio by increasing the number of experts does not alter the fan-in or fan-out of individual experts beyond the effect of width scaling. Moreover, running the proxy search at lower sparsity and transferring to a higher-sparsity target keeps the search itself computationally efficient while still targeting the desired sparse regime.

The second step tackles the problem that even with µP, direct learning rate search over trillions of tokens is prohibitive. The solution is to identify the optimal learning rate within a constrained token budget and extrapolate to the full target horizon. The framework deliberately decouples learning rate optimization from batch size, treating batch size as a system-level variable fixed to maximize GPU throughput. This makes the learning rate scaling law robust to the chosen batch size.

For short-token estimation, the authors train small proxy models over limited token horizons using a Warmup-Stable-Decay (WSD) scheduler as the long-horizon setup. To avoid the computational cost and bias of introducing decay early, proxy runs are terminated during the stable phase without decay, and Exponential Moving Average (EMA) is applied to the model weights. EMA smooths parameter trajectories and mimics the effect of learning rate decay, which is particularly effective in the large-batch regime used here (global batch size of 32M tokens). The EMA parameters $\theta_{\mathrm{EMA}}^{(t)}$θEMA(t)​ are updated as

with a smoothing factor $\alpha = 0.6$α=0.6, ensuring that the most recent 20B tokens retain more than 1% influence. Checkpoints are saved via EMA at approximately 2B-token intervals, and those at every 10B tokens are used for analysis. For each token budget $B$B, the validation loss $\mathcal{L}$L is modeled as a second-order polynomial in log-transformed learning rate:

The vertex $\log \eta^* = -b / 2a$logη∗=−b/2a gives the estimated optimal learning rate $\eta^* = \exp(-b / 2a)$η∗=exp(−b/2a) for that budget.

To extrapolate to long token horizons, the authors fit a linear regression in log-log space between the estimated optimal learning rate $\eta^*$η∗ and the token budget $B$B:

This regression model allows predicting the optimal learning rate for target budgets (e.g., 10T tokens) without additional computation. When batch size scheduling is used, only data points collected after the training dynamics stabilize following a batch size increase are included in the regression, ensuring a consistent optimization regime.

Experiment

The evaluation uses µP, the Muon optimizer, and a WSD scheduler with batch size scheduling on NVIDIA H200 GPUs. Experiments first validate that µP allows optimal learning rates to transfer across MoE model widths, and that this transfer persists across token budgets, with the optimal learning rate decreasing slightly as more tokens are processed. By extrapolating a linear fit from small proxy runs, the optimal learning rate for a 155B-parameter MoE model trained on 10T tokens is predicted and applied, resulting in stable training loss and a model that lies on the Pareto frontier, outperforming comparable open-weight models at similar or lower compute.

Matrix-like parameters in dense and MoE models receive learning rate scaling based on the ratio of fan-in to a base proxy model, while vector-like parameters are only subject to µP initialization. For multi-head latent attention (MLA), the low-rank projection dimensions remain fixed during width scaling, so the learning rate scaling factor for the corresponding up-projection matrices becomes one and has no effect. When scaling MoE models, increasing the total number of experts while keeping active experts per token and intermediate dimension fixed couples sparsity with width, preserving compatibility with µP-based hyperparameter transfer. Learning rate scaling is applied only to matrix-like hidden weights, with factors proportional to the ratio of fan-in to the base proxy model's fan-in. In MLA, fixed query and key-value low-rank dimensions cause the learning rate scaling for up-projection matrices to reduce to 1, nullifying its effect. MoE scaling increases total experts and hidden dimension while holding active experts per token and MoE intermediate dimension constant, coupling sparsity and width under the same µP rules. Vector-like parameters (e.g., embeddings, biases, expert FC2 weights) receive µP initialization but no learning rate scaling during width expansion.

In maximal update parameterization for mixture-of-experts models, parameters are classified by how their shape scales with model width. Matrix-like parameters, such as router weights and expert FC1 weights, use initialization variance and learning rates that scale inversely with fan-in, while vector-like parameters, including embeddings, biases, and expert FC2 weights, use fixed initialization and no learning rate scaling. This classification enables stable hyperparameter transfer when scaling width and sparsity jointly. Matrix-like parameters (router weights, expert FC1) receive both µP initialization and learning rate scaling inversely proportional to fan-in. Vector-like parameters (embeddings, biases, expert FC2) use a fixed initialization variance and a learning rate scaling factor of 1, meaning no adjustment with model width.

The evaluation studies width and sparsity scaling under maximal update parameterization for dense and mixture-of-experts models, including multi-head latent attention. Matrix-like parameters such as router weights and expert FC1 weights receive initialization and learning rate scaling based on fan-in, while vector-like parameters such as embeddings, biases, and expert FC2 weights use fixed initialization and no learning rate scaling. In MLA, fixed low-rank projection dimensions make the up-projection learning rate factor equal to one, and MoE scaling couples sparsity with width by increasing total experts while keeping active experts and intermediate dimensions fixed.