这次我们来看一个优化算法层面的结论Muon 优化器在 Stiefel 流形上的更新存在精确的闭式解不再需要依赖 Newton-Schulz 迭代去近似。如果你关心大模型训练、正交权重约束或者想弄懂优化器底层到底在算什么这个结论值得认真拆一遍。Muon 优化器之所以这几年反复被讨论核心在于它更新完权重后会把矩阵拉回“列正交”的约束集合。这个约束集合在数学上就叫 Stiefel 流形。过去工程实现普遍用 Newton-Schulz 迭代做正交化逼近而标题里的结果指出向 Stiefel 流形的最优投影存在一个闭式公式本质是极分解落到数值计算上就是一次 SVD。这篇文章会展开讲四件事Muon 为什么和 Stiefel 流形绑定在一起闭式更新公式如何从投影问题推导出来PyTorch 参考实现怎么写闭式解和 Newton-Schulz 迭代相比有哪些取舍。适合关心优化器原理的算法工程师、研究生以及想在自有训练框架里替换正交化步骤的开发者。1. 核心能力速览能力项说明主题类型优化算法 / 矩阵流形更新理论核心对象Muon 优化器在 Stiefel 流形上的投影更新关键结论存在精确闭式更新等价于极分解或 SVD 的 $U V^T$解决的问题替代 Newton-Schulz 迭代近似消除正交化误差与迭代步数超参数前置知识梯度下降、正交矩阵、奇异值分解实现依赖PyTorch 或 NumPy不需要额外依赖库计算特点单次 SVD复杂度约 $O(n p^2)$具体取决于矩阵形状批量支持PyTorch 的批量 SVD 可直接处理 batch天然支持批量任务硬件要求无特定门槛GPU 上 SVD 性能需要按矩阵尺寸实测适合场景大模型预训练、正交 RNN、流形约束网络、论文复现、教学实验不适合场景超大矩阵且对延迟极度敏感的在线推理场景这个结论本身是数学性质不依赖特定显卡、CUDA 版本或某个深度学习框架。真正需要工程判断的问题是闭式解和 Newton-Schulz 迭代在给定矩阵尺寸下谁更快、谁更稳。这不是一句“闭式解一定全面胜出”能概括的后面会专门对比。2. Muon 优化器为什么需要 Stiefel 流形Muon 优化器原始的更新逻辑可以拆成两步第一步和传统动量法类似用历史梯度信息算出当前权重的更新方向第二步是特征所在更新完成后再对权重矩阵施加正交约束。为什么要加正交约束深层网络里大量参数以矩阵乘法形式参与计算权重矩阵连乘后奇异值的累积会直接放大或缩小信号。极端情况下要么梯度爆炸要么梯度消失。正交矩阵的奇异值全部等于 1对应线性变换不会改变向量的长度是保持信息传递最理想的约束对象。在注意力矩阵、FFN 权重、embedding 这类参与频繁矩阵乘法的参数上施加正交性训练稳定性和收敛行为都有可能改善。正交约束不是 Muon 独有的思路。很多归一化方法、正交正则化、谱归一化也围绕类似动机展开。区别在于 Muon 把正交性作为优化过程的硬约束每次更新后直接投影回满足 $W^T W I$ 的流形而不是像正则项那样“软性拉回”。从社区公开实现来看常见的正交化手段有三种QR 分解、Newton-Schulz 迭代、极分解。QR 分解实现简单但它求出来的是列空间的一组标准正交基并不保证是 Frobenius 范数意义下的最优投影Newton-Schulz 迭代避免了 SVD适合超大矩阵但需要设定迭代次数并且对输入矩阵的谱范数缩放敏感极分解或 SVD 能给出精确的最优投影代价是 SVD 本身的数值开销。标题里的“Exact Closed-Form Update”指的就是第三种思路具备了严格的理论保证。3. Stiefel 流形与正交投影问题的形式化Stiefel 流形是满足列正交条件的所有矩阵构成的集合$$St(n, p) { W \in \mathbb{R}^{n \times p} : W^T W I_p }$$当 $p 1$ 时它退化为 $\mathbb{R}^n$ 中的单位球面当 $p n$ 时它就是正交群 $O(n)$。大模型里常见的权重矩阵比如 $n \times p$ 的线性层参数只要 $n \geq p$ 且要求列正交就属于这个流形。Muon 的每次迭代可以抽象成下面的流程先根据动量与梯度计算一个“原始”更新量得到脱离开流形的矩阵 $A W_k \mu M_k$其中 $M_k$ 是某种动量项$\mu$ 是步长然后把 $A$ 投影回 Stiefel 流形得到下一轮权重$$W_{k1} \operatorname{proj}_{St(n,p)}(A)$$投影问题的严格定义是在所有满足 $Q^T Q I_p$ 的矩阵中找一个与 $A$ 的 Frobenius 距离最近的 $Q$$$\min_{Q \in St(n,p)} | A - Q |_F^2$$这个优化问题的解就是闭式更新的核心。展开 Frobenius 范数平方$$| A - Q |_F^2 |A|_F^2 |Q|_F^2 - 2 \operatorname{tr}(A^T Q)$$因为 $Q$ 是列正交矩阵$|Q|_F^2 p$ 是常数$|A|_F^2$ 与 $Q$ 无关所以最小化问题等价于最大化 $\operatorname{tr}(A^T Q)$。这一步非常关键只要能在 Stiefel 流形上找到让内积迹最大的矩阵就找到了最优投影。接下来需要回答的是这个最大化问题的解到底是什么。4. 闭式更新的推导极分解与 SVD对任意矩阵 $A \in \mathbb{R}^{n \times p}$假设 $n \geq p$ 且列满秩可以做奇异值分解$$A U \Sigma V^T$$其中 $U \in \mathbb{R}^{n \times p}$ 是列正交矩阵$V \in \mathbb{R}^{p \times p}$ 是正交矩阵$\Sigma \operatorname{diag}(\sigma_1, \dots, \sigma_p)$ 包含全部奇异值。现在取 $Q U V^T$。验证一下它是否落在 Stiefel 流形上$$Q^T Q V U^T U V^T V V^T I_p$$所以 $U V^T$ 确实是合法的列正交矩阵。接下来证明它是最优的。根据 von Neumann 迹不等式对任意满足 $Q^T Q I_p$ 的矩阵都有$$\operatorname{tr}(A^T Q) \leq \sum_{i1}^{p} \sigma_i(A)$$等号成立当且仅当 $Q$ 的左右奇异向量与 $A$ 的左右奇异向量对齐。取 $Q U V^T$ 时迹正好等于奇异值之和达到上界。因此闭式解就是$$\operatorname{proj}_{St(n,p)}(A) U V^T$$这也是矩阵极分解的正交因子。$A$ 的极分解写作 $A Q H$其中 $Q$ 是列正交矩阵$H V \Sigma V^T$ 是半正定矩阵。SVD 直接把 $Q$ 解耦出来丢掉 $\Sigma$只保留 $U V^T$。这个结论有两个值得强调的性质。第一它是 Frobenius 范数意义上的全局最优不是迭代近似投影误差在理论上可以精确到浮点精度。第二它对任意列满秩矩阵成立不要求 $A$ 是方阵因此可以覆盖线性层、注意力投影、embedding 等多种形状的权重。还有一个数值层面容易忽略的点当 $A$ 接近秩亏缺时SVD 的奇异值中有若干接近零$U V^T$ 依然有定义但解的稳定性会变差。训练中如果梯度爆炸导致 $A$ 的元素很大最好先做梯度裁剪或缩放再做投影。5. PyTorch 参考实现与 Newton-Schulz 对比5.1 极分解投影实现PyTorch 的torch.linalg.svd天然支持批量计算因此这个闭式投影可以同时处理多个矩阵非常契合批量训练场景import torch def stiefel_projection(A): 将 A 投影到 Stiefel 流形上。 参数: A: 形状为 (..., n, p) 的张量n p 返回: Q: 与 A 同形状的张量满足 Q^T Q ≈ I_p 且 ||A - Q||_F 达到最小。 U, _, Vh torch.linalg.svd(A, full_matricesFalse) return U Vh这里full_matricesFalse返回的是经济型 SVD$U$ 的形状为 $(..., n, p)$$Vh$ 的形状为 $(..., p, p)$乘积 $U Vh$ 保持原矩阵形状。5.2 Muon 更新循环示意把投影函数放进 Muon 风格的更新循环里整体逻辑如下。动量参数、动量系数、学习率只作为示意具体数值以原论文或仓库实现为准def muon_step(W, m, g, lr0.1, beta0.95, nesterov_alpha0.05): Muon 风格的一步更新。 参数: W: 当前权重矩阵 m: 动量缓冲 g: 当前梯度 lr: 学习率 beta: 动量衰减系数 nesterov_alpha: Nesterov 风格的当前梯度权重 返回: W_new: 投影后的新权重 m_new: 更新后的动量缓冲 m_new beta * m g update m_new nesterov_alpha * g A W lr * update W_new stiefel_projection(A) return W_new, m_new注意上面这段代码的重点是演示“更新后投影”这一结构不承担数值调参任务。实际使用中动量系数、学习率和 Nesterov 权重都需要根据任务重新验证。5.3 Newton-Schulz 迭代实现Newton-Schulz 迭代作为对比实现同样很短它通过若干次矩阵乘法逼近正交矩阵def newton_schulz(Q, n_steps5): 通过 Newton-Schulz 迭代将 Q 近似正交化。 迭代不收敛时通常需要先对 Q 做谱范数缩放。 for _ in range(n_steps): Q 1.5 * Q - 0.5 * Q (Q.transpose(-2, -1) Q) return Q这段迭代的数学原理是基于奇异值映射 $f(\sigma) 1.5\sigma - 0.5\sigma^3$。奇异值大于 $\sqrt{3}$ 时映射会不稳定因此使用前需要把矩阵除以谱范数或乘一个缩放系数。这也是 Newton-Schulz 在实际工程里最容易被忽视的坑。5.4 数值验证换完投影函数后第一件事是检查正交误差和投影误差def check_projection(A): Q stiefel_projection(A) orth_error torch.norm(Q.transpose(-2, -1) Q - torch.eye(Q.shape[-1], deviceQ.device)) proj_error torch.norm(A - Q) print(f正交误差: {orth_error.item():.2e}) print(f投影误差: {proj_error.item():.2e})正交误差应该接近浮点精度投影误差就是理论上界 $\sqrt{|A|_F^2 - \sum \sigma_i^2}$ 的数值表达。如果正交误差明显大于 $10^{-5}$多半是 SVD 返回结果异常或者输入矩阵包含 NaN。6. 闭式更新与 Newton-Schulz 的取舍闭式更新和 Newton-Schulz 迭代不是简单的替代关系它们的差异集中在三个维度精度、计算开销、超参数数量。精度上闭式解是全局最优Newton-Schulz 只经过有限步迭代理论上存在残余误差。这个误差不会让优化算法立刻崩溃但会随迭代步数累积最终训练日志里的正交误差指标会持续偏高。闭式解能给出确定性的精确投影这在复现实验时会友好很多。计算开销上SVD 对 $n \times p$ 矩阵的复杂度约为 $O(n p^2)$而 Newton-Schulz 每一步需要两次矩阵乘法复杂度同样在 $O(n p^2)$ 量级但常数因子小得多。5 次 Newton-Schulz 迭代通常比一次 SVD 快这是工程上选择 Newton-Schulz 的主要原因。矩阵尺寸越大SVD 的常数因子劣势越明显尤其是在 GPU 上常规 SVD 实现未必比矩阵乘法高效。超参数数量上闭式解没有新增超参数Newton-Schulz 至少需要设定迭代次数有时还需要谱范数缩放系数。结论是对于尺寸在数百到数千的权重矩阵闭式解完全可行对于上万维的巨型矩阵可以保留 Newton-Schulz或者对两者做实测对比再决定。另一个务实思路是混合使用把尺寸较小的矩阵交给闭式投影尺寸很大的矩阵继续走迭代方案。这样既能拿到精确投影的优势又不会被 SVD 的耗时拖住。7. 资源占用与性能观察方法7.1 复杂度与显存SVD 除了计算时间还会引入额外的显存开销。torch.linalg.svd在经济型分解下会生成 $U$、$S$、$Vh$ 三部分显存占用约为输入的 2 到 3 倍。对于一个 $n \times p$ 的批量张量额外显存大概在 $O(Bnp Bp^2)$ 量级。训练中如果显存紧张可以逐层投影避免一次性投影整个参数列表。7.2 性能对比脚本在自己的机器上做一次简单 benchmark比任何经验数字都可靠。以下脚本对比闭式投影和 5 步 Newton-Schulz 的耗时import time import torch def bench(fn, A, iters20): # warmup for _ in range(3): fn(A) if torch.cuda.is_available(): torch.cuda.synchronize() t0 time.time() for _ in range(iters): fn(A) if torch.cuda.is_available(): torch.cuda.synchronize() return (time.time() - t0) / iters device cuda if torch.cuda.is_available() else cpu A torch.randn(1024, 1024, devicedevice) svd_time bench(stiefel_projection, A) ns_time bench(lambda x: newton_schulz(x, 5), A) print(f闭式解平均耗时: {svd_time * 1000:.2f} ms) print(fNewton-Schulz 平均耗时: {ns_time * 1000:.2f} ms)观察重点是三类现象的平衡投影耗时、正交误差、训练收敛速度。有的场景下投影快但正交误差大训练反而更不稳闭式解耗时略长但省去了调迭代步数的时间总成本可能更低。建议以自己模型的实际训练曲线为准不要只看单次投影耗时的对比。7.3 训练中如何观察正交质量推荐在训练日志里加一项正交误差指标def log_orth_error(W): I torch.eye(W.shape[-1], deviceW.device, dtypeW.dtype) return torch.norm(W.T W - I).item()每个 checkpoint 记录一次即可。如果这个指标从 $10^{-6}$ 量级漂移到 $10^{-2}$ 量级说明投影没有真正生效或者模型结构里存在不受约束的权重副本。8. 常见问题与排查方法问题现象可能原因排查方式解决方案投影后正交误差很大使用了 QR 分解但没注意它不是最优投影计算 $|Q^T Q - I|_F$改用 SVD 极分解方案SVD 结果出现 NaN输入矩阵含 NaN 或梯度爆炸检查输入张量中的有限值加梯度裁剪或先缩放到合理范围投影后的权重与预期不符权重是方阵时误把 $V^T$ 当成 $V$检查 SVD 返回的 Vh 是否转置保持一致统一用 $U Vh$GPU 上 SVD 很慢矩阵尺寸过大或 GPU 上 SVD 内核效率低用 benchmark 脚本记录耗时换 CPU 投影或混合使用 Newton-SchulzNewton-Schulz 发散输入矩阵谱范数过大迭代映射不稳定打印迭代前后谱范数先除以谱范数或用闭式解替代训练不稳定学习率过大投影前更新量过大观察投影前 A 的范数降低学习率加入 warmup批量投影显存溢出一次性投影过大 batch观察峰值显存逐层或分块投影排查思路上记住一个原则先确认输入数据正常再确认投影算法正确最后才怀疑训练超参数。绝大多数异常都能通过打印张量形状、范数和是否有 NaN 快速定位。9. 最佳实践与使用建议在自有框架里引入闭式投影时建议按下面的顺序推进。第一先用小矩阵验证正确性。固定一个随机矩阵分别用 SVD 投影、Newton-Schulz 迭代、PyTorch 官方正交化接口计算投影结果对比正交误差和与原始矩阵的距离。这一步能快速发现符号、转置和 shape 错误。第二建立最小可运行配置。把 Muon 更新循环封装成独立函数输入输出只依赖权重和梯度不依赖全局状态。这样后续切换投影算法只需要替换函数内部实现。第三分目录管理模型文件、训练脚本、日志和检查点。投影算法切换后训练曲线可能变化保留原始实验配置能帮助回溯差异来源。第四对批量任务做好日志和失败重试。批量训练中如果某个 batch 的梯度包含 NaNSVD 会把这个 NaN 扩散到整个矩阵。建议在投影前检查梯度有限性异常时跳过当前 batch 并记录。第五涉及商业模型训练时注意训练数据授权和模型权重的开源协议。优化器本身是数学算法但训练数据和模型发布仍然受数据和代码许可约束发布前需要确认合规边界。第六保持一个端到端可复现的实验脚本。不要只保存代码还要保存随机种子、数据版本、优化器超参数和投影方式。闭式解的优势之一就是参数少这会让复现成本明显降低。10. 总结与下一步这个结论最值得尝试的点是把 Muon 里原本需要调迭代步数的正交化步骤替换成一个无超参数、有最优性保证的 SVD 投影。相比 Newton-Schulz它减少了正交误差也减少了需要调的超参数。第一次尝试时建议先验证两件事一是在小规模矩阵上确认投影误差和正交误差符合理论预期二是把闭式投影接入现有训练循环观察训练损失曲线是否和 Newton-Schulz 版本趋势一致。最容易踩的坑是拿超大矩阵直接跑 SVD耗时无法接受却没有做分块和混合策略。后续可以继续扩展的方向有两个。一个是性能优化调研批量化 SVD、低秩近似投影、或者把闭式解嵌入到更大的训练框架里做端到端 benchmark。另一个是理论延伸把极分解投影推广到带权重衰减或带约束的变体验证在更多模型结构上的稳定性。先把小规模验证跑通再决定是否替换现有生产环境里的正交化实现。
Muon优化器在Stiefel流形上的闭式更新：极分解与SVD的精确解法
这次我们来看一个优化算法层面的结论Muon 优化器在 Stiefel 流形上的更新存在精确的闭式解不再需要依赖 Newton-Schulz 迭代去近似。如果你关心大模型训练、正交权重约束或者想弄懂优化器底层到底在算什么这个结论值得认真拆一遍。Muon 优化器之所以这几年反复被讨论核心在于它更新完权重后会把矩阵拉回“列正交”的约束集合。这个约束集合在数学上就叫 Stiefel 流形。过去工程实现普遍用 Newton-Schulz 迭代做正交化逼近而标题里的结果指出向 Stiefel 流形的最优投影存在一个闭式公式本质是极分解落到数值计算上就是一次 SVD。这篇文章会展开讲四件事Muon 为什么和 Stiefel 流形绑定在一起闭式更新公式如何从投影问题推导出来PyTorch 参考实现怎么写闭式解和 Newton-Schulz 迭代相比有哪些取舍。适合关心优化器原理的算法工程师、研究生以及想在自有训练框架里替换正交化步骤的开发者。1. 核心能力速览能力项说明主题类型优化算法 / 矩阵流形更新理论核心对象Muon 优化器在 Stiefel 流形上的投影更新关键结论存在精确闭式更新等价于极分解或 SVD 的 $U V^T$解决的问题替代 Newton-Schulz 迭代近似消除正交化误差与迭代步数超参数前置知识梯度下降、正交矩阵、奇异值分解实现依赖PyTorch 或 NumPy不需要额外依赖库计算特点单次 SVD复杂度约 $O(n p^2)$具体取决于矩阵形状批量支持PyTorch 的批量 SVD 可直接处理 batch天然支持批量任务硬件要求无特定门槛GPU 上 SVD 性能需要按矩阵尺寸实测适合场景大模型预训练、正交 RNN、流形约束网络、论文复现、教学实验不适合场景超大矩阵且对延迟极度敏感的在线推理场景这个结论本身是数学性质不依赖特定显卡、CUDA 版本或某个深度学习框架。真正需要工程判断的问题是闭式解和 Newton-Schulz 迭代在给定矩阵尺寸下谁更快、谁更稳。这不是一句“闭式解一定全面胜出”能概括的后面会专门对比。2. Muon 优化器为什么需要 Stiefel 流形Muon 优化器原始的更新逻辑可以拆成两步第一步和传统动量法类似用历史梯度信息算出当前权重的更新方向第二步是特征所在更新完成后再对权重矩阵施加正交约束。为什么要加正交约束深层网络里大量参数以矩阵乘法形式参与计算权重矩阵连乘后奇异值的累积会直接放大或缩小信号。极端情况下要么梯度爆炸要么梯度消失。正交矩阵的奇异值全部等于 1对应线性变换不会改变向量的长度是保持信息传递最理想的约束对象。在注意力矩阵、FFN 权重、embedding 这类参与频繁矩阵乘法的参数上施加正交性训练稳定性和收敛行为都有可能改善。正交约束不是 Muon 独有的思路。很多归一化方法、正交正则化、谱归一化也围绕类似动机展开。区别在于 Muon 把正交性作为优化过程的硬约束每次更新后直接投影回满足 $W^T W I$ 的流形而不是像正则项那样“软性拉回”。从社区公开实现来看常见的正交化手段有三种QR 分解、Newton-Schulz 迭代、极分解。QR 分解实现简单但它求出来的是列空间的一组标准正交基并不保证是 Frobenius 范数意义下的最优投影Newton-Schulz 迭代避免了 SVD适合超大矩阵但需要设定迭代次数并且对输入矩阵的谱范数缩放敏感极分解或 SVD 能给出精确的最优投影代价是 SVD 本身的数值开销。标题里的“Exact Closed-Form Update”指的就是第三种思路具备了严格的理论保证。3. Stiefel 流形与正交投影问题的形式化Stiefel 流形是满足列正交条件的所有矩阵构成的集合$$St(n, p) { W \in \mathbb{R}^{n \times p} : W^T W I_p }$$当 $p 1$ 时它退化为 $\mathbb{R}^n$ 中的单位球面当 $p n$ 时它就是正交群 $O(n)$。大模型里常见的权重矩阵比如 $n \times p$ 的线性层参数只要 $n \geq p$ 且要求列正交就属于这个流形。Muon 的每次迭代可以抽象成下面的流程先根据动量与梯度计算一个“原始”更新量得到脱离开流形的矩阵 $A W_k \mu M_k$其中 $M_k$ 是某种动量项$\mu$ 是步长然后把 $A$ 投影回 Stiefel 流形得到下一轮权重$$W_{k1} \operatorname{proj}_{St(n,p)}(A)$$投影问题的严格定义是在所有满足 $Q^T Q I_p$ 的矩阵中找一个与 $A$ 的 Frobenius 距离最近的 $Q$$$\min_{Q \in St(n,p)} | A - Q |_F^2$$这个优化问题的解就是闭式更新的核心。展开 Frobenius 范数平方$$| A - Q |_F^2 |A|_F^2 |Q|_F^2 - 2 \operatorname{tr}(A^T Q)$$因为 $Q$ 是列正交矩阵$|Q|_F^2 p$ 是常数$|A|_F^2$ 与 $Q$ 无关所以最小化问题等价于最大化 $\operatorname{tr}(A^T Q)$。这一步非常关键只要能在 Stiefel 流形上找到让内积迹最大的矩阵就找到了最优投影。接下来需要回答的是这个最大化问题的解到底是什么。4. 闭式更新的推导极分解与 SVD对任意矩阵 $A \in \mathbb{R}^{n \times p}$假设 $n \geq p$ 且列满秩可以做奇异值分解$$A U \Sigma V^T$$其中 $U \in \mathbb{R}^{n \times p}$ 是列正交矩阵$V \in \mathbb{R}^{p \times p}$ 是正交矩阵$\Sigma \operatorname{diag}(\sigma_1, \dots, \sigma_p)$ 包含全部奇异值。现在取 $Q U V^T$。验证一下它是否落在 Stiefel 流形上$$Q^T Q V U^T U V^T V V^T I_p$$所以 $U V^T$ 确实是合法的列正交矩阵。接下来证明它是最优的。根据 von Neumann 迹不等式对任意满足 $Q^T Q I_p$ 的矩阵都有$$\operatorname{tr}(A^T Q) \leq \sum_{i1}^{p} \sigma_i(A)$$等号成立当且仅当 $Q$ 的左右奇异向量与 $A$ 的左右奇异向量对齐。取 $Q U V^T$ 时迹正好等于奇异值之和达到上界。因此闭式解就是$$\operatorname{proj}_{St(n,p)}(A) U V^T$$这也是矩阵极分解的正交因子。$A$ 的极分解写作 $A Q H$其中 $Q$ 是列正交矩阵$H V \Sigma V^T$ 是半正定矩阵。SVD 直接把 $Q$ 解耦出来丢掉 $\Sigma$只保留 $U V^T$。这个结论有两个值得强调的性质。第一它是 Frobenius 范数意义上的全局最优不是迭代近似投影误差在理论上可以精确到浮点精度。第二它对任意列满秩矩阵成立不要求 $A$ 是方阵因此可以覆盖线性层、注意力投影、embedding 等多种形状的权重。还有一个数值层面容易忽略的点当 $A$ 接近秩亏缺时SVD 的奇异值中有若干接近零$U V^T$ 依然有定义但解的稳定性会变差。训练中如果梯度爆炸导致 $A$ 的元素很大最好先做梯度裁剪或缩放再做投影。5. PyTorch 参考实现与 Newton-Schulz 对比5.1 极分解投影实现PyTorch 的torch.linalg.svd天然支持批量计算因此这个闭式投影可以同时处理多个矩阵非常契合批量训练场景import torch def stiefel_projection(A): 将 A 投影到 Stiefel 流形上。 参数: A: 形状为 (..., n, p) 的张量n p 返回: Q: 与 A 同形状的张量满足 Q^T Q ≈ I_p 且 ||A - Q||_F 达到最小。 U, _, Vh torch.linalg.svd(A, full_matricesFalse) return U Vh这里full_matricesFalse返回的是经济型 SVD$U$ 的形状为 $(..., n, p)$$Vh$ 的形状为 $(..., p, p)$乘积 $U Vh$ 保持原矩阵形状。5.2 Muon 更新循环示意把投影函数放进 Muon 风格的更新循环里整体逻辑如下。动量参数、动量系数、学习率只作为示意具体数值以原论文或仓库实现为准def muon_step(W, m, g, lr0.1, beta0.95, nesterov_alpha0.05): Muon 风格的一步更新。 参数: W: 当前权重矩阵 m: 动量缓冲 g: 当前梯度 lr: 学习率 beta: 动量衰减系数 nesterov_alpha: Nesterov 风格的当前梯度权重 返回: W_new: 投影后的新权重 m_new: 更新后的动量缓冲 m_new beta * m g update m_new nesterov_alpha * g A W lr * update W_new stiefel_projection(A) return W_new, m_new注意上面这段代码的重点是演示“更新后投影”这一结构不承担数值调参任务。实际使用中动量系数、学习率和 Nesterov 权重都需要根据任务重新验证。5.3 Newton-Schulz 迭代实现Newton-Schulz 迭代作为对比实现同样很短它通过若干次矩阵乘法逼近正交矩阵def newton_schulz(Q, n_steps5): 通过 Newton-Schulz 迭代将 Q 近似正交化。 迭代不收敛时通常需要先对 Q 做谱范数缩放。 for _ in range(n_steps): Q 1.5 * Q - 0.5 * Q (Q.transpose(-2, -1) Q) return Q这段迭代的数学原理是基于奇异值映射 $f(\sigma) 1.5\sigma - 0.5\sigma^3$。奇异值大于 $\sqrt{3}$ 时映射会不稳定因此使用前需要把矩阵除以谱范数或乘一个缩放系数。这也是 Newton-Schulz 在实际工程里最容易被忽视的坑。5.4 数值验证换完投影函数后第一件事是检查正交误差和投影误差def check_projection(A): Q stiefel_projection(A) orth_error torch.norm(Q.transpose(-2, -1) Q - torch.eye(Q.shape[-1], deviceQ.device)) proj_error torch.norm(A - Q) print(f正交误差: {orth_error.item():.2e}) print(f投影误差: {proj_error.item():.2e})正交误差应该接近浮点精度投影误差就是理论上界 $\sqrt{|A|_F^2 - \sum \sigma_i^2}$ 的数值表达。如果正交误差明显大于 $10^{-5}$多半是 SVD 返回结果异常或者输入矩阵包含 NaN。6. 闭式更新与 Newton-Schulz 的取舍闭式更新和 Newton-Schulz 迭代不是简单的替代关系它们的差异集中在三个维度精度、计算开销、超参数数量。精度上闭式解是全局最优Newton-Schulz 只经过有限步迭代理论上存在残余误差。这个误差不会让优化算法立刻崩溃但会随迭代步数累积最终训练日志里的正交误差指标会持续偏高。闭式解能给出确定性的精确投影这在复现实验时会友好很多。计算开销上SVD 对 $n \times p$ 矩阵的复杂度约为 $O(n p^2)$而 Newton-Schulz 每一步需要两次矩阵乘法复杂度同样在 $O(n p^2)$ 量级但常数因子小得多。5 次 Newton-Schulz 迭代通常比一次 SVD 快这是工程上选择 Newton-Schulz 的主要原因。矩阵尺寸越大SVD 的常数因子劣势越明显尤其是在 GPU 上常规 SVD 实现未必比矩阵乘法高效。超参数数量上闭式解没有新增超参数Newton-Schulz 至少需要设定迭代次数有时还需要谱范数缩放系数。结论是对于尺寸在数百到数千的权重矩阵闭式解完全可行对于上万维的巨型矩阵可以保留 Newton-Schulz或者对两者做实测对比再决定。另一个务实思路是混合使用把尺寸较小的矩阵交给闭式投影尺寸很大的矩阵继续走迭代方案。这样既能拿到精确投影的优势又不会被 SVD 的耗时拖住。7. 资源占用与性能观察方法7.1 复杂度与显存SVD 除了计算时间还会引入额外的显存开销。torch.linalg.svd在经济型分解下会生成 $U$、$S$、$Vh$ 三部分显存占用约为输入的 2 到 3 倍。对于一个 $n \times p$ 的批量张量额外显存大概在 $O(Bnp Bp^2)$ 量级。训练中如果显存紧张可以逐层投影避免一次性投影整个参数列表。7.2 性能对比脚本在自己的机器上做一次简单 benchmark比任何经验数字都可靠。以下脚本对比闭式投影和 5 步 Newton-Schulz 的耗时import time import torch def bench(fn, A, iters20): # warmup for _ in range(3): fn(A) if torch.cuda.is_available(): torch.cuda.synchronize() t0 time.time() for _ in range(iters): fn(A) if torch.cuda.is_available(): torch.cuda.synchronize() return (time.time() - t0) / iters device cuda if torch.cuda.is_available() else cpu A torch.randn(1024, 1024, devicedevice) svd_time bench(stiefel_projection, A) ns_time bench(lambda x: newton_schulz(x, 5), A) print(f闭式解平均耗时: {svd_time * 1000:.2f} ms) print(fNewton-Schulz 平均耗时: {ns_time * 1000:.2f} ms)观察重点是三类现象的平衡投影耗时、正交误差、训练收敛速度。有的场景下投影快但正交误差大训练反而更不稳闭式解耗时略长但省去了调迭代步数的时间总成本可能更低。建议以自己模型的实际训练曲线为准不要只看单次投影耗时的对比。7.3 训练中如何观察正交质量推荐在训练日志里加一项正交误差指标def log_orth_error(W): I torch.eye(W.shape[-1], deviceW.device, dtypeW.dtype) return torch.norm(W.T W - I).item()每个 checkpoint 记录一次即可。如果这个指标从 $10^{-6}$ 量级漂移到 $10^{-2}$ 量级说明投影没有真正生效或者模型结构里存在不受约束的权重副本。8. 常见问题与排查方法问题现象可能原因排查方式解决方案投影后正交误差很大使用了 QR 分解但没注意它不是最优投影计算 $|Q^T Q - I|_F$改用 SVD 极分解方案SVD 结果出现 NaN输入矩阵含 NaN 或梯度爆炸检查输入张量中的有限值加梯度裁剪或先缩放到合理范围投影后的权重与预期不符权重是方阵时误把 $V^T$ 当成 $V$检查 SVD 返回的 Vh 是否转置保持一致统一用 $U Vh$GPU 上 SVD 很慢矩阵尺寸过大或 GPU 上 SVD 内核效率低用 benchmark 脚本记录耗时换 CPU 投影或混合使用 Newton-SchulzNewton-Schulz 发散输入矩阵谱范数过大迭代映射不稳定打印迭代前后谱范数先除以谱范数或用闭式解替代训练不稳定学习率过大投影前更新量过大观察投影前 A 的范数降低学习率加入 warmup批量投影显存溢出一次性投影过大 batch观察峰值显存逐层或分块投影排查思路上记住一个原则先确认输入数据正常再确认投影算法正确最后才怀疑训练超参数。绝大多数异常都能通过打印张量形状、范数和是否有 NaN 快速定位。9. 最佳实践与使用建议在自有框架里引入闭式投影时建议按下面的顺序推进。第一先用小矩阵验证正确性。固定一个随机矩阵分别用 SVD 投影、Newton-Schulz 迭代、PyTorch 官方正交化接口计算投影结果对比正交误差和与原始矩阵的距离。这一步能快速发现符号、转置和 shape 错误。第二建立最小可运行配置。把 Muon 更新循环封装成独立函数输入输出只依赖权重和梯度不依赖全局状态。这样后续切换投影算法只需要替换函数内部实现。第三分目录管理模型文件、训练脚本、日志和检查点。投影算法切换后训练曲线可能变化保留原始实验配置能帮助回溯差异来源。第四对批量任务做好日志和失败重试。批量训练中如果某个 batch 的梯度包含 NaNSVD 会把这个 NaN 扩散到整个矩阵。建议在投影前检查梯度有限性异常时跳过当前 batch 并记录。第五涉及商业模型训练时注意训练数据授权和模型权重的开源协议。优化器本身是数学算法但训练数据和模型发布仍然受数据和代码许可约束发布前需要确认合规边界。第六保持一个端到端可复现的实验脚本。不要只保存代码还要保存随机种子、数据版本、优化器超参数和投影方式。闭式解的优势之一就是参数少这会让复现成本明显降低。10. 总结与下一步这个结论最值得尝试的点是把 Muon 里原本需要调迭代步数的正交化步骤替换成一个无超参数、有最优性保证的 SVD 投影。相比 Newton-Schulz它减少了正交误差也减少了需要调的超参数。第一次尝试时建议先验证两件事一是在小规模矩阵上确认投影误差和正交误差符合理论预期二是把闭式投影接入现有训练循环观察训练损失曲线是否和 Newton-Schulz 版本趋势一致。最容易踩的坑是拿超大矩阵直接跑 SVD耗时无法接受却没有做分块和混合策略。后续可以继续扩展的方向有两个。一个是性能优化调研批量化 SVD、低秩近似投影、或者把闭式解嵌入到更大的训练框架里做端到端 benchmark。另一个是理论延伸把极分解投影推广到带权重衰减或带约束的变体验证在更多模型结构上的稳定性。先把小规模验证跑通再决定是否替换现有生产环境里的正交化实现。
