Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Practical Efficiency of Muon for Pretraining

Essential AI, :, Ishaan Shah, Anthony M. Polloreno, Karl Stratos, Philip Monk, Adarsh Chaluvaraju, Andrew Hojel, Andrew Ma, Anil Thomas, Ashish Tanwer, Darsh J Shah, Khoi Nguyen, Kurt Smith, Michael Callahan, Michael Pust, Mohit Parmar, Peter Rushton, Platon Mazarakis, Ritvik Kapila, Saurabh Srivastava, Somanshu Singla, Tim Romanski, Yash Vanjani, Ashish Vaswani

TL;DR

The paper demonstrates that Muon, a second-order optimizer, expands AdamW's compute-time Pareto frontier, especially in data-efficient training with large batches. It introduces muP as a principled hyperparameter-transfer framework and a Telescoping Protocol to control finite-width and grid-sampling errors across model widths, enabling efficient tuning. The authors provide theoretical analyses of error sources and validate their approach on transformer models up to 3.7B parameters, achieving competitive final losses and showing robust width-based transfer of hyperparameters. Collectively, the work offers a practical methodology for scalable, economical pretraining with second-order optimization and width-parametric transfer.

Abstract

We demonstrate that Muon, the simplest instantiation of a second-order optimizer, explicitly expands the Pareto frontier over AdamW on the compute-time tradeoff. We find that Muon is more effective than AdamW in retaining data efficiency at large batch sizes, far beyond the so-called critical batch size, while remaining computationally efficient, thus enabling more economical training. We study the combination of Muon and the maximal update parameterization (muP) for efficient hyperparameter transfer and present a simple telescoping algorithm that accounts for all sources of error in muP while introducing only a modest overhead in resources. We validate our findings through extensive experiments with model sizes up to four billion parameters and ablations on the data distribution and architecture.

Practical Efficiency of Muon for Pretraining

TL;DR

The paper demonstrates that Muon, a second-order optimizer, expands AdamW's compute-time Pareto frontier, especially in data-efficient training with large batches. It introduces muP as a principled hyperparameter-transfer framework and a Telescoping Protocol to control finite-width and grid-sampling errors across model widths, enabling efficient tuning. The authors provide theoretical analyses of error sources and validate their approach on transformer models up to 3.7B parameters, achieving competitive final losses and showing robust width-based transfer of hyperparameters. Collectively, the work offers a practical methodology for scalable, economical pretraining with second-order optimization and width-parametric transfer.

Abstract

We demonstrate that Muon, the simplest instantiation of a second-order optimizer, explicitly expands the Pareto frontier over AdamW on the compute-time tradeoff. We find that Muon is more effective than AdamW in retaining data efficiency at large batch sizes, far beyond the so-called critical batch size, while remaining computationally efficient, thus enabling more economical training. We study the combination of Muon and the maximal update parameterization (muP) for efficient hyperparameter transfer and present a simple telescoping algorithm that accounts for all sources of error in muP while introducing only a modest overhead in resources. We validate our findings through extensive experiments with model sizes up to four billion parameters and ablations on the data distribution and architecture.
Paper Structure (16 sections, 16 equations, 5 figures, 1 algorithm)

This paper contains 16 sections, 16 equations, 5 figures, 1 algorithm.

Table of Contents

  1. Related Work
  2. Dethroning AdamW.
  3. Impact of the batch size.
  4. Practical aspects of muP.
  5. Choosing Hyperparameters for Muon
  6. The Maximal Update Parameterization (muP)
  7. Sources of Error
  8. The "Telescoping" Protocol
  9. Experiments
  10. Appendix
  11. Hyperparameter Shift Under $\mu$P
  12. Perturbative expansion of the minimizing hyperparameter.
  13. Result.
  14. $\mu P$
  15. Distribution of Losses
  16. ...and 1 more sections

Figures (5)

  • Figure 1: muP hyperparameter scaling rules, reproduced from yang2022tensor. Here, $\mathrm{fan\_in}$ and $\mathrm{fan\_out}$ denote the input and output dimensions, respectively, for a given weight matrix, e.g. for the weight matrix in the $\ell^{\text{th}}$ layer, $W_\ell \in \mathbb{R}^{m\times n}$$\mathrm{fan\_in}$ is $n$ and $\mathrm{fan\_out}$ is $m$.
  • Figure 2: Telescoping algorithm applied to weight decay and learning rate. The loss for each model is shifted to a small offset and then presented on a logarithmic scale to exaggerate the minimum for visualization. The red vertical lines highlight the optimal hyperparameters for comparison. The narrowest model (smallest $n$) has the most grid points (the topmost surface) and the coarsest mesh with cell size $\epsilon = 1/(m-1)$ (see Eq. \ref{['eq:epsilon']}), with $m$ being the number of sampled points for each hyperparameter. Submeshes at each subsequent doubling of width (lower surfaces) shrink in size, highlighting how the minima of each surface remain close across widths. Yellow regions indicate higher losses; color scales differ per stage for clarity. The shift in the red lines, marking the minima at each of the four layers, is proportional to $\alpha/n$ in Eq. \ref{['eq:opt-shift']}. The largest grid is $8\times 8$, followed by $5\times 5$, $3\times 3$ and $1\times 1$, approximately following our geometric schedule in Alg. \ref{['alg:tele']}. The third level, with cell size $\epsilon^{1/4}$ has extra points plotted to demonstrate that the true minimum is contained in the internal $3\times3$ subgrid. The final layer of the telescope, before training the final model, with cell size $\epsilon^{1/8}$ is $3\times 3$ rather than $1\times 1$ to demonstrate that the selected point is a true minimum (is locally flat).
  • Figure 3: $\mu$P hyperparameter scaling rules, reproduced from yang2022tensor. Here, $\mathrm{fan\_in}$ and $\mathrm{fan\_out}$ denote the input and output dimensions, respectively, for a given weight matrix.
  • Figure 4: The distribution of loss values at each level of the telescope in Fig. \ref{['fig:telescope']}, truncated at 2.5 to omit clearly bad data. We note that as the telescope levels increase, the distribution becomes tighter around its mean, demonstrating that we are producing a precise estimate, and the loss continues to decrease demonstrating that we maintain accuracy.
  • Figure 5: $N_c$ is the calibration level of the telescope - the deepest level we perform the telescoping up to. Two figures of merit: a) the percentage of compute saved using telescoping $\mu$P as opposed to performing a full grid search on the full model size, for $G$ grid points, and $N$ levels the base model width. Savings are largest for large grid sweeps, large models, and early truncation ratios. b) The percentage of the total compute spent on the final model run. Because we perform the same amount of compute at each stage of the telescope, this is maximized for small $N_c$, following the procedure suggested by yang2022tensor. This measures a form of exploration-exploitation trade off.