NorMuon 的自适应缩放为何退化,DGA-Muon 又如何修正?

NorMuon 的行级自适应性源自正交化几何而非梯度,DGA-Muon 将其解耦并按矩阵形状重新对齐。

直接答案

NorMuon 的行级自适应缩放很大程度上是正交化几何的副产物,而非真正的优化信号。在精确正交化下,它对方形和宽矩阵退化为单一全局标量;在近似正交化下,其唯一的行间差异来自正交性残差,因此正交化越精确,自适应性越弱。DGA-Muon 改为直接从原始梯度计算缩放因子,并让缩放轴与极因子单侧正交结构对齐;论文报告其在 NanoChat 上的验证 BPB 为 0.6910,而 NorMuon 为 0.6924。

1篇文献引用

本文由 WisPaper 驱动的搜索和论文分析生成。

NorMuon 缩放正交化动量的各行

Muon 沿正交化动量的方向更新二维权重矩阵,而不是沿原始动量方向。给定动量 M_t,它由薄 SVD M_t = UΣV^T 构造 Orth(M_t) = UV^T,把所有奇异值置为一,只保留方向结构;该矩阵 UV^T 即极因子。实践中 Muon 用 Newton–Schulz 迭代近似这一操作,而不做完整 SVD。

NorMuon 增加了一个位于 R^m 的二阶矩估计 v_t,跟踪正交化动量的行级均方,并以衰减率 β2 做指数移动平均更新。参数更新随后把正交化动量除以 v_t 的平方根加一个小常数 δ,按行广播。关键设计选择在于该信号的来源:与 AdaGrad、Adam 或 Shampoo 不同,NorMuon 从正交化动量而非历史梯度导出缩放,这使自适应性紧密耦合于正交化步骤。

精确正交化把缩放塌缩为单一标量

对于方形和宽矩阵(m ≤ n),论文用 Gram 矩阵表示 O_t 逐元素平方的行级均值:mean_cols(O_t ⊙ O_t) = (1/n) diag(O_t O_t^T)。假设 M_t 满行秩,SVD 给出 U 正交、V 列正交归一,因此精确正交化得到 O_t O_t^T = U(V^T V)U^T = UU^T = I_m。矩阵因此是行正交归一的,行级均方变为 (1/n)·1_m,在所有行上完全相同。

把它代入 EMA 递推并以 v_0 = 0 展开,得到 v_t = (1 − β2^t)·(1/n)·1_m,它在各行上一致,仅通过一个全局标量随时间变化。论文指出,若施加 NorMuon 省略的标准偏差校正,会完全抵消该时间相关因子,恰好留下 (1/n)·1_m。无论哪种情况,缩放向量 1/√(v_t + δ) 都给每一行分配相同系数,因此原本设想的神经元级自适应性退化为一个平凡的全局标量。

正交性残差是行间差异的唯一来源

真实实现使用有限次 Newton–Schulz 迭代,因此近似正交化后的矩阵满足 O_t O_t^T = I_m + E_t,其中 E_t 是正交性残差。行级均方随之变为 (1/n)(1_m + diag(E_t)),展开 EMA 得到 v_t = (1 − β2^t)(1/n)1_m + R_t,其中 R_t = ((1 − β2)/n) Σ_{i=1}^{t} β2^{t−i} diag(E_i)。第一项正是精确情形下的行不变项;R_t 是行间差异的唯一来源。

假设残差在谱范数上被一个小量 ε 界定,则对角元满足 |(E_t)_jj| ≤ ε,从而 ‖R_t‖_∞ ≤ ε/n。因此 v_t 任意两个元素之差至多为 2ε/n,由中值定理,缩放因子 1/√(v_t + δ) 的相应差异被某个常数 C 界定为 Cε/n。由于 Newton–Schulz 在实践中把 ε 压得很小,行级缩放因子只有可忽略的变化。论文的解释是:NorMuon 表面上的自适应性是不精确正交化的副产物,而非对梯度统计或曲率信息的响应。

正交化越精确,NorMuon 自适应性越弱

该分析得出作者所称的正交化—自适应性悖论:NorMuon 的自适应性与正交化质量反向耦合。正交化越精确,行间差异越弱;在方形和宽矩阵的精确极限下它完全消失。正交化越不精确,允许的差异越大,表面自适应性越强。

论文将此解读为 NorMuon 的自适应机制在补偿正交化的不完美,而非提供优化驱动的自适应性。这与 Muon 研究中另一条旨在让正交化更精确或更高效的路线形成张力,因为该方向上的改进会缩小 NorMuon 缩放所依赖的那份残差。该分析假设满行秩;论文把秩亏视为不利的边缘情形,而非一般设定。

高矩阵暴露行级缩放的几何错配

