Adam 改进 Muon：带正交化动量的自适应矩估计 Adam Improves Muon: Adaptive Moment Estimation with Orthogonalized Momentum
提出 NAMO/NAMO-D 优化器，将 Adam 的噪声自适应机制融入 Muon 的正交化动量框架
前置知识
Adam 优化器
Adam（Adaptive Moment Estimation）是一种广泛使用的自适应学习率优化算法。它维护梯度的一阶矩估计 $m_t = ?eta_1 m_{t-1} + (1-?eta_1)g_t$ 和二阶矩估计 $v_t = ?eta_2 v_{t-1} + (1-?eta_2)g_t^2$，通过比值 $\hat{m}_t / \sqrt{\hat{v}_t}$ 实现坐标级别的自适应步长。当梯度噪声大或接近驻点时，有效步长自动减小，从而提升训练稳定性。AdamW 是其解耦权重衰减的变体。
本文的核心思想是将 Adam 的噪声自适应机制与正交化动量结合，理解 Adam 的工作原理是理解本文动机的关键
Muon 优化器
Muon 是一种专门为矩阵参数设计的优化器，它对动量矩阵 $M_t$ 进行正交化处理 $\text{Orth}(M_t) = UV^T$（其中 $M = U\Sigma V^T$ 是 SVD），得到最近正交矩阵作为更新方向。在实践中通过 Newton-Schulz 迭代近似计算正交化。Muon 在大语言模型训练中展现出优于 Adam 的性能。
NAMO 和 NAMO-D 是在 Muon 基础上引入噪声自适应机制，Muon 是本文的直接基线和改进对象
正交化（Orthogonalization）
对于矩阵 $M \in \mathbb{R}^{m \times n}$，其正交化定义为 $\text{Orth}(M) := UV^T$，其中 $M = U\Sigma V^T$ 是约化 SVD。正交化矩阵是极分解中的正交因子，也是 Frobenius 范数下最接近 $M$ 的正交矩阵。正交化是无界操作，可能放大噪声矩阵的影响，这正是本文要解决的核心问题。
正交化是 Muon 系列算法的核心操作，理解其数学性质（特别是无界性带来的噪声放大问题）是理解本文动机的基础
核范数与谱范数
核范数 $\|\cdot\|_*$ 是矩阵奇异值之和，谱范数 $\|\cdot\|_2$ 是最大奇异值。论文在核范数-谱范数对偶框架下建立收敛理论，因为正交化梯度下降可解释为谱范数下的最速下降方向。这两种范数在分析矩阵结构优化问题中比 Frobenius 范数更自然。
论文的核心理论贡献建立在核范数-谱范数的对偶关系上，理解这两种范数是读懂收敛证明的前提
偏差校正（Bias Correction）
由于动量估计从零初始化，初始阶段存在系统性偏差。Adam 通过 $\hat{m}_t = m_t / (1-\beta_1^t)$ 和 $\hat{v}_t = v_t / (1-\beta_2^t)$ 进行偏差校正。NAMO 同样采用这种技术，将动量矩阵和范数估计校正为无偏估计。论文中的凸组合表示 $\hat{M}_t = \sum_{\tau=1}^{t} w_{1,t,\tau} \nabla L(\Theta_{\tau-1})$ 将校正后的动量表示为历史梯度的加权平均。
偏差校正是连接理论分析与实际算法的关键技术，论文的收敛证明大量使用凸组合表示来控制偏差项
Hessian 的近似块对角结构
神经网络的 Hessian 矩阵通常呈现近似块对角结构，即不同神经元（或参数块）之间的二阶交互较弱。这一观察启发了 Adam-mini 等按块分配学习率的方法，也是 NAMO-D 采用列级别（即神经元级别）自适应缩放的理论依据。
这一结构先验是 NAMO-D 设计的核心动机——列级别的对角缩放与 Hessian 的块对角结构对齐，从而实现更细粒度的噪声自适应
研究动机
当前主流优化器面临一个根本性的两难困境。Adam/AdamW 通过坐标级别的自适应步长提升训练稳定性，但其符号下降分量会导致泛化性能下降（Wilson et al., 2017 的实验和 Balles & Hennig, 2018 的理论分析都证实了这一点）。Muon 利用矩阵结构的正交化动量作为更新方向，在大语言模型训练中展现出优越性能，但正交化是一个无界操作（Higham, 1986），可能放大动量矩阵中的噪声，导致训练不稳定（He et al., 2025）和超参数敏感性增加（Crawshaw et al., 2025）。现有尝试将 Adam 与 Muon 结合的工作（如 AdaMuon、NorMuon）缺乏理论收敛保证；PRISM 的计算开销较高；DeVA 需要维护 Kronecker 预条件器和周期性特征分解，带来高昂的计算和内存开销。
本文的目标是本文的目标是设计一个理论上严谨的优化算法，将正交化更新方向与 Adam 类型的方差自适应机制进行原则性整合。具体而言，NAMO 通过单个自适应标量缩放正交化动量，保持更新方向的正交性，同时以可忽略的 $O(mn)$ 额外计算成本和零额外内存开销提升 Muon 的性能。NAMO-D 进一步通过右乘对角矩阵实现神经元级别的自适应步长，与神经网络常见的近似块对角 Hessian 结构对齐。两个算法都追求在确定性和随机环境下达到最优收敛速率。
与已有工作不同的是，本文的独特切入角度在于将优化器设计分解为两个正交的组成部分：（i）在无噪声情况下表现良好的更新方向选择规则，和（ii）通过衰减梯度不确定性下的更新来稳定迭代的自适应步长机制。基于正交化将矩阵更新的方向和幅度解耦的性质（Bernstein & Newhouse, 2024, Proposition 5），论文选择基于范数的重缩放作为自适应机制，而非 Adam 的坐标级别缩放。这种设计既保留了正交化方向的结构优势，又引入了噪声自适应能力，是首次为正交化更新提供理论收敛保证的 Adam 类型自适应方法。
核心方法
NAMO 的整体思路可以这样理解：Muon 的正交化动量提供了一个优秀的更新方向（谱范数下的最速下降方向），但这个方向在随机梯度噪声下可能不稳定。借鉴 Adam 的自适应机制，NAMO 不是像 Adam 那样对每个坐标分别缩放（这会破坏正交化方向的结构），而是用一个基于 Frobenius 范数的自适应标量来缩放整个正交化动量。技术路线上，NAMO 维护梯度矩阵 Frobenius 范数的二阶矩估计 $v_t = \mu_2 v_{t-1} + (1-\mu_2)\|G_t\|_F^2$，通过比值 $\alpha_t = \|\hat{M}_t\|_F / \sqrt{\hat{v}_t + \epsilon_t}$ 作为自适应缩放因子。当梯度噪声大或接近驻点时，$\alpha_t$ 自动减小，实现稳定收敛。NAMO-D 则更进一步，对每个神经元（矩阵的每一列）分别计算自适应步长，并通过截断机制保证更新方向的条件数有界。
本文的核心创新在于认识到正交化将矩阵更新的方向和幅度解耦（Bernstein & Newhouse, 2024 的命题 5），因此基于范数的重缩放是自然的设计选择。与已有方法的本质区别体现在三个层面：第一，NAMO 使用单个自适应标量而非 Adam 的坐标级别缩放，从而严格保持正交化方向——这是 AdaMuon/NorMuon 未能做到的理论保证。第二，NAMO-D 的列级别对角缩放与神经网络 Hessian 的近似块对角结构对齐，这是基于对深度学习优化景观结构的深刻理解。第三，NAMO-D 的截断机制 $\tilde{d}_t = \min(\max(d_t, c\bar{d}_t \mathbf{1}), \bar{d}_t/c)$ 通过参数 $c \in (0,1]$ 平衡两个竞争目标：保持良好条件的更新方向和利用细粒度噪声自适应。这是首次在理论上证明这种截断机制能保证最优收敛速率。
方法步骤详情
NAMO 算法的具体步骤如下：输入学习率 $\eta$、动量系数 $\mu_1, \mu_2 \in [0,1)$、批大小 $b$、小常数 $\epsilon > 0$。初始化参数 $\Theta_0$、动量矩阵 $M_0 = 0$、二阶矩标量 $v_0 = 0$。在每次迭代 $t$ 中：(1) 采样大小为 $b$ 的小批量，计算随机梯度 $G_t = \nabla L_t(\Theta_{t-1})$；(2) 更新动量矩阵 $M_t \leftarrow \mu_1 M_{t-1} + (1-\mu_1) G_t$；(3) 更新二阶矩估计 $v_t \leftarrow \mu_2 v_{t-1} + (1-\mu_2) \|G_t\|_F^2$；(4) 计算正交化动量 $O_t \leftarrow \text{Orth}(M_t)$；(5) 计算自适应缩放因子 $\alpha_t \leftarrow \frac{1-\mu_2^t}{1-\mu_1^t} \frac{\|M_t\|_F}{\sqrt{v_t} + \epsilon}$；(6) 更新参数 $\Theta_t \leftarrow \Theta_{t-1} - \eta \alpha_t O_t$。NAMO-D 的额外步骤包括：对每列计算范数 $[d_t]_j = [\hat{M}_t]_{:j} / \sqrt{[\hat{v}_t]_j + \epsilon_t}$，计算平均值 $\bar{d}_t = \|d_t\|_1 / n$，截断 $\tilde{d}_t = \min(\max(d_t, c\bar{d}_t), \bar{d}_t/c)$，构建对角矩阵 $D_t = \text{diag}(\tilde{d}_t)$，更新 $\Theta_t \leftarrow \Theta_{t-1} - \eta O_t D_t$。
技术新颖性
本文的技术新颖性体现在多个方面。首先，收敛分析框架是全新的——论文在核范数-谱范数对偶框架下建立收敛理论，利用正交化下降不等式（Zhang et al., 2025, Lemma B.1）将正交化更新的分析转化为核范数下的分析。其次，论文发展了一套处理偏差校正动量估计的技术：通过凸组合表示 $\hat{M}_t = \sum_{\tau=1}^{t} w_{1,t,\tau} \nabla L(\Theta_{\tau-1})$，将偏差项的控制转化为几何级数估计和伸缩求和。第三，NAMO-D 的截断机制分析是原创的——论文证明截断后 $d_{t,\min} \geq c^2 d_{t,\max}$，从而控制 $D_t$ 的条件数，这保证了更新方向的良好条件性。第四，论文的随机收敛分析展示了噪声自适应性：收敛速率分解为优化项 $O(T^{-1/4})$ 和方差项 $O(\sigma b^{-1/4} T^{-1/8})$，当批大小 $b = \Omega(\sigma^2 \sqrt{T})$ 时恢复最优速率。
实验结果
实验在 GPT-2 预训练上验证了 NAMO 和 NAMO-D 的有效性。在 GPT-2 (124M) 模型上，通过 10K 步的网格搜索确定最优学习率后训练 50K 步，NAMO 的训练损失为 2.9272、验证损失为 3.0351，NAMO-D 的训练损失为 2.9167、验证损失为 3.0246，均优于 AdamW（训练 3.0456、验证 3.0643）和 Muon（训练 3.0265、验证 3.0435）。在 GPT-2 (355M) 模型上训练 10K 步，NAMO 的训练损失为 2.9359、验证损失为 2.9516，NAMO-D 的训练损失为 2.9351、验证损失为 2.9507，同样优于 AdamW（训练 2.9760、验证 2.9914）和 Muon（训练 2.9524、验证 2.9684）。值得注意的是，NAMO 和 NAMO-D 的最优学习率（0.007-0.012）比 AdamW 和 Muon 的最优学习率（0.0009-0.0013）大约大一个数量级，这表明自适应缩放机制确实改变了优化景观的尺度。此外，超参数扫描结果显示 NAMO 和 NAMO-D 在更宽的学习率范围内都能保持较低的损失，展示了更好的调参鲁棒性。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| GPT-2 (124M) 预训练 50K 步 | 验证损失 (Validation Loss) | NAMO: 3.0351, NAMO-D: 3.0246 | AdamW: 3.0643, Muon: 3.0435 | NAMO-D 相比 Muon 降低 0.0189，相比 AdamW 降低 0.0397 |
| GPT-2 (355M) 预训练 10K 步 | 验证损失 (Validation Loss) | NAMO: 2.9516, NAMO-D: 2.9507 | AdamW: 2.9914, Muon: 2.9684 | NAMO-D 相比 Muon 降低 0.0177，相比 AdamW 降低 0.0407 |
局限与改进
论文存在几个明显的局限性。首先，实验仅在 GPT-2 的两个规模（124M 和 355M）上进行，远未达到当前大语言模型的规模（7B、13B 甚至更大），因此无法判断这些改进在更大规模模型上是否能保持。其次，实验仅在 OpenWebText 数据集上进行，缺乏在多语言、代码、数学等多样化任务上的验证。第三，NAMO-D 引入了额外的超参数 $c$（截断参数），虽然论文展示了其影响，但最优值在不同模型规模间变化较大（124M 用 $c=0.1$，355M 用 $c=0.9$），这可能限制其在更大规模上的实用性。第四，论文使用 Newton-Schulz 迭代近似正交化，但未详细分析近似精度对性能的影响。第五，与 PRISM 和 DeVA 的直接对比缺失，使得难以全面评估 NAMO 在同类方法中的相对优势。最后，理论分析假设精确正交化，与实践中的近似正交化存在差距。
独立分析的弱点
独立分析来看，NAMO 和 NAMO-D 存在以下弱点。第一，计算开销分析不够充分——虽然论文声称 NAMO 的额外成本仅为 $O(mn)$，但正交化本身的计算成本（Newton-Schulz 迭代）是主要瓶颈，NAMO-D 的列范数计算和对角矩阵构建也需要额外开销，论文未提供完整的 wall-clock 时间对比。第二，NAMO-D 的截断机制虽然理论上保证了条件数有界，但最优截断参数 $c$ 在不同模型规模间差异巨大，缺乏自适应选择策略。第三，论文的收敛理论假设梯度的 Lipschitz 连续性和有界方差，这些假设在深层网络中可能过于理想化。改进方向包括：开发自适应截断机制以减少 NAMO-D 的超参数调优需求；在更大的模型和更多任务上进行系统性评估；分析近似正交化精度与收敛性的关系；与更多同类方法（如 PRISM、DeVA）进行直接对比。
未来方向
论文作者提出了三个未来研究方向：在更大的大语言模型上评估 NAMO 和 NAMO-D、开发 NAMO-D 的轻量调参变体、以及进一步研究噪声自适应缩放对正交化更新的理论和实践优势。基于本文的成果，还可以延伸以下方向：第一，将 NAMO 的范数自适应机制扩展到其他结构化更新方向（如 Kronecker 分解的预条件器），可能在保持计算效率的同时进一步提升性能。第二，探索 NAMO-D 的截断参数 $c$ 的自适应选择策略，例如基于梯度统计量动态调整 $c$。第三，将 NAMO 的框架与学习率调度策略（如 cosine annealing、WSD）结合，研究自适应缩放与调度的交互效应。第四，将 NAMO 应用于其他模态的大型模型（如视觉 Transformer、多模态模型），验证其泛化性。第五，研究 NAMO 在微调和强化学习场景中的表现，这些场景的噪声特性与预训练不同。
复现评估
论文的复现条件较好。代码已开源在 GitHub（https://github.com/minxin-zhg/namo），这大大降低了复现门槛。实验基于 nanoGPT 实现，这是一个广泛使用的 GPT-2 训练代码库。数据集 OpenWebText 是公开可用的。实验在 4× NVIDIA H100 GPU 上进行，这对大多数研究机构来说是可获得的硬件资源。论文提供了详细的超参数设置（学习率、动量系数、权重衰减等）和训练配置（批大小、梯度累积步数、上下文长度等）。然而，复现的难度主要在于：(1) NAMO-D 的最优截断参数 $c$ 需要针对不同模型规模进行网格搜索；(2) 论文未提供完整的训练脚本和配置文件；(3) 在不同硬件上的性能可能有差异。总体而言，复现难度中等，有经验的研究者应该能够在相似硬件上复现主要结果。
论文图表