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

Source-linked AI summary

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

arXiv:2608.20061v1Aug 20, 2026cs.LGcs.AIcs.CL

TL;DR

Large-scale MoE hyperparameter sweeps are computationally prohibitive, especially when both model width and token horizon vary. The paper transfers learning rates across MoE model scales with µP and extrapolates them across token budgets using a predictive scaling law from proxy runs. The approach predicts a 3.85 × 10^-4 learning rate for 10T-token training and supports stable pretraining of a 155B-total, 17B-active-parameter model, while its demonstrated scope is limited to MLA-based MoEs with the Muon optimizer.

Problem
Hyperparameter search for MoE models exceeding 100B parameters is computationally prohibitive, and reliable transfer across width and sparsity scaling remains unclear.
Method
The framework combines MoE-specific µP transfer across width and sparsity with a token-budget scaling law fitted from small proxy runs.
Results
3.85 × 10^-4 was predicted as the optimal learning rate for 10T-token training, with stable loss during foundation-model pretraining.
Takeaways & Limitations
Small proxy experiments can reliably predict learning rates for trillion-token large-scale MoE pretraining without exhaustive sweeps.
Takeaways & Limitations
The study focuses on MoE architectures using MLA and the Muon optimizer, leaving broader architectural and optimizer generalization for future work.

Abstract

from arXiv · show

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.

1 Introduction

MoE models expand capacity efficiently, but extreme-scale hyperparameter sweeps remain prohibitively expensive. The paper proposes a two-step framework that transfers learning rates across model widths and extrapolates them across token budgets using small proxy runs.

Motivation: MoE architectures increase effective capacity without proportional computational cost, supporting stronger performance without additional training cost.They also enable model specialization and cost-efficient inference.
Motivation: Existing hyperparameter-transfer methods were designed for dense models, while MoE scaling adds sparsity as an additional axis and makes width-only expansion increasingly impractical.Increasing total experts while fixing active experts offers a more practical capacity-scaling path, but its transfer behavior remains unclear.
Approach: The proposed framework decouples model and token scales, using µP for model-scale transfer and a token scaling law for direct extrapolation from proxy runs.This replaces exhaustive 2D sweeps with a lightweight search along the token dimension.
Approach: The study adapts µP to MoE architectures while examining zero-shot transfer across model width, total expert count, and sparsity expansion with the Muon optimizer.This extends the empirical scope of µP-based MoE scaling beyond the previously described dense-model setting.
Approach: The framework extrapolates optimal learning rates to unseen long token horizons using cost-effective small-scale proxy runs, reducing hyperparameter-tuning overhead.The method targets extreme scales spanning both massive model widths and long training horizons.
Validation: The authors validate the framework by pretraining a 155B-total, 17B-active-parameter MoE foundation model from scratch over 10 trillion tokens.This large-scale validation is presented as evidence of practical effectiveness at scale.

2 Methods

The method transfers learning rates across width- and sparsity-scaled MoE models using an MoE-specific µP formulation, then extrapolates optimal rates to long token horizons with log-log regression. It reduces large-scale tuning to short proxy runs while accounting for EMA-based evaluation and stabilized training regimes.

µP adaptation for MoE: µP adapts hyperparameter transfer to MoE width and sparsity scaling by classifying parameters according to their shape behavior.Router and expert FC1 weights are matrix-like, whereas expert FC2 weights are vector-like because active experts and the MoE intermediate dimension remain fixed.
µP adaptation for MoE: Learning-rate scaling is applied only to matrix-like hidden weights while depth remains fixed because depth scaling is unstable for µP transfer.For MLA, fixed low-rank query and key-value projection dimensions make the corresponding up-projection scaling factors equal to 1.
Token-scale extrapolation: Direct searches for optimal learning rates remain computationally prohibitive at trillion-token horizons, motivating extrapolation from limited-budget proxy experiments.The framework uses small models trained over short token horizons across a range of learning rates, while the target setup uses Warmup-Stable-Decay.
Evaluation and fitting details: EMA checkpoints use α = 0.6 and are analyzed every 10B tokens, while batch-scheduled regressions retain only points after optimization dynamics stabilize.The EMA configuration gives the most recent 20B tokens more than 1% influence on merged weights.
Width transfer: Under µP, the optimal learning rate transfers across 2×, 4×, and 8× width-scaled MLA MoE models, unlike Standard Parameterization.The optimal rate is defined as the rate yielding the lowest training loss; the tested models range from a 0.6B-total proxy to 30.7B total parameters.
Token-scale extrapolation: The optimal learning rate for each token budget is estimated from the vertex of a fitted quadratic, then extrapolated with a log-log linear regression.The regression relates log(η*) to log(B), with coefficients β and γ, and can predict rates for horizons such as 10T tokens without additional computational cost.

