Before You Spend 10 Trillion Tokens: Separate Width From Horizon

TL;DR for operators

A training team preparing a multi-trillion-token MoE run usually cannot afford to test several full-scale learning rates. Kim et al. show a way to reduce that search before the expensive run begins: transfer the learning-rate optimum across model width, then estimate separately how that optimum moves as the token budget grows.¹

The strongest validation is not the final 10T-token run itself. A token-scaling model fitted only on roughly 255B–350B tokens predicted optimal learning rates at five unseen budgets from 462.7B to 502.3B tokens with about 4.4% average discrepancy from independently estimated optima.

The method was then used to select $3.85\times10^{-4}$ for a 155B-total, 17B-active MoE trained for 10T tokens. Training remained stable without reported loss spikes. That demonstrates practical viability under the studied setup; it does not establish that $3.85\times10^{-4}$ is the globally optimal full-scale rate, because no exhaustive 10T-token learning-rate sweep was performed.

The tuning decision arrives before the expensive evidence

A large pretraining run creates an awkward sequence of decisions. The team must choose its learning rate before it has evidence from the run it actually cares about. If the rate is poorly chosen, the cost accumulates over trillions of tokens. Testing several alternatives at target scale may cost enough to defeat the purpose of tuning.

The usual proxy strategy has two sources of uncertainty. A learning rate selected on a smaller model may stop being appropriate when width and expert count increase. Even if it transfers across model size, the best rate may shift as training continues for much longer.

The paper’s central move is to treat these as two different scaling problems rather than search them jointly.

μP removes model width from the repeated search

The first step addresses model size. Maximal Update Parameterization, or μP, is designed to keep feature-learning and update behavior comparable as width changes. In this paper’s implementation, matrix-like hidden parameters receive fan-in-dependent initialization and learning-rate scaling, while other parameter classes follow different scaling rules.

The comparison with Standard Parameterization is the relevant evidence. Under the standard setup, the learning rate minimizing training loss shifts as MoE width increases. Under μP, the optimum selected on a 0.6B-total, 0.3B-active proxy remains aligned as the models scale through 30.7B total and 3.6B active parameters.

This is more than ordinary dense-width transfer. Along the studied MoE path, hidden dimension and total expert count increase together, from 16 to 128 experts, while the number of active experts per token and expert intermediate dimension remain fixed. Dense MLA models from 0.24B to 8.02B show similar transfer behavior.

The operational implication is specific: when those μP assumptions hold, a team may not need to repeat a complete learning-rate sweep at every larger model size.

It does not establish that sparsity scaling itself is independently harmless. Width and total expert count change together in these experiments, so their individual contributions cannot be separated.

Training longer still changes the optimum

Removing width from the search does not make the learning rate constant over training horizon.

The authors test several learning rates at a fixed token budget and fit validation loss as a quadratic function of log learning rate:

The vertex gives the estimated optimum,

Repeating that procedure at different token budgets produces a sequence of estimated optimal rates. The paper then models their movement as

where $B$ is the training-token budget.

For the 10.8B-total, 3.3B-active proxy, the estimated optimum declines modestly as the horizon grows. Over stabilized checkpoints from roughly 255B to 502B tokens, the log-log relation reaches $R^2=0.95$.

That fit should be interpreted as evidence that the token dependence is regular enough to model over the observed proxy range. An $R^2$ value by itself says little about whether extrapolation remains accurate well beyond that range.

The held-out 500B test is the more consequential result

The authors therefore perform a retrospective out-of-range test. They fit the token-scaling relation using only approximately 255B–350B tokens, then predict optimal learning rates at five unseen horizons near 500B.

The average discrepancy is about 4.4%.

This experiment is best read as a robustness test of the extrapolation rule rather than a second contribution. It asks whether a relation estimated on shorter runs retains predictive value outside its fitting range. Within the tested extension—from the mid-300B range toward 500B—the answer is favorable.

For a training organization, this suggests a practical validation gate: do not jump directly from a fitted scaling law to a multi-trillion-token commitment. First test whether the rule predicts several withheld horizons that were not used to fit it.

The 10T run demonstrates viability, not optimality

Using the proxy fit, the authors extrapolate an optimal learning rate of $3.85\times10^{-4}$ at 10T tokens. They then train a 155B-total, 17B-active MoE from scratch for the full 10T-token horizon.

The reported training-loss trajectory is stable, with no reported loss spikes, and the model achieves competitive downstream results. The paper also reports it on the Pareto frontier of estimated training compute versus MMLU-Pro accuracy among the evaluated open-weight MoE base models.

That is substantial application evidence: a rate derived from much smaller experiments supported an actual extreme-scale training run.

But stability is a weaker claim than optimality. The authors do not train several 155B models for 10T tokens at neighboring learning rates. The experiment therefore shows that the extrapolated rate was workable and competitive under the chosen configuration, not that every alternative would have been worse.

The business value is a staged commitment process

Cognaptus’s inference is that the method is most valuable as a decision process for allocating training compute, not merely as another scaling-law result.

A team can structure the decision in four stages:

1. Establish width transfer. Verify that μP keeps the learning-rate optimum aligned across progressively larger versions of the intended architecture.
2. Measure horizon dependence. Run cheaper learning-rate sweeps on a proxy model at several token budgets and estimate each local optimum.
3. Withhold later horizons. Test whether the fitted token relation predicts learning rates beyond the range used to estimate it.
4. Commit to full scale only after those checks pass.

The paper’s compute accounting illustrates why this matters. It estimates 64.8 ZFLOPs for the proxy runs, while an illustrated model-scale search needed for a two-dimensional strategy would add 240.3 ZFLOPs. Those figures are specific to the study’s setup, but the organizational consequence is broader: eliminating a repeated model-size sweep can materially change the economics of pretraining experimentation.

Revalidate when the training regime changes

The strongest evidence is confined to MLA-based MoEs trained with Muon along the scaling path studied here. Several changes should reopen the transfer question rather than inherit the reported rule automatically.

A different optimizer may alter update scaling. A different MoE architecture or routing design may change how expert parameters behave. Changing width and sparsity independently would move outside the paper’s demonstrated scaling path. Batch-size transfer is explicitly excluded, and per-expert learning-rate adaptation is not studied even though top-$k$ routing produces heterogeneous effective batch sizes across experts.

The use of EMA-smoothed proxy checkpoints is also an implementation choice worth retaining in view: the authors use these checkpoints as a computational substitute for repeatedly completing the decay phase at every candidate horizon.

The correct operational conclusion is therefore narrower than “learning rates can be predicted from small models.” Under a validated parameterization and training regime, model width can become a transferable dimension, leaving token horizon as the main quantity to model explicitly.

For organizations facing a pretraining run too expensive to sweep directly, that is enough to change the tuning workflow.

Cognaptus: Automate the Present, Incubate the Future.