Dion3：全栈正交更新 Dion3: Full-Stack Orthogonal Updates
从数学、内核、更新规则到通信全栈加速 Muon 正交化，优化器步时间最多降低约 6 倍
前置知识
Muon 优化器
一种基于谱范数最速下降的优化器：维护动量矩阵 $M \leftarrow \mu M + G$，然后按 $W \leftarrow W - \eta\,\mathrm{polar}(M)$ 更新权重。极分解把更新矩阵的奇异值全部压到 1，使每步更新在所有方向上幅度均衡、保持满数值秩。Kimi K2、GLM-5 等前沿大模型已采用 Muon 训练。
本文的全部工作都是围绕削减 Muon 中正交化步骤的计算与通信成本展开，必须先理解其更新规则与开销来源。
Newton-Schulz 迭代
用矩阵多项式迭代逼近极分解的方法：$X_{t+1} = a_t X_t + b_t X_t X_t^\top X_t + c_t (X_t X_t^\top)^2 X_t$，作用是把奇异值逐步推向 1。Muon 通常迭代 $T=5$ 次，复杂度 $O(n^3)$，共 15 次 GEMM，是优化器步的主要计算瓶颈。
Dion3 的 Gram 牛顿-舒尔茨是对该迭代的数学重构，理解其多项式视角（$p_t(x)=x\,h_t(x^2)$）是读懂第 4 节的前提。
极分解（正交化）
若矩阵的 SVD 为 $X = U\Sigma V^\top$，则 $\mathrm{polar}(X) = UV^\top$。它保留奇异向量、抹平奇异值差异，把任意矩阵投影到正交矩阵。它无法精确计算（需做 SVD），实践中只能用牛顿-舒尔茨等迭代法逼近。
正交化是 Muon 更新的核心操作，本文四项优化全部围绕如何更便宜地近似它。
谱范数最速下降
Muon 可解释为在谱范数度量下的最速下降：$\mathrm{polar}(M)$ 是谱范数约束下与 $-M$ 内积最大的方向。由 Frobenius 范数可算出更新规模 $\|\eta\,\mathrm{polar}(M)\|_F = \eta\sqrt{n}$（$n$ 为较小维度），这一量纲分析直接导出 Dion3 的学习率迁移规则 $\eta' = \eta/\sqrt{f}$。
理解更新规模与行抽样比例 $f$ 的关系，才能明白为何 Dion3 换 $f$ 时要按 $\eta\sqrt{f}\approx 0.01$ 重调学习率。
FSDP 与 all-to-all 通信
全分片数据并行（FSDP）把权重和优化器状态切片存到各 GPU。元素级优化器可各自更新分片，但 Muon 必须先 all-to-all 聚合完整动量矩阵做正交化、再分发结果，朴素实现的通信轮数为 $O(N/\text{world\_size})$，且逐参数的小消息打不满互联带宽。
论文的行抽样削减通信量、megabatching 削减通信轮数，都是针对这一分布式瓶颈，是理解第 6、7 节的背景。
半精度浮点与数值稳定性
bfloat16/float16 用 16 位存储以换取速度，但舍入误差会破坏 Gram 矩阵 $XX^\top$ 的半正定性，引入理论上不存在的负特征值（实测可低至 $-4\times 10^{-4}$）。逆平方根迭代对负输入会发散，这曾导致朴素版 Gram 牛顿-舒尔茨训练崩溃。
Gram 牛顿-舒尔茨的“重启”稳定化技巧正是为解决该问题，是第 4.2 节和附录 B 的关键内容。
对称 GEMM 内核
对称矩阵乘法 $AB$（输出对称）只需计算下三角和对角线，再转置复制到上三角，可省约一半 FLOPs。作者用 CuteDSL 为 Hopper/Blackwell GPU 编写了“三角调度器 + 转置写回 epilogue”的专用内核，大矩阵下比 cuBLAS 快约 2 倍。
Gram 牛顿-舒尔茨比标准版本使用更多对称乘法，内核收益因此被放大，这是“全栈”各层协同的核心一环。
误差反馈（Error Feedback）
当只用部分行近似动量矩阵时，常规衰减 $M \leftarrow \mu M$ 会永久丢弃未选中行的信息。误差反馈改为 $M \leftarrow \mu\hat{M} + (M - \hat{M})$：只把选中行衰减 $\mu$ 倍，未选中行原样保留。被忽略的梯度会随时间在 $M$ 中累积，最终被选中并参与更新。
这是 Dion3 行抽样不损害、反而略改善训练质量的关键机制，继承自 Dion 并简化。
研究动机
Muon 正被 Kimi K2、GLM-5 等前沿大模型采用，它达到给定损失所需步数少于 AdamW，但每步本身贵得多：正交化依赖 $T=5$ 次牛顿-舒尔茨迭代，对 $n\times m$ 矩阵的代价是三次方的 $2T(2\alpha+1)n^3$ FLOPs（$\alpha=m/n\geq 1$ 为长宽比）。实测 7B 模型在四张 GH200 上，标准 Muon 的优化器步时间是 AdamW 的 26 倍。分布式场景雪上加霜：FSDP 把权重切片到各卡，正交化前必须 all-to-all 聚合完整矩阵、算完再分发，朴素实现每步需 $O(N/\text{world\_size})$ 轮通信，且逐参数小消息打不满互联带宽。更麻烦的是，最新 MoE 架构走向更细粒度、更稀疏的专家，隐层与中间维之比 $\alpha$ 不断增大，矩形乘法占比随之升高。Moonshot 能把 Muon 训到万亿参数，靠的是旧版框架的连续 flat buffer 分片、单稠密层细粒度 MoE、流水线加专家并行三个因素偶然对齐——而第一个因素在新版框架中已被弃用。这种脆弱的适配使 Muon 难以成为 AdamW 的通用替代品。
本文的目标是作者的目标是给出一个全栈、通用的解决方案，让 Muon 及其变体（NorMuon）在各种模型规模、架构（稠密与 MoE）、集群规模和并行策略（FSDP2、DDP、混合分片）下都能以低开销运行，且不牺牲训练质量。具体拆成四个相互叠加的子目标：(1) 从数学上重构牛顿-舒尔茨，降低正交化 FLOP 成本且输出与标准版本完全等价，保证与既有 Muon 改进兼容；(2) 编写利用对称结构的 CuteDSL GPU 内核，把理论 FLOP 节省变成真实墙钟收益；(3) 设计一个比 Dion 更简单的压缩式更新规则——只正交化动量矩阵的一部分，同时削减计算量与通信量；(4) 用 megabatching 把每步通信轮数降到与模型深度无关的常数级。最终打包为 dion 与 gram-newton-schulz 两个 pip 包，作为 Muon 的即插即用替代品，目标是将优化器步时间相对标准 Muon 最多降低 6 倍。
与已有工作不同的是，已有改进走两条互不相干的路。一条改进牛顿-舒尔茨本身：优化多项式系数 $(a_t,b_t,c_t)$ 或归一化方式，但都保留“在大矩形矩阵上直接迭代”的原始形式；Amsel 等人附录 J 虽提出过类似 Gram 思路，却未研究对称内核的放大效应、数值稳定性与实用配方。另一条缩小被正交化的矩阵：Dion 用幂迭代构造低秩基 $\tilde{V}$，Trion 改用 DCT 矩阵的列，但都把牛顿-舒尔茨当黑盒，且通信问题被单独处理。Dion3 的独特之处在于把四个层级打通并让它们互相放大：Gram 形式比标准形式用更多对称乘法，因而自定义内核加速更显著；行抽样把长宽比 $\alpha$ 放大 $1/f$ 倍，恰好落在 Gram 形式最占优的区域；压缩同时削减通信量，与 megabatching 削减通信轮数正交互补。此外它把“低秩近似”推到最简——直接选 $k=\lceil fn\rceil$ 行、不做任何变换，配合误差反馈，比 Dion 和 Trion 都更简单、更快。
核心方法
直觉上，Muon 的全部额外开销来自 $\mathrm{polar}(M)$，Dion3 因此从四个层面同时攻击它。数学层提出 Gram 牛顿-舒尔茨：利用 $\mathrm{polar}(X)=(XX^\top)^{-1/2}X$，把迭代从 $n\times m$ 矩形矩阵搬到 $n\times n$ 对称 Gram 矩阵上进行，全程只需开头 $XX^\top$ 和结尾 $Q_5X$ 两次矩形乘法，$T=5,\alpha=4$ 时省 55%–68% FLOPs。硬件层用 CuteDSL 写对称 GEMM 内核：只算输出下三角再镜像写回，省约一半乘法，大矩阵下比 cuBLAS 快约 2 倍。算法层提出行抽样更新规则：每步只挑动量矩阵中 $\ell_1$ 范数最大的 $k=\lceil fn\rceil$ 行做正交化（推荐 $f=1/4$ 或 $1/8$），配合误差反馈保留未选中行信息，正交化成本至少降 $1/f^2$、通信量降 $1/f$。系统层用 megabatching 把同形状的所有矩阵打包进一次 all-to-all，通信轮数降为与模型深度无关的常数。四者叠加，7B 模型四卡 GH200 上优化器步时间从标准 Muon 的 26× AdamW 降到 4×。
核心数学洞察是定理 2：任何奇次多项式可写成 $p(x)=x\,h(x^2)$，把牛顿-舒尔茨每个 $p_t$ 如此改写后，复合多项式满足 $(p_T\circ\cdots\circ p_1)(x)=q_T x$，其中 $q_T$ 由标量迭代 $r_0=x^2,q_0=1$，$z_t=h_t(r_{t-1})$，$r_t=r_{t-1}z_t^2$，$q_t=q_{t-1}z_t$ 给出。把它提升到矩阵上（这些运算都保持奇异向量不变，各奇异值独立演化），即得直接逼近 Gram 矩阵逆平方根 $Q_T\approx(XX^\top)^{-1/2}$ 的迭代，最后乘回 $Q_T X$ 得到与标准牛顿-舒尔茨逐位相同的结果——数学严格等价，但主导成本从 $O(T\alpha n^3)$ 降到 $O((T+\alpha)n^3)$。第二个核心是行抽样：不做任何低秩构造（Dion 的幂迭代、Trion 的 DCT 列），直接选 $\ell_1$ 范数 top-$k$ 行正交化，配合误差反馈 $M\leftarrow\mu\hat{M}+(M-\hat{M})$，让被忽略的梯度随时间累积、日后仍会被选中。与 Dion 的本质区别是用原始子矩阵替代低秩投影，省掉基构造等全部开销，实验显示质量反而略优。
方法步骤详情
全流程分四块。(1) 稳定化 Gram 牛顿-舒尔茨（算法 3）：输入 $X\in\mathbb{R}^{n\times m}$ 按 Frobenius 范数归一化并转 float16；算 $R_0=XX^\top$、$Q_0=I$；循环 5 次：$Z_t=b_tR_{t-1}+c_tR_{t-1}^2$，$Q_t=Q_{t-1}Z_t+a_tQ_{t-1}$，$R_t=Z_tR_{t-1}Z_t+a_t(\cdot)$；在 $t=3$ 时重启——先算 $X\leftarrow Q_2X$、重建 $R_2=XX^\top$、$Q_2=I$，把半精度伪负特征值归零；输出 $Q_5X$，总成本 $(4T+3\alpha-3)n^3$ FLOPs。(2) 对称 GEMM 内核：三角调度器只把下三角 256×256 瓦片分给线程块簇，epilogue 把结果转置写入上三角。(3) Dion3 更新规则（算法 4）：$M\leftarrow M+G$；选 $\ell_1$ 范数 top-$k=\lceil fn\rceil$ 行得索引集 $S$；用 Gram NS 正交化子矩阵 $M[S,:]$；可选 NorMuon 归一化（须在 float32）；自定义 Triton 内核完成权重衰减 $W\leftarrow(1-\eta\cdot\text{wd})W$ 与行更新 $W[S,:]\leftarrow W[S,:]-\eta O$；误差反馈只衰减选中行 $M[S,:]\leftarrow\mu M[S,:]$。(4) megabatching：按形状打包一次 all-to-all，组装、批量正交化、统一分发；小规模用 CUDA graph 消除内核启动开销。
技术新颖性
新颖性有四点。第一，Gram 牛顿-舒尔茨是数学等价重构而非启发式：输出与标准版逐位相同，因此与几乎所有 Muon 变体及改进系数（如 Polar Express）即插即用兼容；相比 Amsel 等人附录 J 的早期版本，本文给出更稳定的公式、首次揭示对称内核的协同效应并系统分析重启策略；其迭代核还推广并加速了 Lakić 的逆平方根方法。第二，CuteDSL 对称内核针对 Hopper/Blackwell 的张量内存与线程块簇定制，大矩阵下比 cuBLAS 快约 2 倍；关键论断是只有与 Gram 形式结合时其潜力才被完全释放，因为 Gram 形式需要 $\alpha AB+\beta C$ 这类更一般的对称融合运算。第三，行抽样把压缩正交化推到极简：比 Trion 的 DCT 列更直接，分布式下只需每分片本地选行，避免全局同步与 ragged all-to-all；配套 Triton 内核修复了行选择破坏编译器自动融合、引发多轮 bfloat16 转换的精度问题。第四，稳定性分析严谨：指出伪负特征值可低至 $-4\times10^{-4}$、逆平方根迭代对负输入发散，用重启 + float16 彻底解决，并给出条件数界驱动的重启位置扫描方法。
实验结果
质量：1B 模型 100B token 上，调优后 NorMuon（$\eta=0.01,\mu=0.9$）损失 2.194，Dion3 在 $f=1/8,\eta=0.02$ 时达 2.181（$f=1/2$ 为 2.185，$f=1/16$ 为 2.186），全程低于基线约 0.01；学习率迁移规则 $\eta\sqrt{f}\approx0.01$ 在对数坐标上清晰成立。扩到 3B–14B、10B token：Dion3（$f=1/4,\eta=0.02$）验证损失全胜（3B −0.012、4B −0.011、7B −0.014、14B −0.027），14B 为 2.162 对 2.189，14B 下游 12 基准宏平均准确率 58.1% 对 57.4%（+0.7pp），四规模中三项占优。速度：对称内核 + Gram NS 相对标准 Muon 合计提速 1.5× 以上（$\alpha=8$ 的 Gemma/MoE 类架构达 2×）；分数更新 $f=1/2$、$f=1/4$ 再分别带来约 2× 和 3.7×，总提速 3.6× 和 6.5×；7B 模型四卡 GH200 上优化器步时间从 26× AdamW 降至 4×。通信：1B+8 FSDP 分片时 megabatching 把步时间从 80.7ms 降至 52.1ms（−35%），其余场景 −2% 至 −6%。FLOP 理论：$T=5,\alpha=4$ 时 Gram 形式比对称内核版标准 NS 省 55%，比普通实现省 68%。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| 优化器步时间（7B 模型，4×GH200） | 相对 AdamW 的倍数（越低越好） | Dion3 全栈：4× AdamW | 标准 Muon：26× AdamW | 降低约 6.5× |
| 正交化 FLOP 成本（$T=5, \alpha=4$） | 总 FLOPs | $(4T+3\alpha-3)n^3$（Gram NS + 对称内核） | $T(3\alpha+1)n^3$（对称内核版标准 NS）/ $(20\alpha+10)n^3$（普通实现） | 节省 55% / 68% |
| 对称矩阵乘法内核 | 大矩阵下相对 cuBLAS 的加速比 | ≈2×（Hopper 与 Blackwell） | cuBLAS 标准 GEMM | ~2× |
| 1B 语言模型预训练（100B token，ClimbMix） | 验证交叉熵损失（越低越好） | 2.181（Dion3, $f=1/8, \eta=0.02$） | 2.194（NorMuon, $\eta=0.01$） | −0.013 |
| 14B 语言模型预训练（10B token，ClimbMix） | 验证交叉熵损失 | 2.162（Dion3, $f=1/4$） | 2.189（NorMuon） | −0.027 |
| 14B 下游能力评测 | 12 项基准宏平均准确率 | 58.1% | 57.4%（NorMuon） | +0.7 个百分点 |
| 分布式优化器通信（1B 模型，8 FSDP 分片） | 每 GPU 优化器步时间（中位数） | 52.1 ms（megabatching） | 80.7 ms（普通 batching） | −35% |
局限与改进
作者承认的局限：全栈优化后 Dion3 的优化器步时间仍是 AdamW 的约 4 倍，只是差距缩小；分数更新在 $f<1$ 时改变优化轨迹，其质量提升的机理不明——作者明说这是设计"廉价近似 Muon"之外的意外收获，仅引用 Joo 等人随机掩码块可改进 SGD 动量的结果作旁证，泛化性待验证；重启策略额外花费 $3(\alpha-1)n^3$ FLOPs，$\alpha=1$ 时需回退到核化标准 NS；分布式下真全局 top-$k$ 行选择需额外同步和 ragged all-to-all，默认只做分片内选择。我的补充观察：质量实验最高 14B、训练量仅 10B–100B token，远短于前沿训练长度，长程行为未知；下游评测只有 12 项标准基准，未涉及 RLHF 或指令微调；数据集单一（ClimbMix）；学习率规则 $\eta\propto 1/\sqrt{f}$ 是经验拟合而非理论；对称内核依赖 Hopper/Blackwell 特性，其他硬件拿不到 2× 收益；$f=1$ 时与 Muon 仍有四个实现细节差异，等价性靠收敛曲线而非理论确认。
独立分析的弱点
独立分析几个弱点。其一，行选择用 $\ell_1$ 范数 top-$k$ 是相当粗糙的准则，分布式下还退化为每分片独立选择、缺乏全局视角；作者自己发现随机选择效果相差不多，说明该准则远非最优——可探索基于曲率信息或轻量可学习选择器的方案。其二，换用任何 $f$ 都要按 $\eta'=\eta/\sqrt{f}$ 重调学习率，目前依赖经验扫描，与"即插即用"的卖点存在张力——可研究把该缩放内化的自适应步长机制。其三，数值稳定性依赖重启 + float16 + Polar Express 系数的特定组合，朴素版本会发散，换系数序列时重启位置需重新扫描，工程上易踩坑——可研究对系数不敏感的自动稳定格式。其四，质量提升只在稠密 transformer 上验证（附录 A 的 Gemma/MoE 实验只测速度），MoE、多模态、RL 微调场景表现未知。其五，速度实验用合成数据测孤立的优化器步时间，未系统报告真实集群计算-通信重叠下的端到端收益分布；14B 质量实验仅 10B token，对更长训练的外推性存疑。
未来方向
作者提出的方向：确定分数更新质量提升的适用范围并给出理论解释，尤其是与 Joo 等人“随机掩码块可改进动量 SGD 轨迹”的联系；把压缩数据并行推广到跨 pod 的数据中心网络场景——当行集 $S$ 可无需先同步 $M$ 而选出（如随机选择）时，每步只需同步选中的子矩阵 $M[S,:]$，通信量降 $1/f$ 倍，这对 ICI 内模型并行 + 跨 DCN 纯数据并行的常见部署特别有价值。基于本文成果可延伸：把 Gram 牛顿-舒尔茨与 MuonBP 等块正交化方法组合（双方都把 NS 当黑盒，天然可叠加）；在 μP 迁移框架下研究 $f$ 与 $\eta$ 的联合缩放律并给出理论支撑；把行抽样思想移植到 Shampoo 等二阶优化器的预条件子计算；探索动量驱动的自适应选择策略取代静态 $\ell_1$ top-$k$；把三角调度内核推广到 AMD 及更早的 NVIDIA 架构，或在编译器层面自动生成；研究 $f$ 的动态调度——训练早期用 $f=1$ 保证探索、损失平台期后降 $f$ 省算力——以及重启位置的自适应选择。
复现评估
复现条件较好。代码完全开源且分成两个互操作包：microsoft/dion 实现优化器更新规则（Muon、NorMuon 及 Dion3 两种风味）、megabatching、FSDP2/DDP/混合分片的分布式后端与 CUDA graph 优化；Dao-AILab/gram-newton-schulz 是牛顿-舒尔茨的即插即用替代，内置 Polar Express 系数与重启稳定化逻辑。超参有明确推荐：$f=1/4$ 或 $1/8$、$\eta\sqrt{f}\approx0.01$、重启位置 $t=3$、NorMuon 归一化用 float32；附录 D.1 给出架构与训练细节（GQA、滑动窗口注意力、MXFP8 权重、ClimbMix、GH200 集群）。主要门槛是算力：质量实验最高 14B 模型、10B token（多节点 GH200），附录 E 的端到端案例更大；小规模验证（1B、100B token）单节点 8 卡可承受。内核部分需要 Hopper/Blackwell 硬件与 CuteDSL 工具链，附录 C 提供实现策略细节。总体而言中小规模复现难度中等，完整复现 14B 级别质量结论需要可观的 GPU 预算。
论文图表
