NorMuon: Making Muon more efficient and scalable — part 1/3

Cleaned-text code points [0, 15963). Bibliography omitted here; source.md includes it.

# NorMuon: Making Muon more efficient and scalable

###### Abstract

The choice of optimizer significantly impacts the training efficiency and computational costs of large language models (LLMs). Recently, the Muon optimizer has demonstrated promising results by orthogonalizing parameter updates, improving optimization geometry through better conditioning. Despite Muon’s emergence as a candidate successor to Adam, the potential for jointly leveraging their strengths—has not been systematically explored. In this work, we bridge this gap by proposing NorMuon (Neuron-wise Normalized Muon), an optimizer that synergistically combines orthogonalization with neuron-level adaptive learning rates. Our analysis reveals that while Muon effectively reduces condition numbers, the resulting updates exhibit highly non-uniform neuron norms, causing certain neurons to dominate the optimization process. NorMuon addresses this imbalance by maintaining second-order momentum statistics for each neuron and applying row-wise normalization after orthogonalization, ensuring balanced parameter utilization while preserving Muon’s conditioning benefits. To enable practical deployment at scale, we develop an efficient distributed implementation under the FSDP2 framework that strategically distributes orthogonalization computations across devices. Experiments across multiple model scales demonstrate that NorMuon consistently outperforms both Adam and Muon, achieving 21.74% better training efficiency than Adam and 11.31% improvement over Muon on 1.1B pretraining setting, while maintaining a comparable memory footprint to Muon. Our findings suggest that orthogonalization and adaptive learning rates are complementary rather than competing approaches, opening new avenues for optimizer design in large-scale deep learning. We open-source our implementation at https://github.com/zichongli5/NorMuon.git .

## 1 Introduction

Training efficiency remains a central challenge in scaling large language models (LLMs) Sachdeva et al. (2024) ; Wan et al. (2023) , where optimizer choice directly impacts convergence speed, computational requirements, and ultimately, the feasibility of training at scale Jordan et al. (2024b) ; Wen et al. (2025) . The community standard, Adam ( Kingma and Ba, 2015 ) , achieves robust performance through coordinate-wise preconditioning: dynamically adjusting learning rates for each parameter based on the second moment of its gradient history. While this per-coordinate adaptivity is computationally efficient and generally stable, it suffers from a fundamental limitation—it treats each parameter independently, ignoring the rich geometric structure and cross-coordinate dependencies inherent in neural network layers.

Recent advances have sought to address this limitation through various approaches to capturing cross-coordinate structure. Adam-mini ( Zhang et al., 2025 ) exploits the near-block-diagonal Hessian structure of neural networks by applying adaptive learning rates to parameter blocks (e.g. each neuron) rather than individual coordinates. More ambitious second-order methods, such as Shampoo ( Gupta et al., 2018 ) and SOAP ( Vyas et al., 2024 ) , employ full-matrix preconditioning through singular value decomposition to capture curvature information and parameter interdependencies. However, these approaches incur substantial memory and communication overhead, while introducing hyperparameter sensitivity that limits their practical adoption at scale.

Recently, Muon ( Jordan et al., 2024b ) has emerged as a compelling middle ground, applying truncated Newton-Schulz iterations to approximate the orthogonal polar factor of momentum matrices. This approach yields matrix-wise orthogonalized updates that improve conditioning while maintaining modest computational overhead and approximately half the memory consumption of Adam, demonstrating promising results in LLM training.

These optimizers fundamentally differ in their preconditioning granularity and objectives. Adam and Adam-mini, when considered without exponential moving averages (EMAs), apply L 2 L^{2} normalization at the per-coordinate and per-neuron levels respectively, adjusting learning rates while preserving update signs. In contrast, the idealized version of Shampoo and Muon operate at the per-matrix level, actively orthogonalizing parameter updates.

The varied preconditioning strategies employed by these optimizers raise an important question: Are different forms of preconditioning inherently conflicting, or can they be combined in a way that yields complementary benefits?

To investigate this, we analyzed key properties of the update matrices from different optimizers during pretraining of a 1.1B-parameter Transformer model, examining both singular value distributions and per-neuron norms. As illustrated in Figure 1(a) , raw momentum accumulates updates with extremely high condition numbers, indicating that certain directions dominate while leaving other parameters underutilized. AdamW produces moderately more balanced singular values, though the improvement remains limited. In contrast, Muon’s approximate orthogonalization successfully addresses this conditioning issue, yielding well-balanced singular values across the spectrum. However, examining per-neuron update norms (Figure 1(b) ) reveals a complementary perspective. AdamW demonstrates superior performance in reducing variance across per-neuron update norms compared to SGD momentum. Conversely, while Muon’s orthogonalization effectively improves matrix-level conditioning, the per-neuron update norms exhibit high variance, with some neurons receiving disproportionately large updates relative to others.

