Muon 优化器是 **Keller Jordan** 提出的一个优化器，全称是MomentUm Orthogonalized by Newton-schulz，这是一个用于隐藏层的优化器。Momentum 是指动量法，就是用来优化梯度下降的那个动量法。Orthogonalized 是正交的意思，可看出跟矩阵的正交化有关。Newton-schulz 是指牛顿-舒尔茨迭代法。
Muon 的作者 **Keller Jordan** 在博文《Muon: An optimizer for hidden layers in neural networks》中提出了这一优化器，凭着 muon 优化器，他进入了 OpenAI 工作。有趣的是，当时他参加了 OpenAI 和 xAI 的面试，xAI 的面试官认为他的想法都是错的然后拒绝了他。
随着月之暗面的 Kimi2 的发布，证明了 muon 优化器在万亿参数模型上的有效性，今天我们来学习一下这个优化器。
首先来看一下 Muon 优化器的更新规则
$$ \begin{align} M_t &= \beta M_{t-1} + G_t \\ W_t &= W_{t-1} - \eta \left[ \mathrm{msign}(M_t) + \lambda W_{t-1} \right] \end{align}
$$
$M$是动量，和常见的 Momentum 方法一样。公式2中的$\lambda W_{t-1}$是weight decay，可以先忽略。这里唯一陌生的是 msign 方法。
msign是矩阵符号函数，它是sign函数的矩阵化推广。首先来看sign函数，如下：
$$ \operatorname{sign}(x) = \begin{cases} 1 & \text{if } x > 0 \\ 0 & \text{if } x = 0 \\ -1 & \text{if } x < 0 \end{cases} $$
sign 函数是一个将实数压缩到其符号上的函数，用于表示方向。
设有矩阵$M \in \mathbb{R}^{n \times m}$，$U \in \mathbb{R}^{n \times n}，\Sigma \in \mathbb{R}^{n \times m}，V \in \mathbb{R}^{m \times m}$，$r$ 是$M$的秩，根据奇异值分解有
$$ U, \Sigma, V^\top = \mathrm{SVD}(M) \;\Rightarrow\;\mathrm{msign}(M) = U_{[:,\, :r]} V^\top_{[:,\, :r]} $$
也就是说将sign函数作用到对角矩阵$\Sigma$上，因为对角矩阵的奇异值是非负数，所以sign计算后都为1。
可以这么理解，sign 函数对实数的处理是，保留方向，舍弃“幅度”，msign 同样如此。在几何意义上，U 和 V 代表旋转变换，对角阵$\Sigma$是伸缩变换，旋转不变，使用 sign 函数处理伸缩变换。
把 $U$ 和 $V$ 写成外积和展开式来理解可能更容易，$\sigma_{i}$代表“幅度”，$u_{i}v_{i}^\top$代表“方向”：
$$ \begin{align}\mathrm{SVD}(M) & = U\Sigma V^\top \\ & = \left [ u_{1},u_{2},\dots,u_{n} \right ] \Sigma \left [ v_{1},v_{2},\dots,v_{m} \right ]^\top \\& = \left [ u_{1},u_{2},\dots,u_{n} \right ] \left[\begin{array}{cccccc}\sigma_{1} & & & & & \\& \ddots & & & & \\& & \sigma_{r} & & & \\& & & 0 & & \\& & & & \ddots & \\& & & & & 0\end{array}\right]\begin{bmatrix} v_{1}^\top \\ v_{2}^\top \\ \vdots \\ v_{m}^\top\end{bmatrix} \\& = \sigma_{1}u_{1}v_{1}^\top + \dots + \sigma_{r}u_{r}v_{r}^\top \end{align} $$
用 sign 函数对代表“幅度”的$\sigma_{i}$处理有：
$$ \operatorname{msign}(M) = \sum_{i=1}^r \operatorname{sign}(\sigma_i)\, \mathbf{u}_i \mathbf{v}_i^\top
$$