提出线性时间优化器AuON,解决传统正交动量方法计算慢、内存高问题。
AuON: A Linear-time Alternative to Orthogonal Momentum Updates
- 通过非线性归一化实现单位范数动量更新,避免近似正交矩阵
- 计算复杂度降至线性,内存使用减少3倍,训练更稳定
- 适合大模型训练,尤其在语言建模任务中超越Muon
正交动量梯度更新被提出以克服向量型优化器(如Adam)的局限性:高内存开销和病态条件下的动量更新。传统正交方法(如SVD/QR分解)虽能改善性能,但计算与内存开销高,且表现不如调优后的SGD+动量。近期的Muon通过在正交化前应用动量,并用Newton-Schulz迭代近似正交矩阵,提升了GPU利用率,达到高TFLOPS,内存降低达3倍。然而,Muon(原版)存在注意力日志爆炸问题,且计算复杂度为立方级。本文深入分析正交动量更新的核心特性,提出AuON(基于归一化非线性缩放的替代单位范数动量更新),一种线性时间优化器,无需近似正交矩阵,同时保持结构对齐与病态更新重调节能力。AuON内置自动“紧急制动”机制,防止注意力日志爆炸。进一步提出混合变体Hybrid-AuON,结合Newton-Schulz线性变换,在语言建模任务中优于Muon。代码已公开于https://github.com/ryyzn9/AuON。
原文摘要 · Abstract (English)
Orthogonal momentum gradient updates have emerged to overcome the limitations of vector-based optimizers like Adam. The vector-based optimizer Adam suffers from high memory costs and ill-conditioned momentum gradient updates. However, traditional Orthogonal momentum approaches, such as SVD/QR decomposition, suffer from high computational and memory costs and underperform compared to well-tuned SGD with momentum. Recent advances, such as Muon, improve efficiency by applying momentum before orthogonalization and approximate orthogonal matrices via Newton-Schulz iterations, which gives better GPU utilization, active high TFLOPS, and reduces memory usage by up to 3x. Nevertheless, Muon(Vanilla) suffers from exploding attention logits and has cubic computation complexity. In this paper, we deep dive into orthogonal momentum gradient updates to find the main properties that help Muon achieve remarkable performance. We propose AuON (Alternative Unit-norm momentum updates by Normalized nonlinear scaling), a linear-time optimizer that achieves strong performance without approximate orthogonal matrices, while preserving structural alignment and reconditioning ill-posed updates. AuON has an automatic "emergency brake" to handle exploding attention logits. We further introduce a hybrid variant, Hybrid-AuON, that applies the linear transformations with Newton-Schulz iterations, which outperforms Muon in the language modeling tasks. Code is available at: https://github.com/ryyzn9/AuON
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。