CMuon：通过分块动量正交化加速并稳定扩散Transformer训练

摘要

扩散Transformer（DiTs）在视觉生成模型中取得了最先进的（SOTA）性能，但其训练仍然非常耗时。虽然最近提出的动量正交化（Muon）优化器为AdamW提供了一个有希望的替代方案，但将其直接应用于DiTs导致其在晚期阶段的收敛效果不佳。在这篇论文中，我们识别了这一瓶颈的根本原因：标准DiT架构将功能不同的权重（例如，AdaLN和QKV层中的权重）融合成统一的张量以提高计算效率。将Muon应用于这些融合的张量无意中诱导了隐性的子空间耦合，这扭曲了更新方向并降低了全局优化。为了解决这个问题，我们引入了分块动量正交化（CMuon），这是一种简单而高度有效的策略，通过将这些矩阵分割成独立的子组件来避免动量的直接应用。广泛的实验结果表明，使用CMuon训练的675M参数的DiT在ImageNet 256上在200个epoch后实现了1.18的FID。这比使用AdamW的训练速度提高了2倍，同时有效地克服了基本动量（Muon）在晚期阶段的收敛平顶问题。

CMuon：通过分块动量正交化加速并稳定扩散Transformer训练：论文配图 — 图1 ： CMuon的概述和实证验证。左图： CMuon通过在扩散变换器（ DiT ）中分块选定的权重矩阵，以去耦的方式应用动量正交化，用每块正交化替换单个全局正交化。这种设计减轻了跨块干扰，提高了相对于香草Muon的优化稳定性和效率。中：我们提出的CMuon的插图。右图：在ImageNet上， 256 × 256256\ times 256 ， CMuon在训练675M参数DiT模型时比原生Muon产生了可观的性能提升。