Aurora：一种针对矩形矩阵的杠杆感知优化器

# Aurora：一种感知杠杆率的矩形矩阵优化器 | Tilde 来源：https://blog.tilderesearch.com/blog/aurora 返回 (https://tilderesearch.com/research) 5.05.2026 Alec Dewulf*, Dhruv Pai*, Li Yang*, Ashley Zhang*, Ben Keigwin ## 摘要 - 我们发现 Muon 的更新在 Tall（高瘦）矩阵上继承了行范数各向异性，这会导致 MLP 层中相当一部分神经元永久死亡。对更新进行行归一化可以修复此问题，但代价是牺牲正交性。 - 我们制定了在行范数均匀性和正交性联合约束下的最速下降问题，并提出 **Aurora** 优化器作为解决方案。 - 我们使用 Aurora 训练了一个 1.1B 模型，该模型在开源互联网数据上实现了 **100 倍的数据效率**，并在 HellaSwag 等通用评估中优于更大的模型。 - 我们向 modded-nanoGPT speedrun 提交了一个 PR，其中 Aurora **优于当前的 SOTA（最佳）运行**。 - 未调优的 Aurora 仅比传统 Muon 多出 6% 的开销，且可作为直接替换方案。 - 我们开源了黎曼和原生 Aurora 的代码：github.com/tilde-research/aurora-release (https://github.com/tilde-research/aurora-release)。 加载中... **大纲** 1. 引言 (https://blog.tilderesearch.com/blog/aurora#1-introduction) — Muon 的成功与 NorMuon 谜题 2. 行归一化带来的张力 (https://blog.tilderesearch.com/blog/aurora#2-tension-due-to-row-normalization) — 在 Tall 矩阵上，行归一化与正交性相互竞争 3. 归一化防止神经元死亡 (https://blog.tilderesearch.com/blog/aurora#3-normalization-prevents-neuron-death) — Muon 导致 MLP 神经元死亡；杠杆率分数解释了原因 4. 增强 Muon 目标 (https://blog.tilderesearch.com/blog/aurora#4-augmenting-the-muon-objective) — 从联合约束推导 Aurora 5. 结果 (https://blog.tilderesearch.com/blog/aurora#5-results) — nanoGPT speedrun SOTA，1B 预训练及下游评估 6. 拐点 (https://blog.tilderesearch.com/blog/aurora#6-inflection) — 为什么这很重要 ## 1: 引言 Muon 优化器 [1] 最初因其 nanoGPT speedrun 竞赛 [2] 中的成功而备受关注，尽管每一步所需的计算量更大，但其收敛所需的墙钟时间优于 AdamW [3]。此后，人们投入了大量精力通过改进分布式实现 [4]、更精确的正交化例程 [5][6] 以及减少计算和通信开销的方法 [4] 来扩展 Muon 的规模。结果是，Muon（包括裁剪版和每头版）已成为训练前沿规模模型 [7][8][18] 的日益流行的选择。 最近，推动开放研究基准前沿水平的 Muon 变体层出不穷 [9][10]。一个例子是 NorMuon [9]，它目前在 modded-nanoGPT speedrun 中处于 SOTA 地位。NorMuon 在 Muon 的基础上增加了一个步骤，通过其逆 RMS 范数对每一行进行缩放，类似于 Adam 的每参数归一化。值得注意的是，这种行归一化步骤会使 Muon 的更新远离真正的正交梯度。NorMuon 仍然成功的事实表明，Muon 的公式化中可能存在被行归一化所解决的差距。 我们研究了行归一化的影响，发现 Muon 可能导致 MLP 层中的 *神经元死亡*，即某些神经元在训练早期持续接收较小的更新并无法恢复。我们表明，通过重新分配 MLP 层中 up 和 gate 投影更新的质量，使其均匀分布在各行，可以避免这种失败模式。受此观察的启发，我们提出了 **Aurora**，它使用这种机制来防止神经元死亡，同时不牺牲梯度正交化的精度。我们证明，Aurora 在 1B 规模上取得了比 Muon 和 NorMuon 更大的增益，并在 modded-nanoGPT 优化赛道上优于当前的 SOTA。此外，我们发现 Aurora 的性能增益随 MLP 宽度扩展，这使我们相信 Aurora 对于具有大 MLP 扩展因子的网络特别有效。 ## 2: 行归一化带来的张力 Muon 中的核心算法组件是一种用于计算矩阵极分解因子的迭代算法。对于具有薄奇异值分解 (SVD) $G \in \mathbb{R}^{m \times n}$ 的 $G = U\Sigma V^\top$，*极因子* $\operatorname{polar}(G) = UV^\top$ 是 Frobenius 范数下最接近 $G$ 的半正交矩阵。如果 $G$ 是参数 $W$ 的梯度，则 Muon 形成 $\operatorname{polar}(G) = UV^\top$ 并更新 $W \gets W - \eta\, UV^\top$，其中 $\eta$ 为学习率。 存在仅使用 matmul 的迭代算法来计算 $\operatorname{polar}(G)$ [5][6][1]，这 largely 使得 Muon 在大规模下可行。我们在这里重点关注两种这样的算法：来自 modded-nanoGPT 的五次迭代 [21] 和具有 8 步的 Polar Express (PE-8) [5]。我们在附录 A (https://blog.tilderesearch.com/blog/aurora#a-comparing-newton-schulz-iterations) 中对计算极因子的不同算法进行了更全面的概述。 我们发现，在我们的训练设置中（如附录 B (https://blog.tilderesearch.com/blog/aurora#b-training-setting) 所述），Muon 在更精确的正交化下实现了单调更优的下游损失；即在更准确地近似 $\operatorname{polar}(G)$ 的迭代算法下。我们在图 1 (https://blog.tilderesearch.com/blog/aurora#figure-1) 中绘制了使用我们测试的最不精确迭代（五次）和最精确迭代（PE-8）的 Muon 运行结果，以说明这一差距。这一结果证实，极因子精度是强大 Muon 实现中的重要组成部分。 加载中... **图 1.** Muon 随着更精确的正交化例程实现单调降低的损失。在 1B 规模上，PE-8 大幅优于五次迭代。NorMuon 在 Muon 的基础上增加了一个步骤，将极因子的每一行缩放为单位 RMS 范数。 算法：NorMuon 需要：初始权重 $W_0 \in \mathbb{R}^{m \times n}$，损失 $L$，学习率 $\eta$，动量 $(\beta_1, \beta_2)$，扰动 $\varepsilon$，权重衰减 $\lambda$ 初始化 $M_0 \in \mathbb{R}^{m \times n} \leftarrow 0, v_0 \in \mathbb{R}^m \leftarrow 0$ 对于 $t=1, 2, \dots$ 执行 $G_t \leftarrow \nabla_W L(W_t)$ $M_t \leftarrow \beta_1 M_{t-1} + (1-\beta_1) G_t$ $O_t \leftarrow \operatorname{polar}(M_t)$ $v_t \leftarrow \beta_2 v_{t-1} + (1-\beta_2) \operatorname{mean\_cols}(O_t \odot O_t)$ $V_t \leftarrow \operatorname{ExpandRows}(v_t)$ ($V_t \in \mathbb{R}^{m \times n}$) $\hat{O}_t \leftarrow O_t \oslash (V_t + \varepsilon)$ $\hat{\eta} = 0.2 \eta \sqrt{mn} / \|\hat{O}_t\|_F$ $W_{t+1} \leftarrow W_t - \eta \lambda W_t - \hat{\eta} \hat{O}_t$ 结束循环 我们表明，对于 Tall 矩阵，这种行归一化会显著降低极因子精度，这在理论上以及鉴于我们之前的结果（图 1 (https://blog.tilderesearch.com/blog/aurora#figure-1)）都是不希望的。 **命题 1。** Tall 矩阵不能同时具备列正交性和均匀的单位行范数。也就是说，如果 $M \in \mathbb{R}^{m \times n}$ 且 $m > n$ 满足 $M^\top M = I_n$，则 $M$ 不能拥有所有行范数均为一。 *证明.* 假设反证地，$M$ 的所有行范数均为 1。则 $\operatorname{tr}(M^\top M) = \operatorname{tr}(I_n) = n$ 且 $\operatorname{tr}(M^\top M) = \operatorname{tr}(MM^\top) = \sum_{i=1}^m \|M_{i, \cdot}\|_2^2 = m$ 综上所述，这些事实意味着 $m = n$，这与我们的假设 $M$ 是 Tall 矩阵相矛盾。$M$ 中行的平均范数为 $\sqrt{n/m}$。□ 因此，NorMuon 中的行归一化必然会在正交化例程中引入精度缺陷。我们发现，对于行范数不均匀的矩阵，这种极精度缺陷可能非常大。特别是，引入的误差幅度可能远大于相同设置下 PE-8 相对于五次迭代的精度改进（请参阅附录 A (https://blog.tilderesearch.com/blog/aurora#a-comparing-newton-schulz-iterations) 以获取不同 Newton-Schulz 迭代的完整比较）。 行归一化精度 **图 2.** (a) 我们采样了 $n=1000$ 个随机 $512 \times 128$ 矩阵，其行范数标准差分别为 $\{1, 2, 3, 4, 5\}$，并使用不同的极因子算法（PE-8, CANS-10, CANS-12 和五次）对其进行正交化；图中显示了所得极因子的精度（虚线）。然后我们应用单位行归一化并绘制更新后的精度（深色线）。在行范数标准差为 3 时，正交性缺陷峰值达到 0.06。(b) 绘制了每个行标准范数偏差的平均缺陷幅度。事实上，当这些算法与 NorMuon 在同一设置下使用时，PE-8 和五次迭代之间的性能差距消失了。然而，两者的运行都优于我们的 Muon+PE-8 基线，表明行归一化可以独立发挥作用。 加载中... **图 3.** 在 1B 规模下，NorMuon 并未从更精确的正交化中受益。PE-8 和五次迭代实现了相同的下游损失，这与我们关于行归一化降低极因子精度的结果一致。 为了缓解行归一化引入的缺陷，我们可以简单地将 Tall 矩阵归一化，使其行范数为 $\sqrt{n/m}$ 而不是 1。我们将此变体称为 **U-NorMuon**，并发现它在我们的 340M 设置中优于 NorMuon（图 4 (https://blog.tilderesearch.com/blog/aurora#figure-4)；设置描述见附录 B (https://blog.tilderesearch.com/blog/aurora#b-training-setting)）。 算法：U-NorMuon 需要：初始权重 $W_0 \in \mathbb{R}^{m \times n}$，损失 $L$，学习率 $\eta$，动量 $(\beta_1, \beta_2)$，扰动 $\varepsilon$，权重衰减 $\lambda$ 初始化 $M_0 \in \mathbb{R}^{m \times n} \leftarrow 0, v_0 \in \mathbb{R}^m \leftarrow 0$ 对于 $t=1, 2, \dots$ 执行 $G_t \leftarrow \nabla_W L(W_t)$ $M_t \leftarrow \beta_1 M_{t-1} + (1-\beta_1) G_t$ $O_t \leftarrow \operatorname{polar}(M_t)$ $v_t \leftarrow \beta_2 v_{t-1} + (1-\beta_2) \operatorname{mean\_cols}(O_t \odot O_t)$ $V_t \leftarrow \operatorname{ExpandRows}(v_t)$ ($V_t \in \mathbb{R}^{m \times n}$) $\hat{O}_t \leftarrow O_t \oslash (V_t + \varepsilon)$ $\hat{\eta} = 0.2 \eta n / \|\hat{O}_t\|_F$ $W_{t+1} \leftarrow W_t - \eta \lambda W_t - \hat{\eta} \hat{O}_t$ 结束循环 现在我们将分析转向宽矩阵和方阵，它们在 NorMuon 下也接收单位行范数更新。左正交的宽矩阵或方阵必然具有所有行范数均为 1，因此在这种情况下，行范数均匀性是由正交性 *隐含* 的。 因此，我们预计在使用 PE-8 等精确正交化例程时，行归一化是不必要的，甚至可能是有害的。在训练早期，当原始梯度往往条件较差时，迭代极因子算法可能无法精确收敛，导致行范数统计量的积累迅速变得陈旧。在这种情况下，行归一化会使行变得 *更不* 均匀，直到统计量得到纠正。我们在 340M 规模上发现了这一现象的证据，其中仅对 Tall 矩阵进行行归一化以微小但非平凡的幅度优于 NorMuon 和 U-NorMuon（图 4 (https://blog.tilderesearch.com/blog/aurora#figure-4)）。 加载中... **图 4.** 在 NorMuon 下，down 投影的列归一化与 up 和 gate 行归一化表现非常相似。仅应用于 Tall 矩阵的 NorMuon 优于所有 NorMuon 变体和 Muon。 我们将展示，Muon 可能允许 MLP 层中的大部分神经元实际上死亡，但这一病理现象被 (U-)NorMuon 所缓解。基于文献中的思想，我们提出这一作为在我们所有设置中发现的 Muon 和 NorMuon 之间性能差距的解释。然后我们推导了 Aurora，它有效地对 Tall 参数的更新进行行归一化，而不牺牲极因子的精度。 ## 3: 归一化防止神经元死亡 我们将 *死亡模型组件* 定义为模型参数的一组子集，在训练的最早阶段之后持续接收微小的学习信号。我们使用以下三个标准来识别死亡的模型组件： 1. **低有效梯度范数。** 对于特定批次，接收低有效梯度范数的模型组件必须对应于损失景观中相对平坦的方向。具有持续小梯度范数的组件因此将接收很少的学习信号。在归因和剪枝文献中，使用梯度范数作为特征重要性的代理有着强烈的先例 [11][12][13]。 2. **低有效 *更新* 范数。** 对于使用 Muon 等优化器训练的网络，小的梯度范数本身并不足以识别死亡的网络组件。特别是，正交化步骤在应用之前可以剧烈改变梯度，包括放大小幅度方向。因此，我们要求有效更新范数，它衡量实际