当前位置：首页 > news >正文

Muon优化器与Stiefel流形：正交约束的闭式更新与工程实践

news 2026/8/30 16:34:14

正交约束在深度学习里一直是个“既重要又麻烦”的话题。一方面，很多模型希望权重保持正交性，用来缓解梯度消失/爆炸、增强表示稳定性；另一方面，正交化过程往往需要额外计算，比如经典的 Newton-Schulz 迭代或者 QR 分解。最近关于 Muon 优化器的讨论，让“正交化”又一次成为热点：Muon 在训练大模型时表现出色，因为它会在每个 step 里把动量矩阵拉回到正交方向上。更让人感兴趣的是，当权重本身落在 Stiefel 流形上时，Muon 的正交更新可以写成精确的闭式解，不需要迭代也能严格保持约束。

本文将围绕这个结论展开，先把 Stiefel 流形、Muon 更新规则、Cayley 变换这些基础概念讲清楚，再从数学上推一推“为什么存在闭式更新”，最后给出 PyTorch 示例和工程建议。读完你可以理解 Muon 与正交流形之间的关系，也能在代码里实现一套比 Newton-Schulz 更干净的更新方式。

1. 背景：为什么 Muon 和 Stiefel 流形会被放在一起讨论

1.1 正交约束在模型中的价值

正交性不是新概念。在 RNN 里，循环权重如果接近正交，可以避免梯度在时间维上指数级放大或衰减；在 Transformer 里，也有不少工作会对注意力矩阵或前馈层施加正交正则化，用 Lipschitz 约束提升鲁棒性；在自监督学习和子空间方法里，我们干脆希望网络学到的特征落在某个低维正交子空间里。

于是“把权重约束成正交矩阵”变成一个常见的训练需求。严格来说，若一个矩阵 (W \in \mathbb{R}^{n \times p}) 满足：

[ W^T W = I_p ]

则 (W) 的列是单位正交向量组。这种矩阵的集合叫作 Stiefel 流形，记作：

[ \mathrm{St}(n,p)={W\in\mathbb{R}^{n\times p} : W^T W=I_p} ]

当 (p=1) 时，它就是单位球面；当 (n=p) 时，它其实是正交群 (O(n))。所以 Stiefel 流形可以看作“列正交矩阵”的泛化空间。

问题在于：普通梯度下降得到的新权重几乎不可能还满足 (W^T W=I)，所以我们必须想办法把每次更新后的矩阵拉回流形上。这个“拉回”操作，在流形优化里叫作 retraction。

1.2 Muon 优化器的出现

Muon 是最近比较受关注的一种优化器。它的核心思想不是从头开始做流形优化，而是保持一个动量矩阵 (M)，然后把 (M) 正交化后作用到当前权重上：

[ W_{t+1} = O_t W_t ]

其中 (O_t) 是一个正交矩阵，通常由 (M_t) 经过某些正交化步骤得到。为什么这样做有效？一种直觉解释是：如果每一步都用正交矩阵左乘权重，那么权重矩阵的行内积结构会保持不变。如果初始权重接近正交，那么训练过程中它也能一直保持一种“稳定的几何结构”，这对信号传播和梯度回传都比较友好。

传统实现里，(O_t) 的计算要靠 Newton-Schulz 迭代或者 QR 分解。QR 分解本身是精确的，但每次 step 都做一次分解，计算开销不小；Newton-Schulz 则是迭代逼近，迭代次数不够时，(O_t) 并不是严格正交矩阵。

1.3 一个关键观察

如果权重 (W_t) 本身落在 Stiefel 流形上，并且我们使用的更新方向能够保持正交性，那么实际上存在一个精确的闭式更新，不需要迭代。这个闭式更新和 Cayley 变换密切相关。把这个结论展开来推一遍，就能明白 Muon 在 Stiefel 流形上为什么可以“甩掉” Newton-Schulz。

2. Stiefel 流形：定义、切空间与黎曼梯度

2.1 流形直觉

对初学者来说，“流形”听起来很抽象，其实可以把它理解成一个带约束的曲面。单位球面就是一个流形：你可以在球面上移动，但移动后的点必须仍然在球面上。Stiefel 流形就是“列正交矩阵”组成的曲面，表面上每个点都是一个满足 (W^T W=I) 的矩阵。

