上篇文章介绍了Muon等新兴深度学习优化器背后的原理，即约束参数矩阵的诱导范数下得到新的更新方向。
在Muon对参数更新方向$-\boldsymbol{U}\boldsymbol{V}^\top$的计算中用到了Newton-Schulz迭代方法，本质上是在寻找这样一个多项式函数
$$ f(x)=ax+bx^3+cx^5+\ldots $$
使其满足对任意$x\in(0, 1]$，对$x$应用多次$f(\cdot)$，都能收敛到1附近。这里我们尝试设计一个能work的参数组合。
我的一个简单的想法是，设计一个多项式函数，使$x=1$是它的一个吸引不动点：
定义1（不动点）当
$x_0$被函数$f(\cdot)$映射到自身，即$f(x_0)=x_0$时，称$x_0$是函数$f(\cdot)$的一个不动点。
定义2（吸引不动点）
$f$的吸引不动点是$f$的不动点$x_0$使得，对在足够接近$x_0$的定义域中的任何$x$值而言，迭代函数序列$x,f(x),f(f(x)),f(f(f(x))),\ldots$收敛于$x_0$。
要令$x=1$是$f(x)$的一个吸引不动点，要满足如下的必要条件：
$f(1)=1$$|f'(1)|<1$
使用这两个条件是无法确定具体的参数值$a,b,\ldots$的，但是对于三阶（参数包括$a,b$两个）或者五阶（参数包括$a,b,c$三个）的Newton-Schulz迭代，可以大大缩小搜索的空间。下面展开看下。
三阶迭代
先讨论三阶迭代的形式
$$ f(x)=ax+bx^3 $$
代入上面的两个必要条件：
$$ \begin{split} f(1)=a+b = 1\\ -1 < f'(1)=a+3b < 1\\ \end{split} $$
根据第一个条件，可以把$b$用$1-a$重参数化，然后就有可行的条件
$$ 1 < a < 2 $$
我们记五次迭代后的函数$\phi(x)=f(f(f(f(f(x)))))$，可视化看一下不同$a$取值下对应的情况（理想情况下，对于$(0,1]$区间内的$x$，曲线要尽可能接近$y=1$）
注意到在$a$接近1的时候，$\phi(x)$收敛到1附近的邻域是比较窄的，随着$a\to 2$，收敛到1附近的「邻域」范围逐渐拓宽，但在$a=2$附近，曲线开始出现一定的抖动。对于优化器而言，这样的局部近似的方差是可以容忍的，因此我们可以选取一个比较接近2的值作为$a$的参数，例如$a=1.99,b=-0.99$。
作为对比，在Bernstein & Newhouse 2024.中，作者给出的参数是$a=3/2,b=-1/2$。可以在下图中对照两种设定下的$\phi(x)$.
可以看到Bernstein给出的参数虽然更平滑地收敛于1，但是对于在0附近的初始$x$，普遍无法收敛到1。也就是说对于较小的奇异值对应的$\boldsymbol{u}_i, \boldsymbol{v}_i$，倾向于在更新中被忽略。
在$x=0$附近$\phi(x)$能否快速接近1，主要取决于参数$a$的大小，这是因为$\phi'(0)=a^5$。所以应该在尽可能保证$\phi(x)\approx 1,\forall x\in(0,1]$的同时，让$a$尽可能大。
五阶迭代
现在来考虑五阶迭代的形式
$$ f(x)=ax+bx^3+cx^5 $$
代入上面的两个必要条件：
$$ \begin{split} f(1)=a+b+c = 1\\ -1 < f'(1)=a+3b+5c < 1\\ \end{split} $$
这里的参数比条件多，所以需要给更多的假设（e.g., 固定一个参数）。
将$c=1-a-b$代入第2个条件，得到:
$$ -1 < 5-4a-2b < 1\\ $$
得到
$$ 2-2a<b<3-2a $$
不妨固定$b=2.99-2a$，则$c=a-1.99$，这样我们可以变化$a$的值看下$\phi(x)$曲线长什么样：
类似三阶迭代中的情况，随着$a$变大，收敛到1附近的$x$范围逐步扩大，但是$a$超过一定阈值的时候，曲线就发散了。这是因为当$a$过大时，f(x)在$x=1$右侧的增长率越来越大，则$f(x)$会在迭代中落到超出$(0,1]$区间特别远的地方。
我们可以选取一个互联网人的幸运数字$a=3.25$，这样$b=3.51,c=1.26$，对比Keller Jordan给出的参数$(3.4445, -4.7750, 2.0315)$如下：
可以看到我们按照简单的规则选取的参数在快速收敛到1这个性质上和Jordan的版本差不多，方差更小（但是需要注意在$x=0$附近这里的增长率是小于Jordan的版本的，因为这个增长率与$a$正相关）。
根据笔者的实验，微调一下学习率的情况下，这组参数跟Jordan的参数收敛性能做到基本一致。
参考：