宾夕法尼亚大学：为什么给AI训练的每一层都配上"专属教练"，效果会更好？

这项研究来自宾夕法尼亚大学，以预印本形式发布于2026年5月，论文编号为arXiv:2605.18106。有兴趣深入了解的读者可以通过该编号在arXiv平台上查阅完整论文。

深度学习已经成为现代科技的基础设施，而训练这些模型的"优化器"——也就是那个告诉模型"参数该往哪个方向调整、调多少"的核心算法——几十年来几乎没有从根本上改变过。Adam及其变体统治了整个领域，它的逻辑非常简单粗暴：把模型里所有的参数拼成一条超长的数字串，然后对每个数字独立地做调整。这种"人人平等、各自为政"的方式虽然好用，却忽略了一件至关重要的事——模型里不同层的参数，其实生活在完全不同的"几何世界"里，它们天然地具有不同的结构和对称性。

用一个更直观的比方来说：假设你在管理一支乐队，里面有钢琴手、吉他手和鼓手。Adam的做法相当于给所有人发同一本完全相同的练习手册，钢琴手和吉他手都按同样的方式练习指法。而这篇论文的核心主张是：钢琴手应该按照钢琴的逻辑练习，吉他手应该按照吉他的逻辑练习，鼓手应该按照打击乐器的逻辑练习——每种乐器有其独特的对称性和规律，优化算法也理应尊重这种结构差异。

这个听起来颇为直觉性的洞察，被研究团队发展成了一套严谨的理论框架，并在多个真实的大语言模型预训练实验中得到了验证。

一、 为什么Adam不够用——坐标系的烦恼

要理解这篇论文解决的问题，先从一个简单的思想实验出发。假设你有一张地图，上面标着从A城到B城的最短路径。现在，有人把整张地图旋转了90度，但路径本身没有变化——A城和B城的相对位置、最短路径的长度，这些东西都是客观存在的，不会因为地图怎么摆放而改变。

神经网络的权重矩阵也有类似的性质。一个线性层的权重矩阵W，本质上描述的是一个从输入空间到输出空间的线性变换。如果你对输入空间做一个正交变换（相当于换一组基向量，类似于旋转坐标系），对输出空间做另一个正交变换，那么表示同一个线性变换的矩阵会相应地变化，但它所代表的数学对象没有本质改变。换句话说，权重矩阵有一种内禀的"双正交不变性"——它不依赖于坐标系的选择。

Adam完全不知道这件事。Adam看到的是一个矩阵的每个元素，独立地估计每个元素的梯度大小，然后对每个元素做自适应的缩放更新。如果你旋转了坐标系，矩阵的元素会变，Adam的更新方向也会变——即使两个坐标系描述的是完全相同的神经网络，Adam也会给出不同的训练轨迹。这意味着Adam的更新方式依赖于我们如何"命名"权重矩阵的行和列，而不仅仅取决于权重矩阵所代表的线性变换本身。

研究团队把这种现象称为"几何不匹配"：优化器的更新规则，与它正在优化的参数所天然具备的几何结构，并不一致。更糟糕的是，Adam的更新方向还会把梯度中的低秩结构（也就是梯度中真正重要的方向）淹没在大量高秩的坐标噪声里。

近年来，一类被称为"谱优化器"的方法开始引起关注，其中最具代表性的是Muon算法。Muon的核心思想是：不要直接用梯度矩阵G来更新权重，而是用G的"正交极因子"polar(G)来更新。所谓正交极因子，就是对G做奇异值分解G=UΣV?，然后取UV?——相当于把G的"方向信息"提取出来，把所有奇异值都压缩成1，只保留梯度告诉我们"该往哪个方向走"的信息，而不关心"该走多远"。

这种更新方式有一个美妙的性质：如果你对梯度做双正交变换G→PGQ?，那么polar(PGQ?)=P·polar(G)·Q?，更新方向也以完全相同的方式变换。这意味着Muon的更新规则尊重了权重矩阵的几何结构，与坐标系的选择无关。研究团队把这种性质称为"双正交等变性"。

然而，Muon在实践中有一个公认的局限：它不适合用于词嵌入矩阵（embedding matrix）和语言模型头（LM head），通常这两个部分仍然使用Adam。这件事听起来很奇怪——既然Muon那么好，为什么不直接用在所有矩阵上？这正是这篇论文要回答的核心问题。

