HIERARCHICAL MUON: TILED NEWTON–SCHULZ UPDATES FOR EFFICIENT MUON OPTIMIZATION

Muon-type optimizers improve neural network training by applying a finite Newton-Schulz (NS) iteration to momentum-gradient matrices. While effective, the full-matrix NS update is computationally expensive, requiring O(r^2sK) operations and creating global spectral coupling that limits parallelization. The authors introduce Hierarchical Muon (HiMuon), which replaces the global NS update with a tiled version. By partitioning matrices into T x T tiles and applying the NS map independently to each, HiMuon transforms a large, monolithic matrix operation into a series of independent, small-matrix computations.

HiMuon reduces the leading computational cost of the NS update to O(HWTK), where T is the tile size. Because the tilewise operations are independent, the method is highly amenable to batched GPU execution. The authors demonstrate that this structure allows for tile-size-dependent kernels, cross-layer batching, and memory-bounded chunking. Empirical results on transformer training show that HiMuon achieves training behavior very close to the full-matrix Muon baseline, validating that local spectral interactions are often sufficient for effective optimization.

As neural networks grow in size, the computational overhead of sophisticated optimizers becomes a bottleneck. HiMuon provides a scalable alternative that bridges the gap between coordinate-wise methods and full-matrix updates. By making the tile size a tunable parameter, researchers can explicitly control the trade-off between spectral coupling (numerical quality) and computational efficiency (hardware throughput), facilitating more efficient training for large-scale models on distributed hardware.