We analyze Muon as originally proposed and used in practice -- using the momentum orthogonalization with a few Newton-Schulz steps. The prior theoretical results replace this key step in Muon with an exact SVD-based polar factor. We prove that Muon with Newton-Schulz converges to a stationary point at the same rate as the SVD-polar idealization, up to a constant factor for a given number $q$ of Newton-Schulz steps. We further analyze this constant factor and prove that it converges to 1 doubly exponentially in $q$ and improves with the degree of the polynomial used in Newton-Schulz for approximating the orthogonalization direction. We also prove that Muon removes the typical square-root-of-rank loss compared to its vector-based counterpart, SGD with momentum. Our results explain why Muon with a few low-degree Newton-Schulz steps matches exact-polar (SVD) behavior at a much faster wall-clock time and explain how much momentum matrix orthogonalization via Newton-Schulz benefits over the vector-based optimizer. Overall, our theory justifies the practical Newton-Schulz design of Muon, narrowing its practice-theory gap.


翻译:我们分析了Muon算法最初提出及实际应用时的形式——即采用动量正交化并结合若干步Newton-Schulz迭代。先前的理论研究成果将Muon中的这一关键步骤替换为基于精确奇异值分解的极因子计算。我们证明了采用Newton-Schulz迭代的Muon算法能以与理想化SVD极分解版本相同的收敛速率达到稳定点,其收敛常数仅取决于给定的Newton-Schulz迭代步数$q$。我们进一步分析了该常数因子,证明其以$q$的双指数速率收敛于1,且收敛性能随着Newton-Schulz算法中用于逼近正交化方向的多项式阶数提高而改善。我们还证明了相较于其基于向量的对应方法——带动量的随机梯度下降,Muon能够消除典型的秩平方根损失。我们的研究结果解释了为何采用少量低阶Newton-Schulz迭代的Muon算法能在更短的壁钟时间内达到与精确极分解相同的效果,并阐明了通过Newton-Schulz实现动量矩阵正交化相较于基于向量的优化器所具有的优势。总体而言,我们的理论为Muon采用Newton-Schulz算法的实际设计提供了依据,从而缩小了其实践与理论之间的差距。

0
下载
关闭预览

相关内容

在数学和计算机科学之中,算法(Algorithm)为一个计算的具体步骤,常用于计算、数据处理和自动推理。精确而言,算法是一个表示为有限长列表的有效方法。算法应包含清晰定义的指令用于计算函数。 来自维基百科: 算法
【新书】机器学习算法深度剖析
专知会员服务
75+阅读 · 2024年7月28日
《资源分配博弈中的收敛率》
专知会员服务
41+阅读 · 2023年3月10日
专知会员服务
37+阅读 · 2021年9月12日
伊利诺伊《算法》书籍,集20年之大成,附472页pdf
专知会员服务
65+阅读 · 2020年9月27日
【AAAI2021】对比聚类,Contrastive Clustering
专知
26+阅读 · 2021年1月30日
机器学习的可解释性:因果推理和稳定学习
DataFunTalk
13+阅读 · 2020年3月3日
机器学习中的最优化算法总结
人工智能前沿讲习班
22+阅读 · 2019年3月22日
数据分析师应该知道的16种回归方法:泊松回归
数萃大数据
35+阅读 · 2018年9月13日
最新|深度离散哈希算法,可用于图像检索!
全球人工智能
14+阅读 · 2017年12月15日
国家自然科学基金
6+阅读 · 2015年12月31日
国家自然科学基金
0+阅读 · 2014年12月31日
国家自然科学基金
3+阅读 · 2014年12月31日
国家自然科学基金
0+阅读 · 2014年12月31日
国家自然科学基金
8+阅读 · 2014年12月31日
国家自然科学基金
0+阅读 · 2014年12月31日
国家自然科学基金
0+阅读 · 2014年12月31日
国家自然科学基金
1+阅读 · 2014年12月31日
国家自然科学基金
0+阅读 · 2014年12月31日
国家自然科学基金
2+阅读 · 2014年12月31日
Arxiv
0+阅读 · 2月12日
VIP会员
相关基金
国家自然科学基金
6+阅读 · 2015年12月31日
国家自然科学基金
0+阅读 · 2014年12月31日
国家自然科学基金
3+阅读 · 2014年12月31日
国家自然科学基金
0+阅读 · 2014年12月31日
国家自然科学基金
8+阅读 · 2014年12月31日
国家自然科学基金
0+阅读 · 2014年12月31日
国家自然科学基金
0+阅读 · 2014年12月31日
国家自然科学基金
1+阅读 · 2014年12月31日
国家自然科学基金
0+阅读 · 2014年12月31日
国家自然科学基金
2+阅读 · 2014年12月31日
Top
微信扫码咨询专知VIP会员