在这个曲面上做优化，我们需要回答三个问题：

从当前点出发，哪些方向是合法的移动方向？
梯度在约束下应该如何修正？
修正后的点如何落回流形上？

这三个问题分别对应切空间、黎曼梯度和 retraction。

2.2 切空间

设 (W \in \mathrm{St}(n,p))，切空间是满足如下条件的矩阵 (V) 的集合：

[ W^T V + V^T W = 0 ]

这个条件其实就是对约束 (W^T W=I) 求导得到的。如果写成矩阵分块形式，任意切向量 (V) 都可以分解为：

[ V = W A + W_{\perp} B ]

其中 (A) 是 (p \times p) 的反对称矩阵，(B) 是任意矩阵，(W_{\perp}) 是 (W) 的正交补。这个分解非常重要：(A) 描述的是“在已有列张成空间内的旋转”，(B) 描述的是“向外扩张的分量”。

2.3 黎曼梯度

假设在欧几里得空间中，损失函数对 (W) 的梯度是 (G=\nabla f(W))。直接拿 (G) 做更新，显然会破坏正交约束。在标准内积下，把 (G) 投影到切空间，得到黎曼梯度：

[ \operatorname{grad}_{\mathrm{St}} f

G - W \operatorname{sym}(W^T G) ]

其中：

[ \operatorname{sym}(A)=\frac{A+A^T}{2} ]

另一种等价写法是：

[ \operatorname{grad}_{\mathrm{St}} f

(I - W W^T) G + W \operatorname{skew}(W^T G) ]

其中：

[ \operatorname{skew}(A)=\frac{A-A^T}{2} ]

这个形式把梯度分成两部分：

((I - W W^T)G)：垂直于 (W) 列空间的外部分量；
(W \operatorname{skew}(W^T G))：在 (W) 列空间内的反对称旋转分量。

实际实现时，我们通常用第一种投影公式，因为它只需要矩阵乘法和转置，不涉及 (W_{\perp}) 的显式构造。

3. Muon 的经典更新流程

3.1 Muon 的基本伪代码

Muon 的经典步骤可以简化如下：

计算当前梯度 (G_t)。
更新动量： [ M_t = \beta M_{t-1} + (1-\beta) G_t ]
将 (M_t) 正交化，得到 (O_t)。
更新权重： [ W_{t+1} = O_t W_t ]

其中第 3 步是关键。早期实现常用 QR 分解：

import torch def ortho_qr(X): # 把 X 分解为 X = QR，取 Q 作为正交矩阵 Q, _ = torch.linalg.qr(X) return Q

也可以用 Newton-Schulz 迭代逼近：

def ortho_newton_schulz(X, iters=5): # 让 X 更快地接近正交矩阵 # 这里仅示意，实际系数需要根据矩阵范数调整 X = X / (torch.norm(X) + 1e-8) I = torch.eye(X.shape[-1], dtype=X.dtype, device=X.device) for _ in range(iters): X = X @ (3 * I - X.T @ X) / 2 return X

Newton-Schulz 的好处是不用显式分解，在 GPU 上比较友好；缺点是迭代步数有限时，结果并不是严格正交。

3.2 Muon 为什么有效

Muon 的一个出发点是：很多神经网络层的权重矩阵并不需要在整个一般线性空间里自由移动。若权重在正交矩阵族附近活动，那么用正交矩阵左乘来更新，既保留了权重的“尺度”，又能提供稳定的旋转方向。

从优化角度看，Muon 相当于给梯度动量做了一个“正交投影”，这个投影保留的是改变的“方向”，而把许多随机的缩放信息压缩掉了。因此在训练深层 Transformer 时，Muon 有时比 AdamW 收敛更快，尤其在大规模并行训练中表现出色。

3.3 Muon 与流形优化的关系

如果你把 Muon 看作“先算动量，再投影到正交矩阵附近”，会发现它天然适合定义在 Stiefel 流形上的优化问题。因为 (W_{t+1}=O_tW_t) 这个更新，只要 (O_t) 是正交矩阵，那么：

[ W_{t+1}^T W_{t+1}

W_t^T O_t^T O_t W_t

W_t^T W_t

I ]

