改进优化器让扩散模型训练更快更高效。
Scaling Muon for Diffusion Transformers
- 每K步做一次高开销更新，其余步用轻量级局部更新
- 在1.3B到15B模型上生成质量比AdamW提升12.9%~19.1%
- 降低50%以上计算通信开销，训练时间减少超三分之一
矩阵感知优化器Muon通过平衡奇异方向的更新，提升大模型训练效果，但其在大规模扩散变换器（DiTs）上的扩展性与端到端效率仍不明确。我们首次在1.3B至15B参数的DiTs上研究了Muon的缩放行为，发现其在优化与生成质量上对AdamW的优势在不同规模下持续存在。然而，在大规模下，每步执行5步Newton--Schulz迭代（NS5）并全动量显式存储带来显著计算与通信开销，抵消了步级效率优势。为此，我们提出周期性行向量Muon：每K步进行一次完整NS5谱更新，其余步骤采用低开销行约束更新。同时设计分布式实现，非刷新步直接在分片动量上操作，并通过桶化all-gather和通信-计算重叠加速谱刷新。在所有规模下，Muon相比AdamW生成质量提升12.9%–19.1%。与原始Muon相比，周期性行向量Muon在1.3B–4B模型上质量损失小于0.5%，在9B模型上提升4.5%；优化器时间减少46.9%–54.3%，端到端步时减少15.7%–24.3%，逻辑通信量降低66.7%，达到最优生成质量所需活跃训练时间减少33.7%–64.8%。结果表明，周期性行向量Muon在保持生成质量优势的同时，将效率优势转化为大规模DiTs的端到端训练效率。
原文摘要 · Abstract (English)
The matrix-aware optimizer Muon improves large model training by balancing updates across singular directions, yet its scaling behavior and end-to-end efficiency on large Diffusion Transformers (DiTs) remain unclear. We first establish Muon's scaling behavior on DiTs from 1.3B to 15B parameters, showing that its optimization and generative quality advantages over AdamW persist across model scales. However, at scale, the 5-step Newton--Schulz iteration (NS5) performed at every optimization step, together with full-momentum materialization, introduces substantial computation and communication overhead that can offset Muon's step-efficiency advantage. We introduce \emph{Periodic Row-wise Muon}, which performs a full NS5 spectral update once every \(K\) steps and applies a low compute and communication cost row-wise constrained update based on the current momentum at the remaining steps. We further co-design a distributed implementation that operates directly on sharded momentum during non-refresh steps and accelerates spectral refreshes through bucketed all-gather and communication--computation overlap. Across all scales, Muon improves the best observed generative quality over AdamW by 12.9--19.1\%. Compared with vanilla Muon, Periodic Row-wise Muon remains within 0.5\% in best generative quality on the 1.3B--4B models and improves it by 4.5\% at 9B. It reduces optimizer time by 46.9--54.3\%, end-to-end step time by 15.7--24.3\%, and logical communication volume by 66.7\%, while reaching its respective best generative quality with 33.7--64.8\% less active training time. These results show that Periodic Row-wise Muon preserves Muon's generative quality advantage while translating it into end-to-end training efficiency for large DiTs.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。