MuonMax: Robust Neural Network Optimization

• MuonMax is a gradient-based optimization algorithm that uses non-Euclidean steepest descent over product spaces to update neural network parameters.
• It employs spectral norms for weight matrices and adaptive norms for other parameters, aggregating these via a hybrid ℓ₂-product norm to enhance robustness.
• Empirical results show that MuonMax, especially when combined with model-based momentum (Momo), tolerates a wide range of hyperparameter values while matching or outperforming baselines.

MuonMax is a gradient-based optimization algorithm for training neural networks, developed within the general framework of non-Euclidean steepest descent on product spaces. It is a variant of the Muon optimizer that replaces constrained (normalized) updates with regularized updates under a hybrid ℓ₂-product norm, resulting in improved robustness to learning rate selection. MuonMax is particularly distinguished by its use of spectral and adaptive norms on different parameter blocks and aggregation via an ℓ₂ combination, as well as scaling update magnitudes using the sum of nuclear norms across all matrix parameter blocks. Empirical evaluation demonstrates that MuonMax, especially when combined with model-based momentum (Momo), offers significantly increased tolerance to hyperparameter choices, reducing the need for extensive tuning, and matches or slightly outperforms baseline optimizers on autoregressive language modeling tasks (Crawshaw et al., 10 Oct 2025).

1. Non-Euclidean Steepest Descent Framework

Optimization of neural networks in the MuonMax paradigm is formulated as steepest descent over a partitioned parameter space. Let $\Theta = (W^1, ..., W^L, \theta)$ denote the full set of network parameters, where each $W^\ell$ is a weight matrix and $\theta$ encompasses all other parameters (biases, embeddings, layer-norm gains, etc.). Each block is assigned a norm tailored to its structure: spectral norm $\|W^\ell\|_{2 \rightarrow 2}$ for matrices and an adaptive vector norm (such as Ada$_2$ or Ada$_\infty$) for $\theta$, with the adaptive norms utilizing Adam first- and second-moment estimates.

To aggregate block-wise norms into a unified norm on parameter space, a product norm is constructed:

$$h(x^1,...,x^L,y) = f(g_1(x^1), ..., g_L(x^L), g_{L+1}(y)),$$

where $g_i$ are the block norms and $f$ is typically an $W^\ell$0 norm, with possible choices such as $W^\ell$1, $W^\ell$2, or a hybrid combination for special purposes.

Steepest-descent updates are implemented via either:

• Constrained SD: normalizes each update to have a fixed norm.
• Regularized SD: scales the step by the dual norm $W^\ell$3, adapting to gradient history magnitude.

The direction of steepest descent under the norm $W^\ell$4 is given by computing $W^\ell$5.

2. Origin and Distinction from Muon

The Muon optimizer performs constrained steepest descent (SD) using a layer-wise spectral norm for matrices and an Ada$W^\ell$6 adaptive norm for non-matrices, aggregating with $W^\ell$7 across all blocks. Muon updates adopt distinct step sizes ($W^\ell$8 for $W^\ell$9, $\theta$0 for $\theta$1) and are given by:

• $\theta$2,
• $\theta$3, where $\theta$4 is a momentum average and $\theta$5, $\theta$6 are Adam statistics.

MuonMax alters two key components:

• It employs regularized SD, so update magnitude is proportional to the dual norm rather than forcibly normalized.
• It aggregates block-wise spectral/Ada$\theta$7 norms using an $\theta$8 combination (hybrid norm), promoting robustness and automatic step-size scaling.

3. MuonMax Algorithmic Definition

MuonMax defines its norm as

$\theta$9

with $\|W^\ell\|_{2 \rightarrow 2}$0.

The dual norm (for scaling the update direction) is

$\|W^\ell\|_{2 \rightarrow 2}$1

where $\|W^\ell\|_{2 \rightarrow 2}$2 denotes the nuclear norm.

The update equations are:

