正交约束在深度学习里一直是个“既重要又麻烦”的话题。一方面，很多模型希望权重保持正交性，用来缓解梯度消失/爆炸、增强表示稳定性；另一方面，正交化过程往往需要额外计算，比如经典的 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 XNewton-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.WCayley 版本只允许权重在现有列空间内旋转，训练时相当于在 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 初始化权重 |
6.1 为什么正交误差仍然存在
理论上极分解和 Cayley 都是精确保持正交的，但浮点运算会引入微小的误差。如果 (W) 的维度较大，误差会缓慢累积。常用做法是每隔若干步强制做一次 QR 或 SVD 修正，或者直接使用更高精度计算。
6.2 Cayley 更新是否丢失梯度信息
是的，Cayley 更新只保留列空间内的旋转分量。如果模型期望权重“扩张”到新的子空间，只做旋转可能不够。这时可以使用极分解更新，它会保留垂直分量，让权重能够离开当前列空间。
6.3 动量参数怎么选
Muon 的动量系数通常可以用 0.9 到 0.99。过大会导致收敛迟缓，过小会让梯度噪声变大。在流形优化中，建议同时配合学习率调度，尤其是 Cayley 更新的平方根求逆对步长比较敏感。
6.4 用 bfloat16 训练时的注意事项
CUDA 上的 bfloat16 矩阵求逆可能不稳定。如果使用 Cayley 更新，建议在 float32 下完成求逆，再转回 bfloat16。SVD 分解也有类似问题，最好在更高精度下执行。
7. 工程建议与最佳实践
7.1 先确定“你是否真的需要正交流形”
并不是所有场景都需要把权重严格限制在 Stiefel 流形上。普通全连接层用 AdamW 也能训练得很好。正交约束通常用于子空间学习、正则化约束、稳定的循环结构等特定需求。如果只是跟风使用 Muon，却没有一个清晰的目标，很容易陷入调参泥潭。
7.2 区分“正交化衰减”和“正交硬约束”
有些实现会在损失函数里加入正则项：
[ \mathcal{L}_{reg} = \lambda |W^T W - I|^2 ]
这是“软约束”，权重不会严格正交。Stiefel 流形上的闭式更新则是“硬约束”，每个 step 后都严格满足 (W^T W=I)。两者各有适用场景：
- 软约束适合训练早期，对优化器改动小；
- 硬约束适合对权重正交性要求极高的任务，比如某些低秩子空间模型。
7.3 推荐的做法
在综合实验时，我比较推荐下面这一套组合：
- 对需要正交约束的层，用 Stiefel-Muon；
- 对其他层，继续使用 AdamW；
- 权重初始化时用 QR 或 SVD 保证初始点一定在流形上；
- 每 N 个 step 校验一次正交误差；
- 把学习率和动量分成两组，避免不同层使用同一套超参；
- 生产环境里加上正交误差监控，一旦超过阈值就自动切换到惩罚项更新。
7.4 性能建议
Stiefel-Muon 的主要计算开销会出现在 SVD 或求逆上。如果矩阵较大，优先选择 Cayley 更新，因为它只涉及矩阵乘法和一次 (p \times p) 求逆。对于 (n \times p) 的权重，只要 (p) 远小于 (n)，Cayley 的计算量会明显小于极分解。
另一方面，牛顿-舒尔茨迭代虽然每轮只有矩阵乘法，但迭代次数一多，成本并不低。闭式更新的最大优势不是“省掉所有计算”，而是“不需要为了精确度多次迭代”。在 GPU 上，一次矩阵求逆比多次迭代更容易做并行和梯度裁剪。
8. 进一步思考
Muon 在 Stiefel 流形上的闭式更新揭示了一个更一般的思想：优化器和几何结构可以深度绑定。传统的优化器只关注梯度方向和学习率，而流形优化关注的是“每一步更新后，参数仍然合法”。当模型有强约束时，闭式 retraction 往往比惩罚项和迭代投影更稳健。
如果你对 orthogonal optimization 感兴趣，下一步可以学这几个方向：
- 李群与李代数：理解 Cayley 变换、指数映射、对数映射；
- 黎曼流形上的随机梯度下降：把
torch.optim.SGD改成流形版本； - 分布式训练中的正交化：如何在大规模并行中减少通信又能保持正交性；
- Muon 与 Transformer 的结合：在真实 NLP 任务里对比 AdamW 与 Muon 的收敛曲线。
流形优化不是银弹，但它为“带约束的深度学习训练”提供了一条非常优雅的路径。与其在每次 step 后被动地把权重拉回流形，不如直接用几何上正确的更新，让每一步都天然落在约束空间里。这种思路无论对新一代优化器设计，还是对模型结构的理解，都很有帮助。
如果你在自己的实验中用到了 Stiefel-Muon 或 Cayley 更新，欢迎收藏本文备用。实际运行中遇到数值问题，也可以按第 6 节的表格逐项排查。希望这篇内容能帮你在正交优化这条路上少踩一些坑。
