面向优化器设计的对称性兼容原理：嵌入层、LM Head、SwiGLU MLP 与 MoE 路由器 Symmetry-Compatible Principle for Optimizer Design: Embeddings, LM Heads, SwiGLU MLPs, and MoE Routers
提出按层对称性为每类矩阵参数定制优化器的层式优化框架
前置知识
等变性（Equivariance）
等变性指函数满足 $\mathcal{U}(g \cdot x) = g \cdot \mathcal{U}(x)$，即输入被群 $g$ 变换时输出同步变换。本文把该原则作用到优化器更新映射：梯度在坐标变换下变为 $PGQ^\top$ 时，更新也必须做相同变换。不变性是其特例。
论文的核心论点是优化器应与参数块对称性等变，这是把架构等变性从模型层迁移到优化器层的关键概念
谱范数最速下降与极分解（Polar Decomposition）
谱范数酉变换下不变，其最速下降方向是矩阵的极分解 $\text{polar}(D) = UV^\top$（$D = U\Sigma V^\top$）。Muon 类优化器以 Newton–Schulz 与 Polar Express 多项式在 GPU 上近似 $UV^\top$，常用动量优先变体。
Muon、SSD、Scion、PolarGrad 这些谱优化器本质上是双正交等变更新；理解极分解才能读懂 Theorem 3.7 的完整刻画
AdamW 与逐坐标自适应方法
AdamW 把参数展平成向量，逐元素维护一阶动量 $m_t$ 与二阶动量 $v_t$ 缩放更新，权重衰减与梯度解耦。其前提是每个权重分量是独立坐标，与权重矩阵几何相冲突。
论文把 AdamW 视为对照基线，证明其在大型嵌入、LM Head、路由器上与参数化几何不兼容，但实际中却被默认使用
混合专家（MoE）与路由器对称性
MoE 路由器 W 输出 softmax(Wx)，有专家置换与共享对数偏移不变性两类对称。更新须满足 U(PD + 1ea^T) = PU(D) 且水平性 1e^T U(D) = 0。
路由器对称性决定 centered row-norm、centered left-spectral 与 hybrid 路由器优化器，是论文新引入的模块。
SwiGLU 与中间神经元置换对称
SwiGLU 块 = W_down(sigma(W_gate x) ⊙ (W_up x))，其中 sigma 逐元素。它对中间神经元置换不变，但因 sigma 逐元素性，不能放宽为正交。
SwiGLU 的中间神经元置换对称性决定了其门控/上投影矩阵（行方向）和下投影矩阵（列方向）需要不同的等变优化器
研究动机
现代大语言模型训练几乎被 Adam、AdamW、RMSprop 这类逐坐标自适应优化器主导，它们把所有参数展平成一个长向量后逐元素独立更新，隐含假设权重矩阵的每个 entry 都是高维空间中的独立坐标。然而这一假设几乎从未被质疑，却严重影响了 LLM 的训练动力学——例如逐坐标优化会破坏线性层、注意力投影、SwiGLU 中间投影、MoE 路由器、嵌入、LM Head 等不同层类型各自的矩阵几何结构。Muon 等新近矩阵优化器以谱范数最速下降或极分解为原理启发了这一波研究，但其原始推导紧密绑定谱范数，未必能自然推广到嵌入、LM Head、SwiGLU、MoE 路由器等具有不同对称类别的矩阵参数块。更糟的是，已有的优化器基准（如 AlgoPerf）评测的 workload 远小于现代 LLM 预训练规模，且不包含语言建模任务，因此对不同优化器族的比较严重依赖经验性 speedrunning（modded-nanogpt）和大规模基准，理论层面几乎空白。Muon 在实践中甚至经常推荐对嵌入和 LM Head 退回使用 AdamW，这种权宜之计暴露出当前优化器几何与模型几何严重不匹配的问题。
本文的目标是本文建立一个统一的逐层（layerwise）等变优化器设计原理：根据参数块自身的对称群（而不是全局统一的优化器族）为每类矩阵参数定制更新映射。具体目标可量化为：(1) 在 dense（Qwen3-0.6B-style 625M、Gemma 3 1B-style 1.087B）和 sparse MoE（OLMoE-style 2.82B、downsized gpt-oss 3.47B）端到端预训练中，用与层对称性兼容的优化器替换 AdamW 一致性地降低最终验证损失；(2) 把 Muon、Scion、PolarGrad、SSD 全部纳入谱优化器这一更大的对称类，使这些方法有统一的理论根基；(3) 为嵌入、LM Head、SwiGLU MLP 投影、MoE 路由器分别提供工程可用的 RowNormM、RightPolarGradM、LeftPolarGradM、HybridPolarGradM 四个动量变体。
与已有工作不同的是，已有矩阵优化器视角通常把谱范数或某具体矩阵范数当作出发点，把 Muon 这种算法看作 normalizer 或 preconditioner 启发式。和这些工作相比，本文的关键切入角度是把优化器设计问题重新表述为等变性问题（equivariance principle for optimizer design）：优化的核心不是某个特定矩阵范数，而是参数块本身的对称群。决定更新形状的不是「找一个 unitarily invariant 的矩阵范数做 steepest descent」这一选择，而是「该层的对称群是什么」——左置换/右正交（嵌入、LM Head、SwiGLU gate/up）、中间神经元置换（SwiGLU）、专家置换加共享对数偏移不变性（MoE 路由器）、双正交（普通线性层）。相应地，更新映射必须是群等变的，并在参数化具有对称冗余方向时去除之。这一视角把优化器几何、矩阵几何、架构几何三件过去相互割裂的事情统一在「symmetry-to-optimizer」原理之下。
核心方法
直观上「给不同零件配不同扳手」：旋转矩阵/普通线性层/注意力投影两侧都是连续特征基底（对称群 O_m × O_n），需要 bi-orthogonal 等变更新，即 Muon 这类谱优化器做的事；嵌入矩阵行是离散的 token ID（旋转无意义），只有置换才不被注意（P_v × O_d），允许 RowNormM 按行缩放；LM Head 多 softmax 共享对数偏移不变性，更新须投影到水平商空间；SwiGLU 中间神经元置换对称性推到 gate/up 行方向与 down 列方向；MoE 路由器叠加 shared-logit-shift 不变性，定义 centered row-norm 与 centered left-spectral。技术路线：先建等变框架（§3.1），对普通层推导出 bi-orthogonal 等变→谱算子等价（Theorem 3.7），对四个特殊层分别推导优化器类，最后以 Polar Express / Gram Newton–Schulz 高效近似极分解并在四个端到端预训练中验证。
本文最本质的创新是把优化器更新映射的等变性需求与参数块的对称群严格对齐：参数 W 在群作用 g·W 下不变时，更新方向 D_k（动量）按 D_k → g·D_k 变换，更新映射 U 必须满足 U(g·D_k) = g·U(D_k)；参数化存在对称冗余方向时（如 LM Head 共享 vocab 偏移），实际更新必须投影到水平商空间。和已有方法的区别：Muon/Scion 等矩阵优化器只从谱范数或谱诱导范数出发，本文把等变性本身作为原理，由此不仅得到谱优化器这一更大对称类（Theorem 3.7），还得到针对非双正交对称层的全新优化器族。MoE 路由器的 centred row-norm 更新需要先去掉 D 的共享行方向 Π⊥D、按行范数缩放、再投影回水平——这一三步法（先中心化、行处理、再投影）是 LM Head 与路由器共用的关键技术，由 Π⊥ = I − 1v 1v^T/v 显式给出。
方法步骤详情
方法分四步。第一步「参数-对称识别」：线性/注意力 O(d_out)×O(d_in)，嵌入/LM Head P_v×O_d（LM Head 还商 1v R^{1×d}），SwiGLU W_gate/W_up 为 P_{dff}×O_{dmodel}、W_down 转置，MoE 路由器 P_e×(1e R^{1×d}) 商去 shared-logit-shift。第二步「更新构造」：bi-orthogonal 用极分解 U(G)=UV^T；LPRO 行范数 U_row=Diag(η(‖D_i‖))D、右谱 U_R=D(D^TD+εI)^(-1/2)、hybrid U_R∘U_row；MoE 路由器先 D_c=Π⊥D，输出再投回水平。第三步「动量实现」：M_k=βM_{k-1}+(1-β)G_k 应用更新，谱部分由 Polar Express/Gram Newton–Schulz 在 float32 近似 (M^TM+εI)^(-1/2)。第四步「step-size 校正」：ν_k=max(tr(C_k R_k),ε) 防小 trace 爆炸，W_{k+1}=(1−γ_kλ)W_k−γ_kν_k^{−α}M_k R_k。
技术新颖性
技术新颖性有三。第一，完整刻画双正交等变更新映射与谱算子的等价（Theorem 3.7），给 Muon、SSD、Scion、PolarGrad 统一的 symmetry-based 理论解释，明确「谱优化器 = bi-orthogonal 等变类」——首次把矩阵优化器合法性从某个具体范数提升到几何原理。第二，提出 left-permutation right-orthogonal equivariance 这一介于 bi-orthogonal 与完全坐标化之间的中间概念，证明其合法更新类是右谱、行范数、hybrid 并集（Proposition 3.3 封闭性），突破过去只有谱范数一种选择的局限。第三，把 softmax shared-logit-shift invariance 翻译成 Π⊥ 水平投影，使 LM Head 与 MoE 路由器有「先中心化、行列操作、再投回水平」三步法 U_LM_row(D)=Π⊥^v Diag(η(‖D_{c,1:}‖)) D_c。这些构造把 Aurora、SCALE、RMNP、REG、Nora、NorMuon、Muon+、MuonEq 统一到同一框架。
实验结果
四个端到端预训练一致验证逐层对称兼容优化器有效。Qwen3-0.6B-style（625.8M，词表 151936）：SwiGLU 用 Muon 时 RowNormM/Hybrid/AdamW 验证损失 4.2017/4.2050/4.2084；改用 HybridPolarGradM 后 4.1950/4.1955/4.2046，AdamW 早期领先、后期被反超。Gemma 3 1B-style（1.087B，词表 262144）Muon 下三组 4.0702/4.0655/4.1046，Hybrid 下 4.0516/4.0435/4.0862，验证「词表行维度越大、几何不匹配被放大」。OLMoE-1B-7B-style（2.824B）四组 4.0955/4.0815/4.1187/4.1240，对称兼容组击败 AdamW 组；全 AdamW 在 ~2.1B token 出现更明显损失尖峰。downsized gpt-oss（3.468B）四组 4.3080/4.3136/4.3388/4.3687，RowNormM 系三组相对全 AdamW 改进 ~0.05。整体替换 AdamW 一致带来 0.005-0.06 损失下降。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| Qwen3-0.6B-style 预训练验证损失 | Validation loss (10B token, Muon 用于 SwiGLU MLP) | RowNormM = 4.2017；HybridPolarGradM = 4.2050 | AdamW = 4.2084 | RowNormM 比 AdamW 低 0.0067；HybridPolarGradM 比 AdamW 低 0.0034 |
| Qwen3-0.6B-style 预训练（Hybrid 用于 SwiGLU MLP） | Validation loss | RowNormM = 4.1950；HybridPolarGradM = 4.1955 | AdamW = 4.2046 | RowNormM 低 0.0096；HybridPolarGradM 低 0.0091 |
| Gemma 3 1B-style 预训练（Muon 用于 SwiGLU MLP） | Validation loss (10B token, 词表 262144) | RowNormM = 4.0702；HybridPolarGradM = 4.0655 | AdamW = 4.1046 | RowNormM 低 0.0344；HybridPolarGradM 低 0.0391 |
| Gemma 3 1B-style 预训练（Hybrid 用于 SwiGLU MLP） | Validation loss | RowNormM = 4.0516；HybridPolarGradM = 4.0435 | AdamW = 4.0862 | RowNormM 低 0.0346；HybridPolarGradM 低 0.0427 |
| OLMoE-1B-7B-style 稀疏 MoE 预训练 | Validation loss (10B token, 32 专家) | (i) RowNormM 全部 = 4.0955；(ii) RowNormM+LeftPolarGradM 路由 = 4.0815 | (iii) RowNormM+AdamW 路由 = 4.1187；(iv) 全 AdamW = 4.1240 | (ii) 比 (iv) 低 0.0425；(ii) 比 (iii) 低 0.0372，且训练损失尖峰更少 |
| downsized gpt-oss 稀疏 MoE 预训练 | Validation loss (4B token, 词表 201088, 含 bias) | (i) RowNormM 全部 = 4.3080；(ii) RowNormM+LeftPolarGradM = 4.3136 | (iii) RowNormM+AdamW 路由 = 4.3388；(iv) 全 AdamW = 4.3687 | (i) 比 (iv) 低 0.0607；RowNormM 系三组相对全 AdamW 平均改进 ~0.05 |
局限与改进
作者承认论文不是 optimizer benchmark，只是给出 design principle 在代表性任务验证。所有模型只在 10B token（dense）或 4B token（gpt-oss）训练，远低于 scaling-law 预算，收益「modest but visible」。HybridPolarGradM 比 RowNormM 引入更高算力开销，全文未给出每 step 的 wall-clock 或 GPU-hour 数字。独立观察：(1) 没做 LR sweep——AdamW 在合理 sweep 后可能接近，须报告 LR sensitivity；(2) 实验只覆盖 dense Qwen3/Gemma 与 MoE OLMoE/gpt-oss，attention、CNN、SSM、ViT 未触及；(3) MoE 路由器 (i)(ii) 差距很小且不稳定（OLMoE 0.014、gpt-oss 反向），可能与 RowNormM LR 或极分解 oracle 精度有关；(4) 与 Shampoo/SOAP/Newton-Muon 差异只理论未定实验对比。
独立分析的弱点
三大弱点。第一，最大缺口是 LR 与超参网格：Gemma 3 1B-style 上 AdamW 早期领先、之后被反超，未明示是否做了匹配 LR sweep。若 AdamW 在合理 sweep 后接近，本结论会弱化。改进：报告每 config 的 LR sensitivity，并通过 Mu-Transfer / μP 类超参迁移规则做公平 LR 选择。第二，hybrid 对数值精度敏感——论文指出「即使 1^T D_c=0，按行缩放会重新引入共享行方向」，必须再投影；若与 GS / Newton–Schulz 内迭代耦合，bfloat16 下可能 drift。改进：HybridPolarGradM 算法显式记录「行向非线性后立即 Π⊥」并加可选 post-step renormalization 防 error accumulation。第三，架构覆盖太窄：多模态视觉 embedding、MoE 每 expert 共享参数化、Conv/ViT/SSM 矩阵类未处理；Theorem 3.7 仅刻画 memoryless 等变算子，stateful 类对称性刻画缺失。
未来方向
作者在 Remark 3.7 提 projected/proximal 扩展：用 prox_{γ_k h} 或 proj_C 配合 G-invariant 正则 h 或约束集 C 保留等变性，引出 ProxPolarGrad。Section 5 与 G 指认：探索非 elementwise 优化器的分布式系统挑战（Dion、Disco、Distributed Muon），与 tensor/pipeline/sequence/expert parallelism 集成；探索更稳定的 inexact polar oracles（Polar Express、PRISM、Turbo-Muon、Gram Newton–Schulz）。基于成果可延伸：(a) 把 symmetry-to-optimizer 泛化到卷积核（translation group）、state-space（对角+low-rank）；(b) 把 RowNormM 移植到 LAMB/LARS；(c) 与 μP 结合得跨模型尺寸可迁移 stack；(d) 从路由器 z-loss 反推在 RLHF/DPO 对齐阶段的稳定性。
复现评估
代码已在 https://github.com/timlautk/equivariant_optimizers 开源，配套 Polar Express、Gram Newton–Schulz 等 GPU 友好极分解近似。实验细节在 Section G 附录：FineWeb-Edu 10B token，context length 1024。算力：四个模型在消费或中型 GPU 集群完成；Qwen3 dense 约 7-8h，OLMoE ~10h，gpt-oss ~20h wall-clock。复现门槛中等：四种变体可直接 modify 一个 modded-nanogpt 训练栈复现，算法伪代码（Algorithm 1, 2）完整给出。MoE 实验处理 fused gate/up 张量、去 auxiliary load-balancing 与 router z-loss，复杂度略高；G.2.3 的 LM head projection 降低 vocab-logit growth 需 logit norm 监控，论文有详细附录。整体 reproducibility 高。
论文图表
该图左右对比两种优化视角：左侧「Coordinate-wise view」把矩阵参数视作长向量，由 Adam/AdaGrad/Adafactor/RMSprop 等逐 entry 自适应更新，会破坏正交等变性、丢弃谱结构、引入不匹配几何；右侧「Symmetry-aware matrix view」识别每个矩阵参数类各自的 layerwise 对称与几何，使用 spectral / one-sided spectral / row-norm / hybrid 等优化器族，构成「architecture–optimizer co-design」。图中用箭头标示从左到右是要「rethink optimizer geometry」。
这是论文的全图总览式图，把核心论点「优化器几何应匹配参数块对称性」以视觉对比呈现，去掉则读者无法抓住论文的中心论断