发表机构
The University of Tokyo; RIKEN（东京大学; 理化学研究所）
机构由 AI 辅助整理，请以论文原文为准。AI 中文总结
该研究针对大型语言模型预训练优化器Muon，证明带有限牛顿-舒尔茨迭代的Muon可通过平滑极映射提升非光滑非凸优化的收敛性，其样本复杂度达最优，且优于采用精确极因子更新的Muon。
AI 中文摘要
Muon已成为大型语言模型预训练中针对矩阵值参数的强大优化器，它通过数次牛顿-舒尔茨（Newton-Schulz）迭代近似正交化其动量。现有理论要么将该迭代替换为它所近似的精确极因子，要么将其有限深度视为近似误差，因此Muon实际运行的迭代只会损害理论保证。我们证明有限牛顿-舒尔茨迭代反而对非光滑非凸优化有益。为此，我们通过“在线到非凸转换”分析Muon，该转换将更新规则视为在线学习者，并将其遗憾边界转化为平稳性保证。有限牛顿-舒尔茨迭代将不连续的极映射平滑为奇异值的利普希茨映射，带有限牛顿-舒尔茨迭代的Muon可被视为具有平滑谱势的在线学习者。这种平滑正是转换所需的：我们证明，牛顿-舒尔茨深度仅随目标精度对数增长，即可实现非光滑非凸优化中平稳点的收敛，而采用精确极因子更新的Muon可能无法收敛。所得样本复杂度边界与非光滑非凸优化的最佳已知保证匹配，且对于光滑非凸优化，在依赖问题的因子范围内是最优的。该论证可从牛顿-舒尔茨扩展到具有相同平滑特性的一般谱映射。
英文摘要
Muon has emerged as a strong optimizer for the matrix-valued parameters in large language model pretraining, approximately orthogonalizing its momentum with a few Newton-Schulz iterations. Existing theory either replaces this iteration with the exact polar factor it approximates, or treats its finite depth as an approximation error, and thus the iteration Muon actually runs can only hurt the guarantees. We show that finite Newton-Schulz can instead be beneficial for nonsmooth nonconvex optimization. To this end, we analyze Muon through the online-to-nonconvex conversion, which views the update rule as an online learner and converts its regret bound into a stationarity guarantee. The finite Newton-Schulz iteration smooths the discontinuous polar map into a Lipschitz map of the singular values, and Muon with finite Newton-Schulz can be regarded as an online learner with a smoothed spectral potential. This smoothing is exactly what the conversion needs: we prove that a Newton-Schulz depth growing only logarithmically in the target accuracy suffices for convergence to stationary points in nonsmooth nonconvex optimization, whereas Muon with the exact-polar update may fail to converge. The resulting sample complexity bounds match the best-known guarantees for nonsmooth nonconvex optimization and are optimal for smooth nonconvex optimization up to problem-dependent factors. The argument extends beyond Newton-Schulz to general spectral maps with the same smoothing property.
Comments37 pages, 3 figures