二、 每种"乐器"有其独特的对称性

答案藏在这些特殊矩阵的结构里。词嵌入矩阵E的形状是v×d，其中v是词表大小（通常几万到几十万），d是隐藏层维度（通常几百到几千）。矩阵的每一行对应词表中的一个词，每一列对应隐藏空间的一个特征维度。

关键在于：行和列在这里扮演的角色完全不同。词表中的词是离散的、可以互换顺序的——你可以把第3行和第5行对换，只要同时更新词表的索引，对整个模型没有任何影响。这是一种"左置换对称性"：左乘一个置换矩阵P（把行打乱顺序），不改变模型的任何实质。而列（隐藏特征维度）则不然，特征维度之间是连续的、可以做正交变换的——右乘一个正交矩阵R，相当于在特征空间做了一次旋转，模型的功能同样不变，只要其他层也做了相应的变换。

所以词嵌入矩阵的自然对称群，不是Muon所对应的"双正交群"（左边正交变换×右边正交变换），而是"左置换×右正交"群。研究团队把这类对称性称为"左置换右正交（LPRO）等变性"，并从这一对称性出发推导了适合词嵌入矩阵和LM头矩阵的优化器类。

在LPRO等变的更新映射中，最自然的一类是"行范数更新"（RowNorm）：对梯度矩阵的每一行，按照该行的范数大小进行缩放，再用缩放后的矩阵作为更新方向。这种更新满足LPRO等变性，因为置换行只是重新排列行范数，右正交变换保持行范数不变。另一类是"右谱更新"（RightSpectral）：计算梯度矩阵的右Gram矩阵D?D（这是一个d×d的小矩阵），对其求逆平方根，再右乘到梯度矩阵上。当词表很大（v远大于d）时，这种方法只需要操作d×d的小矩阵，计算效率很高。此外，把这两类更新组合起来，先做一步再做另一步，得到的"混合LPRO更新"（HybridPolarGrad）也是LPRO等变的——研究团队证明了LPRO等变映射在复合运算下是封闭的。

SwiGLU激活函数的MLP层（现代大语言模型普遍采用的前馈网络结构）也有类似的分析。SwiGLU块由门控投影Wgate、上投影Wup和下投影Wdown三个矩阵组成。研究团队证明，如果同时对Wgate和Wup的行做置换，同时对Wdown的列做同样的置换，SwiGLU块的输出完全不变。这是因为中间神经元的顺序并不重要，只要三个矩阵中对应位置的行/列一起重排就行。这种"中间神经元置换对称性"意味着Wgate和Wup的行维度具有置换对称性，因此应该使用LPRO类的优化器；而Wdown的列维度具有置换对称性，相当于对Wdown?做LPRO更新，再转置回来。

值得注意的是，这里的对称群严格比Muon所对应的双正交群更小：中间神经元维度只有置换对称性，而没有完整的正交变换对称性。直觉上的原因是，SwiGLU中有逐元素的非线性激活函数σ，对一般正交矩阵Q，σ(Qz)≠Qσ(z)，所以左乘任意正交矩阵并不构成对称性；但置换是σ的对称操作，因为σ是逐坐标应用的。

混合专家（MoE）模型的路由器矩阵则有另一种完全不同的对称性。路由器矩阵W的形状是e×d，其中e是专家数量，d是隐藏维度。路由器的功能是计算softmax(Wx)来决定每个token该被哪些专家处理。由于专家之间是可以互换编号的（把第3号专家和第5号专家对调，同时调整路由器矩阵的对应行，模型功能不变），路由器有"专家置换对称性"。更重要的是，softmax函数有一个鲜为人知的性质：给所有逻辑值都加上同一个常数c，softmax的输出不变。这意味着路由器矩阵的每一行都加上同一个向量a?，对路由结果没有任何影响——这就是"共享logit平移不变性"。

这两种对称性加在一起，告诉我们路由器优化器应该满足：对专家置换等变，对共享行平移不变。满足这两个条件的更新映射，应该先把梯度矩阵"中心化"（减去各行的均值，消除共享分量），然后对中心化后的矩阵做左谱更新（LeftSpectral）或行范数更新（CenteredRowNorm）。

三、 理论推导：从对称性到具体算法

