在深度学习的训练过程中,优化算法的选择直接影响模型的收敛速度和最终性能。传统的自适应优化器(如Adam)通过动量估计和逐参数学习率调整提升了训练效率,而正交化更新方向(如Muon)则通过约束梯度方向的正交性改善了优化轨迹的稳定性。NAMO(Noise-Adaptive Momentum with Orthogonalization)创新性地将这两种思路结合起来,提出了一个理论完备的优化框架。
NAMO的核心创新在于将正交化更新方向与基于范数的自适应矩估计相结合。具体来说,它通过以下两个关键机制实现:
这种设计使得NAMO在保持Muon正交性优势的同时,能够自适应不同噪声环境,在GPT-2等大规模语言模型预训练中展现出优于AdamW和Muon的性能。
技术细节:NAMO的正交化更新通过计算梯度的极分解实现,即对于梯度矩阵G,其正交部分O=UV^T(其中G=UΣV^T是G的SVD分解)。这种分解保证了更新方向的稳定性。
NAMO-D是NAMO的扩展版本,通过引入对角矩阵实现更细粒度的噪声适应:
实验表明,NAMO-D在GPT-2(355M)模型上相比NAMO有进一步性能提升,特别是在训练初期收敛速度方面。
在满足标准光滑性和有界方差噪声假设的条件下,NAMO在确定性优化场景中可以达到最优收敛速率:
证明中特别处理了自适应步长与正交化更新之间的耦合效应,这是传统分析中未考虑的难点。
当使用随机梯度时,NAMO展现出噪声自适应特性:
这种特性使NAMO特别适合大规模分布式训练,其中梯度噪声水平可能随worker数量变化。
虽然NAMO引入了额外计算(正交化+自适应缩放),但通过以下设计保持高效:
实测表明,NAMO相比Muon仅增加约5-8%的计算开销,远低于其带来的收敛加速收益。
基于论文实验,推荐以下配置:
| 超参数 | 推荐值 | 作用说明 |
|---|---|---|
| μ1 (动量) | 0.9-0.99 | 控制历史梯度信息的衰减率 |
| μ2 (二阶矩) | 0.999 | 控制梯度平方的衰减率 |
| 初始步长η | 1e-4~5e-4 | 需要根据模型规模调整 |
| clamping参数c | 0.1-0.3 | 仅NAMO-D需要,控制对角矩阵范围 |
实际经验:在LLM训练中,建议初始使用较小η,并在训练中期适当提高,以平衡收敛速度与稳定性。
在OpenWebText语料上的对比实验显示:
| 优化器 | 参数量 | 最终perplexity | 达到基准的步数 |
|---|---|---|---|
| AdamW | 124M | 18.7 | 100% (基准) |
| Muon | 124M | 18.3 (-2.1%) | 85% |
| NAMO | 124M | 17.9 (-4.3%) | 70% |
| NAMO-D | 124M | 17.6 (-5.9%) | 65% |
更大规模的355M参数模型上,NAMO-D相比AdamW节省约30%的训练计算量。
通过控制变量实验验证各组件贡献:
特别地,NAMO-D的神经元级适应在深层transformer中效果显著,可能因为不同层的梯度统计特性差异较大。
NAMO系列特别适合以下场景:
对于计算资源严格受限的场景,可能仍需权衡其额外开销。
训练初期震荡:
收敛后期停滞:
内存占用过高:
NAMO的技术思路可延伸至:
我在实际项目中发现,将NAMO-D应用于图像生成任务时,适当调整c值可以显著改善细节生成质量,这可能与不同尺度特征的噪声特性差异有关。