为扩散变换器扩展 Muon

# 面向扩散Transformer的Muon规模化扩展  
来源：https://arxiv.org/html/2608.20818  
Chenghao Li1,2, Xiao Han2, Xinxin Huang2, Wei Liu2, Boyang Li2, Bing Xiao2, Heran Zhang2, Juanma Perez Rua2, Ke Xu2, Kangning Liu2, Linjun Kuang2, Na Li2, Tan Wang2, Tian Xie2, Wei Peng2, Yang Pei2, Yifan Xu2, Yuanhao Zhai2, Yuwei Lin2, Zhe Wang2, Zihao He2, Daniel Li2, Junbiao Tang2, Ziyang Jiang2, Dake Chen2  
1南加州大学 2Meta  

###### 摘要  
矩阵感知优化器Muon通过平衡奇异方向的更新来提升大模型训练效果，但其在大规模扩散Transformer（DiTs）上的规模化行为与端到端效率尚不明确。我们首次建立了Muon在参数规模从1.3B到15B的DiTs上的规模化行为，表明其相对于AdamW的优化与生成质量优势在不同模型规模下均得以保持。然而，在大规模场景中，每个优化步骤执行的5步Newton-Schulz迭代（NS5）以及全动量物化会引入显著的计算与通信开销，可能抵消Muon的步效率优势。为此，我们提出*周期性逐行Muon*方法，该方法每K步执行一次完整的NS5谱更新，并在其余步骤基于当前动量应用低计算与通信成本的逐行约束更新。我们进一步协同设计了分布式实现方案，在非刷新步骤直接对分片动量进行操作，并通过分桶全聚集与通信-计算重叠加速谱刷新。在所有规模下，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在保留Muon生成质量优势的同时，将其转化为面向大规模DiTs的端到端训练效率。

## 1 引言  
扩散Transformer（DiTs）的持续规模化扩展提升了模型容量与生成质量，同时也使训练效率成为核心关注点。先前工作已表征了增加训练计算量带来的可预测改进，但通常将优化器视为固定不变。然而在十亿参数规模下，优化器效率不能仅通过达到更低损失所需的步骤数来评判。在分布式训练中，更少的优化步骤必须转化为更少的GPU小时数与更短的挂钟时间。因此评估优化器可扩展性需同时考虑其优化优势及其在分布式训练中实际实现的计算与通信成本。  
Muon为提升大规模训练效率提供了有前景的替代方案。与应用坐标级自适应更新的AdamW不同，Muon将二维权重张量视为矩阵，并对其动量应用有限步Newton-Schulz（NS）变换，生成具有全局谱结构的更新方向。近期研究表明，在适当校准下，Muon可在大型语言模型预训练与后训练中以显著更少的训练浮点运算量匹配AdamW。来自较小扩散模型的证据同样显示了尽管单步计算与通信成本更高，优化性能仍有所改善，但同时也表明损失、生成质量与运行时长可能对优化器的评价产生不同排序。因此尚不明确当DiTs扩展到百亿参数时，Muon的优化优势是否持续存在，以及其分布式开销能否充分降低以提升端到端训练效率。回答这些问题需要联合考虑优化与计算通信成本。  
在本研究的分片执行路径中，Muon的NS5变换引入了额外矩阵乘法、全动量通信、同步与临时物化。这些成本促使我们降低完整谱变换的频率与暴露开销，同时保留Muon的优化优势。我们首先通过在GPIC数据集上训练1.3B到15B参数的DiTs，确认Muon的优化优势在规模化下持续存在。在所有四个规模下，Muon均取得更低的验证损失，并将最佳观测FD-DINO较AdamW提升12.9%–19.1%。然而，其额外计算与通信使原始Muon的单步成本显著更高。这些结果证明了Muon在大规模DiTs上的算法可扩展性，同时暴露出高效实现其优化优势的系统瓶颈。  
为解决此瓶颈，我们重新审视是否必须在每个优化步骤强制施加全局谱几何结构，并引入*周期性逐行Muon*。该方法每K步执行一次完整NS5更新，其余步骤应用行归一化算子（RowNorm）。该方法在两种互补的矩阵几何结构间交替：周期性谱刷新步骤提供行与奇异方向间的全局耦合；而低计算通信成本的逐行约束步骤提供局部尺度控制，维持连续谱刷新间的稳定优化行为。我们协同设计了分布式执行方案与周期更新机制：在非刷新步骤，RowNorm直接对分片动量操作；当归一化行跨越排名时，仅进行范数统计的全约减而非整个矩阵；在刷新步骤，我们对即将处理的矩阵进行分桶动量全聚集，并与已可用桶上的NS计算形成流水线，以降低暴露的通信延迟。该设计移除了非刷新步骤的全动量物化，并减少了剩余谱刷新的暴露开销。  
在1.3B–4B模型上，周期性逐行Muon的最佳观测FD-DINO与原始Muon保持0.5%以内差异；在9B模型上优于其4.5%，在15B模型上优于2.7%。相对于原始Muon，其优化器时间减少46.9%–54.3%，端到端步时间缩短15.7%–24.3%，逻辑优化器通信量降低66.7%。我们的主要贡献如下：  
(1) 表征了Muon在1.3B–15B DiTs上的表现，证明其相对于AdamW的验证损失与生成质量优势在规模化下持续存在，同时识别出NS5计算与全动量通信为主要系统瓶颈；  
(2) 提出周期性逐行Muon，用基于当前动量的低计算通信成本RowNorm更新替代大部分NS5变换；  
(3) 开发了分布式实现方案，在非刷新步骤对分片动量操作，在刷新步骤将分桶全聚集与NS5计算形成流水线，在保持与原始Muon相当的生成质量的同时大幅降低系统开销。