也就是说，(W_{t+1}) 仍然在 Stiefel 流形上。所以问题的核心变成了：怎样从动量 (M_t) 构造出一个正交矩阵 (O_t)，最好还是精确闭式解。

4. 精确闭式更新：Cayley 变换与极分解

4.1 李代数视角

正交矩阵的切空间由反对称矩阵构成。对任意 (p \times p) 反对称矩阵 (A)，有：

[ A^T = -A ]

它对应的正交矩阵可以通过 Cayley 变换得到：

[ Q(A)

\left(I - \frac{\eta}{2} A\right)^{-1} \left(I + \frac{\eta}{2} A\right) ]

其中 (\eta) 是步长参数。可以验证：(Q(A)^T Q(A)=I)。因为反对称矩阵的转置与其相反，Cayley 变换天然把反对称矩阵映射到正交群。

这个公式最大优势是精确且闭式。它只需要一次矩阵求逆和两次矩阵乘法，不需要迭代，也不需要 SVD。

4.2 对 Stiefel 流形上的权重使用 Cayley 更新

考虑一个 (p \times p) 的反对称矩阵 (A)，我们可以用当前权重 (W) 和动量 (M) 构造它。一个常见做法是取动量在“列空间内旋转分量”的反对称部分：

[ A = W^T M - M^T W ]

由于 (A^T=-A)，它天然反对称。然后用 Cayley 变换得到正交矩阵：

[ Q = \left(I - \frac{\eta}{2} A\right)^{-1} \left(I + \frac{\eta}{2} A\right) ]

最后更新：

[ W_{t+1} = W_t Q ]

因为 (Q) 是正交矩阵，所以 (W_{t+1}) 依旧落在 Stiefel 流形上。这种更新的优点非常明显：它只处理列空间内部的旋转，计算量小，并且数学上是精确正交，而不是近似正交。

4.3 完整闭式更新：极分解收缩

上面的 Cayley 更新只使用动量在 (W) 列空间内的反对称旋转分量，忽略了垂直分量 ((I-WW^T)M)。如果要保留完整梯度信息，可以使用极分解定义 retraction：

[ W_{t+1}

(W_t + \eta V_t) \left( I + \eta^2 V_t^T V_t \right)^{-1/2} ]

其中 (V_t) 是动量的黎曼梯度：

[ V_t

M_t - W_t \operatorname{sym}(W_t^T M_t) ]

这个公式同样在数学上是闭式解。它可以看作对“欧几里得点 (W_t+\eta V_t)”做一次到 Stiefel 流形的极分解投影。计算时通常先对矩阵做 SVD，然后取 (U V^T) 作为投影结果：

def polar_retraction(W, V, eta=1.0): Z = W + eta * V U, _, Vt = torch.linalg.svd(Z, full_matrices=False) return U @ Vt

极分解的优点是有完整梯度信息，对非方阵也能工作；缺点是 SVD 的计算成本比 Cayley 高。Cayley 更新更像是一个“快速且精确”的替代方案，适合那些希望把正交约束嵌入优化器且不太关心垂直分量的场景。

4.4 两种闭式更新的关系

用一句话总结：极分解是“投影到流形上”的闭式精确解，Cayley 变换是“沿流形内旋转方向移动”的闭式精确解。两者都回答了一个问题：更新后如何保持 (W^T W=I)？

经典 Muon 的 Newton-Schulz 迭代，本质上是在逼近某种正交化投影。而 Stiefel 流形上的闭式更新告诉我们：如果我们的问题本身就限定在正交性空间里，那就不需要迭代逼近，直接使用 Cayley 或极分解即可。

5. 实战代码：从优化器到玩具实验

5.1 PyTorch 中的 Cayley 工具函数

我们先把 Cayley 变换实现为通用函数。它输入一个任意矩阵，内部自动取反对称部分，再返回正交矩阵。

import torch def cayley_orthogonal(A: torch.Tensor, eta: float = 1.0) -> torch.Tensor: """ 输入 A: (..., p, p) 张量 返回 Q: (..., p, p) 正交矩阵 """ A_skew = A - A.transpose(-1, -2) A_skew = A_skew / 2.0 I = torch.eye(A_skew.shape[-1], dtype=A_skew.dtype, device=A_skew.device) # Q = (I + eta / 2 * A) * inv(I - eta / 2 * A) left = I + (eta / 2.0) * A_skew right = I - (eta / 2.0) * A_skew return torch.linalg.solve(right, left)

