NAMO与NAMO-D：正交动量与自适应优化的融合算法

1. 正交动量与自适应优化的融合：NAMO与NAMO-D算法解析

在深度学习优化算法的演进历程中，我们见证了两种截然不同但各具优势的技术路线。一方面，以Adam为代表的坐标自适应方法通过调整每个参数的更新步长，显著提升了训练的稳定性和收敛速度；另一方面，Muon等正交化动量方法则充分利用神经网络权重矩阵的结构特性，通过正交化操作加速收敛过程。这两种技术路线看似平行发展，但本文介绍的NAMO（Norm-based Adaptive Moment estimation with Orthogonalized momentum）和NAMO-D（Diagonal extension of NAMO）算法，首次实现了二者的有机融合。

1.1 算法设计背景与核心思想

现代大规模学习中的随机优化面临一个关键挑战：如何在噪声环境下平衡两个看似矛盾的目标——选择在无噪声情况下表现良好的更新方向，同时又能适应随机扰动的影响。传统Adam优化器通过坐标自适应的步长机制（即对每个参数单独调整步长）来应对梯度噪声，但其更新方向可能不利于长期收敛；而Muon优化器虽然通过正交化动量获得了结构上的优势，却缺乏对噪声的自适应能力。

NAMO系列算法的核心创新点在于：

• 正交动量保留：继承Muon的正交化动量更新，保持矩阵参数的结构特性
• 自适应噪声调节：引入Adam式的范数自适应机制，根据梯度噪声水平动态调整步长
• 灵活扩展性：NAMO-D进一步通过右乘对角矩阵实现神经元级的细粒度自适应

这种融合不是简单的功能叠加，而是基于对优化问题本质的深刻理解。从数学角度看，正交化可以将矩阵更新的方向和大小解耦（通过范数对偶性表征），这使得基于范数的正交化更新重新缩放成为一个自然的设计选择。

1.2 算法实现细节

NAMO基础算法

NAMO的核心更新规则可表示为：
Θₜ = Θₜ₋₁ - ηαₜOₜ

其中关键组件包括：

1. 正交化动量Oₜ：通过牛顿-舒尔茨迭代近似计算Mₜ的正交因子
2. 自适应缩放因子αₜ：αₜ = (1-μ₂ᵗ)^(1/2)/(1-μ₁ᵗ) * ||Mₜ||_F/√(vₜ + ϵ)
3. 二阶矩估计vₜ：vₜ = μ₂vₜ₋₁ + (1-μ₂)||Gₜ||_F²

这种设计保留了Muon的正交化方向特性，同时通过单个自适应步长调整更新幅度。当随机梯度噪声较大或接近驻点时，αₜ会自动减小，从而保证训练稳定性。

NAMO-D扩展算法

NAMO-D在NAMO基础上引入了更精细的缩放策略：
Θₜ = Θₜ₋₁ - ηOₜDₜ

其中Dₜ是一个对角矩阵，其对角线元素通过以下方式计算：

1. 对每列（对应单个神经元）计算自适应步长：dₜ = Nc(M̂ₜ)⊘(√v̂ₜ + ϵₜ)
2. 计算平均值d̄ₜ = ||dₜ||₁/n
3. 应用钳位操作：d̃ₜ = min(max(dₜ, c d̄ₜ1), d̄ₜ/c)

这种设计带来了三个关键优势：

1. 神经元级自适应：为每个神经元分配独立的步长，符合神经网络Hessian矩阵常见的近块对角结构
2. 条件数控制：通过钳位参数c保证dₜₘᵢₙ ≥ c²dₜₘₐₓ，维持更新方向良好的条件数
3. 噪声鲁棒性：在保持正交化方向优势的同时，实现对梯度噪声的细粒度适应

实际实现提示：在大型模型训练中，建议使用牛顿-舒尔茨迭代进行近似正交化计算，虽然理论分析假设精确正交化，但实践中3-5次迭代通常就能达到满意精度，且计算开销可接受。

2. 理论收敛性分析

2.1 确定性场景下的收敛保证

在精确梯度（无噪声）的情况下，NAMO和NAMO-D都达到了光滑非凸优化的一阶方法的最优收敛速率。具体而言：

定理1（NAMO确定性收敛）：
在满足Lipschitz光滑假设（假设1）下，当选择η=O(T^(-1/2))，μ₁=Θ(1)，μ₂=Θ(1)时，NAMO满足：
1/T ∑_{t=1}^T ||∇L(Θₜ₋₁)||_F ≤ O(T^(-1/2))

定理2（NAMO-D确定性收敛）：
在相同假设下，当选择η=O(T^(-1/2))，ϵ=O(T^(-1/2)n^(-1))，c=Θ(1)时，NAMO-D也达到相同的收敛速率。

这些结果表明，正交化更新与自适应步长的结合不会损害算法的理论收敛性能。证明的关键在于：

1. 通过正交化下降不等式关联梯度范数与更新量
2. 利用偏置校正后的矩估计的凸组合表示
3. 控制自适应缩放因子的上下界

2.2 随机场景下的噪声自适应

在有噪声的随机梯度环境下，两种算法展现出更强的适应性：