研究团队将上述直觉发展成了精确的数学定理。核心结论是：对于一个矩阵值的更新映射，它满足双正交等变性当且仅当它是一个谱算子——也就是说，它保持更新方向的奇异向量不变，只对奇异值做变换。具体地，如果梯度矩阵G=U·Diag(σ)·V?，那么等变更新一定具有U·Diag(ψ(σ))·V?的形式，其中ψ是一个绝对对称函数（对奇异值做某种变换）。

这个定理的意义在于，它给出了一个完整的刻画：满足双正交等变性的更新规则恰好就是谱优化器这个大类，不多也不少。简单梯度下降（ψ是恒等函数）、Muon（ψ把所有奇异值压成1）、核范数缩放的PolarGrad（ψ把所有奇异值乘以核范数然后压成1）、SSD等，都是这个大类的特例。

对于词嵌入矩阵和LM头，研究团队同样做了类似的分析，得出行范数更新和右谱更新是LPRO等变更新映射的自然子类这一结论。这两个子类都严格包含于LPRO等变映射的全集，而它们的复合（混合更新）也是LPRO等变的。

在实际算法设计上，研究团队将这些更新规则与动量（momentum）结合，提出了几个具体的实用算法。RowNormM是最简单的：维护一个梯度的指数移动平均（动量），然后对动量的每一行按行范数归一化，用归一化后的方向更新权重。RightPolarGradM则更复杂一些：先计算动量，再计算动量的右Gram矩阵M?M（d×d大小），求其逆平方根，再右乘到动量矩阵上，得到更新方向。LeftPolarGradM是路由器的版本：先对动量做中心化，计算左Gram矩阵MM?（e×e大小，e是专家数），求逆平方根后左乘到中心化动量上。HybridPolarGradM则将行归一化和单边谱更新两步结合，按两种顺序各有一个变体：先右谱再行归一化，或先行归一化再右谱。

逆平方根的计算是这些算法在工程上的核心挑战。研究团队使用了"Polar Express"算法的多项式系数（这是一种专门为GPU优化的矩阵符号函数迭代算法）以及Gram Newton-Schulz迭代来完成这一计算，在float32精度下进行以确保数值稳定性。

四、 动量的等变性：为什么"先动量再极因子"更好

研究团队还解决了一个实践中悬而未决的问题：在结合动量和极因子更新时，应该"先算动量，再取极因子"（momentum-first），还是"先取极因子，再算动量"（polar-first）？

答案是前者更好，原论文中的Muon也采用前者，而研究团队给出了理论上的解释。结论是：梯度的指数移动平均（动量）具有和梯度完全相同的双正交等变性——如果梯度按G→PGQ?变换，那么动量也按同样的方式变换。因此，动量本身是一个合法的双正交等变方向，对它取极因子是有意义的，并且所得到的更新也会以正确的方式变换。

而polar-first的方案则不同：先对每步梯度取极因子，再对极因子做指数移动平均。问题在于极因子是非线性的，polar(βM+G)≠β·polar(M)+(1-β)·polar(G)，所以polar-first的动量并不以等变的方式变换，从等变性的角度讲是"不干净"的。直觉上，momentum-first先对带噪声的梯度做平滑，再从平滑后的信号中提取方向，更稳定；而polar-first把每步噪声梯度都压成等范数的方向，再对这些方向做平均，会引入更多噪声。

五、 在真实模型上的验证：改变优化器栈，结果如何？

理论推导固然重要，但最终的检验在于实验。研究团队在四个不同规格的语言模型上做了预训练实验，使用FineWeb-Edu数据集的100亿token子集，覆盖了从6亿参数的密集模型到35亿参数的稀疏混合专家模型。

第一组实验基于Qwen3-0.6B风格的密集模型，词表大小151936，隐藏维度1024。在这里，研究团队比较了三种对词嵌入矩阵和LM头的处理方式：使用RowNormM、使用HybridPolarGradM（行归一化+右谱），以及使用AdamW作为基线。对于SwiGLU MLP投影矩阵，则比较了纯Muon风格更新和HybridPolarGradM（行感知+右谱）两种方案。结果显示，RowNormM和HybridPolarGradM在最终验证损失上均优于AdamW基线，差距虽然不大但稳定存在。同时，对SwiGLU投影矩阵使用HybridPolarGradM（而非Muon）可以进一步改善所有三种词嵌入优化器的验证损失。