这里有两点需要注意：

我们对输入矩阵主动做了一次反对称化，即使调用时传入非反对称矩阵，结果也会保持正交。
torch.linalg.solve对批量矩阵也有效，所以可以直接处理多个矩阵同时更新的情况。

5.2 Stiefel Muon 优化器示例

接下来实现一个简单的 Stiefel-Muon 优化器。为了更直观，我把权重限制为一个列正交矩阵 (W \in \mathbb{R}^{n \times p})。优化器维护一个与 (W) 同尺寸的动量 (M)，并在每个 step 里使用闭式更新。

class StiefelMuon: def __init__(self, W, lr=0.01, momentum=0.9): self.W = W self.lr = lr self.beta = momentum self.M = torch.zeros_like(W) def step(self, grad): # 1. 动量更新 self.M = self.beta * self.M + (1 - self.beta) * grad # 2. 计算黎曼梯度 sym = (self.W.T @ self.M + self.M.T @ self.W) / 2.0 V = self.M - self.W @ sym # 3. 使用极分解闭式更新，包含完整梯度信息 Z = self.W + self.lr * V U, _, Vt = torch.linalg.svd(Z, full_matrices=False) self.W = U @ Vt return self.W

如果你更希望使用 Cayley 更新，可以这样写：

def step_cayley(self, grad): self.M = self.beta * self.M + (1 - self.beta) * grad A = self.W.T @ self.M - self.M.T @ self.W # p x p 反对称 Q = cayley_orthogonal(A, eta=self.lr) self.W = self.W @ Q return self.W

Cayley 版本只允许权重在现有列空间内旋转，训练时相当于在 Stiefel 流形上做“原地旋转”。如果问题本身只需要旋转分量，它的计算效率最高。

5.3 一个最小验证：单位球面上的瑞利商

为了检验更新是否真的保持约束，我们来做一个最简单的实验：在单位球面上优化瑞利商。

[ \min_{w \in \mathbb{R}^d, |w|=1} -w^T A w ]

这里 (A) 是一个对称矩阵，最优解是最大特征值对应的特征向量。单位球面其实是 (\mathrm{St}(d,1)) 的特例。我们用一个极分解闭式更新来验证：

torch.manual_seed(0) d = 8 A = torch.randn(d, d) A = (A + A.T) / 2 w = torch.randn(d) w = w / torch.norm(w) M = torch.zeros(d) lr = 0.5 beta = 0.9 for step in range(200): grad = -2 * A @ w M = beta * M + (1 - beta) * grad # 球面上的黎曼梯度 V = M - (w @ M) * w # 闭式更新：归一化到单位球面 w = w + lr * V w = w / torch.norm(w) if step % 20 == 0: loss = -w @ A @ w norm_error = torch.norm(torch.dot(w, w) - 1.0).item() print(f"step={step:3d} loss={loss:.4f} norm_err={norm_error:.2e}")

训练过程中你会发现 (w^T w=1) 始终成立，误差只来自浮点数舍入。

5.4 记录正交误差

实践中最简单的检查方法是每个 step 都计算：

def orth_error(W): I = torch.eye(W.shape[1], dtype=W.dtype, device=W.device) return torch.norm(W.T @ W - I).item()

如果这个误差在 (1e-5) 量级，说明约束保持得很好；如果误差逐步变大，就要检查更新逻辑或数值稳定性。

6. 常见问题与排查思路

在实际使用中，闭式更新并不是“写上就能跑”的，很多细节会影响效果。

问题现象	常见原因	解决思路
正交误差增大	学习率太大，更新后矩阵离流形太远	降低学习率，或者改用更精细的闭式投影
Cayley 求逆报错	矩阵接近奇异，通常是学习率过大或权重退化	增加单位矩阵扰动，或者检查权重是否满秩
训练不收敛	动量尺度不合适，导致更新方向被放大	归一化梯度，或者对动量做层归一化
SVD 太慢	权重矩阵尺寸过大	使用 Cayley 更新只做列内旋转，减少计算量
模型初始权重不是正交矩阵	初始化时没有做约束	用 QR 分解或 SVD 初始化权重