• $\|W^\ell\|_{2 \rightarrow 2}$3,
• $\|W^\ell\|_{2 \rightarrow 2}$4.

Polar denotes the spectral-norm maximizer, i.e., the leading singular vector direction.

4. Practical Implementation and Pseudocode

A simplified pseudocode for MuonMax (excluding stale approximation and Momo) consists of the following steps:

• Initialize $\|W^\ell\|_{2 \rightarrow 2}$5, $\|W^\ell\|_{2 \rightarrow 2}$6, and $\|W^\ell\|_{2 \rightarrow 2}$7 for all blocks.
• At each iteration:
  • Compute stochastic gradients for all $\|W^\ell\|_{2 \rightarrow 2}$8 and $\|W^\ell\|_{2 \rightarrow 2}$9.
  • Update momentum (matrix blocks) and Adam moments (non-matrix block).
  • Compute $_2$0 and the Ada$_2$1-normalized non-matrix gradient.
  • Aggregate into $_2$2.
  • Update $_2$3 and $_2$4 as above.

Practical recommendations include default hyperparameters: $_2$5, $_2$6, $_2$7, and for Momo-enhanced variants $_2$8 (FineWeb1B) or $_2$9 (SlimPajama). The coefficient ratio $_\infty$0 can usually be set to $_\infty$1 with little performance penalty, and MuonMax is robust to $_\infty$2 in $_\infty$3.

5. Empirical Evaluation

Experiments were conducted in the context of autoregressive language modeling, focusing on two benchmarks:

• FineWeb1B (124M-parameter GPT2-Small, 1B tokens)
• SlimPajama (774M-parameter GPT2-Large, up to 6B tokens)

Baselines included MuonAdam (Muon + Adam), Scion (Muon + SignSGD), MuonAdam+Momo, and MuonMax+Momo. Key findings:

• MuonMax+Momo matches or slightly outperforms MuonAdam in terms of final validation loss on both tasks.
• Robustness: While MuonAdam and Scion required learning rate tuning within a narrow window ($_\infty$4 order of magnitude) to avoid divergence or high loss, MuonMax+Momo attained near-optimal loss over $_\infty$5–$_\infty$6 orders of magnitude variation in either $_\infty$7 or $_\infty$8.
• Ablation studies showed that using stale nuclear-norm sums (rather than freshly computed) changes final loss by $_\infty$9 while saving $\theta$0 wall-clock time per step. The choice of $\theta$1 in Momo was not critical if learning rates were reasonable.

6. Theoretical Motivation and Advantages

The principal theoretical argument for MuonMax’s robustness centers on adaptive magnitude adjustment via the dual norm $\theta$2: the step size automatically increases when gradient magnitudes are small (e.g., early in training), and decreases when gradients are large, providing self-adjustment lacking in constrained SD. By contrast, Muon requires precise tuning to prevent steps from being too small or excessively large.

The hybrid aggregation norm in MuonMax—max spectral norm across weight matrices and Ada$\theta$3 for other parameters, combined via $\theta$4—allows for a more balanced, global assessment of the geometry of the parameter space. This approach reduces sensitivity to the relative scaling of different parameter blocks.

7. Practical Guidelines and Recommendations

For practitioners, the following summarizes best practices for MuonMax:

• Use spectral norm on weight matrices and Ada$\theta$5 on non-matrix parameters.
• Aggregate with $\theta$6 (MuonMax) rather than $\theta$7 (Muon).
• Set default learning rates to $\theta$8 and optimize within $\theta$9.
• For new tasks where hyperparameters are unknown, MuonMax (especially with Momo) minimizes the need for expensive hyperparameter sweeps.

In summary, MuonMax generalizes the Muon optimizer by adopting regularized SD with a hybrid ℓ₂ product norm, providing improved learning-rate robustness and reduced sensitivity to the layer-base learning rate ratio, with minimal loss performance compromise (Crawshaw et al., 10 Oct 2025).