## 2 面向DiTs的Muon  
### 2.1 DiTs训练中的Muon规模化扩展  
给定干净数据样本x、扩散时间步a与噪声ε，令xₐ表示根据指定噪声调度构建的含噪输入。我们通过最小化损失函数来训练DiT预测对应的速度目标v⋆：  
\( \mathcal{L}(\theta) = \mathbb{E}_{x,a,\epsilon} \left[ \left\lVert v_\theta(x_a,a) - v^\star(x_a,a) \right\rVert_2^2 \right] \)  
为简化表示，省略了文本嵌入等其他条件变量。AdamW基线对所有可训练参数应用AdamW。Muon则对二维隐层权重矩阵应用矩阵感知更新，而偏置、归一化参数等非矩阵参数仍由AdamW更新。  
考虑训练步t处的权重矩阵 \( W_t \in \mathbb{R}^{m \times n} \)，其梯度 \( G_t = \nabla_{W_t} \mathcal{L}(\theta_t) \)。Muon对其动量 \( M_t \) 应用由 \( J_{\mathrm{NS}}=5 \) 次Newton-Schulz迭代构成的有限谱变换 \( \Phi_5 \)，并执行解耦权重衰减更新：  
\( \mathcal{P}_t = \Phi_5(M_t), \qquad W_{t+1} = (1 - \eta_t \lambda) W_t - \eta_t \, s(W_t) \mathcal{P}_t \)  
其中 \( \eta_t \) 为学习率，λ为权重衰减系数，\( s(W_t) \) 为依赖形状的更新尺度。有限映射 \( \Phi_5 \) 受矩阵极分解启发，将动量的奇异值重塑为全局耦合的谱方向 \( \mathcal{P}_t \)。重要的是，NS5不等同于计算精确SVD极分解，且在5次迭代后通常不会生成精确正交矩阵。因此我们的算法定义、理论分析与实验均采用已实现的有限映射 \( \Phi_5 \)，而非假设其收敛至精确极分解。  
我们从算法与系统两个维度刻画Muon在大规模DiTs训练中的可扩展性：算法可扩展性回答Muon是否随模型规模增长保持稳定优化行为，及其相对于AdamW的验证损失与生成质量优势是否在不同模型规模下持续存在；系统可扩展性回答这些算法优势能否在规模化分布式训练中转化为实际效率提升。

