通过分块计算提升神经网络优化器效率，保持训练效果。
Hierarchical Muon: Tiled Newton-Schulz Updates for Efficient Muon Optimization
- 将梯度动量矩阵分块，独立计算每块的牛顿-舒尔兹更新。
- 计算复杂度从O(r²sK)降至O(HWT K)，支持高效并行和内存调度。
- 适合大规模Transformer训练，可灵活调整分块大小以优化性能。
Muon类优化器通过有限牛顿-舒尔兹映射对密集神经网络权重的动量-梯度矩阵构造更新方向。对于一个H×W的矩阵，其中r=min{H,W}，s=max{H,W}，K步全矩阵牛顿-舒尔兹更新需要O(r²sK)的计算量，并通过重复的格拉姆矩阵乘积耦合所有行和列。本文提出层级式Muon（HiMuon），一种用于Muon类优化的分块牛顿-舒尔兹方案。HiMuon将每个动量-梯度矩阵划分为T×T的子块，对每个子块独立应用相同的有限牛顿-舒尔兹映射，并重新组合结果。当T为低于矩阵维度的固定值时，HiMuon定义了一个局部矩阵函数映射，而非全矩阵更新的收敛近似：仅保留块内谱交互，丢弃块间交互。在固定有限的T下，主导牛顿-舒尔兹工作量降至O(HWT K)，且计算可分解为独立的小型密集矩阵操作。该结构支持基于块大小的GPU核、跨层批处理、内存受限的分块以及运行时块大小调度。在Transformer训练和受控矩阵函数诊断实验中，结果显示HiMuon提升了优化步骤效率，同时在测试范围内保持了与全矩阵Muon相近的训练行为。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。