3 Experiments and Analysis

Experiments validate µP-based learning-rate transfer across width-scaled MoE models and use token-scale regression to predict learning rates for 10T-token training. The predicted configuration supports stable pretraining of a 155B-total, 17B-active-parameter foundation model.

3.2 µP and Learning Rate Transfer with MoE Models: Under µP, width-scaled MoE models share optimal learning rates, unlike models using Standard Parameterization.The optimal learning rate is defined as the rate yielding the lowest training loss.
3.3.1 Learning Rate Dynamics Across Token Horizons: At each token scale, proxy and width-scaled models show consistent quadratic curvature and vertex locations, indicating reliable width transfer.The tested proxy has 5.6B total and 1.8B active parameters; the held-out model has 20.7B total and 3.8B active parameters.
3.3.2 Predicting Optimal Learning Rate for Extremely Long Token Horizon: R2 = 0.95 supports predicting an optimal learning rate of 3.85 × 10^-4 for 10T-token training from proxy runs.The regression uses optimal learning rates estimated across token scales after 255B tokens.
3.3.3 Applying to Our Foundation Model: The 155B-total, 17B-active-parameter model maintains highly stable Stage 1 training loss without spikes under the extrapolated learning rate.Stage 1 pretraining uses a 10T-token horizon and evaluates English, multilingual, math, and code domains.
3.3.3 Applying to Our Foundation Model: The foundation model lies on the Pareto frontier, achieving higher MMLU-Pro accuracy than dots.llm1 and GLM-4.5-Air at comparable or lower estimated compute.Compute is estimated with 6ND, using active parameters for N and reported tokens for D.

4 Related Works

Prior work enables hyperparameter transfer for dense models and studies scaling laws for MoE architectures, but reliable transfer for large, fine-grained-expert MoEs remains open.

Hyperparameter Transfer: µP and µ-Transfer enable zero-shot transfer of hyperparameters across model widths by modifying parameterization and tensor-specific scaling.These methods were originally developed for dense models.
Hyperparameter Transfer: A prior Switch Transformer study was limited to AdamW, fixed total and active expert counts, and hidden-dimension scaling, restricting generalization to fine-grained MoEs.The paper identifies reliable transfer for MoEs with many experts as an open problem.
Scaling Laws: Existing MoE scaling research largely studies optimal architectural configurations as parameter count increases, while token horizons also affect optimization dynamics.The related-work discussion motivates studying learning-rate scaling across both model and token dimensions.

5 Conclusion

The paper proposes compute-efficient learning-rate transfer for large-scale MoE pretraining and validates it on a 155B-total, 17B-active-parameter model. Stable loss and competitive benchmarks support the extrapolated configuration, while broader architectural and optimizer coverage remains future work.

Conclusion: The two-step framework combines µP-based width transfer with a linear scaling law across token budgets to avoid exhaustive sweeps.Small proxy experiments are used to estimate learning rates for trillion-token MoE pretraining.
Conclusion: The framework is validated by pretraining a 155B-total, 17B-active-parameter foundation model from scratch.The target training horizon is 10T tokens.
Conclusion: Exhaustive full-scale sweeps cannot definitively verify the predicted learning rate, so stable loss and benchmark results provide empirical rather than exhaustive validation.The predicted learning rate is 3.85 × 10^-4.
Conclusion: The study focuses on MoEs using Multi-head Latent Attention and the Muon optimizer, leaving diverse architectures and optimizers for future work.Per-expert learning-rate adaptation is also left for future research because of routing, data-mixture, and implementation complexities.

A Training Details

Training details specify optimizer normalization, fixed regularization coefficients, and batch-size schedules used across the experiments.

Optimizer and Regularization: Muon’s update RMS is matched to AdamW by scaling the learning rate by 0.2 · p_max(A, B).A and B are the fan-in and fan-out dimensions of the full-rank matrix parameter.
Optimizer and Regularization: Weight decay is fixed at 1 × 10^-1 and auxiliary z-loss at 5 × 10^-6 throughout the experiments.These coefficients remain constant across the reported experiments.
Batch Scheduling: Experiments use 0.5M-token batches in Section 3.2 and 8M-token batches in Section 3.3.1.Section 3.3.2 increases the batch size from 8M to 32M tokens after 200B tokens.

