Hierarchical Muon: Tiled Newton-Schulz Updates for Efficient Muon Optimization

2026-06-25

AI Summary

This work addresses the computational inefficiency of the Muon optimizer in large-scale neural networks, where full-matrix Newton–Schulz updates incur prohibitive costs and poor parallelizability. To overcome this limitation, the authors propose Hierarchical Muon (HiMuon), which partitions the momentum gradient matrix into fixed-size blocks and applies a limited number of Newton–Schulz iterations independently within each block. This block-wise local matrix function approximation preserves intra-block spectral structure while neglecting inter-block coupling, thereby maintaining optimization behavior with substantially improved efficiency. By integrating GPU-customized kernels, cross-layer batching, memory-aware blocking, and dynamic block-size scheduling, HiMuon achieves significant acceleration of per-step optimization in Transformer training and matrix function experiments, with negligible loss in optimization performance.

Abstract

Muon-type optimizers construct update directions for dense neural-network weights by applying a finite Newton-Schulz map to momentum-gradient matrices. For an $H \times W$ matrix, with $r=\min\{H,W\}$ and $s=\max\{H,W\}$, $K$ steps of the full-matrix Newton-Schulz update require $O(r^2 s K)$ work and couple all rows and columns through repeated Gram matrix products. We introduce Hierarchical Muon (HiMuon), a tiled Newton-Schulz scheme for Muon-type optimization. HiMuon partitions each momentum-gradient matrix into $T \times T$ tiles, applies the same finite Newton-Schulz map independently to each tile, and reassembles the results. For finite $T$ below the matrix dimensions, HiMuon defines a local matrix-function map rather than a convergent approximation to the full-matrix update: spectral interactions are preserved within tiles and discarded across tile boundaries. For fixed finite $T$, the leading Newton-Schulz work decreases to $O(H W T K)$, and the computation decomposes into independent small dense matrix operations. This structure enables tile-size-dependent GPU kernels, cross-layer batching, memory-bounded chunking, and runtime tile-size schedules. Experiments on transformer training and controlled matrix-function diagnostics show that HiMuon improves optimizer-step efficiency while keeping training behavior close to full-matrix Muon in the tested regimes.

Problem

• Muon optimization
• Newton-Schulz iteration
• matrix function
• computational efficiency
• large-scale optimization

Innovation

• Hierarchical Muon
• tiled Newton-Schulz
• matrix-function approximation
• efficient optimization
• block-wise update