NorMuon: Neuron-Adaptive Optimizer

• NorMuon is a neuron-wise normalized variant of the Muon optimizer that integrates matrix-level orthogonalization with adaptive scaling to mitigate the shortcomings of Adam and standard Muon.
• It delivers significant convergence gains in LLM training, reducing iterations by up to 15% while maintaining nearly the same state memory as AdamW.
• The optimizer supports distributed implementations and schedule-free training, making it a practical and efficient choice for large-scale deep learning applications.

NorMuon is a neuron-wise normalized variant of the Muon optimizer, designed to improve the efficiency and scalability of large-scale deep learning, particularly in LLM pretraining. NorMuon synergistically combines matrix-level momentum orthogonalization (as in Muon) with neuron-wise adaptive step-size normalization, mitigating the pathologies of both coordinate-wise adaptive optimizers (e.g., Adam) and matrix-orthogonal optimizers. Empirical benchmarks demonstrate that NorMuon delivers substantial convergence and efficiency gains over both Adam and plain Muon, with minimal increases in computational and memory overhead (Li et al., 7 Oct 2025, Grigorev, 25 Jan 2026). This innovative class of optimizers has also been adapted to new regimes, including schedule-free, anytime training (Apte et al., 21 May 2026).

1. Motivation: Limitations of Adam and Muon

Adam and AdamW apply element-wise (diagonal) normalization based on second-moment estimates, yielding exceptional stability but failing to account for correlations in 2D parameter matrices. This limitation results in pathological conditioning for whole layers or blocks. The Muon optimizer mitigates this by projecting the first-order momentum update onto (semi-)orthogonal matrices—improving the singular value spectrum and accelerating convergence. However, Muon’s orthogonalization can lead to substantial row-norm (i.e., neuron-norm) variability, causing certain neurons to dominate the update and leading to inefficient parameter utilization and slow convergence (Li et al., 7 Oct 2025, Grigorev, 25 Jan 2026).

2. NorMuon: Algorithmic Structure

NorMuon’s update consists of three principal stages: matrix-level orthogonalization, neuron-wise adaptive scaling, and learning-rate rescaling. The key steps for each 2D parameter $W_t \in \mathbb{R}^{m \times n}$ are:

1. First-order momentum: Maintain $M_t = \beta_1 M_{t-1} + (1-\beta_1) G_t$.
2. Matrix orthogonalization: Apply Newton–Schulz (NS) iteration to approximate the polar decomposition, $X_0 = M_t/\|M_t\|_F$, $$X_k = a X_{k-1} + b (X_{k-1} X_{k-1}^\top) X_{k-1} + c (X_{k-1} X_{k-1}^\top)^2 X_{k-1}, k=1...N$$. Set $O_t = X_N$.
3. Neuron-wise normalization: Maintain per-row second-moment $v_t \in \mathbb{R}^m$, $r_t^{(i)} = \frac{1}{n} \sum_{j=1}^n O_t(i,j)^2$, $v_t = \beta_2 v_{t-1} + (1-\beta_2) r_t$. Normalize: $$\widehat{O}_t = O_t \oslash \sqrt{V_t + \varepsilon}$$, $V_t$ is $M_t = \beta_1 M_{t-1} + (1-\beta_1) G_t$0 broadcast over columns.
4. Learning-rate scaling: $M_t = \beta_1 M_{t-1} + (1-\beta_1) G_t$1, $M_t = \beta_1 M_{t-1} + (1-\beta_1) G_t$2 (calibrates RMS step size).
5. Update: $M_t = \beta_1 M_{t-1} + (1-\beta_1) G_t$3.

For 1D parameters (biases, LayerNorm), standard AdamW is used, optionally with a smaller learning rate (Grigorev, 25 Jan 2026).

3. Empirical Efficiency and Scaling

Extensive experiments on models ranging from 124M to 5.4B parameters demonstrate that NorMuon consistently outperforms baseline optimizers:

• On 1.1B parameter LLM pretraining, NorMuon reduces the number of steps to reach Adam’s final loss by 21.74%, doubling the gain relative to plain Muon (≈10.4% step reduction).
• For 5.4B models, NorMuon yields 13.9% step reduction over Adam (vs. 6.1% for Muon).
• For sub-billon models, NorMuon consistently yields 6–15% fewer iterations than Muon (Li et al., 7 Oct 2025).
• In targeted ablations, switching from AdamW to NorMuon reduces perplexity by ≈8.5%, and adding cautious weight decay (“CWD”) results in an additional ≈3% perplexity improvement, cumulating to a ≈5.2% reduction in loss (Grigorev, 25 Jan 2026).

