Scaling Muon for Diffusion Transformers
Chenghao Li, Xiao Han, Xinxin Huang, Wei Liu, Boyang Li, Bing Xiao, Heran Zhang, Juanma Perez Rua, Ke Xu, Kangning Liu, Linjun Kuang, Na Li, Tan Wang, Tian Xie, Wei Peng, Yang Pei, Yifan Xu, Yuanhao Zhai, Yuwei Lin, Zhe Wang, Zihao He, Daniel Li, Junbiao Tang, Ziyang Jiang, Dake Chen
cs.LG, cs.AI, cs.CV
2026-08-21
Muon在1.3B到15B的DiT上,最优FD-DINO相对AdamW改进12.9到19.1%。Periodic版每3步才做一次NS5,生成质量相当,优化器时间砍掉一半,到达最优质量少花33.7到64.8%训练时间。
Muon把二维权重当成矩阵来更新:对动量做五步Newton–Schulz(NS5),把各奇异方向上的更新幅度拉齐。语言模型预训练里它已经能用更少的FLOPs追上AdamW。扩散模型的小规模实验也提示过优化更好,但loss、生成质量和墙钟时间会给优化器排出不同名次。十亿参数级的Diffusion Transformer(DiT)上,这件事还没人量过。
十亿参数以后,优化器不能只比步数。Muon的NS5在每一步都做全矩阵乘法,而且在FSDP分片下要把完整动量all-gather出来,每个rank上都跑一遍谱变换。步效率再好,墙钟上也可能被AdamW追上。这篇要回答两件事:Muon的质量优势能不能跟到15B,以及这笔系统和通信开销能不能砍下来。
实验从零训练text-to-image MMDiT,数据是GPIC-Full的1亿图文对(Flickr与Wikimedia,Qwen3-VL打标),分辨率512×512。四个规模约1.3B、4B、9B、15B,双流骨干相同,图像走冻结的FLUX.1-schnell VAE,文本走冻结的CLIP-L、CLIP-G和T5-XXL。Muon只改Transformer块里的二维权重,约占可训参数的99%;bias、normalization和块外参数仍走AdamW。三组优化器都跑6万步、全局batch 4096,256张H100,PyTorch FSDP2。
Periodic Row-wise Muon把每一步都做的NS5拆成两种几何。
设计理由是极因子方向在远离秩退化时对动量的小变化是局部稳定的,所以不必每步重做全局谱变换。K管刷新频率,γ管RowNorm分支相对NS5的步长。两种归一化对应不同可行集,直接复用Muon学习率等于设γ=1,没有理论依据。超参在1.3B上用16张H100、3万步扫过:K取2/3/4,γ取0.10/0.15/0.25/0.35,取20k、25k、30k的平均FD-DINO。K从2提到3,NS5次数少33%,晚期平均FD-DINO只差2.8%;再提到4能再省25%的NS5,FD-DINO却差8.2%。最终固定K=3、γ=0.15,四个规模都不重调。
分布式实现跟算法绑在一起。非刷新步RowNorm直接在分片动量上算:扁矩阵的行在本地,零通信;高矩阵只all-reduce每列的范数标量,不必all-gather整块动量。刷新步把矩阵打成bucket,下一桶的all-gather和当前桶的NS5重叠。K=3时,优化器侧逻辑通信量大约变成vanilla Muon的三分之一。
Muon对AdamW的优势在四个规模上都在。验证loss全程更低,最优FD-DINO(DINOv2特征空间的Fréchet距离,越低越好)相对AdamW好12.9–19.1%。但把loss对墙钟重画,AdamW会先到达同一中间loss:每步的系统和通信开销把步效率吃掉了。
Periodic Row-wise Muon把质量优势留住,并把开销变成墙钟收益。6万步终点相对AdamW的FD-DINO改进是11.1–17.8%;最优checkpoint相对vanilla Muon,1.3B和4B差不超过0.5%,9B好约4.5%,15B好约2.7%。到达各自最优FD-DINO所用有效训练时间分别少33.7%、36.1%、64.8%、57.4%。
| 规模 | AdamW终点FD-DINO | Muon | Periodic |
| 1.3B | 53.93 | 46.08 | 45.65 |
| 4B | 51.39 | 41.62 | 42.25 |
| 9B | 41.26 | 33.97 | 36.70 |
| 15B | 40.23 | 33.51 | 33.99 |
这是60k终点,不是训练中见过的最好点。9B上Periodic的GenEval2算术平均从Muon的51.77升到57.33,几何平均从11.35升到15.97;15B上Coverage和Density更好,GenEval2却从57.93掉到52.26。指标并不同步。
相对vanilla Muon,优化器时间砍46.9–54.3%,端到端步时砍15.7–24.3%,逻辑通信量固定少66.7%。15B上优化器时间从2.884降到1.317(以1.3B AdamW平均步时为1),步时从6.527降到5.002。
消融说明两件事缺一不可。1.3B上每步只做RowNorm,最优FD-DINO是51.322,几乎退回AdamW的51.56;γ=1的周期版是44.294,校准后的γ=0.15才到41.913,贴近vanilla Muon的41.733。系统侧,15B朴素周期实现每步仍物化完整动量,步时5.685;分片RowNorm把通信量从28.62降到9.54 GiB/rank/step,再加bucket和重叠,步时到5.002。
给正在训大DiT、已经想上Muon的团队一个能落地的折中:质量跟得上每步NS5,墙钟明显更便宜。K和γ在1.3B上选定后直接搬到15B,说明这组默认值至少在这一族MMDiT上可迁移。
这不是新优化器原理。谱刷新降频、行归一、分片Muon运行时都有前人工作。这篇的贡献是把「隔步RowNorm加分片统计通信」在1.3B到15B的真实DiT训练里跑通,并同时报生成质量和系统开销。只看step-wise loss会高估vanilla Muon;只看墙钟又会低估它。两边一起看,才轮得到Periodic这个解。
作者自己写了:只测了一族MMDiT、一套GPIC数据、一个512分辨率、一套32节点H100加FSDP2。视频DiT、更高分辨率、不同并行策略都还没有。K和γ是全局固定,没有layer-wise或自适应。刷新步仍然要完整动量通信和物化,加速比会随拓扑和并行策略变。
另外几处读下来站不稳。评测主指标用「训练过程中见过的最优FD-DINO」,Muon和Periodic进入最佳质量的步数不同,checkpoint之间波动也不小:4B的Periodic在45k曾跳到62.89,15B在50k跳到52.25,都比邻近checkpoint差一截。终点表里9B的Periodic FD-DINO不如Muon(36.70对33.97),15B的GenEval2也不如。把「最优checkpoint」和「终点」混着宣传时,需要对着表看。学习率两边也不对称:Muon主学习率大约是AdamW的十几倍,辅助AdamW组另有1e-4,公平性依赖这套调参。没有公开代码或kernel,数字全部来自内部256卡集群。