Glonce

New method predicts optimal learning rates for MoE models without costly sweeps

researchThe Glonce editorsHugging Face🔥80Glonce rating80Impact65Novelty70Relevance100Surprise45Players35Credibility80
New method predicts optimal learning rates for MoE models without costly sweeps

Mixture-of-Experts (MoE) architectures let a model grow much larger without a matching jump in the compute needed to train it, because only a subset of the parameters fire for each input. That efficiency has a catch: finding the right hyperparameters, especially the learning rate, normally means sweeping across many training runs, and at the model sizes and token budgets used for today's largest models, the authors say that sweep itself becomes computationally prohibitive.

A new paper proposes a two-step framework to get around this. The first step adapts Maximal Update Parameterization (μP), a technique for keeping training hyperparameters stable as a model's width scales up, to MoE architectures that use Multi-head Latent Attention (MLA) and the Muon optimizer. The authors show that with this adaptation, an optimal learning rate found on a narrow, cheap model transfers consistently to wider versions of the same model.

The second step extends that transfer along a different axis: the number of training tokens. The authors fit a scaling law by linear regression, using the optimal learning rates measured on small proxy models trained on limited token budgets, then extrapolate that curve out to far larger training horizons, for example 10 trillion tokens. The extrapolation holds with an R-squared of 0.95, meaning the predicted learning rate fits the actual optimal value with high fidelity even at that scale.

Together, the two steps mean that training small, cheap proxy models is enough to determine the learning rate for a large-scale MoE run, without ever sweeping at the target size. The authors test this directly: they pretrain a foundation model with 155 billion total parameters, 17 billion of which are active per token, from scratch, applying the proposed methodology. Training was stable, and the evaluation results validated that the optimal configuration for a full-scale target model can be accurately predicted, at minimal additional ablation cost.

Key facts

  • The framework works in two steps: first, a Maximal Update Parameterization (μP) adaptation for MoE architectures using Multi-head Latent Attention (MLA) and the Muon optimizer shows that optimal learning rates transfer consistently as model width scales up.
  • Second, a scaling law fit by linear regression on small proxy models extrapolates the optimal learning rate to far larger training horizons, for example 10 trillion tokens, with an R-squared of 0.95.
  • The authors conclude that proxy training on small models is sufficient to determine the optimal learning rate for the extensive training of large-scale MoEs, without a full-scale sweep.
  • They apply the method to pretrain a foundation model with 155 billion total parameters and 17 billion active parameters from scratch.
  • Training was stable, and the evaluation results validated that the optimal configuration for the full-scale model can be accurately predicted, at minimal additional ablation cost.

Why it matters

Training a large MoE model from scratch is expensive, and hyperparameter sweeping at the model sizes and token budgets used for today's largest models is, in the authors' words, computationally prohibitive. This work offers a cheaper alternative: tune small, inexpensive proxy models across widths and token budgets, then use a scaling law to extrapolate straight to the learning rate a much larger model needs, before committing the compute to train it. The authors' own conclusion is direct: proxy training on small models is sufficient to determine the optimal learning rate for the extensive training of large-scale MoEs, so the search that used to require full-scale runs can be pushed down to proxy scale instead.

Who it affects

Researchers and engineers who pretrain large Mixture-of-Experts models from scratch, particularly at the scale where a full hyperparameter sweep on the target-size model would be prohibitively expensive. The specific setup the authors validate combines Maximal Update Parameterization with Multi-head Latent Attention and the Muon optimizer, the same combination used in their own 155-billion-parameter, 17-billion-active-parameter foundation model.

How to use it

The recipe described has two parts. First, adapt Maximal Update Parameterization (μP) for an MoE architecture that uses Multi-head Latent Attention (MLA) and the Muon optimizer, and confirm that the optimal learning rate transfers as the model's width scales up. Second, train a handful of small proxy models on limited token budgets, record each one's optimal learning rate, fit a linear regression across those values, and extrapolate the resulting curve out to the token budget of the intended large-scale run. The text does not say whether the authors released code, model weights, or the 155-billion-parameter, 17-billion-active-parameter model itself, so applying the method means implementing it from the paper's own description.

How solid is it

The extrapolation itself is backed by a quantitative fit: an R-squared of 0.95 for the linear regression that projects the optimal learning rate out to massive training horizons. Beyond that curve fit, the authors also ran an end-to-end test at the scale the method targets: they used the methodology to pretrain an actual 155-billion-parameter, 17-billion-active-parameter foundation model from scratch, and report that training was stable and the evaluation results validated the predicted configuration.

Risks and caveats

The paper gives no numeric comparison to any prior hyperparameter-transfer or learning-rate-search method, only the general claim that sweeping remains computationally prohibitive, so the text alone does not show how much better this approach is. It also does not say which benchmarks or evaluation tasks the stable training and evaluation results cover, and it states no absolute compute savings, in GPU-hours or cost, from skipping a full sweep. The 10-trillion-token figure is an illustrative example in the scaling-law discussion, not a stated fact about how many tokens the 155-billion-parameter model was actually trained on. The text also gives no author names, institutional affiliation, submission date, model name, or release status for the model, code or weights.