Aurora: 一种杠杆感知的谱优化器

# 1 引言  
来源：https://arxiv.org/html/2606.27715  
![[未加标题的图片]](https://arxiv.org/html/2606.27715v1/figures/tilde_logo.png)  
2026年6月26日  
Aurora: 一种杠杆感知的谱优化器  
Alec Dewulf¹¹通信邮箱：alec@tilderesearch\.com, Dhruv Pai, Li Yang, Ashley Zhang, Ben Keigwin  
Tilde Research  

###### 摘要  
我们表明，对于高矩阵参数（如MLP层中的投影矩阵），Muon更新的行范数可能任意不均匀。这可能导致一个自我强化的反馈循环，其中神经元持续接收很小的更新，最终无法对网络输出做出有意义的贡献。通过额外的行归一化步骤可以有效缓解这一问题，但当前方法在这样做时会偏离Muon更新几何中动量矩阵的极因子，而我们发现这是不可取的。我们提出Aurora优化器，它在保留Muon极因子几何的同时，强制矩阵参数更新的行均匀性。在预训练实验中，Aurora优于Muon，并且与现有方法结合时，在modded-nanoGPT速度竞赛的优化器赛道中取得了谱优化器中的最新最优性能。此外，我们发现Aurora相对于Muon的实证收益随MLP扩展因子增大而增加，这表明Aurora可能允许有效训练非常宽的MLP层。  

Muon\[14 (https://arxiv.org/html/2606.27715#bib.bib14)\]已成为训练至前沿规模的语言模型时越来越流行的优化器\[7 (https://arxiv.org/html/2606.27715#bib.bib7),17 (https://arxiv.org/html/2606.27715#bib.bib17),37 (https://arxiv.org/html/2606.27715#bib.bib37),11 (https://arxiv.org/html/2606.27715#bib.bib11)\]。Muon的核心操作是一个正交化例程，它近似地将矩阵更新替换为其极因子，保留更新的奇异向量同时去除奇异值尺度。Muon属于一个不断壮大的谱优化器家族，这些优化器对更新矩阵的谱应用某种显式变换。大多数谱优化器使用正交化，但最近的工作探索了使用更柔和的奇异值映射的可能性，即对大奇异值进行不那么激进的收缩\[19 (https://arxiv.org/html/2606.27715#bib.bib19)\]。其他谱优化器增加了条件步骤、归一化和自适应重缩放，某些变体在社区研究基准（如modded-nanoGPT\[15 (https://arxiv.org/html/2606.27715#bib.bib15),4 (https://arxiv.org/html/2606.27715#bib.bib4),36 (https://arxiv.org/html/2606.27715#bib.bib36)\]）上显示出显著加速的收敛速度。其中一个例子是NorMuon\[20 (https://arxiv.org/html/2606.27715#bib.bib20)\]，它将极因子的每一行按其逆RMS范数重新缩放。提出此行归一化步骤是为了控制矩形矩阵更新行范数的方差，在Muon下该方差可能很大。我们表明，这种方差可能导致MLP层中参数的严重利用不足，并且在某些情况下神经元可能变得实际上死亡。我们发现行归一化有效解决了这个问题，但代价是修改了Muon更新的几何，这在实证上和Muon理论公式下都是不可取的。我们提出Aurora优化器，作为一种在不牺牲极因子计算精度的情况下控制神经元更新行范数方差的方法。Aurora更新规则专门针对MLP层中的up和gate投影矩阵。Aurora在我们所有实验中有效防止了神经元死亡，并且与Muon和其他行归一化谱优化器相比，收敛速度更快。此外，我们发现Aurora相对于Muon的实证收益随MLP扩展因子增大而增加，这表明Aurora可能允许有效训练具有异常宽MLP层的网络。  

#### 贡献。  
我们的贡献如下。  
1. 1\. **归一化谱优化器的分析**。我们表明，事后单位行归一化可能使Muon更新几何远离真实极因子的几何；此外，我们表明对于精确正交化下的方阵或宽参数，不需要行归一化。这导致了U-NorMuon的设计，它是NorMuon的一个简单修改，将行归一化限制为仅应用于高矩阵，并去除了额外的行范数缓冲区。在我们的实验中，U-NorMuon优于Muon，这为行归一化和精确正交化同时可取这一假设提供了证据。  
2. 2\. **Muon下神经元死亡的研究**。我们发现，使用Muon训练的网络可能在MLP层中出现显著的神经元死亡。我们提出了一种机制解释，说明神经元如何在Muon下死亡，并表明U-NorMuon有效防止了这一问题。为验证结果，我们研究了使用Muon训练的开源权重模型，并在密集MLP层中发现了神经元死亡的强烈证据。  
3. 3\. **Aurora和Riemannian-Aurora**。我们提出Aurora，作为一种在严格遵守Muon理论框架的同时强制统一行范数的方法。我们还推导出Riemannian-Aurora作为参考求解器。Riemannian-Aurora在Stiefel流形和等行范数流形交集的切空间中形成更新，以更高代价更忠实地求解这一双重约束问题。我们发现实际实例化Aurora使用的迭代解法在实践中通常也非常精确。  
4. 4\. **广泛的预训练验证**。在340M和1.1B规模下，使用Aurora训练的模型在下游验证损失和基准分数上均优于Muon和NorMuon。值得注意的是，我们使用Aurora训练的1.1B模型在MMLU上比Muon基线高出9.1分。在modded-nanoGPT速度竞赛（优化器赛道）中，Aurora改进了NorMuon基线，并与现有方法结合，在撰写本文时创下了最新最优性能。我们还发现其相对于Muon的性能增益随MLP宽度增加而增大。  

#### 组织结构。  
在第2节 (https://arxiv.org/html/2606.27715#S2)中，我们提供Muon和行归一化谱优化器的背景。我们还讨论了计算Muon中极因子的算法相关近期研究。在第3节 (https://arxiv.org/html/2606.27715#S3)中，我们表明极因子精度和统一行范数同时可取。我们提出U-NorMuon作为通向Aurora的中间步骤，它仅对高矩阵参数应用无状态归一化。在第4节 (https://arxiv.org/html/2606.27715#S4)中，我们形式化了双重约束问题，即同时强制半正交性和行范数均匀性，并给出Aurora和Riemannian-Aurora作为解决方案。最后，在第5节 (https://arxiv.org/html/2606.27715#S5)中，我们展示Aurora的实证验证，然后在第6节 (https://arxiv.org/html/2606.27715#S6)中讨论我们工作的更广泛含义及未来方向。  

## 2 预备知识与相关工作  
符号说明。 对于矩阵 \(A\in\mathbb{R}^{m\times n}\)，我们用 \(\|A\|_F\) 表示其Frobenius范数，用 \(\|A\|_2\) 表示其谱范数。\(A\) 的奇异值记为 \(\sigma_1(A)\geq\cdots\geq\sigma_{\min\{m,n\}}(A)\)。Hadamard积记为 \(\odot\)。对于方阵 \(S\)，\(\operatorname{diag}(S)\) 表示由 \(S\) 的对角元素构成的向量；而对于向量 \(x\)，\(\operatorname{diag}(x)\) 表示以 \(x\) 为对角线的对角矩阵。我们记 \(\operatorname{sym}(S)=\frac{1}{2}(S+S^\top)\)。对于具有瘦SVD分解 \(A=U_r\Sigma V^\top\) 的 \(A\)，定义 \(\operatorname{polar}(A)=U_rV^\top\)。我们用 \(d_{\mathrm{model}}\) 表示模型残差流的维度。  

### 2.1 Muon及其归一化谱优化器  
Muon是一种用于隐藏层权重矩阵的优化器。其更新可以看作是在谱范数信任区域中最陡的一阶下降方向，该方向由梯度动量的极因子给出。对于权重矩阵 \(W\in\mathbb{R}^{m\times n}\)，设 \(M_t\) 为第 \(t\) 步的梯度动量，Muon形成更新：  
\[\Delta W_t = -\eta \operatorname{polar}(M_t) = -\eta UV^\top, \quad M_t = U\Sigma V^\top.\]  
其中 \(\eta\) 是常数学习率。另一个视角是，Muon对矩阵参数执行关于谱范数的最陡下降。Adam(W) 通常用于不自然表示欧几里得空间之间线性映射的参数，例如RMSNorm缩放因子和嵌入/反嵌入矩阵。  

最近，一套“归一化”的Muon变体在社区研究基准中变得流行。这些优化器在Muon算法中添加了一个显式的行或列归一化步骤，位置在极因子计算之前或之后。例如，NorMuon维护每行的二阶矩统计量，并用它们来重新缩放正交化更新矩阵的行。NorMuon的动机来自于观察到Muon更新中高矩阵的行范数可能不均匀。我们将基于这一结果，表明这种更新中的不均匀性可能导致MLP层中参数的严重利用不足。MuonEq则在正交化之前重新缩放动量矩阵，这可以改善Newton-Schulz迭代的收敛性。Aurora类似于MuonEq-R的迭代版本；然而，Aurora仅应用于高矩阵参数，而MuonEq-R应用于所有矩阵参数。我们发现这两个差异——仅归一化高参数的更新，以及迭代细化更新——各自都能独立提升下游性能。一个相关但不同的工作路线使用归一化来提高优化器的内存和计算效率，将无状态的行/列重缩放应用于梯度，作为自适应预条件或正交化本身的轻量级替代\[23 (https://arxiv.org/html/2606.27715#bib.bib23),33 (https://arxiv.org/html/2606.27715#bib.bib33),10 (https://arxiv.org/html/2606.27715#bib.bib10),41 (https://arxiv.org/html/2606.27715#bib.bib41),42 (https://arxiv.org/html/2606.27715#bib.bib42),8 (https://arxiv.org/html/2606.27715#bib.bib8)\]。Aurora的行归一化步骤在机械上与其中几种方法相似，但动机不同。Aurora对更新同时进行行归一化和正交化，实现了精确的行均匀极因子。如我们在第3.1节 (https://arxiv.org/html/2606.27715#S3.SS1)中所展示的，对于高矩阵，行归一化和正交化之间存在天然的张力；解决这种张力是Aurora要处理的核心设计问题。  

### 2.2 计算极因子的算法  
Newton-Schulz (NS) 迭代是一种不动点递推，形式为 \(X_{t+1}=p(X_t)\)，其中 \(p\) 是某个矩阵多项式。实际的Muon实现通常使用NS类型的迭代来近似极因子，通常是通过对Gram矩阵隐式应用平方根逆迭代来实现的。对于矩阵 \(G\in\mathbb{R}^{m\times n}\)，这些方法在 \(T\) 次迭代后产生一个近似 \(\widehat{P}_T(G) \approx \operatorname{polar}(G)\)。Muon之所以能在大规模下可行，很大程度上是因为存在仅需要矩阵乘法的算法来计算极因子。先前的工作发现，更精确的NS迭代（即产生更接近 \(\operatorname{polar}(G)\) 的 \(\widehat{P}_T(G)\) 的迭代）能带来更好的下游性能，但收益在饱和点之后很快递减\[16 (https://arxiv.org/html/2606.27715#bib.bib16)\]。在本文中，我们将考虑三种常用的算法：Jordan的五次多项式五步版本 (quintic-5)\[14 (https://arxiv.org/html/2606.27715#bib.bib14)\]、Polar Express八步版本 (PE-8)\[2 (https://arxiv.org/html/2606.27715#bib.bib2)\] 和CANS三度十二步版本 (CANS-12)\[12 (https://arxiv.org/html/2606.27715#bib.bib12)\]。CANS-12和PE-8都需要24次矩阵乘法，而quintic-5只需15次。我们发现这种计算差异在我们分布式Muon设置中引起的挂钟开销可以忽略不计。  

参照图注  
图1：绘制CANS-12、quintic-5和PE-8在奇异值介于0和1之间的表现。Quinctic-5误差最大，CANS-12对小奇异值的收敛性略优于PE-8。  

直观上，更精确的方法平均将奇异值映射得更接近1。CANS-12和PE-8对大奇异值都非常精确，但对于小奇异值，CANS-12可以明显更精确（图1 (https://arxiv.org/html/2606.27715#S2.F1)）。在第3.1节 (https://arxiv.org/html/2606.27715#S3.SS1)中，我们将实证测量这种精度差距在我们训练设置中的程度。  

## 3 极精度与神经元更新不平衡  
我们在设置中研究了不同的NS迭代，并发现对于更高的极因子精度，验证损失单调降低。受此结果启发，我们研究了归一化对极因子的影响，并表明当前的逐行/逐列归一化方案必然会分别降低高矩阵和宽矩阵参数的极精度（第3.1节 (https://arxiv.org/html/2606.27715#S3.SS1)）。然后，我们研究了使用Muon训练的网络中的参数效率，并发现了MLP层中显著神经元死亡的证据。我们表明，对Muon更新的行归一化有效缓解了这种病态（第3.2节 (https://arxiv.org/html/2606.27715#S3.SS2)），从而改善了下游评估结果。我们通过检查开源权重模型中的死亡神经元来验证结果的普遍性（第3.3节 (https://arxiv.org/html/2606.27715#S3.SS3)）。在所有情况中，我们发现了相当数量的死亡神经元群体，尤其在早期层中集中。与我们提出的机制解释一致，我们在所检查模型的细粒度专家中没有发现死亡神经元。  

### 3.1 分析归一化方案  
很难对不同NS迭代进行绝对比较，因为它们的收敛性很大程度上取决于输入的形状和条件。为了本研究的目的，我们将使用以下实证误差度量来比较这些方法的准确性。  

###### 定义3.1 (极近似误差，PAE)。  
对于矩阵 \(X\in\mathbb{R}^{m\times n}\) 和Newton–Schulz迭代 \(\operatorname{A}\)，设 \(\operatorname{A}_k(X)\) 表示将 \(\operatorname{A}\) 应用于 \(X\) 进行 \(k\) 步后的结果。我们定义 \(\operatorname{A}\) 在 \(X\) 上进行 \(k\) 次迭代后的极近似误差为：  
\[\varepsilon_{k,\operatorname{A}}(X) \coloneqq \frac{\left\|\operatorname{A}_k(X) - \operatorname{polar}(X)\right\|_{\mathrm{F}}}{\left\|\operatorname{polar}(X)\right\|_{\mathrm{F}}}.\]  
我们使用不同NS迭代的Muon训练一个普通的340M transformer，并在训练的每一步计算所有被Muon更新的矩阵参数的平均极近似误差（完整训练细节见第5节 (https://arxiv.org/html/2606.27715#S5)）。直观上，该指标捕捉了近似正交化引入的平均误差。我们还使用奇异值分解运行Muon。