定理3（NAMO随机收敛）：
在假设1-2下，当η=O(T^(-3/4))，1-μ₁=Θ(T^(-1/2))，1-μ₂=Θ(T^(-1/2))时：
1/T ∑_{t=1}^T E[||∇L(Θₜ₋₁)||_F] ≤ O(T^(-1/4) + σb^(-1/4)T^(-1/8))

定理4（NAMO-D随机收敛）：
在相同参数设置下，NAMO-D也达到相同阶数的收敛速率。

这些结果揭示了几个重要特性：

1. 噪声自适应：收敛界自然地分解为优化项O(T^(-1/4))和显式的方差相关项O(σb^(-1/4)T^(-1/8))
2. 批量大小影响：当b=Ω(σ²√T)时，方差项被主导，恢复最优O(T^(-1/4))速率
3. 结构保持：即使引入噪声，算法仍能保持正交化更新的结构优势

收敛分析的技术路线统一依赖于：

1. 偏置校正矩估计的凸组合表示
2. 将误差分解为噪声项和漂移项分别控制
3. 使用正交化下降不等式关联各种量
4. 通过Cauchy-Schwarz等不等式最终导出梯度范数界

3. GPT-2预训练实验验证

3.1 实验设置与基准比较

我们在GPT-2模型预训练任务上验证NAMO系列算法的实际表现，对比基线包括：

• AdamW：目前LLM训练的事实标准，使用坐标自适应步长
• Muon：纯正交化动量方法，缺乏噪声自适应机制

实验配置要点：

• 模型规模：测试124M和355M两种参数量的GPT-2
• 数据集：OpenWebText（约90亿训练token）
• 硬件：4×NVIDIA H100 GPU
• 超参数：通过网格搜索为每个优化器选择最优学习率

3.2 结果分析与讨论

学习率鲁棒性测试（124M模型）

图1展示了各优化器在不同学习率下的表现：

• NAMO和NAMO-D在更广的学习率范围内保持较低的训练/验证损失
• 传统优化器对学习率选择更敏感，尤其在高学习率时性能下降明显
• 这表明自适应机制确实提升了算法的调参鲁棒性

长期训练表现（124M模型，50K步）

使用最优学习率进行延长训练后（图2）：

• NAMO-D最终训练损失2.9167，显著优于Muon的3.0265和AdamW的3.0456
• 验证损失也呈现相同趋势，说明改进非过拟合所致
• NAMO-D相对NAMO的优势随训练进行逐渐显现，表明细粒度自适应在后期更重要

大规模模型验证（355M模型）

表2中355M模型的结果显示：

• NAMO-D在最优钳位参数c=0.9时表现最佳（训练损失2.9351）
• 性能排序为：NAMO-D > NAMO > Muon > AdamW
• c的选择需要平衡方向保持与噪声适应，过大或过小都会损害性能

关键发现总结

1. 结构优势：正交化更新方向确实能加速收敛，验证了Muon的设计理念
2. 自适应价值：NAMO相比Muon的改进证明噪声适应机制的必要性
3. 细粒度收益：NAMO-D的进一步优势说明神经元级自适应与神经网络结构匹配
4. 超参数影响：钳位参数c在大型模型中作用更显著，需要仔细调整

4. 实践指导与扩展思考

4.1 实际应用建议

基于理论分析和实验结果，我们给出以下实践建议：

学习率选择：

• NAMO通常可比Muon使用更大学习率（实验中约大一个数量级）
• 初始尝试范围：NAMO 5e-3到1e-2，NAMO-D稍保守些
• 使用线性预热（2000步左右）有助于稳定初期训练

动量参数设置：

• 推荐μ₁=0.95，μ₂=0.99作为起点
• 更小的1-μ₂值（即μ₂更接近1）会使自适应更"保守"
• 在非常嘈杂的任务中可适当增大1-μ₂

NAMO-D钳位参数：

• 小模型（<1亿参数）：c=0.1左右足够
• 大模型：需要更精细调整，0.5-0.9可能更合适
• 过小的c会限制自适应能力，过大则可能损害方向保持

计算效率考量：

• 相比原始Muon，NAMO仅增加O(mn)计算量（范数计算）
• NAMO-D因对角矩阵计算增加O(n)开销
• 实际测得的额外时间开销<5%，内存占用几乎无增加

4.2 潜在改进方向

虽然NAMO系列表现出色，但仍有若干值得探索的方向：

自适应钳位策略：
当前c为固定超参数，可研究根据训练动态自动调整c的策略，如：

• 基于梯度噪声水平的自适应
• 分层设置不同c值
• 随训练进程衰减或增强

混合精度训练：
正交化操作对数值精度较敏感，需要特别关注：

• 牛顿-舒尔茨迭代在混合精度下的稳定性
• 自适应缩放因子的数值范围控制
• 梯度裁剪与正交化的交互影响

理论扩展：

• 非光滑或非Lipschitz场景下的收敛性
• 更一般的矩阵分解方式（如低秩正交化）
• 与其他自适应方法（如LAMB）的结合可能性

正交动量与自适应优化的融合为深度学习优化算法设计开辟了新思路。NAMO和NAMO-D的成功表明，充分挖掘参数结构特性并与噪声适应机制有机结合，确实能带来实质性的性能提升。这一范式也可能启发更多结合几何特性与统计适应的新型优化器设计。