📊 statistics

Symmetry-Compatible Principle for Optimizer Design: Embeddings, LM Heads, SwiGLU MLPs, and MoE Routers

本文提出了一种与对称性相容的优化器设计原则，该原则使梯度更新规则与不同参数块（如嵌入层、语言模型头、SwiGLU 多层感知机以及混合专家路由器）的特定等变性结构相一致，并通过大量预训练实验证明，这些定制化优化器在最终验证损失和训练稳定性方面始终优于 AdamW。

原作者： Tim Tsz-Kit Lau, Weijie Su

发布于 2026-05-19

📖 1 分钟阅读☕ 轻松阅读

CC BY 4.0

原作者： Tim Tsz-Kit Lau, Weijie Su

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

想象一下，你正在尝试教一个巨大且复杂的机器人学习一门新语言。这个机器人由许多不同的部分组成：一个词汇词典、一个用于处理句子的“大脑”，以及一个用于决定咨询哪位专家的决策系统。

多年来，教导这个机器人的标准方法就像使用一把通用的、一刀切的画笔。你将画笔蘸上“学习颜料”，然后不管机器人的各个部分实际功能如何，都一视同仁地涂抹在它的每一个部位。这种方法（称为AdamW）将机器人“大脑”中的每一个微小数值都视为独立、孤立的点。这种方法虽然勉强可行，但它忽略了一个事实：机器人的某些部分实际上是结构化的网格、表格或相互连接的网络。

本文认为，我们应该停止用同一把画笔处理所有事情。相反，我们应该使用专用工具，这些工具要匹配机器人每个部分的特定形状和对称性。作者将这一理念称为**“对称兼容原则”**。

以下是他们利用简单类比进行的分解说明：

1. 问题：“坐标-wise"的错误

大多数现有的优化器将机器人的“大脑”视为一长串扁平的数字列表。想象你有一张城市的照片。一种“坐标-wise"的优化器试图通过独立地、逐个地调整每个像素的亮度来修复这张照片。它没有意识到建筑物是一个连接的结构，或者天空具有平滑的渐变。它忽略了数据的几何结构（形状和结构）。

2. 解决方案：让工具匹配工作

作者指出，神经网络的不同部分具有不同的“对称性”（关于它们如何在不被破坏的情况下进行重排的规则）。优化器应当尊重这些规则。

词典（嵌入层与语言模型头部）：
- 形状： 将词汇表想象成一个巨大的电子表格，其中行是单词，列是特征。你可以打乱单词（行）的顺序而不改变含义，但你不能随意旋转特征（列）。
- 旧方法： 通用优化器将每个“单词 - 特征”对视为一个独立的点。
- 新方法： 本文建议使用**“行范数”或“右谱”**更新。
- 类比： 与其给每个像素上色，不如根据每个单词的活跃程度来调整整行（单词）的亮度。如果一个单词很少被使用，就给它一个轻微的推动；如果它经常被使用，就给它一个更强的推动。这尊重了单词作为列表中独立项的事实。
大脑的隐藏层（SwiGLU MLPs）：
- 形状： 这些层拥有可以相互交换的“神经元”。如果你将神经元 A 与神经元 B 交换，数学运算依然以同样的方式工作。
- 新方法： 本文建议使用**“行感知”或“列感知”**更新。
- 类比： 想象一个工人团队。如果你交换两名工人的名字，团队依然能正常运作。优化器意识到了这一点，并一起调整整个团队的努力，而不是将每个工人视为孤立的个体。
专家交换机（MoE 路由器）：
- 形状： 在“专家混合”模型中，路由器决定呼叫哪位专家。专家是可互换的（如果你交换专家 1 和专家 2 的名字，它们是一样的），并且路由器具有“共享偏移”（给所有分数加上一个常数不会改变选择）。
- 新方法： 本文建议使用**“中心化行范数”或“左谱”**更新。
- 类比： 想象一位经理给一群相同的专家分配任务。经理并不在乎谁是“专家 1"或“专家 2"。新的优化器确保经理将团队作为一个整体来对待，调整他们的集体工作量，而不会被任意的标签所迷惑。

3. 结果：一个更好的机器人

作者通过训练几种不同类型的语言模型（如 Qwen、Gemma 和 OLMoE）来测试这一想法。他们用这些针对模型特定部分的专用、对称感知的工具，替换了通用的“一刀切”优化器。

