如果你最近在训练扩散 Transformer（Diffusion 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 和训练步数，观察中后期是否更稳定。
尤其要注意“采样质量”这个指标。优化器对比不能只看 loss，loss 更低不代表模型生成的图像质量更好。在扩散模型里，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 稳定性 | 已经很稳定 | 明显震荡或频繁 spike |
| batch size | 中等或较小 | 很大，梯度噪声被压低 |
| 采样质量 | 满意，只是提速问题 | 不满意，且怀疑优化过程有问题 |
| 调试资源 | 想开箱即用 | 愿意做额外超参搜索 |
如果表格里大多数观察项落在右侧，才值得为 CMuon 投入时间。如果只是觉得“新优化器可能更好”，建议先等更多实验数据，不要贸然迁移。
6. 一个可复用的判断框架：要不要尝试 CMuon
6.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 最终能不能成为主流，你都能建立起属于自己的训练判断力。