第二组实验基于Gemma 3 1B风格的密集模型，词表大小更大，达到262144，隐藏维度1152，这意味着词嵌入矩阵和LM头矩阵更大也更可能出现梯度的各向异性或病态条件。在这里，HybridPolarGradM对词嵌入矩阵的改善更加显著，超过了RowNormM，而两者都大幅优于AdamW基线。具体来说，最终验证损失从AdamW的约4.10降到了RowNormM的约4.07和HybridPolarGradM的约4.05，差距在这个量级下是相当可观的。同样，对SwiGLU投影矩阵使用HybridPolarGradM也带来了额外改善。研究团队还做了学习率扫描和三个随机种子的重复实验，证明这些提升是稳定的，不是偶然现象。

第三组实验基于OLMoE-1B-7B风格的稀疏MoE模型，词表大小50304，隐藏维度2048，每层有32个专家，激活8个。相比密集模型，这里的路由器矩阵成为了一个新的实验变量。研究团队比较了四种配置：词嵌入和LM头使用RowNormM，路由器分别使用RowNormM、LeftPolarGradM或AdamW；以及全部使用AdamW的基线。结果显示，全AdamW基线在前500步左右更新最快，但随后被RowNormM+RowNormM路由和RowNormM+LeftPolarGradM路由两个配置超越，差距在学习率衰减阶段持续拉大。在路由器优化器的比较中，LeftPolarGradM略优于RowNormM，而两者都优于AdamW。研究团队还观察到，全AdamW配置在约21亿token处出现了更明显的训练损失尖峰，而使用等变路由器优化器的配置则更稳定。

第四组实验基于缩小版gpt-oss模型，词表大小201088，参数量约35亿，这是所有实验中词表最大的架构。结果再次证实了相同的规律：使用RowNormM更新词嵌入和LM头矩阵，相比全AdamW基线，验证损失有约0.06的改善；在此基础上对路由器也使用等变更新，能进一步改善；而在词嵌入和LM头使用RowNormM的情况下，路由器使用AdamW则略逊于使用RowNormM或LeftPolarGradM。

综合四组实验，研究团队观察到一个一致的规律：词表越大（词嵌入矩阵越"高瘦"），等变更新相比AdamW的收益越大。这与理论预期一致——词表越大，词嵌入矩阵的梯度中低维特征子空间的结构越重要，坐标方向上的更新越容易与这种结构产生错配。

六、 层次优化器栈：一张为每层定制的菜单

研究团队的最终贡献是将上述所有分析整合成一个系统性的设计原则，并给出一张"层次优化器栈"的完整菜单。普通线性层和注意力投影矩阵使用完整的谱优化器（双正交等变，具体实现为Muon/PolarGrad）；词嵌入矩阵和LM头矩阵使用LPRO等变的更新（行范数、右谱或混合）；SwiGLU MLP的门控投影和上投影使用同样的LPRO更新，下投影则对其转置使用LPRO更新再转置回来；MoE路由器使用中心化后的行范数更新或左谱更新；标量和向量参数（如层归一化的权重和偏置）继续使用AdamW。

这张菜单并不是一个全有或全无的替换方案，而是一个精细化的搭配策略：根据每类参数的对称群，匹配相应等变性的优化器。研究团队提出了一个四步设计原则：首先，识别参数化的对称群；其次，判断对称性作用在左边、右边、两侧还是只在某个商空间上；第三，选择相应的等变优化器类；第四，使用与该对称性相容的最小Gram矩阵或不变统计量进行计算。

与单纯替换单个优化器的做法不同，这种端到端的层次化设计确保了每个主要参数类都接受与其几何结构相匹配的更新，这也是研究团队所强调的"架构-优化器协同设计"理念的核心体现。

七、 这不是"又一篇谱优化器论文"——理论视角的统一

在介绍了大量新结果之后，研究团队也用这个等变性框架为现有的谱优化器提供了统一的理论解释。Muon、SSD（随机谱下降）、Scion、PolarGrad这些表面上看起来各不相同的方法，在这个框架下都是同一个基本原理的不同实例：它们都是双正交等变的谱算子，区别只在于选择了不同的奇异值变换函数ψ。

