AuON：正交动量更新的一个线性时间替代方案
机器学习
2025-12-19 v4 机器学习
摘要
正交动量梯度更新已涌现出来，旨�克服基于向量的优化器如Adam的局限性。基于向量的优化器Adam存在高内存成本和不良条件动量梯度更新的问题。然而，传统的正交动量方法，如SVD/QR分解，因计算和内存成本高，性能远不如经过精调的带动量SGD。最近的进展，如Muon，通过在正交化之前应用动量并通过Newton-Schulz迭代近似正交矩阵，从而实现更好的GPU利用率、活跃的高TFLOPS，并将内存使用降低最高可达3倍。然而，Muon(Vanilla)存在注意力logit爆炸问题，并具有三次方计算复杂度。本文深入研究正交动量梯度更新，以发现帮助Muon实现卓越性能的主要属性。我们提出AuON（Alternative Unit-norm momentum updates by Normalized nonlinear scaling），一种无需近似正交矩阵即可实现强大性能的线性时间优化器，同时保持结构对齐并重构不良条件的更新。AuON内置自动的“紧急刹车”以处理注意力logit的爆炸。我们进一步引入了混合变体Hybrid-AuON，该变体应用线性变换与Newton-Schulz迭代，在语言建模任务中超越了Muon。代码已公开：https://github.com/ryyzn9/AuON
关键词
引用
@article{arxiv.2509.24320,
title = {AuON: A Linear-time Alternative to Orthogonal Momentum Updates},
author = {Dipan Maity},
journal= {arXiv preprint arXiv:2509.24320},
year = {2025}
}