AI 中文总结
该论文首次系统分析NorMuon的自适应机制,揭示其源于正交化几何而非优化信息,并提出解耦几何对齐的DGA-Muon优化器,通过解耦缩放与正交化、按形状对齐缩放方向,实现优于NorMuon的性能。
AI 中文摘要
尽管NorMuon通过增强Muon的行式自适应缩放,在大规模预训练中取得了强劲的经验性能,但其潜在的自适应机制仍鲜为人知。在本工作中,我们首次对NorMuon的自适应性进行了系统分析,揭示其主要源于正交化诱导的几何特性,而非真正的优化相关信息。在精确正交化下,自适应缩放因子对于方阵和宽矩阵退化为单一的全局标量,而对于高矩阵,其变化源于正交化后行能量的非均匀分布。在近似正交化下,正交化残差为缩放因子引入了额外的变化,导致“正交化-自适应性悖论”:更精确的正交化削弱了自适应性。我们进一步表明,NorMuon刚性的行式缩放与高矩阵的单侧正交结构在几何上不匹配。基于对这些局限性的分析,我们提出了Muon理想自适应机制应满足的两个核心设计原则。首先,自适应缩放应与正交化解耦,缩放因子直接从原始梯度计算。其次,自适应缩放应与极分解因子的形状相关正交结构对齐,对宽矩阵使用行式缩放,对高矩阵使用列式缩放。通过纳入其他设计考虑,包括基于求和的一阶矩估计、偏差校正和缩放因子的自适应裁剪,我们得到了解耦几何对齐Muon(DGA-Muon)优化器。我们为DGA-Muon建立了收敛保证,并从经验上验证了我们对NorMuon缩放退化的理论刻画以及DGA-Muon相对于NorMuon的优越性。
英文摘要
While NorMuon has achieved strong empirical performance in pretraining, its underlying adaptive mechanism remains largely heuristic and poorly understood. In this work, we provide the first systematic theoretical analysis of NorMuon's adaptivity, revealing that it primarily arises from orthogonalization-induced geometry rather than genuine optimization dynamics, serving to offset the resulting geometric non-uniformity. Under exact orthogonalization, the adaptive scaling factors degenerate into a single global scalar for square and wide matrices, while for tall matrices their variation results from unevenly distributed row energy after orthogonalization. Under approximate orthogonalization, the orthogonality residuals introduce additional variation into the scaling, giving rise to a counterintuitive Orthogonalization--Adaptivity Paradox: more accurate orthogonalization weakens adaptivity. We further show that NorMuon's rigid row-wise scaling is geometrically misaligned with the one-sided orthogonal structure of tall matrices by distorting column orthogonality. Motivated by these limitations, we propose two core design principles that a desirable adaptive mechanism for Muon should satisfy. First, adaptive scaling should be decoupled from orthogonalization, with the scaling factors computed directly from raw gradients. Second, adaptive scaling should be aligned with the shape-dependent orthogonal structure of the polar factor, using row-wise scaling for wide matrices and column-wise scaling for tall matrices. We prove that this geometry-aligned scaling preserves the orthogonal structure of the update. By incorporating several other techniques, we obtain Decoupled Geometry-Aligned Muon (DGA-Muon). We establish convergence guarantees for DGA-Muon and empirically validate both our theoretical characterization of NorMuon's scaling degeneration and the superiority of DGA-Muon.
Commentsabstract and some details refined