PaperlayerView exact revision on arXiv

Source-linked AI summary

NorMuon: Making Muon more efficient and scalable

Zichong Li, Liming Liu, Chen Liang, Weizhu Chen, Tuo Zhao

arXiv:2510.05491v1cs.LGcs.CL

TL;DR

The paper addresses the limited exploration of combining Muon’s orthogonalization with adaptive learning rates. It proposes NorMuon, which adds neuron-wise normalization to Muon, and reports improved training outcomes while preserving the central benefits of orthogonalization.

  • Problem

    The potential for jointly leveraging Muon’s orthogonalization and adaptive learning rates has not been systematically explored, despite optimizer choice affecting LLM training efficiency.

  • Method

    NorMuon combines Muon’s orthogonalization with neuron-wise adaptive learning rates and applies normalization after orthogonalization.

  • Results

    The experiments report notable improvement over Muon while NorMuon maintains uniform neuron norms and Muon’s low condition number.

  • Takeaways & Limitations

    Orthogonalization and adaptive scaling need not conflict, supporting their combination in optimizer design for large-scale deep learning.

  • Takeaways & Limitations

    Coordinate-wise adaptive methods ignore the geometric structure and cross-coordinate dependencies inherent in neural network layers.

Abstract

from arXiv · show

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.1 B 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.

1 Introduction

NorMuon combines Muon’s matrix orthogonalization with neuron-wise adaptive scaling to address complementary optimization problems. The paper also presents an FSDP2 implementation designed to distribute orthogonalization efficiently for larger-scale training.

  • Motivation: Adam-style methods use adaptive learning rates but ignore cross-coordinate structure, while full-matrix methods incur substantial memory and communication overhead.These trade-offs motivate combining geometric conditioning with coarser-grained adaptivity.
  • Research question: The paper asks whether orthogonalization and adaptive preconditioning can be combined rather than treated as conflicting strategies.This question is investigated through singular-value distributions and per-neuron norms during 1.1B-parameter Transformer pretraining.
  • Optimization geometry: Muon improves update conditioning through approximate orthogonalization, but its per-neuron update norms remain highly variable.Some neurons receive disproportionately large updates despite Muon’s improved matrix-level conditioning.
  • NorMuon: NorMuon adds neuron-wise adaptive learning rates based on second-order momentum statistics to Muon’s orthogonalized updates.The resulting updates target both low condition numbers and more uniform neuron norms.
  • NorMuon: NorMuon combines Muon’s conditioning with balanced neuron contributions, supporting more even utilization of representational capacity.The paper characterizes orthogonalization and blockwise adaptive learning rates as complementary.
  • Distributed implementation: The distributed implementation uses FSDP2 sharding to distribute orthogonalization across devices and perform neuron-wise normalization without additional communication overhead.This addresses replicated computation arising because FSDP2 shards nearly all parameters across devices.

2 Related Works and Background

Prior optimizers trade off adaptive learning, geometric conditioning, memory, and distributed efficiency. Muon addresses some of these costs through momentum orthogonalization, but is primarily designed for matrix-valued parameters.

  • Adaptive Gradient Methods: Adam-family optimizers use first- and second-moment estimates for coordinate-wise step sizes, but storing two extra tensors per parameter creates high memory overhead.This motivates reduced-state methods such as AdaFactor.
  • Second-order Methods: Shampoo and SOAP capture parameter correlations through structured preconditioners, but their memory and computational costs hinder applicability at LLM scale.
  • Orthogonal Update Methods: Muon applies Newton–Schulz iterations to momentum, extracting an approximately orthogonal component without storing second-order momentum.Its design targets 2D weight matrices in neural-network hidden layers.
  • Orthogonal Update Methods: Muon improves convergence and memory efficiency compared to Adam, supporting its potential for scaling model pretraining.
  • Orthogonal Update Methods: Dion extends orthogonal updates for distributed settings by using low-rank amortized power iteration and decoupled momentum buffers across devices.
  • Orthogonal Update Methods: Muon is typically applied only to 2D hidden-layer weights, while scalar parameters, biases, embeddings, and unembedding layers use standard optimizers.

3 Method

NorMuon combines Muon’s orthogonalization with neuron-wise adaptive normalization, then distributes orthogonalization under FSDP2 to reduce redundant computation. It retains Muon-like memory efficiency while adding modest communication overhead.

  • NorMuon: NorMuon combines Muon’s orthogonalization with block-wise adaptive learning rates motivated by high variance in orthogonalized update norms across neurons.
  • NorMuon: NorMuon tracks each neuron’s squared update magnitude with a second-order momentum and applies row-wise normalization after orthogonalization.The statistic treats each neuron, or row, as a block.
  • Memory Overhead: The additional neuron-wise statistics require only m scalars for an m × n weight matrix, making NorMuon’s memory overhead over Muon negligible.NorMuon uses m(n + 1) optimizer-state units versus mn for Muon and 2mn for Adam.
  • NorMuon: Row-wise normalization reduces variance in update magnitudes across neurons while preserving favorable conditioning, although it increases the resulting direction’s norm.
  • NorMuon: NorMuon’s normalization remains useful for m ≤ n matrices in practice because orthogonalization is approximated rather than exact.
  • Distributed NorMuon: FSDP2 distributes orthogonalization by assigning parameter tensors across devices, gathering each full momentum matrix only on its assigned device.This avoids replicated orthogonalization while maintaining load balance.
  • Distributed NorMuon: The assigned device scatters the orthogonalized result back, while row normalization operates locally on complete row-wise shards without additional communication.

4 Experiments

