Abstract:Preconditioned adaptive methods have gained significant attention for training deep neural networks, as they capture rich curvature information of the loss landscape. The central challenge in this field lies in balancing preconditioning effectiveness with computational efficiency of implementing the preconditioner. Among recent advances, Muon stands out by using Newton-Schulz iteration to obtain preconditioned updates without explicitly constructing the preconditioning matrix. Despite its advantages, the efficiency of Muon still leaves room for further improvement. In this paper, we introduce RMNP (Row Momentum Normalized Preconditioning), an optimizer that replaces Newton-Schulz iteration with a simple row-wise ($d_{\text{in}}$) $\ell_2$ normalization operation, motivated by the empirically observed diagonal block structure of the Transformer layerwise Hessian. We empirically verified that orthogonalization and row-wise (on input dim) $\ell_2$ normalization are asymptotically equivalent in the case of the transformer. This substitution reduces the per-iteration computational complexity from ${O}(mn\cdot\min(m,n))$ to ${O}(mn)$ for an $m\times n$ weight matrix while maintaining comparable optimization performance. Theoretically, we establish convergence guarantees for RMNP in the non-convex setting that match recent results for Muon optimizers, achieving the minimax optimal complexity. Extensive experiments on large language model pretraining show that RMNP delivers competitive optimization performance compared with Muon while substantially reducing preconditioning wall-clock time. Our code is available at this https URL.
| Comments: | The 43rd International Conference on Machine Learning (ICML 2026) |
| Subjects: | Machine Learning (cs.LG) |
| Cite as: | arXiv:2603.20527 [cs.LG] |
| (or arXiv:2603.20527v3 [cs.LG] for this version) | |
| https://doi.org/10.48550/arXiv.2603.20527 arXiv-issued DOI via DataCite |
From: Shenyang Deng [view email]
[v1]
Fri, 20 Mar 2026 21:55:28 UTC (7,444 KB)
[v2]
Tue, 31 Mar 2026 18:27:08 UTC (4,519 KB)
[v3]
Wed, 13 May 2026 10:41:39 UTC (7,293 KB)
此内容由惯性聚合(RSS阅读器)自动聚合整理，仅供阅读参考。 原文来自 — 版权归原作者所有。