RMNP：基于行动量归一化的预条件化技术用于可扩展矩阵优化

经过预条件化的自适应方法因能捕捉损失景观的丰富曲率信息而受到广泛关注。在此领域的核心挑战在于平衡预条件化效果与实现预条件器的计算效率。近期的进展中，Muon 通过 Newton-Schulz 迭代获取预条件化更新，无需显式构造预条件化矩阵，凸显其优势。尽管如此，Muon 的效率仍有提升空间。本文引入 RMNP（Row Momentum Normalized Preconditioning，行动量归一化预条件化），一种优化器，取代 Newton-Schulz 迭代，使用简单的按行 ($d_{\text{in}}$) $\ell_2$ 归一化操作，基于 Transformer 层级 Hessian 的经验观察到的对角块结构所致。我们经验性地验证了在 Transformer 情形下，对角化和按输入维度的行 $\ell_2$ 归一化在渐近意义上是等价的。该替换将 $m\times n$ 权重矩阵的每迭代计算复杂度从 ${O}(mn\cdot\min(m,n))$ 降低至 ${O}(mn$，同时保持可比的优化性能。理论上，我们在非凸情形下为 RMNP 建立了收敛保证，达到与 Muon 优化器相近的复杂度，实现最小混沌复杂度。大规模语言模型预训练中的广泛实验表明，RMNP 在与 Muon 相当的优化性能下，显著降低了预条件化的 wall-clock 时间。我们的代码已公开于 https://github.com/Dominator-Index/RMNP。