RMNP: Row-Momentum Normalized Preconditioning for Scalable Matrix-Based Optimization

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 $\mathcal{O}(mn\cdot\min(m,n))$ to $\mathcal{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 https://github.com/Dominator-Index/RMNP

Training a large AI model like ChatGPT is essentially a "downhill hike": the model starts from a random point and gradually walks toward the lowest valley of an error landscape, guided at every step by an algorithm called an optimizer. A recent optimizer called MUON trains models very effectively because it can "read" the terrain in fine detail, but doing so requires a costly chain of matrix calculations at every step, which may slows training down. We analyzed the training dynamics of the Transformer (the architecture underlying the vast majority of today's large language models) under MUON and uncovered a striking regularity in its loss geometry: the curvature information that governs training organizes itself into a clean, block-diagonal pattern. This is an intrinsic property of the network, not an assumption we impose. By following this inherent structure, we built a new optimizer called RMNP that replaces MUON's expensive calculation with one extremely simple operation: normalizing each row of the update. In experiments on large language model pretraining, RMNP matches MUON's quality while substantially cutting per-step computation time. We also provide a corresponding non-convex convergence analysis. The result: training powerful AI models with noticeably less time and compute.