缩放 μ 介子进行扩散 Transformer

矩阵感知优化器 Muon 通过平衡单个方向上的更新来改进大型 模型训练，但其在大型扩散 Transformer (DiTs) 上的扩展行为和端到端效率仍不清楚。我们首先在 DiT 上从 1.3B 到 15B 参数建立 Muon 的缩放行为，表明其相对于 AdamW 的优化和生成质量优势在整个模型尺度上持续存在。然而，在规模上，每个优化步骤执行的 5 步牛顿-舒尔茨迭代 (NS5) 以及全动量实现，会引入大量的计算和通信开销，从而抵消 Muon 的步长效率优势。我们引入了 \emph{Periodic Row-wise Muon}，它每 \(K\) 步骤执行一次完整的 NS5 谱更新，并根据剩余步骤的当前动量应用低计算和通信成本的行约束更新。我们进一步共同设计了一种分布式实现，该实现在非刷新步骤期间直接在分片动量上运行，并通过桶式全收集和通信计算重叠来加速频谱刷新。在所有尺度上，Muon 将观察到的最佳生成质量比 AdamW 提高了 12.9--19.1%。与普通 Muon 相比，Periodic Row-wise Muon 在 1.3B--4B 模型上的最佳生成质量保持在 0.5% 以内，并在 9B 模型上提高了 4.5%。它将优化器时间减少了 46.9--54.3\%，端到端步骤时间减少了 15.7--24.3\%，逻辑通信量减少了 66.7\%，同时以 33.7--64.8\% 的主动训练时间减少了 33.7--64.8\% 来达到各自的最佳生成质量。这些结果表明，Periodic Row-wise Muon 保留了 Muon 的生成质量优势，同时将其转化为大型 DiT 的端到端训练效率。