arXiv:2603.20527cs.LG2026-03被引 7

用行归一化替代迭代计算,提升大模型优化效率

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

  • 以行方向ℓ2归一化替代牛顿-舒尔迭代,降低计算开销
  • 对m×n权重矩阵复杂度从O(mn·min(m,n))降至O(mn)
  • 适用于大规模语言模型训练,显著减少预处理耗时

预条件自适应方法因其能捕捉损失曲面的丰富曲率信息而受到广泛关注。当前核心挑战在于预条件效果与实现效率之间的平衡。近期的Muon方法通过牛顿-舒尔迭代获得预条件更新,无需显式构造预条件矩阵。但其效率仍有提升空间。本文提出RMNP(行动量归一化预条件),将牛顿-舒尔迭代替换为简单的行方向(输入维度d_in)ℓ2归一化操作,基于Transformer层间海森矩阵的对角块结构经验观察。实验验证了正交化与行方向ℓ2归一化在Transformer情形下渐近等价。该替换使每轮迭代复杂度从O(mn·min(m,n))降至O(mn),同时保持相当的优化性能。理论上,我们建立了非凸设置下的收敛性保证,达到近期Muon类优化器的最小最大复杂度。大规模语言模型预训练实验表明,RMNP在优化性能上可媲美Muon,且大幅降低预条件耗时。代码已开源:https://github.com/Dominator-Index/RMNP。

原文摘要 · Abstract (English)

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 https://github.com/Dominator-Index/RMNP.

优化器大模型训练预条件Transformer

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。