Experiments evaluate NorMuon against AdamW, Muon, and Dion across model scales, ablations, and efficiency analyses. NorMuon generally improves training efficiency while retaining manageable computational and memory overhead.

  • Experimental setup: NorMuon is evaluated on 124M, 350M, 1.1B, and 5.4B parameter models using separate large- and small-model experimental settings.The experiments include extensive ablations and efficiency analyses.
  • Main results: NorMuon consistently and substantially improves over AdamW, Muon, and Dion across the reported model-scale comparisons.Muon and Dion already outperform AdamW, while NorMuon amplifies the advantage through neuron-wise adaptive learning rates.
  • Main results: 21.74% and 13.91% are NorMuon’s efficiency gains over Adam for the 1.1B and 5.4B models, respectively.The metric is the percentage reduction in training steps required to reach Adam’s final validation loss.
  • Ablation studies: Selective normalization only for m > n matrices underperforms full NorMuon, indicating that normalization for m ≤ n matrices is helpful.The m > n subset corresponds only to MLP up-projection matrices and yields only a marginal improvement over Muon.
  • Efficiency analysis: NorMuon adds only a 2.9% increase in training step time compared with AdamW and has nearly 50% lower memory use than AdamW or Muon + Adam.Its memory efficiency is comparable to Muon, and neuron-wise norm computation adds minimal cost relative to orthogonalization.

5 Conclusion

NorMuon combines Muon’s orthogonalization with neuron-wise adaptive learning rates and uses an efficient FSDP2 implementation for large-scale training. Experiments show notable improvement over Muon, supporting the complementarity of orthogonalization and adaptive scaling.

  • Conclusion: NorMuon integrates Muon’s orthogonalization with neuron-wise adaptive learning rates.The optimizer is presented as a simple yet effective combination of these components.
  • Conclusion: An FSDP2 implementation coordinates momentum gathering and orthogonalization to reduce redundant computation and communication overhead.The implementation targets practical large-scale training.
  • Conclusion: Experiments report notable improvement over Muon when orthogonalization and adaptive scaling are combined.The conclusion characterizes the two approaches as complementary rather than mutually exclusive.

A.1 Adam-mini’s results on optimization geometry analysis

The supplied passage identifies Figure 6 as an analysis of optimization geometry during 1.1B-parameter pretraining for different optimizers.

  • A.1 Adam-mini’s results on optimization geometry analysis: Figure 6 analyzes optimization geometry during 1.1B model pretraining across different optimizers.The supplied passage provides the figure’s subject but no specific geometric outcome.

A.2 Ablation Experiments on Modded-NanoGPT

Ablations test normalization granularity, placement, and matrix-shape selectivity in Modded-NanoGPT. Full NorMuon performs better than simplified normalization variants across the reported comparisons.

  • A.2 Ablation Experiments on Modded-NanoGPT: The ablations compare standard NorMuon, Muon, direct Muon normalization, and a selective NorMuon variant on 350M Modded-NanoGPT.The setting uses FineWeb and labels the variants as NorMuon, Muon, Muon w/ normalization, and NorMuon ablation.
  • A.2 Ablation Experiments on Modded-NanoGPT: Muon with normalization is initially slightly better than Muon but is eventually surpassed, so NorMuon’s effectiveness is not attributed to normalization alone.The comparison concerns the training trajectory.
  • A.2 Ablation Experiments on Modded-NanoGPT: Applying NorMuon only to m > n matrices greatly diminishes its effect, producing only a marginal improvement over Muon.Matrices with m > n correspond only to MLP up-projection matrices, a small portion of the model.

B Implementation Details

Experiments used different hardware configurations by model scale: larger models ran across two nodes, while Modded-NanoGPT models ran on one node. Figure 7 reports ablation studies for the 350M model.

  • 1.1B and 5.4B models were trained on 2 nodes, each equipped with 8 A100 GPUs with 80GB memory.The GPUs were connected via NVLink for optimized inter-GPU communication.
  • Figure 7 presents ablation studies on the 350M model.
  • Training lasted approximately 2 days for the 1.1B model and 7 days for the 5.4B model.
  • 124M and 350M Modded-NanoGPT experiments were performed on a single node with 8 A100 GPUs.

C Experiment Setup of Modded-NanoGPT

The experiments follow Muon’s setup using GPT-2-style models trained on FineWeb, with specified optimizer configurations and a WSD learning-rate schedule. The 124M and 350M models differ in scale, token budget, and selected hyperparameters.

  • Experimental setup: The experiments strictly follow Muon’s experimental setup, with details specified for the 350M model.
  • Models: The model architecture follows GPT-2, with 124M and 350M parameter configurations created by adjusting width and depth.
  • Dataset: All models were trained on FineWeb, using approximately 3.2B tokens for 124M and 4B tokens for 350M.
  • Hyperparameters: Training used batch size 512, sequence length 1024, the WSD schedule, and 6,200 or 7,500 iterations for 124M or 350M.The configurations adopted Muon’s optimized settings except for slight tuning of β1.
  • 124M configuration: For 124M, Adam used a 3.6 × 10^-3 learning rate, while Muon and NorMuon used 3.6 × 10^-4.Adam used (β1,β2) = (0.9,0.95); Muon and NorMuon grid-searched β1 over {0.9,0.95}, with NorMuon β2 = 0.95.
  • 350M configuration: For 350M, Adam used learning rates of 0.3 for embeddings and 3 × 10^-3 for the output layer.Adam used (β1,β2) = (0.8,0.95).
  • 350M configuration: For 350M hidden layers, Muon and NorMuon used a 7.5 × 10^-4 learning rate, selecting β1 from {0.9,0.95} by validation performance.
Loading 2510.05491v1…