Tag: AdamW

  • Muon Optimizer: Scaling Matrix Orthogonalization for Efficient Deep Learning Training

    Muon Optimizer: Scaling Matrix Orthogonalization for Efficient Deep Learning Training

    TL;DR

    • The Muon optimizer trains 2D hidden layer weight matrices via steepest descent under the spectral norm using efficient iterative Newton-Schulz orthogonalization.
    • Muon operates as a hybrid optimizer: it manages 2D hidden-layer weight matrices, while AdamW handles everything else — genuinely 1D parameters (biases, normalization scales) plus the input embedding and output layers, which are excluded from Muon for role-based reasons even though they are typically 2D matrices themselves.
    • Compared to AdamW, Muon cuts training compute — about 1.35x faster wall-clock in Keller Jordan’s NanoGPT speedrun and up to 2x fewer FLOPs under Moonlight’s compute-optimal scaling law — while also cutting VRAM overhead by eliminating the second-moment buffer and reducing the high condition number (near low-rank concentration) of raw weight updates.

    Introduction: Understanding the Muon Optimizer and Orthogonal Updates

    The Muon Optimizer—short for MomentUm Orthogonalized by Newton-Schulz—is an optimization method designed for training the Hidden Layers of neural networks. Unlike conventional first-order optimizers such as SGD, Adam, and AdamW, which treat weight parameters as flattened 1D Parameters and apply coordinate-wise scaling, Muon treats parameters as 2D matrices. It performs Steepest Descent with respect to the Spectral Norm (the matrix operator norm), optimizing both the directional geometry and the magnitude of Weight Updates.

    Muon was introduced in late 2024 by Keller Jordan and collaborators, building on foundational theoretical work by Jeremy Bernstein and Laker Newhouse (‘Old Optimizer, New Norm: An Anthology’), which frames this Steepest Descent step under an RMS-to-RMS operator norm — a normalized formulation that is effectively equivalent to the classical Spectral Norm used elsewhere in this article, differing mainly in how the norm is scaled by matrix dimensions rather than in the underlying geometry. It was first benchmarked at scale in the open-source Modded-NanoGPT speedrun repository. The optimizer has since been expanded through theoretical analyses, such as the Newton–Muon optimizer studied by Du and Su at UPenn, and scaled by Moonshot AI (the Kimi Team) to train the Moonlight 3B/16B Mixture-of-Experts (MoE) model across 5.7 trillion tokens.

    At its core, Muon operates through a multi-step update pipeline:

    • Momentum Collection: Gradient updates are first accumulated using standard Nesterov momentum.
    • Newton-Schulz Orthogonalization: Rather than computing an exact, computationally expensive Singular Value Decomposition (SVD) on GPUs, Muon applies an iterative Newton-Schulz Iteration (typically 5 steps). This quickly drives all singular values of the update matrix close to 1.0, generating an approximate Matrix Orthogonalization equivalent to the matrix-sign function and Polar Decomposition.
    • Aspect-Ratio Scaling: The orthogonalized update is rescaled according to matrix dimensions to preserve stable signal propagation across layers with varying shapes. Keller Jordan’s original implementation uses a ratio-based factor, sqrt(max(1, rows/cols)); the Moonlight paper (Kimi Team) proposes an explicit coefficient form instead, 0.2 * sqrt(max(A, B)), noting the two scalings are equivalent up to a global scale only when the matrices share the same second dimension, and can otherwise diverge.

    Muon is implemented as a hybrid optimizer. It is specialized strictly for 2D weight matrices in Hidden Layers and linear layers—such as Multi-Layer Perceptron (MLP) weights, attention projections in Transformer Architectures, and reshaped convolutional kernels. Everything else is optimized concurrently using a traditional optimizer like AdamW: genuinely 1D parameters such as Biases and Normalization Scales, plus the input embedding and output (unembedding) layers. Those embedding and output layers are typically 2D matrices themselves, but Muon’s authors exclude them for role-based reasons rather than dimensionality—embeddings on the theoretical grounds of the modular norm framework, and output layers on purely empirical grounds.

    This Matrix Orthogonalization scheme provides several key advantages over standard coordinate-wise optimizers:

    • Compute Efficiency and Convergence: Reported efficiency gains come from two distinct sources — Keller Jordan’s NanoGPT speedrun recorded roughly 1.35x faster wall-clock training time, while the Moonlight paper’s compute-optimal scaling law separately found up to 2x fewer training FLOPs versus AdamW to reach the same validation loss.
    • Lower Optimizer Memory Overhead: Muon tracks only a single first-moment momentum buffer per parameter and eliminates the second-moment buffer (v_t) required by AdamW, yielding notable VRAM savings on Hidden Layers.
    • Reducing Update Condition Number: By normalizing singular values, the optimizer prevents Weight Updates from collapsing into a small subset of high-variance singular directions—a near low-rank, high-condition-number pathology of raw gradients—better engaging the full representational capacity of weight matrices.
    • Empirical Performance: In addition to setting speedrun training records on single-node setups for CIFAR-10 and GPT-2 (small and medium), Muon has demonstrated stable scaling in multi-billion-parameter LLM pretraining. The Moonlight paper’s Distributed Muon implementation follows ZeRO-1 to partition optimizer states across the data-parallel dimension — with tensor parallelism usable only as a complementary, non-essential option — combined with adjusted weight decay.

    Sources: Muon: An optimizer for hidden layers in neural networks, GitHub – KellerJordan/Muon: Muon is an optimizer for hidden layers in neural networks · GitHub, Deriving Muon, medium.com/@md.abir1203/muon-optimizer-the-complete-guide-to-2-faster…, [2502.16982] Muon is Scalable for LLM Training

    Closing thoughts

    Taken together, these technical strengths illustrate how shifting from coordinate-wise updates to geometry-aware 2D matrix orthogonalization enables Muon to directly address the structural properties of hidden layers, yielding clear optimization advantages over standard first-order methods. Its hybrid architecture—applying efficient Newton-Schulz iterations to weight matrices while retaining AdamW for 1D parameters and embeddings—strikes a sensible balance between theoretical rigor and implementation practicality.

    Ultimately, with demonstrated FLOP reductions and lower memory overhead validated from single-node speedruns up to multi-trillion-token LLM pretraining, Muon establishes matrix-level spectral descent as a genuinely viable, high-efficiency paradigm for training modern neural networks.

    Frequently Asked Questions

    What is the Muon optimizer?

    Muon (MomentUm Orthogonalized by Newton-Schulz) is an optimization method designed for training hidden layers of neural networks by treating parameters as 2D matrices rather than flattened 1D vectors. It performs steepest descent with respect to the spectral norm to optimize the directional geometry and magnitude of weight updates.

    How does Muon achieve matrix orthogonalization efficiently?

    Instead of computing an expensive Singular Value Decomposition (SVD) on GPUs, Muon applies an iterative Newton-Schulz iteration, typically in 5 steps. This quickly drives update matrix singular values close to 1.0, providing an approximate orthogonalization equivalent to the matrix-sign function and polar decomposition.

    Why is Muon used as a hybrid optimizer alongside AdamW?

    Muon is specialized strictly for 2D weight matrices in hidden layers, such as MLP weights and attention projections. Everything else uses a traditional optimizer like AdamW instead: 1D parameters such as biases and normalization scales, plus the input and output embedding layers. Those embedding layers are typically 2D themselves, but are excluded from Muon for role-based reasons—not because they are non-2D.

    What are the primary benefits of Muon compared to AdamW?

    Muon reduces training compute versus AdamW — about 1.35x faster wall-clock time in Keller Jordan’s NanoGPT speedrun, and up to 2x fewer FLOPs under Moonlight’s compute-optimal scaling law — while also lowering VRAM usage by eliminating AdamW’s second-moment buffer. Additionally, it helps reduce the high condition number of raw weight updates by preventing them from collapsing into a small subset of high-variance singular directions.


    Keep reading: Training & Architectures

    All guides in this series: LLM Training & Architecture Guides: Fine-Tuning, RL & New Models