1/4 新论文与 @weijie444！我们提出一种对称兼容的 LLM 优化器设计原则，并作为副产品，获得了一个端到端的逐层优化器堆栈，其中每个主要的矩阵值参数（嵌入、LM 头、SwiGLU MLP、MoE 路由器）都有自己原则性的更新！http://arxiv.org/abs/2605.18106 http://github.com/timlautk/equivariant_optimizers…

嵌入、LM 头、SwiGLU MLP 和 MoE 路由器

来源：https://arxiv.org/html/2605.18106

MnLargeSymbols’164 MnLargeSymbols’171

优化器设计的对称兼容原则：嵌入、LM 头、SwiGLU MLP 和 MoE 路由器

摘要

深度学习实践中长期存在一种显著的几何差异。现代神经网络架构天然展现出丰富的对称性和等变性，而流行的优化方法（如 Adam 及其变体）本质上是逐坐标操作的，无法尊重参数空间的等变结构。在本文中，我们通过引入优化器设计的对称兼容原则来解决这一差异。具体而言，我们认为梯度更新规则应作用于神经网络的相应权重块时，在对称群的作用下保持等变。遵循这一原则，我们首先为通用矩阵层的双正交等变更新提供了一个统一视角，这种更新已被随机谱下降、Muon、Scion 和极坐标梯度方法所采用。更重要的是，通过从正交群转向置换群和共享移位对称性，我们推导出针对参数块（其对称性与通用矩阵层不同）的新型对称兼容优化器类别：对于嵌入和 LM 头矩阵，左置换和右正交等变性导致单边谱更新、行范数更新以及混合行范数/谱更新；对于 SwiGLU MLP 投影，中间神经元置换对称性激发了行感知和列感知变体；对于 MoE 路由器，专家置换对称性结合共享对数移位不变性，产生了中心化行范数更新和左谱更新。这些构造形成了一个端到端的逐层优化器堆栈，其中每个主要的矩阵值参数类别都被分配了一个更新，其等变性与该参数类别的对称群相匹配。我们通过在密集和稀疏 MoE 语言模型上的大量预训练实验（包括 Qwen3-0.6B 风格、Gemma 3 1B 风格、OLMoE-1B-7B 风格以及缩小版 gpt-oss 架构）验证了这一优化器设计原则。在这些实验中，与相应的 AdamW 更新相比，对称兼容的更新规则一致地改善了最终验证损失，并在多个案例中提升了训练稳定性。

https://github.com/timlautk/equivariant_optimizers

1 引言