发生了什么？

性能提升： 在几乎每一项测试中，这个机器人学习得更快，并且最终比使用通用方法训练的机器人拥有更低的错误率（更好的验证损失）。
稳定性： 训练过程更加平滑，错误值的突然尖峰更少。
可扩展性： 在更大的模型和拥有巨大词汇量的模型中，这些优势尤为明显。

核心结论

本文并不声称旧方法“坏了”或者它将不再被使用。相反，它论证了几何结构至关重要。正如你不会用锤子去拧灯泡一样，你也不应该对复杂神经网络的每一个部分都使用通用的、基于坐标的优化器。

通过设计“学习工具”以匹配其正在更新的大脑部分的特定形状和对称性，我们可以构建更好、更高效、更稳定的 AI 模型。这是一种转变：从将 AI 视为一堆随机数字，转变为将其视为一个结构化的几何对象。

技术摘要：优化器设计的对称性兼容原则

1. 问题陈述

现代深度学习实践中存在一个根本性的几何差异。尽管神经网络架构（嵌入层、语言模型头、SwiGLU 多层感知机、MoE 路由器）拥有丰富的内在对称性和等变性属性，但主流的优化方法（如 Adam、AdamW 和 RMSprop）却是按坐标（coordinate-wise）运行的。这些方法将矩阵参数视为独立坐标的扁平化向量，从而忽略了矩阵结构，破坏了参数空间固有的自然等变性。

这种“几何盲视”的方法无法区分不同层类型的独特几何结构。因此，流行的优化器可能会丢弃谱结构、引入不匹配的几何结构，并无法尊重特定参数块固有的不变性和等变性结构（例如词汇表中的置换不变性或 MoE 中的专家可交换性）。目前，关于按坐标自适应方法与谱优化器之间的理论理解和实证比较仍然有限，往往依赖于大规模基准测试，而缺乏对几何行为的深入分析。

2. 方法论：对称性兼容原则

本文提出了一种用于优化器设计的逐层对称性兼容原则。其核心论点是：参数块 $W$ 的优化器更新规则 $\mathcal{U}$ 在该特定块作用的对称群 $G$ 下应当是等变的。

2.1 理论框架

设 $W \in \mathbb{R}^{m \times n}$ 为参数矩阵。如果输入空间和输出空间的坐标通过可逆矩阵 $P$ 和 $Q$ 进行变换，则线性映射变换为 $\tilde{W} = P W Q^{-1}$ 。梯度相应地按逆变/协变方式变换。优化器更新映射 $\mathcal{U}$ 必须满足：
$\mathcal{U}(g \cdot D) = g \cdot \mathcal{U}(D)$
其中 $D$ 是更新方向（例如梯度或动量）， $g$ 是对称群的一个元素。

本文基于不同层类型的对称群推导出了特定的优化器类别：

普通矩阵层（线性层/注意力层）：
- 对称性： 双正交（ $O_m \times O_n$ ）。
- 优化器类别： 谱优化器（Spectral Optimizers）。更新必须是作用于奇异值同时保持奇异向量的谱算子。这统一了随机谱下降（SSD）、Muon、Scion 和极化梯度（Polar Gradient）方法。更新通常是动量方向的极化正交因子。
嵌入层和语言模型头（LM Heads）：
- 对称性： 左置换（ $P_v$ ）和右正交（ $O_d$ ）。行索引词汇项（离散置换），列索引隐藏特征（连续旋转）。
- 优化器类别： 左置换右正交（LPRO）等变更新。
  - 右谱（Right-Spectral）： 通过右格拉姆矩阵（ $D^\top D$ ）全局作用。
  - 行范数（Row-Norm）： 局部作用于行范数（词汇项）。
  - 混合（Hybrid）： 行范数和右谱更新的组合。
SwiGLU MLP 投影：
- 对称性： 中间神经元置换（ $P_{d_{ff}}$ ）。门控（Gate）和上投影（Up projection）的行对应中间神经元；下投影（Down projection）的列对应它们。
- 优化器类别： 行感知和列感知变体。门控/上投影使用 LPRO 兼容更新（行范数/右谱）；下投影使用转置几何结构。
MoE 路由器：
- 对称性： 专家置换（ $P_e$ ）和共享 Logit 平移不变性（对所有 Logit 添加常数不改变 Softmax 分布）。
- 优化器类别： 路由器兼容更新。这些更新作用于中心化方向 $D_c = \Pi^\perp D$ $D_{c} = Π^{⊥} D$ （移除共享平移）。
  - 中心化左谱（Centered Left-Spectral）： 作用于中心化专家格拉姆矩阵。
  - 中心化行范数（Centered Row-Norm）： 作用于中心化专家行范数。