This observation motivates our key insight: while Muon’s orthogonalization effectively reduces the condition number of updates, the remaining high variance in neuron norms still creates an imbalanced learning dynamic, potentially leading to inefficient parameter usage. Drawing inspiration from Adam-mini’s success Zhang et al. (2025) with per-neuron adaptive learning rates, we propose to incorporate second-order momentum to normalize these disparate scales and ensure more balanced parameter updates. Our method, NorMuon , augments Muon’s orthogonalization with neuron-wise adaptive learning rates computed from accumulated second-order statistics. As demonstrated in our analysis, NorMuon yields updates with both low condition numbers (Figure 1(a) ) and uniform neuron norms (Figure 1(b) ), thereby combining the advantages of Muon and AdamW and achieving more balanced utilization of the network’s representational capacity.

Beyond algorithmic innovation, the distributed implementation of orthogonalization-based optimizers remains relatively underexplored in the literature. To enable training at larger scales, we develop a distributed version of NorMuon compatible with the FSDP2 framework ( Feng et al., 2022 ) . While previous work on distributed Muon ( Liu et al., 2025a ) was implemented using ZeRO-1 with Megatron-LM ( Rajbhandari et al., 2020 ; Shoeybi et al., 2019 ) , FSDP2 can offer greater flexibility and memory efficiency. However, direct adaptation of the previous distributed approach to FSDP2 would result in extensive replicated computation, as FSDP2 shards nearly all parameters across devices. Our implementation addresses this challenge by distributing orthogonalization computation across devices, eliminating redundant calculations while maintaining load balance. Furthermore, we leverage FSDP2’s row-wise parameter sharding to enable efficient neuron-wise normalization without incurring additional communication overhead.

In summary, our contributions are threefold:

∙ \bullet We propose NorMuon, a simple and effective optimizer that combines Muon’s orthogonalization with neuron-wise adaptive learning rates. NorMuon maintains uniform neuron norms to ensure balanced parameter utilization while preserving the low condition number achieved by Muon’s orthogonalization.

∙ \bullet We develop an efficient distributed implementation under the FSDP2 framework. By carefully orchestrating sharded optimizer states, we gather updated momentum and distribute Muon orthogonalization computation uniformly across GPUs, achieving optimal memory efficiency with manageable communication and computational overhead.

∙ \bullet Through extensive experiments across multiple scales of LLM pretraining, we demonstrate that orthogonalization and blockwise adaptive learning rates are complementary rather than conflicting, with their combination yielding superior training dynamics compared to either approach in isolation.

## 2 Related Works and Background

### 2.1 Related Works

Adaptive Gradient Methods . The introduction of per-parameter adaptive learning rates has been instrumental in training deep networks. Optimizers such as AdaGrad ( Duchi et al., 2011 ) , RMSProp ( Hinton, 2012 ) , Adam ( Kingma and Ba, 2015 ) and AdamW ( Loshchilov and Hutter, 2017 ) use first- and second-moment estimates to adjust each weight’s step size individually. This coordinate-wise preconditioning improves stability and convergence in heterogeneous settings, and has become the de facto standard for LLM training. However, treating each weight independently ignores the underlying structure of neural network layers and incurs high memory overhead by storing two extra tensors per parameter. This memory cost motivated techniques like AdaFactor ( Shazeer and Stern, 2018 ) , which factorizes the second-moment accumulator across rows and columns to reduce memory. Similarly, Adam-mini ( Zhang et al., 2025 ) partitions parameters into blocks (e.g. each neuron’s weights) and assigns a single learning rate to each block, matching AdamW’s performance on different model sizes, while halving memory cost. GaLore ( Zhao et al., 2024 ) maintains momentum in a low-rank subspace derived from the SVD of gradients, although its effectiveness diminishes for long sequence lengths ( Liu et al., 2025b ) . Lion ( Chen et al., 2023 ) applies a coordinate-wise signed update, abandoning second-order moment estimates to achieve memory savings.

Second-order Methods . In parallel, other optimizers capture the rich geometry of the loss surface by coupling updates across parameters. K-FAC ( Martens and Grosse, 2015 ) and its variants ( Martens et al., 2018 ; Gao et al., 2021 ) approximate curvature information beyond individual coordinates, capturing correlations across parameters. Shampoo ( Gupta et al., 2018 ) and its distributed variant ( Shi et al., 2023 ) employ Kronecker-factored preconditioners and have demonstrated strong performance in practice ( Dahl et al., 2023 ) . More recently, SOAP ( Vyas et al., 2024 ) establishes a connection between Shampoo and Adafactor and further improves convergence performance. Despite these advances, Shampoo and SOAP incur substantial memory cost and computational overhead, which hinders their applicability at LLM scale.

