---
title: 'NorMuon: Neuron-Adaptive Optimizer'
url: https://www.emergentmind.com/topics/normuon
type: topic
---

# NorMuon: Neuron-Adaptive Optimizer

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 large language model (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 [2510.05491][2602.02522]. This innovative class of optimizers has also been adapted to new regimes, including schedule-free, anytime training [2605.23061].

## 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 [2510.05491][2602.02522].

## 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 $v_t$ broadcast over columns.
4. **Learning-rate scaling**: $\hat\eta_t = \eta \alpha \sqrt{mn}/\|\widehat{O}_t\|_F$, $\alpha=0.2$ (calibrates RMS step size).
5. **Update**: $W_t = W_{t-1} - \eta \lambda W_{t-1} - \hat\eta_t \widehat{O}_t$.

For 1D parameters (biases, LayerNorm), standard AdamW is used, optionally with a smaller learning rate [2602.02522].

## 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 [2510.05491].
- 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 [2602.02522].

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 [2510.05491].

## 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 $\lambda w$, 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 [2602.02522].
- **µ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 $O(T^{-1/4})$ convergence guarantee, and requires no learning-rate schedule tuning across horizons of $1$ to $8\times$ Chinchilla [2605.23061]. 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 ($\sim$10x) for 2D parameters compared to AdamW, as validated in IMU-1 and other recipes [2602.02522].

## 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., $\beta_2$) at a per-layer or per-block level [2510.05491].

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 [2510.05491][2605.23061][2602.02522].

Source: https://www.emergentmind.com/topics/normuon