Search papers, labs, and topics across Lattice.
This paper introduces a compute-efficient two-step hyperparameter transfer framework designed to optimize learning rates for large Mixture-of-Experts (MoE) models. By leveraging a Maximal Update Parameterization (渭P) and Multi-head Latent Attention (MLA) along with the Muon optimizer, the authors demonstrate that optimal learning rates can be effectively transferred across models of varying widths and extrapolated to extensive training scenarios of up to 10 trillion tokens. The results show a high fidelity in predicting ideal learning rates, achieving an R虏 of 0.95, thus confirming that small proxy models can guide the training of significantly larger MoE architectures with minimal computational overhead.
Optimal learning rates for training massive Mixture-of-Experts models can be accurately predicted from small proxy models, achieving high fidelity even at trillion-token scales.
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.