Orthogonal Update Methods : Muon (Momentum Orthogonalized by Newton-Schulz, Jordan et al. (2024b) ) represents a breakthrough that leverages matrix geometry without the full cost of second-order methods. Muon performs an approximate polar decomposition (via Newton–Schulz iterations) on the momentum to extract its orthogonal component and also eliminates the need to store second-order momentum. Muon thus simultaneously improves both convergence and memory efficiency compared to Adam, demonstrating great potential for scaling up model pretraining ( Liu et al., 2025a ; Shah et al., 2025 ) . We provide a detailed description of the Muon algorithm in Section 2.2 . More recently, Dion ( Ahn et al., 2025 ) extends the orthogonal update paradigm to be more communication- and compute-efficient in distributed settings. Dion applies a low-rank orthonormalization scheme via amortized power iteration instead of full Newton–Schulz, and decouples momentum buffers across devices to avoid full gradient synchronization.

### 2.2 Background: Muon optimizer

Muon ( Jordan et al., 2024b ) is an optimizer designed for the 2D weight matrices in neural network hidden layers. The key innovation lies in orthogonalizing the momentum before applying parameter updates, thereby improving the conditioning of the optimization trajectory. Formally, at iteration t t , given weight matrix W t − 1 W_{t-1} , learning rate η t \eta_{t} , and loss function L L , Muon maintains a first-order momentum M t M_{t} and computes updates as: M t \displaystyle M_{t} = μ ​ M t − 1 + ∇ L ​ ( W t − 1 ) , \displaystyle=\mu\,M_{t-1}+\nabla L(W_{t-1}), (1) O t \displaystyle O_{t} = NS5 ⁡ ( M t ) , \displaystyle=\mathrm{NS5}(M_{t}), (2) W t \displaystyle W_{t} = W t − 1 − η t ​ O t , \displaystyle=W_{t-1}-\eta_{t}\,O_{t}, (3) where M 0 = 𝟎 M_{0}=\mathbf{0} and μ \mu is the momentum coefficient. The critical component is the orthogonalization operator NS5 ⁡ ( ⋅ ) \mathrm{NS5}(\cdot) , which aims to approximate the orthogonal projection of the momentum matrix: Ortho ​ ( M ) = arg ⁡ min O ​ { ‖ O − M ‖ F : O ⊤ ​ O = I ​ or ​ O ​ O ⊤ = I } . \displaystyle\text{Ortho}(M)=\arg\min_{O}\{\|O-M\|_{F}:O^{\top}O=I\text{ or }OO^{\top}=I\}. (4) Muon approximates this orthogonalization through a fixed number of Newton-Schulz iterations. Starting with the Frobenius-normalized momentum X 0 = M t / ‖ M t ‖ F X_{0}=M_{t}/\|M_{t}\|_{F} , the algorithm performs N N iterations (typically N = 5 N=5 ): X k = a X k − 1 + b ( X k − 1 X k − 1 ⊤ ) X k − 1 + c ( X k − 1 X k − 1 ⊤ ) 2 X k − 1 , k = 1 , … , N , \displaystyle 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},\quad k=1,\ldots,N, (5) with the final orthogonalized update O t = X N O_{t}=X_{N} . The coefficients ( a , b , c ) (a,b,c) are carefully chosen such that singular values of the update matrix converge toward unity. In practice, Muon is typically applied only to 2D weight matrices in hidden layers, while scalar parameters, bias vectors, embeddings, and unembedding layers continue to use standard optimizers such as Adam.

## 3 Method

In this section, we introduce NorMuon, where we aims to combine Muon’s orthogonalization with block-wise adaptive learning rates based on the observation that the approximated orthogonalized updates can experience a high variance on update directions norm of each neuron.

### 3.1 NorMuon

We present our update rule in Algorithm 1 . The algorithm maintains two momentum states: the standard first-order momentum 𝐌 t ∈ R m × n \mathbf{M}_{t}\in\mathbb{R}^{m\times n} used by Muon (line 5), and an averaged second-order momentum 𝐯 t ∈ R m \mathbf{v}_{t}\in\mathbb{R}^{m} that tracks the squared magnitude of each neuron’s update direction (lines 7). Importantly, 𝐯 t \mathbf{v}_{t} requires minimal additional memory overhead, storing only m m scalars compared to the m × n m\times n first-order momentum.

At each iteration and given the gradient, we first follow the Muon’s update rule that update the first-order momentum and apply Newton-Schulz iteration for orthogonalization (line 4-6), producing 𝐎 t \mathbf{O}_{t} with improved conditioning. Rather than directly using this orthogonalized update, we compute row-wise statistics to capture the per-neuron update magnitudes. Specifically, we calculate the mean squared value across columns for each row of 𝐎 t \mathbf{O}_{t} (line 7). This statistic is accumulated into our averaged second-order momentum 𝐯 t \mathbf{v}_{t} using exponential moving average with decay rate β 2 \beta_{2} . We then apply 𝐯 t \mathbf{v}_{t} for row-wise normalization (line 9). This second-order momentum is similar to Adam-mini’s ( Zhang et al., 2025 ) block-wise reduced-dimensional statistics, where we treat each neuron (i.e. each row) as a block. As illustrated in Figure 1, this normalization reduces the variance in update magnitudes across neurons while preserving the favorable conditioning properties.