Muon with Finite Newton-Schulz: The Smoothing Benefit in Nonsmooth Nonconvex Optimization

Muon 与有限牛顿-舒尔茨迭代：非光滑非凸优化中的平滑增益

Abstract: 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.

摘要： Muon 已成为大语言模型预训练中矩阵值参数的一种强力优化器，它通过少量牛顿-舒尔茨（Newton-Schulz）迭代对其动量进行近似正交化。现有的理论要么用其近似的精确极因子（exact polar factor）来代替这种迭代，要么将其有限深度视为一种近似误差，因此 Muon 实际运行的迭代过程在理论上反而会削弱其收敛保证。

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.

我们证明，有限牛顿-舒尔茨迭代实际上对非光滑非凸优化是有益的。为此，我们通过“在线到非凸”（online-to-nonconvex）转换来分析 Muon，该方法将更新规则视为在线学习器，并将其遗憾界（regret bound）转换为平稳性保证。有限牛顿-舒尔茨迭代将不连续的极映射平滑为奇异值的 Lipschitz 映射，因此带有有限牛顿-舒尔茨迭代的 Muon 可以被视为具有平滑谱势（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.

这种平滑处理正是该转换过程所需要的：我们证明，牛顿-舒尔茨迭代深度只需随目标精度呈对数增长，就足以在非光滑非凸优化中收敛到平稳点，而使用精确极更新的 Muon 则可能无法收敛。由此得出的样本复杂度界限与目前非光滑非凸优化领域已知的最佳保证相匹配，并且在考虑与问题相关的因子后，对于平滑非凸优化也是最优的。该论证不仅适用于牛顿-舒尔茨迭代，还可扩展至具有相同平滑特性的通用谱映射。