对高矩阵而言情况不同。正交化只强制单侧正交性,即列正交归一,而非行正交归一。所有行的平方 ℓ2 范数总和固定,但每一行各自的平方 ℓ2 范数不受约束,因此即使在精确正交化下,行级缩放因子也可以变化。这种变化来自正交化后行能量的非均匀分布,而非梯度信息。

论文进一步论证,NorMuon 僵硬的行级缩放与这种单侧结构在几何上错配。对高矩阵而言,极因子是列正交的,其列定义了正交方向,而行级缩放沿行维度作用,会扭曲列几何,破坏极因子所诱导的正交性。这是结构性错配,而非调参问题。

两条设计原则解耦并重新对齐自适应性

从这些局限出发,论文提出理想 Muon 自适应机制应满足的两条原则。第一,自适应缩放应与正交化解耦,缩放因子直接从原始梯度计算,再作用于极因子。第二,自适应缩放应与极因子依赖形状的正交结构对齐:宽矩阵用行级缩放,高矩阵用列级缩放,方形矩阵两种皆可。论文给出形式化证明,说明这种几何对齐的缩放能保持更新的正交结构。

DGA-Muon 还加入三项设计选择。二阶矩估计由平方项的求和而非均值构成,因此每个行或列梯度向量被视为一个自适应单元,按其累计梯度总能量缩放;累计能量越大的单元得到越小的缩放因子。由于二阶矩估计直接由原始梯度计算,会承受 EMA 常见的初始化偏差,因此对其施加标准偏差校正;而一阶矩无需校正,因为正交化具有尺度不变性。最后,自适应裁剪把缩放因子限制在中位数附近的 c·s_med 与 d·s_med 之间,约束所诱导对角缩放矩阵的条件数,并保留 Muon 支持大步长的能力。

NanoChat 结果把理论与基线区分开

论文报告了两项实证检验。第一,它检验 NorMuon 缩放会在方形矩阵上退化的理论论断:在 NanoChat 中移除 Q/K/V/O 投影矩阵上的自适应缩放,这些矩阵在所有注意力头上作为完整矩阵优化,因此是方形的。训练运行 7,308 步,每 250 步验证一次,使用验证集每字节比特数(BPB),越低越好,且该指标不受分词器词表大小影响。移除缩放后验证 BPB 基本不变,两条曲线几乎重合,作者将其作为退化分析的实证支持。

第二,DGA-Muon 在 24 层、1.38B 参数的 NanoChat Transformer 上与 NorMuon 做端到端比较,数据与参数比从 10 提高到 18。DGA-Muon 最终验证 BPB 为 0.6910,NorMuon 为 0.6924;最终训练损失分别为 2.2708 与 2.2799。在 DCLM CORE 基准上,DGA-Muon 得分为 0.3089,NorMuon 为 0.2942,绝对提升 0.0147,相对提升 5.0%。论文指出,两种优化器因缩放信号不同而需要不同超参数,DGA-Muon 采用更平缓的学习率调度,包括延长预热、增大衰减比例并降低最终学习率比例。

结论在哪里成立,又在哪里止步

退化结论是在满行秩假设下的精确正交化情形中给出的:方形和宽矩阵塌缩为全局标量,高矩阵则保留来自行能量分布的差异。缩放差异的残差界是在正交性残差谱范数被 ε 界定的假设下推导的,论文称这与实际 Newton–Schulz 容差相符。退化的实证验证仅针对 NanoChat 中方形 Q/K/V/O 投影矩阵,端到端比较也只使用一种架构、一个数据与参数比,并为每种优化器分别调参。

论文还把贡献框定得比单一优化器修补更广:它主张矩阵感知优化器的自适应机制应把优化相关信号的提取与矩阵变换分开,并认为行级或列级自适应性是对矩阵方向结构的自然响应,而不只是逐元素自适应的近似。证据支持的实践优先级正是解耦与对齐这一选择本身:在比较自适应 Muon 变体时,该问的是缩放信号是否来自原始梯度,以及其轴是否匹配极因子的单侧正交结构,因为正是这两点把本次评估中 DGA-Muon 的 0.3089 CORE 分数与 NorMuon 的 0.2942 区分开来。

本文引用的文献

1

DGA-Muon: Decoupled Geometry-Aligned Adaptive Scaling for Muon

NorMuon 的行级自适应性源自正交化诱导的几何,而非优化相关信息,在精确正交化下对方形和宽矩阵退化为全局标量;DGA-Muon 将缩放与正交化解耦,并使其与极因子依赖形状的正交结构对齐,在 NanoChat 上达到 0.6910 的验证 BPB,而 NorMuon 为 0.6924,DCLM CORE 得分为 0.3089,NorMuon 为 0.2942。