如果你最近在训练扩散 TransformerDiffusion Transformer简称 DiT这一类模型可能会碰到一个很拧巴的现象——模型结构调得合理数据也没问题但 loss 曲线就是不稳定动不动就飙一下或者训练卡在一个平台期很久都下不去。很多人第一反应是学习率太大或者 batch size 不合适于是反复调整学习率、warmup、梯度裁剪但效果常常有限。最近我注意到一个叫 CMuon 的工作全称是 Chunked Momentum Orthogonalization从名字看它想通过分块动量正交化来加速并稳定扩散 Transformer 的训练。这个方向很有意思因为它没有去改模型结构而是直接针对优化过程开刀。这类工作真正值得关注的地方不在于它又给你一个“更好用的优化器”而在于它把训练稳定性问题重新拉回到了优化算法层面。过去我们遇到 DiT 训练不稳定第一反应是调学习率、换初始化、改 loss但很少会问动量在更新时是不是已经堆积了错误的梯度信息参数矩阵在更新时是不是已经在病态方向上走了太久如果 CMuon 确实像论文标题暗示的那样通过分块、动量正交化来约束更新方向那它解决的不只是“更快收敛”更是“让大规模训练稳定、可控、可重复”。这篇文章我会从训练 DiT 的痛点出发拆解 CMuon 涉及的核心机制再落到工程落地时需要关注的参数、边界和验证方法。因为原始论文的完整实验细节还没有放出来文章里会明确区分哪些是标题和常识层面的判断、哪些是从工程经验推出来的参考做法。1. 先搞清楚扩散 Transformer 训练中的优化难点是什么1.1 为什么 DiT 对优化器更敏感扩散 Transformer 和普通分类模型有个很大的不同它不是在拟合一个确定标签而是在拟合一个逐步去噪的过程。每一步去噪都对应一组输入和输出模型要在不同噪声强度下都保持稳定行为。这意味着 loss 曲面本身就比普通分类任务更复杂不同时间步的梯度方向可能差异很大。从工程经验看DiT 训练初期特别容易出现 loss 突然变大的现象。常见原因包括模型某一层的权重矩阵在更新中不断积累“同向分量”导致特征分布漂移。时间步 embedding 和图像 token 之间的梯度尺度不一致优化器难以同时适配。大规模 batch 下梯度的随机噪声虽然被平均了但高维参数空间中某些方向仍然会被反复放大。这些问题的共性是单靠全局学习率缩放很难处理。ADAM 类优化器通过二阶矩估计给每个参数自适应学习率在大多数任务上很强但在 DiT 这类高维、多尺度、长训练任务上它的自适应机制本身也可能引入不稳定。1.2 大规模训练中的动量失效现象动量是优化器里最容易被忽略但其实非常关键的组件。SGD with Momentum、Adam 都有动量项作用是让更新方向更平滑减少震荡。但在大规模模型训练中动量存在一个很隐蔽的问题当梯度方向在不同 batch 间变化较大时动量缓冲里累积的其实是“历史梯度的加权平均”而高维空间中这些历史梯度可能来自相互冲突的方向。结果就是动量越大更新方向越容易被早期梯度主导一旦学习率偏大模型很容易越过一个足够好的局部区域loss 突然爆掉。反过来动量太小又起不到平滑作用loss 会在一个范围内反复抖动。在扩散 Transformer 里这个问题更明显。因为模型需要同时建模不同噪声尺度的去噪过程不同 time step 的梯度天然存在方向差异。如果动量设计得不够合理训练初期积累的梯度方向可能会对后期更新产生持续的“拖拽效应”。1.3 一种更本质的解决方向传统做法是通过学习率 warmup、梯度裁剪、EMA 等方式“事后补救”这些方法有效但都属于外部约束。CMuon 这类思路选择在优化器内部做文章更新参数时不再只是沿着某个下降方向走一步而是对参数矩阵施加正交化约束让不同维度的更新尽量互不干扰。这个想法其实有相当长的历史。Shampoo、K-FAC 这类二阶优化器也尝试建模参数之间的相关性但因为计算量太大很难在大规模 DiT 训练中直接用。CMuon 的“分块”策略本质是想用较低成本拿到类似的正交化收益。一句话总结DiT 训练不稳定的一个深层原因是优化器没有约束更新方向之间的耦合。CMuon 走的是“分块 动量正交化”的路线目标是让更新更干净、更可控。2. CMuon 的核心设计分块、动量、正交化三者缺一不可2.1 正交化约束到底约束了什么在参数更新中如果权重矩阵为 W常规优化器会计算一个增量 ΔW然后 W ← W - lr * ΔW。正交化约束的意思是对 ΔW 做一个变换使得它在某种度量下接近正交于当前权重的某些方向或者至少让更新矩阵的行/列之间的相关性被削弱。为什么要这么做一个合理的解释是神经网络的权重矩阵往往存在大量冗余维度如果更新方向在这些冗余维度上积累模型参数会逐渐偏向一个病态区域导致 loss 曲面变得非常陡峭。正交化更新相当于给参数更新装了一个“防漂移装置”避免权重矩阵在错误方向上越走越偏。拿生活中的类比来说这有点像整理行李箱。普通优化器是看到哪个缝隙就往里塞东西塞多了箱子会变形正交化更新则是每放一件东西都稍微压一下保持箱子形状稳定。扩展到训练上就是让每一层在更新时尽量保持“形状”的合理性。2.2 分块策略为什么必要完整的正交化计算在高维参数空间里代价极高。如果对几百兆参数的大矩阵做全局正交化每一步都要做 SVD 或 QR 分解训练速度会慢到无法接受。CMuon 取了一个折中把参数矩阵分成若干块每一块单独做正交化处理。分块的好处是清晰且直接计算成本可控尤其适合 GPU 并行。每一块能获得局部正交性虽然不如全局正交化强但已经能约束大部分病态更新。不同块之间天然可以独立处理适合分布式训练。风险也在这里分块大小直接决定了正交化的强度和开销。分块太小正交化效果接近没有分块太大计算成本又上去了。这个平衡是 CMuon 这类算法调优时最需要实验的地方。2.3 动量正交化的组合效应单独看“动量”和“正交化”都不算新东西。CMuon 的关键是把两者结合起来先对梯度做分块每块内积累动量再对这个动量矩阵做正交化变换最后用于更新参数。这个流程的直觉可以这样理解动量解决“单步梯度噪声大”的问题给更新一个更平滑的估计。正交化解决“累积方向漂移”的问题让动量缓冲不会朝着病态方向一路狂奔。分块则让上述两者能在大模型上落地。过去很多优化器要么只做动量平滑要么只做梯度变换很少同时把分块、动量和正交化统一在一个更新规则里。CMuon 如果真能做到这一点那它在 DiT 这种长训练、大模型、高敏感度的场景下确实可能比 AdamW 更稳。3. CMuon 相对已有优化器的差异在哪里3.1 与 AdamW 的差异AdamW 是当前大模型训练的事实标准。它为每个参数维护一阶矩和二阶矩自适应调整每个参数的学习率。这个机制非常通用几乎开箱即用但也有明显短板二阶矩估计在高维空间中很容易受极端梯度影响。它对参数矩阵内部的相关性结构建模很弱。在很长的训练中AdamW 的更新方向容易变得“过于自信”导致 loss 突然飙升。CMuon 与 AdamW 的本质差异不是用不用自适应学习率而是它多了一个“全局结构约束”。AdamW 是逐参数独立处理CMuon 是分块后做矩阵层面的正交化等于把参数之间的相关性也纳入了更新过程。当然这不是说 CMuon 一定能完全替代 AdamW。AdamW 经过行业大规模验证稳定性在大多数任务上都不差。CMuon 更适合的问题是当你遇到 AdamW 解决不了的训练不稳定时它是一个值得试验的方向。3.2 与 Muon / 其他正交化优化器的差异Muon 这类方法已经提出过用正交化约束提升训练效率。CMuon 的差异点大概率落在“Chunked”上。从工程角度看Muon 如果对超大矩阵直接做正交化计算瓶颈会很明显。CMuon 通过分块把这个操作的复杂度降下来让它能真正用于现代大模型训练。不过这里有一个需要冷静看待的点分块会牺牲一部分全局正交性。某些情况下块与块之间的耦合依然会导致训练不稳定。所以 CMuon 并不一定在所有任务上都比原版 Muon 好它更像是在“理论最优”和“工程可行”之间取平衡。我个人更关注的是CMuon 是否会和混合精度训练、梯度累积、分布式通信有良好兼容性。如果它分块后的正交化操作需要额外做多卡通信那在千卡集群上的开销可能比单卡实验大得多。这是工程落地时最需要先验证的。3.3 与学习率调度、梯度裁剪的关系很多人把优化器当成“换个公式就能更好”的魔法模块但实际训练里学习率调度和梯度裁剪对稳定性的影响往往比优化器本身更大。CMuon 如果引入更强的更新方向约束它对学习率敏感度可能会改变。一个合理的推断是使用 CMuon 时学习率可以不用像 AdamW 那样保守。因为正交化削弱了病态方向的更新幅度较大学习率不一定立刻导致发散。但这不代表可以完全去掉 warmup 和梯度裁剪尤其是训练前几千步梯度方差仍然很大。建议的落地顺序是先在原有 AdamW 配置上保持相同学习率上限用 CMuon 替换优化器观察 loss 稳定性如果训练更稳了再逐步尝试提高学习率寻找新的最优区间。不要一上来就大幅调参否则很难判断到底是优化器的功劳还是学习率调对了。注意换优化器不是简单地“替换一行代码”。它的最优学习率、权重衰减、动量系数都可能和 AdamW 不同。建议先用小规模实验扫描学习率再放大到正式训练。4. 工程落地时最需要关心的四个问题4.1 分块大小和计算开销分块大小是 CMuon 引入的一个新超参。对于全连接层和注意力层参数矩阵通常是二维的分块策略可以选择按行分块、按列分块或者按固定大小矩形分块。不同分块方式会影响正交化操作的数值稳定性和并行效率。从工程经验看分块不宜过小也不宜过大。一个参考做法是对每个二维权重矩阵先看 shape。如果两个维度都在 1024 附近可以按固定块大小如 128×128 切分。如果维度很大比如 4096×4096可以尝试 256×256 或 512×512。实际最优值需要看 GPU 上的运行时间和 loss 曲线。另外要关注的是正交化操作在混合精度下可能产生数值误差。建议在计算正交化时保持 FP32或者在关键层做额外的数值检查。4.2 动量项的超参调整CMuon 里的动量术语和 Adam 的 beta1 不完全一样。它很可能是对分块后的梯度做指数移动平均再做正交化。这意味着动量系数如何设置直接影响正交化矩阵的稳定性和更新方向的历史信息占比。参考建议先使用常见的动量系数 0.9 作为起点。如果 loss 震荡明显可以尝试 0.95 到 0.99让历史信息占比更高更新更平滑。如果 loss 过于保守、下降很慢可以降低到 0.8 到 0.9。动量系数不能单独调。它和学习率、batch size 是联动关系。batch size 越大梯度噪声越低动量系数可以适当减小。4.3 与现有训练框架的集成不管 CMuon 的实现是独立 PyTorch 优化器、DeepSpeed 的扩展还是 JAX 实现落地时都要先确认它能和当前框架的以下功能兼容混合精度训练AMP / bf16梯度累积分布式数据并行DDP / FSDP / DeepSpeed Zero学习率调度器模型并行 / 序列并行Checkpoint 保存和恢复中优化器状态的兼容性这些环节里最容易出问题的是 FSDP 和 DeepSpeed 的优化器状态分片因为 CMuon 需要维护分块动量矩阵和正交化状态优化器状态的形状可能与标准 AdamW 不同保存和恢复 checkpoint 时要格外小心。4.4 正确的验证方式不要直接拿一个 7B 模型跑完整训练来验证 CMuon 是否有效。成本太高而且调参周期太长。建议分三阶段小模型 小数据集先确认能跑通loss 下降趋势是否正常显存和耗时是多少。标准参数量 固定步数用同一个模型和数据集分别跑 AdamW 和 CMuon固定相同步数对比 loss 曲线和采样质量。大批量 长训练如果前两步都没问题再放大 batch size 和训练步数观察中后期是否更稳定。尤其要注意“采样质量”这个指标。优化器对比不能只看 lossloss 更低不代表模型生成的图像质量更好。在扩散模型里FID 等生成指标才是最终裁判。如果只优化 loss很可能出现训练稳定了但生成效果没有显著提升的情况。5. 哪些场景适合用哪些场景不适合用5.1 适合大规模 DiT 预训练CMuon 这类优化器的目标场景非常明确大规模、长训练、高维参数空间。DiT 类模型包括视频生成、图像生成都属于典型的大规模预训练任务。这类场景下训练稳定性本身就是第一优先级。如果你遇到的情况是模型参数量超过 1B。训练 batch size 很大比如上万。loss 起伏明显经常需要中断恢复。试过调整学习率、warmup、梯度裁剪但收益有限。那么 CMuon 是一个值得尝试的方向。它通过更新方向的正交化可能在源头上减少了 loss 震荡的诱因。5.2 不适合小模型、微调和快速实验小模型和快速实验场景中CMuon 的优势并不明显。原因很简单小模型的参数空间相对简单AdamW 已经能稳定训练微调场景下预训练权重已经提供了一个很好的初始化更新方向受约束的必要性大幅降低。另外分块正交化会带来额外的计算开销即便它有加速收敛的潜力在小模型上也可能体现为“每步更慢”。如果只是快速验证一个想法继续用 AdamW 反而更合适。5.3 如何判断自己是不是应该切换我建议用下面这张表做一个快速判断观察项更适合 AdamW更适合 CMuon模型规模小于 1B大于 1B且层数深训练时长短几小时内长多天甚至数周loss 稳定性已经很稳定明显震荡或频繁 spikebatch size中等或较小很大梯度噪声被压低采样质量满意只是提速问题不满意且怀疑优化过程有问题调试资源想开箱即用愿意做额外超参搜索如果表格里大多数观察项落在右侧才值得为 CMuon 投入时间。如果只是觉得“新优化器可能更好”建议先等更多实验数据不要贸然迁移。6. 一个可复用的判断框架要不要尝试 CMuon6.1 三步走先复现再对比最后看边际收益面对一个新优化器我一般不会直接在生产任务里切换而是按下面三步走在标准任务上复现论文里的做法。如果没有开源代码就根据标题里的机制自己实现一个简化版先在小模型上跑通。固定同一个模型、数据、训练步数和 AdamW 做对比。不要同时改学习率、模型结构、数据增强变量越少越好。判断边际收益。如果 CMuon 让训练更稳定但最终采样质量没有明显提升说明瓶颈不在优化器如果 loss 更稳定且采样指标更好再考虑在更大规模任务上验证。这个过程看起来慢但比盲目替换优化器要可靠得多。优化器是训练系统里的“地基”地基换了上面所有环节都要重新验证。6.2 排查优化器问题时各层级的检查顺序如果你已经决定排查训练不稳定问题而且怀疑是优化器层面引起的可以按以下顺序检查先看 loss 曲线特征是持续震荡、突然 spike还是长期平台期。再看数据层batch 里有没有异常样本、数据增强是否引入了过大噪声、label 有没有错。再看模型层初始化是否合理、attention 层数值是否溢出、梯度范数有没有异常。再看学习率层warmup 够不够、峰值学习率是否偏高、调度器衰减是否太快。最后再考虑优化器机制是否需要对动量做正交化约束、分块怎么切。很多人一上来就怀疑优化器但其实前四层出问题的概率更高。CMuon 只有在确认前几层都没问题之后才应该进入你的尝试列表。6.3 记录实验时的关键元信息尝试新优化器时记录内容不能只有 loss 和精度还应该包括每步训练耗时显存占用变化梯度范数分布动量缓冲区的数值范围正交化操作的耗时占比不同分块大小下的更新方向变化这些信息能帮你判断CMuon 带来的收益到底来自正交化约束还是来自动量的平滑效应或者只是分块引入的隐式正则化。长期来看理解机制比拿到一个能用的配置更重要。建议每次改优化器设置只改一个变量并保留 base line 的完整日志。没有对照的实验等于什么都没做。7. 回到最开始的问题优化器能解决训练不稳定但不能解决所有问题CMuon、分块动量正交化这类工作的最大价值是给了我们一个不同的思考角度当大规模扩散 Transformer 训练怎么调都不稳时还可以从更新方向的结构性约束入手而不是永远停留在学习率和梯度裁剪里打转。但要冷静看待它这不是万能开关。模型结构、数据质量、并行策略、日志监控、Checkpoint 恢复依然是训练工程里的核心。CMuon 只是让“优化过程”这个环节变得更可控一些。如果你现在正被 DiT 训练曲线的毛刺困扰我的建议是先建好基准实验记录所有可复现的日志再在隔离环境里小规模尝试 CMuon。不要看到一个新优化器名字就直接全量切换。训练稳定性是一场系统工程优化器是其中一块重要的拼图而不是全部。真正的高手做法是把优化器当成可替换的组件用框架判断、用小成本验证、用数据决策。这样无论 CMuon 最终能不能成为主流你都能建立起属于自己的训练判断力。