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

arXiv:2606.27216v1 Announce Type: cross 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 imes 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 imes 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.

Source: arXiv cs.LG | 2026-06-26