特别地，Muon对应ψ把所有奇异值压成1，这等价于对谱范数做最速下降——因为谱范数是酉不变的，其最速下降方向是正交极因子，天然地具有双正交等变性。从等变性的角度来看，这不是Muon偶然具备的性质，而是它从谱范数最速下降这一出发点"天然继承"的结构特性。

研究团队还讨论了几个相关方法的定位。Shampoo维护左右预条件器的指数移动平均，并用其进行预条件化更新，原则上也具有双正交等变性（当其状态变量正确变换时），但它不是当前更新方向的谱算子——预条件器和当前梯度的奇异向量一般并不对齐。SOAP在学习到的特征空间中做坐标方向的自适应更新，通常也不是严格意义上的谱算子，除非特殊情况。行归一化或列归一化的方法（如NorMuon、RMNP等）则一般只具有单边或置换的等变性，适合LPRO类的层，但不适合普通矩阵层的双正交等变优化器。

说到底，这篇论文的核心信息相当简洁：不同类型的参数矩阵在数学上"生活在"不同的对称性世界里，好的优化器应该尊重并利用这种结构差异，而不是把所有矩阵都当作无差别的数字数组来处理。这种理念看起来朴素，但在实际大模型预训练中带来了稳定且可重现的性能提升，并且为未来设计新架构时如何配套设计优化器提供了一个可操作的框架。

当然，研究团队也坦承了若干局限性。首先，他们的实验训练步数并没有达到通常的计算最优规模，因此不能简单外推到更大规模的训练。其次，非元素级别的优化器在工程上带来了新挑战，包括矩阵逆平方根的数值计算、分布式训练中的通信开销等，不过最近的一些工程工作（如Distributed Muon、Flash-Muon等）正在积极解决这些问题。第三，实验结果应被视为对这一设计原则有效性的证据，而非优化器之间的全面基准测试。

归根结底，这项工作的价值或许不只在于它提出的那些具体算法，更在于它提供了一种思维方式：在设计深度学习优化算法时，应该首先问一问"这类参数具有什么样的对称性？"然后再去寻找与之匹配的更新规则。随着基础模型的架构越来越多样化——从密集语言模型到稀疏MoE，从视觉变换器到多模态模型，再到扩散语言模型和状态空间模型——用同一个坐标方向的Adam来统治一切的做法，恐怕会越来越像用同一把螺丝刀来修所有的机器：凑合能用，但总有些地方拧不进去。

---

Q&A

Q1：RowNormM和HybridPolarGradM哪个更适合用在词嵌入矩阵上？

A：两者都比AdamW表现更好。RowNormM计算更简单，只需对动量的每一行按行范数归一化，几乎不增加额外计算量。HybridPolarGradM额外计算了一个右Gram矩阵的逆平方根，计算开销更大，但在词表较大（如Gemma 3风格的26万词表）时效果更好，因为大词表矩阵的梯度更可能存在各向异性，右谱步骤能更好地利用这种结构。实践建议是：计算资源有限时优先RowNormM，追求最佳性能时考虑HybridPolarGradM。

Q2：MoE路由器的对称性和普通矩阵有什么不同？

A：普通矩阵层具有双正交对称性（左右都可以做正交变换），适合Muon这样的完整谱优化器。MoE路由器矩阵则有两种额外的特殊对称性：一是专家可以互换编号（专家置换对称性），二是给所有专家的logit值加同一个常数不影响softmax结果（共享logit平移不变性）。这两种对称性加在一起，要求路由器的优化器先对梯度做中心化（消除共享分量），再对中心化后的方向做左谱或行归一化更新，而不能直接套用普通矩阵的双正交等变更新。

Q3：谱优化器的双正交等变性为什么Adam不具备？

A：Adam对梯度矩阵的每个元素独立估计其历史方差，然后用每个元素的梯度除以其历史方差的平方根来做更新。这个操作本质上依赖于矩阵元素的具体数值，而不是矩阵作为线性算子的内禀几何结构。如果对权重矩阵做一个正交坐标变换（行或列的正交旋转），矩阵元素会变，Adam估计的各元素方差也会变，最终的更新方向会完全不同——即使两个坐标系描述的是完全相同的神经网络。谱优化器的更新则只依赖于梯度矩阵的奇异值和奇异向量，这些量在正交坐标变换下以可预测的等变方式变化，因此更新规则与坐标系的选择无关。