正交约束在深度学习里一直是个“既重要又麻烦”的话题。一方面很多模型希望权重保持正交性用来缓解梯度消失/爆炸、增强表示稳定性另一方面正交化过程往往需要额外计算比如经典的 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 WI_p} ]当 (p1) 时它就是单位球面当 (np) 时它其实是正交群 (O(n))。所以 Stiefel 流形可以看作“列正交矩阵”的泛化空间。问题在于普通梯度下降得到的新权重几乎不可能还满足 (W^T WI)所以我们必须想办法把每次更新后的矩阵拉回流形上。这个“拉回”操作在流形优化里叫作 retraction。1.2 Muon 优化器的出现Muon 是最近比较受关注的一种优化器。它的核心思想不是从头开始做流形优化而是保持一个动量矩阵 (M)然后把 (M) 正交化后作用到当前权重上[ W_{t1} 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 WI) 的矩阵。在这个曲面上做优化我们需要回答三个问题从当前点出发哪些方向是合法的移动方向梯度在约束下应该如何修正修正后的点如何落回流形上这三个问题分别对应切空间、黎曼梯度和 retraction。2.2 切空间设 (W \in \mathrm{St}(n,p))切空间是满足如下条件的矩阵 (V) 的集合[ W^T V V^T W 0 ]这个条件其实就是对约束 (W^T WI) 求导得到的。如果写成矩阵分块形式任意切向量 (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}} fG - W \operatorname{sym}(W^T G) ]其中[ \operatorname{sym}(A)\frac{AA^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_{t1} 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, iters5): # 让 X 更快地接近正交矩阵 # 这里仅示意实际系数需要根据矩阵范数调整 X X / (torch.norm(X) 1e-8) I torch.eye(X.shape[-1], dtypeX.dtype, deviceX.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_{t1}O_tW_t) 这个更新只要 (O_t) 是正交矩阵那么[ W_{t1}^T W_{t1}W_t^T O_t^T O_t W_tW_t^T W_tI ]也就是说(W_{t1}) 仍然在 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_{t1} W_t Q ]因为 (Q) 是正交矩阵所以 (W_{t1}) 依旧落在 Stiefel 流形上。这种更新的优点非常明显它只处理列空间内部的旋转计算量小并且数学上是精确正交而不是近似正交。4.3 完整闭式更新极分解收缩上面的 Cayley 更新只使用动量在 (W) 列空间内的反对称旋转分量忽略了垂直分量 ((I-WW^T)M)。如果要保留完整梯度信息可以使用极分解定义 retraction[ W_{t1}(W_t \eta V_t) \left( I \eta^2 V_t^T V_t \right)^{-1/2} ]其中 (V_t) 是动量的黎曼梯度[ V_tM_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, eta1.0): Z W eta * V U, _, Vt torch.linalg.svd(Z, full_matricesFalse) return U Vt极分解的优点是有完整梯度信息对非方阵也能工作缺点是 SVD 的计算成本比 Cayley 高。Cayley 更新更像是一个“快速且精确”的替代方案适合那些希望把正交约束嵌入优化器且不太关心垂直分量的场景。4.4 两种闭式更新的关系用一句话总结极分解是“投影到流形上”的闭式精确解Cayley 变换是“沿流形内旋转方向移动”的闭式精确解。两者都回答了一个问题更新后如何保持 (W^T WI)经典 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], dtypeA_skew.dtype, deviceA_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, lr0.01, momentum0.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_matricesFalse) 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, etaself.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(fstep{step:3d} loss{loss:.4f} norm_err{norm_error:.2e})训练过程中你会发现 (w^T w1) 始终成立误差只来自浮点数舍入。5.4 记录正交误差实践中最简单的检查方法是每个 step 都计算def orth_error(W): I torch.eye(W.shape[1], dtypeW.dtype, deviceW.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 WI)。两者各有适用场景软约束适合训练早期对优化器改动小硬约束适合对权重正交性要求极高的任务比如某些低秩子空间模型。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 节的表格逐项排查。希望这篇内容能帮你在正交优化这条路上少踩一些坑。 Muon优化器与Stiefel流形：正交约束的闭式更新与工程实践
正交约束在深度学习里一直是个“既重要又麻烦”的话题。一方面很多模型希望权重保持正交性用来缓解梯度消失/爆炸、增强表示稳定性另一方面正交化过程往往需要额外计算比如经典的 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 WI_p} ]当 (p1) 时它就是单位球面当 (np) 时它其实是正交群 (O(n))。所以 Stiefel 流形可以看作“列正交矩阵”的泛化空间。问题在于普通梯度下降得到的新权重几乎不可能还满足 (W^T WI)所以我们必须想办法把每次更新后的矩阵拉回流形上。这个“拉回”操作在流形优化里叫作 retraction。1.2 Muon 优化器的出现Muon 是最近比较受关注的一种优化器。它的核心思想不是从头开始做流形优化而是保持一个动量矩阵 (M)然后把 (M) 正交化后作用到当前权重上[ W_{t1} 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 WI) 的矩阵。在这个曲面上做优化我们需要回答三个问题从当前点出发哪些方向是合法的移动方向梯度在约束下应该如何修正修正后的点如何落回流形上这三个问题分别对应切空间、黎曼梯度和 retraction。2.2 切空间设 (W \in \mathrm{St}(n,p))切空间是满足如下条件的矩阵 (V) 的集合[ W^T V V^T W 0 ]这个条件其实就是对约束 (W^T WI) 求导得到的。如果写成矩阵分块形式任意切向量 (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}} fG - W \operatorname{sym}(W^T G) ]其中[ \operatorname{sym}(A)\frac{AA^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_{t1} 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, iters5): # 让 X 更快地接近正交矩阵 # 这里仅示意实际系数需要根据矩阵范数调整 X X / (torch.norm(X) 1e-8) I torch.eye(X.shape[-1], dtypeX.dtype, deviceX.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_{t1}O_tW_t) 这个更新只要 (O_t) 是正交矩阵那么[ W_{t1}^T W_{t1}W_t^T O_t^T O_t W_tW_t^T W_tI ]也就是说(W_{t1}) 仍然在 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_{t1} W_t Q ]因为 (Q) 是正交矩阵所以 (W_{t1}) 依旧落在 Stiefel 流形上。这种更新的优点非常明显它只处理列空间内部的旋转计算量小并且数学上是精确正交而不是近似正交。4.3 完整闭式更新极分解收缩上面的 Cayley 更新只使用动量在 (W) 列空间内的反对称旋转分量忽略了垂直分量 ((I-WW^T)M)。如果要保留完整梯度信息可以使用极分解定义 retraction[ W_{t1}(W_t \eta V_t) \left( I \eta^2 V_t^T V_t \right)^{-1/2} ]其中 (V_t) 是动量的黎曼梯度[ V_tM_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, eta1.0): Z W eta * V U, _, Vt torch.linalg.svd(Z, full_matricesFalse) return U Vt极分解的优点是有完整梯度信息对非方阵也能工作缺点是 SVD 的计算成本比 Cayley 高。Cayley 更新更像是一个“快速且精确”的替代方案适合那些希望把正交约束嵌入优化器且不太关心垂直分量的场景。4.4 两种闭式更新的关系用一句话总结极分解是“投影到流形上”的闭式精确解Cayley 变换是“沿流形内旋转方向移动”的闭式精确解。两者都回答了一个问题更新后如何保持 (W^T WI)经典 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], dtypeA_skew.dtype, deviceA_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, lr0.01, momentum0.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_matricesFalse) 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, etaself.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(fstep{step:3d} loss{loss:.4f} norm_err{norm_error:.2e})训练过程中你会发现 (w^T w1) 始终成立误差只来自浮点数舍入。5.4 记录正交误差实践中最简单的检查方法是每个 step 都计算def orth_error(W): I torch.eye(W.shape[1], dtypeW.dtype, deviceW.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 WI)。两者各有适用场景软约束适合训练早期对优化器改动小硬约束适合对权重正交性要求极高的任务比如某些低秩子空间模型。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 节的表格逐项排查。希望这篇内容能帮你在正交优化这条路上少踩一些坑。