2.2 实际实现

作者将这些原则实例化为实用的基于动量的优化器：

RightPolarGradM / LeftPolarGradM： 使用牛顿 - 舒尔茨（Newton-Schulz）迭代进行矩阵逆平方根的单侧谱更新。
RowNormM： 逐行归一化更新。
HybridPolarGradM： 行范数和单侧谱步骤的组合。
中心化（Centering）： 对于 MoE 路由器，更新应用于中心化动量 $M_{k,c} = \Pi^\perp M_k$ 。

3. 主要贡献

对称性兼容原则： 正式论证了按坐标优化器对于矩阵值参数在几何上是不匹配的。本文确立了更新规则应与参数块的对称群相匹配。
谱优化器的统一视角： 提供了基于对称性的谱优化器（Muon、SSD 等）解释，将其视为唯一的双正交等变更新映射类。
新的等变优化器类别： 推导了具有非标准对称性层的具体更新规则：
- 嵌入层和语言模型头的 LPRO 等变更新。
- SwiGLU MLP 的行/列感知更新。
- MoE 路由器的中心化专家置换更新。
端到端预训练证据： 构建了一个完整的逐层优化器栈，其中每个主要矩阵类都获得了对称性匹配的更新。这在稠密和稀疏 MoE 语言模型（Qwen3-0.6B、Gemma 3 1B、OLMoE-1B-7B、缩小版 gpt-oss）上进行了评估。

4. 实验结果

作者在 FineWeb-Edu 的 100 亿 token 子集上，针对四种架构进行了预训练实验。

稠密模型（Qwen3-0.6B、Gemma 3 1B）：
- 在词汇索引矩阵（嵌入层/语言模型头）上用 RowNormM 或 HybridPolarGradM 替换 AdamW，一致地改善了最终验证损失。
- 对于较小的 Qwen3 模型，提升幅度适中，但对于较大的 Gemma 3 模型则更为显著，这表明随着词汇矩阵变大且梯度变得更加各向异性，几何感知更新变得更为关键。
- 对 SwiGLU MLP 投影使用 HybridPolarGradM 进一步提升了性能，这可能是由于纠正了高瘦矩阵中的行尺度不平衡。
稀疏 MoE 模型（OLMoE-1B-7B、缩小版 gpt-oss）：
- 使用对称性兼容更新配置嵌入层、语言模型头和路由器（例如 RowNormM + LeftPolarGradM）的方案，优于在这些组件上使用 AdamW 的配置。
- 完全按坐标的基线（处处使用 AdamW）表现最差。
- 对称性兼容的路由器更新优于按坐标的路由器更新，并且在某些情况下减少了训练损失尖峰，表明稳定性有所提高。
总体趋势： 在所有模型中，将大型词汇索引矩阵上的 AdamW 替换为对称性兼容优化器，一致地降低了最终验证损失。这些收益在不同的架构和稀疏度水平下均持续存在。

5. 意义与主张

本文声称在优化器设计上实现从“几何盲视”的按坐标视角到“对称性感知”的矩阵视角的原则性转变。

架构 - 优化器协同设计： 本文论证优化器设计应是一个架构感知的问题。不同的参数类别（嵌入层、路由器、线性层）具有不同的自然对称性，需要不同的优化器几何结构。
理论根基： 它通过等变性提供了谱优化器为何在线性层中有效的严格数学表征，并将此逻辑扩展以推导其他层的新优化器。
实用价值： 结果表明，对称性兼容的分配不仅在理论上站得住脚，而且在实践中有益，能够改善最终损失并可能提高大规模预训练中的训练稳定性。
谦逊声明： 作者明确表示，他们并不声称等变优化器在所有情形下都优于按坐标方法。相反，他们提供了证据表明，逐层对称性兼容原则在代表性的预训练设置中能够产生具有竞争力且通常有益的分配。他们将结果视为该原则有用性的证据，而非详尽的基准测试。

本文 concludes 指出，将对称性和等变性引入优化器设计，为更稳定、参数化一致且具有理论基础的大规模模型预训练开辟了一条道路。