### 2.2 计算复杂度  
Muon与AdamW最直接的计算差异源于对二维矩阵的更新。考虑 \( W \in \mathbb{R}^{m \times n} \)，令 \( r = \min(m,n) \)，\( c = \max(m,n) \)。由于AdamW维护并应用逐元素统计量，该矩阵的算术复杂度为 \( C_{\mathrm{AdamW}}(m,n) = \Theta(mn) = \Theta(rc) \)。相反，有限步Newton-Schulz变换需要重复执行矩阵乘法。即使选择较小的Gram矩阵方向，其复杂度也为 \( C_{\mathrm{NS}}(m,n) = \Theta(J_{\mathrm{NS}} r^2 c) \)。因此NS5为每个矩阵引入额外 \( \Theta(J_{\mathrm{NS}} r) \) 倍的算术运算量。忽略实现常数与硬件吞吐量，单矩阵的算术成本比为 \( \Theta(J_{\mathrm{NS}} r) \)。

### 2.3 规模化分布式训练下的通信成本  
仅局部算术量无法描述Muon在规模化下的成本，还取决于其完整矩阵谱变换如何映射至分片优化器状态。在我们的执行路径中，每个二维参数及其动量沿第零维分片，因此每个排名存储行子集及对应动量分片。AdamW具有坐标可分性，因为每个排名可仅使用本地一阶与二阶矩分片更新本地参数，无需物化完整矩阵。在此执行路径下，Muon不具有分片可分性，因为Newton-Schulz迭代通过Gram矩阵与矩阵乘法耦合行与奇异方向。通常，排名无法仅从其本地动量分片恢复其谱更新部分。因此我们的原始实现将动量分片全聚集至所有排名，在每个排名执行完整矩阵的NS5，保留与本地参数分片对应的输出切片，随后丢弃剩余输出并释放临时完整矩阵缓冲区。  
Muon的系统成本因此包括复制NS5执行的GEMM及相关矩阵操作，以及全动量通信与临时完整矩阵物化。假设分片组包含p个排名，每个动量元素占用b字节。忽略填充与协议开销，全聚集一个 \( m \times n \) 动量矩阵要求每个排名逻辑接收约 \( V_{\mathrm{AG}}(m,n) \approx bmn \frac{p-1}{p} \) 字节。该数量已排除前向与反向传播所需的参数通信。然而仅逻辑通信量不能决定通信时间，还取决于集合通信次数、桶大小、启动延迟、节点间流量比例、网络拓扑与竞争以及通信-计算重叠程度。特别是，许多小桶可能在不改变总通信量的情况下引入显著启动延迟。  
重复的谱计算与全动量通信使Muon的单步成本显著更高，可能抵消其优化优势。这促使我们在保留有效更新的同时降低完整谱变换的频率。

## 3 周期性逐行Muon  
我们提出*周期性逐行Muon*，周期性执行NS5谱更新，并在其余步骤对当前动量应用RowNorm。与将RowNorm视为NS5的数值近似不同，我们的方法在由不同约束几何诱导的两种矩阵更新间交替。随后利用RowNorm的局部性减少动量集合通信与完整矩阵物化。

### 3.1 互补的谱与逐行几何结构  
对于二维动量 \( M_t \in \mathbb{R}^{m \times n} \)，令 \( \widetilde{M}_t = \mathcal{O}(M_t) \) 表示：若 \( m \leq n \) 则为 \( M_t \)，否则为 \( M_t^\top \)，从而 \( \widetilde{M}_t \in \mathbb{R}^{r \times c} \)（其中 \( r = \min(m,n) \leq c = \max(m,n) \)）。所得方向通过 \( \mathcal{O}^{-1} \) 映射回原空间。对于满行秩矩阵 \( X \in \mathbb{R}^{r \times c} \)，理想极分解 \( \mathcal{P}(X) \) 为：  
\( \mathcal{P}(X) = (XX^\top)^{-1/2} X \in \operatorname*{arg\,max}_{\|U\|_2 \leq 1} \langle X,U \rangle, \qquad \langle A,B \rangle = \operatorname{tr}(A^\top B) \)  
谱范数约束全局耦合X的行与奇异方向。实际Muon采用受此谱几何启发的有限NS5映射，而非计算精确极分解。  
先前在对称性兼容优化器设计中研究的RowNorm则独立归一化每行：  
\( \mathcal{R}_\epsilon(X)_{i:} = \frac{X_{i:}}{\max\left(\lVert X_{i:} \rVert_2, \epsilon\right)}, \qquad i=1,\ldots,r \)