深度学习中最广泛使用的优化器，如 Adam [81 (https://arxiv.org/html/2605.18106#bib.bib74)]、Adafactor [136 (https://arxiv.org/html/2605.18106#bib.bib798)]、RMSprop [147 (https://arxiv.org/html/2605.18106#bib.bib78)]、AdaGrad [43 (https://arxiv.org/html/2605.18106#bib.bib75),114 (https://arxiv.org/html/2605.18106#bib.bib640)] 及其变体，都属于逐坐标自适应梯度方法这一大类。这些方法将模型参数视为一个单一的长连接向量，并独立更新每个坐标。尽管它们取得了经验上的成功，但这种设计隐含地假设权重矩阵的每个条目都是高维向量空间中的一个独立坐标。这一假设很少被质疑，却深刻塑造了现代神经网络的训练动态。具体而言，这种忽视几何特性的处理忽略了神经网络参数丰富的矩阵结构，并且未能区分不同层类型（如嵌入层、LM 头、密集线性层、注意力投影、SwiGLU MLP 投影和 MoE 路由器）的几何结构。

与此同时，我们对优化器行为的理论理解在两大与当代大规模训练最相关的家族中仍然有限：逐坐标自适应梯度优化器和谱优化器。特别是在语言模型预训练中，这些优化器家族之间的比较在很大程度上仍然依赖经验，依赖于大规模基准测试练习 [153 (https://arxiv.org/html/2605.18106#bib.bib956),134 (https://arxiv.org/html/2605.18106#bib.bib957)] 和速通 [75 (https://arxiv.org/html/2605.18106#bib.bib781)]，而对其不同几何行为和训练动态的分析相对较少。超参数迁移规则 [160 (https://arxiv.org/html/2605.18106#bib.bib697)] 和缩放定律处方 [77 (https://arxiv.org/html/2605.18106#bib.bib522),67 (https://arxiv.org/html/2605.18106#bib.bib482)] 通常跨优化器应用，尽管它们的原始开发主要与逐坐标自适应方法（特别是 AdamW [109 (https://arxiv.org/html/2605.18106#bib.bib528)]）相关。另一个值得注意的基准测试工作是 AlgoPerf：训练算法 [29 (https://arxiv.org/html/2605.18106#bib.bib527),78 (https://arxiv.org/html/2605.18106#bib.bib833)]，它评估仅通过更改训练算法获得的训练加速，并旨在提供更全面的优化器比较。然而，AlgoPerf 并未包含语言建模任务，其任务规模远小于现代预训练中考虑的语言模型。这种基准测试实践隐含地假设不同的优化器家族可以直接比较并共享相似的训练现象，但这未必是事实。

本文的核心论点是，现代神经网络的优化器设计应当是逐层且对称兼容的。我们不建议对所有参数应用单一的逐坐标优化器，而是提出一个逐层对称兼容的原则：每个主要的矩阵值参数类别应当由一个等变性与该参数类别的对称性相匹配的优化器来更新。这导致了一个广泛的等变优化器家族，其更新法则与所作用的参数块的对称群匹配。图 ˜1 (https://arxiv.org/html/2605.18106#S1.F1) 总结了这一转变。逐坐标视角将矩阵值参数视为独立坐标的向量化集合，导致更新可能丢弃谱结构并破坏自然的等变性。相比之下，对称感知的矩阵视角从每个参数类别的逐层几何结构出发，推导出等变性与该几何结构相匹配的优化器更新。

逐坐标视角： 参数被视为长向量 逐条目自适应更新（Adam, AdaGrad, Adafactor, RMSprop） 破坏矩阵层的正交等变性 丢弃谱结构并引入不匹配的几何

对称感知矩阵视角： 矩阵参数具有逐层对称性和几何 更新映射应与每个参数类别的对称性匹配 谱、单边谱、行范数和混合优化器 用于线性层、SwiGLU MLP、嵌入层、头层和 MoE 路由器的架构-优化器共同设计

重新思考优化器几何

贡献。本工作做出以下贡献。

1. 1. 矩阵梯度优化器设计的对称兼容原则。我们认为流行的逐坐标自适应优化器（如 Adam、AdamW 和 RMSprop）对于矩阵值参数而言在几何上不匹配，因为它们的更新通常未能尊重矩阵层的自然等变性和不变性结构。全连接层、注意力投影、嵌入和 LM 头矩阵、密集和专家 SwiGLU MLP 投影以及 MoE 路由器权重矩阵都拥有非平凡的行、列、置换和谱几何。它们的梯度通常表现出相关性、低秩结构和主导奇异方向，这些并未被逐元素更新显式表示。我们的核心信息是，神经网络权重矩阵所处的几何结构是逐坐标自适应方法无法捕捉的。
2. 2. 谱优化器的统一等变视角。我们表明，由正交等变性支配的优化器更新自然导致了谱优化器类别。该类别包括或提供了对随机谱下降 (SSD) [21 (https://arxiv.org/html/2605.18106#bib.bib795)]、Muon [76 (https://arxiv.org/html/2605.18106#bib.bib775)]、Scion [122 (https://arxiv.org/html/2605.18106#bib.bib892)] 和极坐标梯度方法 (PolarGrad) [89 (https://arxiv.org/html/2605.18106#bib.bib1058)] 的统一解释。这些方法精确或近似地计算更新方向 (D)（如梯度 (G) 或动量 (M)）的正交极因子：

   [ D = U \Sigma V^\top \quad\Rightarrow\quad U_{\mathsf{p}} \coloneqq \mathrm{polar}(D) = UV^\top. ]

   此类更新是双正交等变的，保留了更新方向的奇异向量结构，并自然源于矩阵几何。这一视角给出了 Muon [11 (https://arxiv.org/html/2605.18106#bib.bib769),12 (https://arxiv.org/html/2605.18106#bib.bib768),76 (https://arxiv.org/html/2605.18106#bib.bib775)] 所依据的谱范数最速下降原则的基于对称性的解释：由于谱范数是酉不变的，相应的极坐标更新自然也是双正交等变的。

3. 3. 用于逐层架构-优化器共同设计的等变优化器家族。超越普通矩阵层的全谱优化器，我们为对称性与标准线性映射不同的层推导了等变优化器类别。这些包括单边谱优化器（例如嵌入和 LM 头矩阵的右谱更新、MoE 路由器的左谱更新），以及非谱的行范数优化器和混合行范数/单边谱优化器。我们进一步表明，SwiGLU MLP 投影矩阵具有中间神经元置换几何，这激发了门控和上投影的行感知更新以及下投影的列感知更新。相应的实用动量变体记为 RightPolarGradM、LeftPolarGradM、RowNormM 和 HybridPolarGradM。这些构造实例化了一个基于逐层等变性的架构-优化器共同设计原则。
4. 4. 端到端预训练证据。我们在密集和稀疏 MoE 语言模型预训练实验（第 ˜4 节 (https://arxiv.org/html/2605.18106#S4)）中评估了所提出的等变优化器分配。据我们所知，这些实验实例化了第一个端到端预训练优化器堆栈，其中语言模型中所有主要的矩阵值参数类别都根据其逐层对称性分配了更新。用行范数或混合等变更新替换大词汇表索引矩阵上的 AdamW 一致地改善了最终验证损失。对于较小的 Qwen3-0.6B 风格密集模型，增益虽小但可见；对于较大的 Gemma 3 1B 风格模型，增益更为明显；在基于 OLMoE-1B-7B 和缩小版 gpt-oss 的稀疏 MoE 实验中，增益持续存在（图 ˜2 (https://arxiv.org/html/2605.18106#S1.F2)）。在密集模型中，SwiGLU MLP 投影的混合行范数/谱更新进一步改善了验证损失。在 MoE 设置中，对称兼容的路由器更新优于逐坐标路由器更新，并可能减少训练损失尖峰。作为一个代表性示例，图 ˜2 (https://arxiv.org/html/2605.18106#S1.F2) 显示了在稀疏 MoE 预训练实验中对称兼容分配的效果。

   见图注

   图 2： 缩小版 gpt-oss 预训练的验证损失。不同配置在嵌入、LM 头和路由器矩阵的优化器上有所差异；详见第 ˜4.4 节 (https://arxiv.org/html/2605.18106#S4.SS4)。配置 (i) 和 (ii) 使用基于逐层等变原则推导出的对称兼容优化器，而配置 (iii) 将路由器更新替换为 AdamW，配置 (iv) 对嵌入、LM 头和路由器矩阵使用 AdamW。

范围和局限性。我们的目标并非声称等变优化器在所有机制中都优于逐坐标自适应方法。相反，我们发展了一个针对矩阵值参数的逐层等变原则，并表明它导致了实用的优化器分配，这些分配在代表性预训练设置中具有竞争力且通常有益。经验结果应被视为该原则有用性的证据，而非详尽的大规模优化器基准测试。

组织结构。我们首先在第 ˜2 节 (https://arxiv.org/html/2605.18106#S2) 中介绍符号和密切相关的相关工作。在第 ˜3 节 (https://arxiv.org/html/2605.18106#S3) 中，我们发展逐层对称兼容原则，从矩阵参数的线性算子视图及其导致的无关坐标的等变需求开始。然后我们为嵌入、LM 头、SwiGLU MLP 投影和 MoE 路由器推导等变优化器类别，包括单边谱、行范数和混合变体。在第 ˜3.8 节 (https://arxiv.org/html/2605.18106#S3.SS8) 中，我们确定谱优化器正是与双正交等变性兼容的方向性更新映射。我们在第 ˜4 节 (https://arxiv.org/html/2605.18106#S4) 中介绍密集和 MoE 语言模型预训练实验。最后我们在第 ˜5 节 (https://arxiv.org/html/2605.18106#S5) 中讨论更广泛的影响和未来方向。

2 预备知识与相关工作

在本节中，我们介绍必要的符号和相关工作，以保证自包容性。对于相关工作的扩展综述，请参阅附录 ˜A (https://arxiv.org/html/2605.18106#A1)。

符号。对于任意实值方阵 (S \in \mathbb{R}^{d \times d})，(\mathrm{diag}(S) \in \mathbb{R}^{d}) 表示其对角元素构成的向量，(\operatorname*{Diag}(S) \in \mathbb{R}^{d \times d}) 表示对角线上元素与 (S) 相同的对角矩阵，(\mathrm{tr}(S)) 是其迹。对于任意 (x \in \mathbb{R}^{d})，(\operatorname*{Diag}(x) \in \mathbb{R}^{d \times d}) 是对角线上元素为 (x) 中条目的对角矩阵。对于任意 (m \times n) 的实值矩阵 (A \coloneqq (a_{i,j}){1 \leqslant i \leqslant m, 1 \leqslant j \leqslant n}) 和 (B \coloneqq (b{i,j}){1 \leqslant i \leqslant m, 1 \leqslant j \leqslant n})，我们记 (A) 和 (B) 的 Frobenius 内积为 (\llangle A, B \rrangle{\rm F} \coloneqq \mathrm{tr}(A^\top B) = \sum_{i,j} a_{i,j} b_{i,j})。对于矩阵 (A \in \mathbb{R}^{m \times n})，我们记其 Frobenius 范数为 (\lvert\kern-1.07639pt\lvert\kern-1.07639pt\lvert A \rvert\kern-1.07639pt\rvert\kern-1.07639pt\rvert_{\mathrm{F}} \coloneqq \sqrt{\llangle A, A \rrangle_{\rm F}})，其谱范数为 (\lvert\kern-1.07639pt\lvert\kern-1.07639pt\lvert A \rvert\kern-1.07639pt\rvert\kern-1.07639pt\rvert_{\mathrm{S}} \coloneqq \sup_{x \in \mathbb{R}^{n}, x \neq 0} { |Ax|_2 / |x|_2 })。