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

• URL: http://arxiv.org/abs/2606.27216v1
• Date: Thu, 25 Jun 2026 16:05:55 GMT
• Title: Hierarchical Muon: Tiled Newton-Schulz Updates for Efficient Muon Optimization
• Abstract summary: Muon-types construct update directions for dense neural-network weights by applying a finite Newton-Schulz map to momentum-gradient matrices.<n>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.
• Score: 4.275508336828267
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• 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.

Related papers

• Two-Tower Quantum Matrix Chain Multiplication: Trading Qubits for Depth [73.08853228981701]
  Matrix chain multiplication is a problem in scientific computing, machine learning, and graph analysis.<n>We present emphTwo-Tower Matrix, a quantum subroutine that encodes the product $mathcalW$ of the $K$ matrices into a quantum state in circuit depth.
  arXiv  Detail & Related papers  (2026-07-14T18:42:40Z)
• Muon$^2$: Boosting Muon via Adaptive Second-Moment Preconditioning [18.570226339282296]
  Muon has emerged as a promising computation for large-scale foundation model pre-training.<n>We propose Muon$2$, a memory-efficient factorized variant that preserves most of the gains of Muon$2$ with negligible memory overhead.
  arXiv  Detail & Related papers  (2026-04-11T00:27:40Z)
• The Newton-Muon Optimizer [2.1583596008762935]
  We introduce a surrogate model that sheds new light on the design iteration of Muon.<n>New optimization method, which we refer to as Newton-Muon, shows that standard Muon can be interpreted as an implicit Newton-type method.
  arXiv  Detail & Related papers  (2026-04-01T23:29:08Z)
• UNSO: Unified Newton Schulz Orthogonalization [10.113387496783007]
  The Newton-Schulz (NS) has gained increasing interest for its role in the Muon and Stiefel manifold.<n>We consolidate the iterative structure into a unified framework, named Unified Newton-Schulz Orthogonalization (UNSO)<n>These learnable coefficients are then optimized, and achieve an outstanding performance with stable convergence.
  arXiv  Detail & Related papers  (2026-01-18T17:54:43Z)
• Evolution Strategies at the Hyperscale [57.75314521465674]
  We introduce EGGROLL, an evolution strategies (ES) algorithm designed to scale backprop-free optimization to large population sizes.<n>ES is a set of powerful blackbox optimisation methods that can handle non-differentiable or noisy objectives.<n>EGGROLL overcomes these bottlenecks by generating random matrices $Ain mathbbRmtimes r, Bin mathbbRntimes r$ with $rll min(m,n)$ to form a low-rank matrix perturbation $A Btop$
  arXiv  Detail & Related papers  (2025-11-20T18:56:05Z)
• INC: An Indirect Neural Corrector for Auto-Regressive Hybrid PDE Solvers [61.84396402100827]
  We propose the Indirect Neural Corrector ($mathrmINC$), which integrates learned corrections into the governing equations.<n>$mathrmINC$ reduces the error amplification on the order of $t-1 + L$, where $t$ is the timestep and $L$ the Lipschitz constant.<n>We test $mathrmINC$ in extensive benchmarks, covering numerous differentiable solvers, neural backbones, and test cases ranging from a 1D chaotic system to 3D turbulence.
  arXiv  Detail & Related papers  (2025-11-16T20:14:28Z)
• LREI: A fast numerical solver for quantum Landau-Lifshitz equations [0.0]
  We develop LREI, a memory- and time-efficient scheme for solving quantum Landau-Lifshitz (q-LL) and quantum Landau-Lifshitz-Gilbert (q-LLG) equations.
  arXiv  Detail & Related papers  (2025-08-28T20:33:37Z)
• Compute Better Spent: Replacing Dense Layers with Structured Matrices [77.61728033234233]
  We identify more efficient alternatives to dense matrices, as exemplified by the success of convolutional networks in the image domain. We show that different structures often require drastically different initialization scales and learning rates, which are crucial to performance. We propose a novel matrix family containing Monarch matrices, the Block-Train, which we show performs better than dense for the same compute on multiple tasks.
  arXiv  Detail & Related papers  (2024-06-10T13:25:43Z)
• Conv-Basis: A New Paradigm for Efficient Attention Inference and Gradient Computation in Transformers [16.046186753149]
  Self-attention mechanism is the key to the success of transformers in recent Large Language Models (LLMs) We leverage the convolution-like structure of attention matrices to develop an efficient approximation method for attention using convolution matrices. We hope our new paradigm for accelerating attention computation in transformer models can help their application to longer contexts.
  arXiv  Detail & Related papers  (2024-05-08T17:11:38Z)
• Softmax-free Linear Transformers [90.83157268265654]
  Vision transformers (ViTs) have pushed the state-of-the-art for visual perception tasks. Existing methods are either theoretically flawed or empirically ineffective for visual recognition. We propose a family of Softmax-Free Transformers (SOFT)
  arXiv  Detail & Related papers  (2022-07-05T03:08:27Z)
• Monarch: Expressive Structured Matrices for Efficient and Accurate Training [64.6871423399431]
  Large neural networks excel in many domains, but they are expensive to train and fine-tune. A popular approach to reduce their compute or memory requirements is to replace dense weight matrices with structured ones. We propose a class of matrices (Monarch) that is hardware-efficient.
  arXiv  Detail & Related papers  (2022-04-01T17:37:29Z)
• Newton-LESS: Sparsification without Trade-offs for the Sketched Newton Update [88.73437209862891]
  In second-order optimization, a potential bottleneck can be computing the Hessian matrix of the optimized function at every iteration. We show that the Gaussian sketching matrix can be drastically sparsified, significantly reducing the computational cost of sketching. We prove that Newton-LESS enjoys nearly the same problem-independent local convergence rate as Gaussian embeddings.
  arXiv  Detail & Related papers  (2021-07-15T17:33:05Z)
• Metric Transforms and Low Rank Matrices via Representation Theory of the Real Hyperrectangle [17.808087068037985]
  We show how to compute the eigenvectors and eigenvalues of certain matrices arising from hyperrectangles. We then use our new technique along with these connections to prove several new structural results.
  arXiv  Detail & Related papers  (2020-11-23T16:03:12Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.