NorMuon achieves near-parity with AdamW in total optimizer state memory (m(n+1) vs. 2mn for AdamW) and adds only ≈2.9% per-step runtime overhead.

Optimizer	State Memory (per m×n)	Typical Speed Overhead
AdamW	2mn	Baseline
Muon	mn	+1–2%
NorMuon	m(n+1)	+2.9%

Row-wise normalization contributes negligible compute compared to the Newton–Schulz operation.

4. Distributed Implementation and FSDP2 Support

Scaling NorMuon to model-parallel LLM pretraining required a tailored distributed design. The Newton–Schulz orthogonalization requires access to full momentum matrices, but under FSDP2 (Full Sharded Data Parallel v2) both parameters and state are sharded by rows across GPUs. NorMuon addresses this by:

• Sorting all 2D tensors by size and round-robin assigning them to ranks.
• For each tensor, gathering shards to a single rank, performing NS-orthogonalization, and scattering back.
• Neuron-wise statistics and normalization are computed locally, requiring no extra communication.
• NorMuon increases per-param communication from 12 bytes (FSDP) to 16 bytes (two additional bf16 gather/scatter), a ≈33% increase.

This distributed orthogonalization keeps walltime costs within 2.9% of AdamW, while naive replication would increase step time by 8% or more (Li et al., 7 Oct 2025).

5. Variants: Cautious Weight Decay, µP Parametrization, and Schedule-Free NorMuon

Multiple variants and extensions of NorMuon have emerged:

• Cautious Weight Decay (CWD): Instead of always subtracting $M_t = \beta_1 M_{t-1} + (1-\beta_1) G_t$4, weight decay is only applied for weights whose sign matches the sign of the update. This modification further improves convergence and stability in LLM training (Grigorev, 25 Jan 2026).
• µP Parametrization: NorMuon is compatible with µP meta-parameter scaling, allowing fixed learning rates across model widths and straightforward transfer of hyperparameters.
• Schedule-Free NorMuon: In the schedule-free optimization regime, SF-NorMuon extends NorMuon by incorporating operator-norm steepest descent with schedule-free averaging and warmup. The complete algorithm supports anytime checkpointing, with a theoretical $M_t = \beta_1 M_{t-1} + (1-\beta_1) G_t$5 convergence guarantee, and requires no learning-rate schedule tuning across horizons of $M_t = \beta_1 M_{t-1} + (1-\beta_1) G_t$6 to $M_t = \beta_1 M_{t-1} + (1-\beta_1) G_t$7 Chinchilla (Apte et al., 21 May 2026). SF-NorMuon matches or exceeds horizon-tuned AdamW in both 125M and 772M LLMs.

6. Ablation Studies and Design Rationale

Ablation experiments highlight the criticality of certain design choices:

• Granularity: Matrix-level orthogonalization with neuron-wise normalization (NorMuon) outperforms coordinate-wise normalization (as in “Muon+Adam”) and uses much less memory.
• Normalization Position: Applying neuron-wise normalization after, not before, orthogonalization yields faster convergence.
• Selective vs. Universal: Normalizing only matrices with more rows than columns is suboptimal; normalization should be universal for all 2D tensors.
• Hyperparameter robustness: NorMuon tolerates significantly higher learning rates ($M_t = \beta_1 M_{t-1} + (1-\beta_1) G_t$810x) for 2D parameters compared to AdamW, as validated in IMU-1 and other recipes (Grigorev, 25 Jan 2026).

7. Implications and Future Directions

The NorMuon family illustrates that matrix-level conditioning (via orthogonalization) and block-level adaptivity (via neuron-wise normalization) address distinct and complementary optimization pathologies: spectral flattening and per-neuron magnitude balancing. By integrating both, NorMuon unlocks faster convergence and improved sample efficiency for LLMs, with minimal overhead.

Open research questions include:

• Extending orthogonalization and neuron-wise adaptivity to higher-order block structures (e.g., attention heads).
• More efficient, possibly low-rank, orthonormal approximation schemes that preserve both spectral conditioning and row-norm uniformity.
• Transferability to convolutional kernels and other structured parameterizations.
• Adaptive tuning of normalization statistics (e.g., $M_t = \beta_1 M_{t-1} + (1-\beta_1) G_t$9) at a per-layer or per-block level (Li et al., 7 Oct 2025).

NorMuon opens a new class of optimizers that harmonize global matrix geometry with local adaptability, broadening the optimizer design space for efficient large-scale deep learning (Li et al., 7 Oct 2025, Apte et al., 21 May 2026, Grigorev, 25 Jan 2026).