B Model Configuration Details

The configurations use MLA with controlled depth and fixed MoE expert settings during width scaling, while Figure 9 contrasts learning-rate transfer under SP and µP.

All model configurations employ Multi-head Latent Attention, with model depth held constant across configurations.
During MoE width scaling, the expert intermediate dimension and number of active experts remain fixed, preserving dff = k × dexpert.
Under SP, the optimal learning rate diverges across increasing dense-model widths, whereas µP transfers the 0.24B proxy optimum consistently.

C Learning Rate Transfer with Dense Models

The dense-model baseline validates learning-rate transfer protocols for MLA and Muon before MoE extension, using controlled width-scaled models trained under matched depth and token budgets.

The study first validates hyperparameter transfer on dense models because MLA and Muon have limited prior evaluation under µP learning-rate scaling.
The base dense proxy has dmodel 384, dff 1536, and 4 attention heads, then scales to widths of 768, 1536, and 3072.
The width-scaled dense models contain 0.24B, 0.7B, 2.27B, and 8.02B parameters, with 48 layers and 1.3B training tokens.

D Evaluation Details

Evaluation after Stage 1 pretraining uses a customized framework spanning English, multilingual knowledge, mathematics, and code benchmarks with specified few-shot settings.

The evaluation framework is applied after Stage 1 pretraining on 10T tokens.
English General Knowledge: English general knowledge is evaluated with MMLU, MMLU-Pro, and BBH using five, five, and three shots respectively.
Multilingual General Knowledge: Global-MMLU evaluates multilingual general knowledge across Korean, Japanese, Vietnamese, and Chinese with five shots.
Math: Mathematical reasoning is evaluated with MATH at four shots and GSM8K at eight shots.
Code: Code evaluation uses MBPP and HumanEval with Pass@1, using three shots for MBPP and zero shots for HumanEval.

E Small-Scale Validation of Learning Rate Extrapolation

The token-scale extrapolation is validated retrospectively by fitting on earlier budgets and testing on held-out budgets near 500B tokens. Predicted optimal learning rates closely match independently estimated optima, with an average discrepancy of approximately 4.4%.

Direct validation at the 10T-token target is computationally infeasible because a five-learning-rate sweep would consume roughly half of full-model training compute.
The extrapolation model fits the first 11 budgets from approximately 255B–350B tokens and predicts optimal learning rates for five unseen budgets near 500B.
≈4.4% average discrepancy separates predicted and independently estimated optimal learning rates on the held-out budgets.

F Additional Analysis of Expert Routing Bias in Staged Pretraining

Stage 2 expert-bias choices affect routing balance after the data distribution shifts, but the three settings produce nearly identical training-loss trends.

Stage 2 follows general-domain pretraining with smaller, higher-quality, domain-augmented data, which can destabilize MoE expert routing.
The study compares continuously updated bias, frozen Stage 1 bias, and bias terms removed through zero initialization.
Continuously updating expert bias yields the most balanced expert utilization during Stage 2.
With frozen bias updates, zero initialization balances utilization better than inheriting Stage 1 values, suggesting Stage 1-tuned biases may mismatch Stage 2 data.
The three expert-routing bias settings exhibit nearly identical training-loss trends.

G Layer-wise Routing Balance and Expert Specialization

The analysis separates marginal expert-load balance from domain-dependent specialization, showing that routing can remain balanced while specialization increases with depth and varies by domain.

Metrics: The evaluation reports Aggregate MaxVio, normalized mutual information, and mean pairwise Jensen–Shannon divergence to measure imbalance and specialization.
Metrics: Lower Aggregate MaxVio indicates more balanced routing across experts.
Metrics: Normalized mutual information uses a uniform domain prior, preventing larger domains from dominating the statistic.
Metrics: Mean pairwise JSD averages divergence across domain pairs; zero means identical routing distributions, while higher values indicate more distinct patterns.
Layer-wise findings: Across depth, Aggregate MaxVio stays mostly stable while normalized mutual information and mean pairwise JSD rise toward the deepest MoE layers.
Layer-wise findings: The results show that balanced marginal routing can coexist with domain-dependent expert specialization.
Layer-wise findings: Code routing diverges from the marginal distribution early, whereas multilingual routing remains close until specialization increases sharply in final MoE layers.