香港中文大学研发频谱守护者优化器提升AI训练稳定性

训练大型语言模型，如同在云端构建一座持续生长的知识大厦。随着模型层数不断增加，任何微小的参数偏差都可能被逐层放大，最终导致训练过程失控。如何确保这座大厦在建造过程中始终保持结构稳定，一直是困扰研究人员的核心挑战。

近期，一项由香港中文大学、马克斯·普朗克智能系统研究所和西湖大学联合发布的技术报告，带来了一个极具启发性的解决方案。这项研究提出了一种名为Pion的全新优化器。简而言之，优化器是深度学习中的“总工程师”，负责指导模型如何调整其数百万乃至数十亿的参数。Pion的独特之处在于，它不像传统方法那样直接“增减”参数数值，而是像调整承重柱的方向一样，仅对参数矩阵进行“旋转”操作，确保其内在的“材料总量”（即数学上的奇异值频谱）恒定不变。

这种“只转不增”的思路，直指大模型训练不稳定的一个深层根源：当参数矩阵的频谱在训练过程中发生漂移，模型的行为就会变得难以预测，训练过程如同在钢丝上行走。

一、为什么现有的优化器方案存在局限

在Pion出现之前，深度学习领域主要由两位“总工程师”主导：经典主流的AdamW和后起之秀Muon。AdamW应用广泛，但其更新方式会同时改变参数的“方向”和“幅度”，长此以往，模型的内在特性会悄然改变。Muon通过“正交化更新”使更新步伐更规整，但它同样无法阻止参数矩阵“能量分布”的缓慢漂移。

这种漂移在数学上表现为权重矩阵奇异值谱的偏移。大量实践经验表明，当权重矩阵的谱范数（可理解为“最大能量密度”）保持稳定时，模型的泛化能力——即处理未见过的任务的能力——通常会更强。

另一个关键框架是“最大更新参数化”（μP）。这套业界规范要求，当模型规模（宽度）扩大时，权重和更新的幅度必须按特定比例缩放，才能保持激活值的稳定。Muon天然满足了“更新幅度”的要求，但对“权重本身的幅度”却无能为力。而Pion从设计之初，就将“保持权重幅度”作为核心机制，从另一个维度原生地满足了μP的要求。

二、Pion的核心原理：仅旋转，不伸缩

Pion的灵感来源于一个名为POET的早期方法。POET的策略是将每个权重矩阵拆分为一个固定基座和两个可旋转的“夹具”，通过旋转夹具来间接调整矩阵方向，同时保持其频谱不变。这好比用两个转盘夹住一块钢板，只改变钢板朝向，不改变其厚度和材质。

但POET的问题在于，这种拆解与重组引入了间接性，容易导致训练损失突然跳变，在不同模型架构间迁移也较为繁琐。

Pion采取了更直接的路径。其名称“Pion”正是“POET诱导优化器，无需重参数化”的缩写。它不再拆分矩阵，而是直接对权重矩阵本身施加旋转操作。每次更新，权重矩阵会同时被一个左侧旋转矩阵和一个右侧旋转矩阵相乘，整个过程如同从两侧温和地扭转一个物体，而其“骨架结构”（奇异值频谱）则完好无损。

这一思想的数学基础是李群优化——一个专门研究如何在旋转、反射等变换空间中进行优化的分支。Pion的每一步更新，都严格地在名为“等谱流形”的曲面上滑动，绝不偏离。

三、Pion如何计算更新方向

具体到操作层面，Pion在每一步训练中执行一套精密的计算流程。假设当前权重矩阵为W，损失函数给出的梯度信号为G。

首先，Pion将梯度G与权重W进行两次矩阵乘法，分别得到“输入侧”和“输出侧”的梯度分量。接着，为了将这些梯度转化为描述旋转运动的“语言”，需要对它们进行“反对称化”处理——即用该矩阵减去其自身的转置。得到的反对称矩阵，在数学上称为李代数元素，正是描述旋转运动的恰当形式。

获得这两个旋转指令后，Pion通过矩阵指数映射（一种将“旋转速度”转换为“实际旋转量”的数学工具）生成两个正交矩阵，分别从左、右两侧乘以当前权重矩阵，从而完成一次更新。

可以这样直观理解：想象你双手捏住一张方形卡片的两条对边，然后双手同时向相反方向轻微扭转。卡片的面积和形状没有改变，只是朝向发生了变化。Pion对权重矩阵所做的，正是这样一次受控的双向扭转。

四、确保稳定训练的四大工程改进

然而，仅有基础机制还不够。研究团队在实验中发现，最朴素的Pion更新会引发实际问题，就像一位理论完美的建筑师也需要规范的施工流程。他们系统性地探索了四项关键的工程改进。

第一项：一致性缩放。 如果直接使用基本规则，同一网络层中不同矩阵受到的旋转强度会差异巨大，甚至同一矩阵的左右两侧旋转幅度也会失衡。为解决此问题，团队引入了两项措施：对每个矩阵两侧的旋转指令按其尺寸进行归一化，同时引入一个基于权重参数量的全局缩放系数（用均方根RMS控制）。实验表明，这一改进极大地拓宽了学习率的可用范围，而原始版本在稍大的学习率下就会崩溃。

第二项：动量设计。 动量是优化器中提供“方向惯性”的机制。在弯曲的等谱流形上，动量的设计比在平坦空间中更微妙。团队比较了三种方案：概念最严谨但计算适中的“平行传输动量”、最节省资源但存在几何偏差的“普通环境空间动量”，以及几何上最严格但需要额外存储的“李代数动量”。实验结果显示，“两侧均采用李代数动量”的组合表现最佳。

第三项：交替更新。 原始Pion每一步都进行双侧旋转，效果最好但计算成本高。团队提出了一个简洁变体：奇数步只进行输入侧旋转，偶数步只进行输出侧旋转。实验显示，交替更新的最终损失仅比双侧同步更新高出约0.23%，在训练早期甚至更快，是追求性价比的绝佳选择。

第四项：矩阵指数的高效近似。 精确计算矩阵指数开销很大。幸运的是，Pion每次都是从“零旋转”状态开始计算，误差不会累积。团队对比了多种近似方案，发现二阶泰勒展开（保留到平方项）在精度和效率间取得了最佳平衡。

五、Pion的收敛性保证

除了直觉和实验，研究团队还为Pion提供了严格的数学收敛性保证。这好比为建筑方案出具了结构安全计算书。

在“目标函数足够光滑且下有界”以及“随机梯度噪声有界”这两个标准假设下，团队证明：随着训练步数T趋向无穷，Pion找到的最佳“停止点”处，其旋转梯度的平方和会以O(1/√T)的速率收敛到零。这意味着训练越久，Pion越能找到等谱流形上的稳定点，且收敛速度符合非凸随机优化的经典标准。

六、与μP框架的兼容性：实现超参数跨规模迁移

“最大更新参数化”（μP）是一个极具实用价值的工程框架。其核心承诺是：在小模型上调试好的最优学习率，在按比例放大后的大模型上依然最优，无需重新调参。这对于动辄耗费数百万美元算力的大模型训练而言，意味着巨大的成本节约。

μP对优化器有两条要求：权重矩阵的谱范数需按特定比例缩放；更新步长的谱范数也需满足同样比例。Pion天然满足第一条，因为它根本不改变频谱。对于第二条，团队提出了两种方案：一是对两侧的李代数梯度矩阵的谱范数进行归一化；二是借鉴Muon，对梯度矩阵做正交化处理。实验证实，两种方案均能有效实现超参数的跨规模迁移。

七、实战检验：Pion在三大任务场景中的表现

研究团队在三大类真实任务上对Pion进行了全面测评，以AdamW和Muon作为主要基准。

大规模预训练： 使用13亿参数的LLaMA架构模型，在5400亿词元的C4数据集上训练。结果显示，Pion在八个主流基准测试上的平均得分最高（47.69），验证集损失（2.7350）与Muon（2.7225）接近，均优于AdamW（2.7700）。更重要的是稳定性指标：Pion使得注意力权重、激活值等关键信号的范数在整个训练过程中几乎保持水平，而AdamW和Muon下这些指标持续攀升。频谱变化图更直观地显示，Pion训练后的奇异值分布与初始化时几乎重合。

极端压力测试： 团队进行了两组“魔鬼”实验。一是彻底移除所有归一化层（如LayerNorm）。在没有这个“稳定器”的情况下，AdamW和Muon均在早期因梯度溢出而失败，而Pion全程平稳收敛。二是将网络深度暴力增至200层。在这个极易出现梯度消失的场景下，Pion的训练损失波动最小，下降速度最快，且各层的表达能力分布最为均匀。

有监督微调与强化学习微调： 在数学推理和代码生成任务的微调中，Pion在代码任务上表现最佳，在数学任务上则显示出更强的抗灾难性遗忘能力。在强化学习微调中，Pion在两个基础模型上都取得了最高平均分，且验证准确率上升最快。有趣的是，在强化学习场景中，交替更新策略反而优于双侧同步更新，推测原因是其引入了更多探索性变化，有助于避免过早锁定次优策略。

八、计算代价分析：是否物有所值

天下没有免费的午餐。Pion的额外开销主要来自三部分：构建两侧李代数梯度、RMS缩放计算以及矩阵指数近似。在实际的大批次训练中，由于前向/反向传播的计算量与批次大小成正比，而Pion的额外矩阵运算与批次大小无关，因此其相对开销会被摊薄。

实测数据显示，完整版Pion的峰值显存占用比AdamW高约16%，主要源于额外的动量缓冲区。但如果去掉二阶动量，显存占用反而会低于AdamW。每步训练时间上，Pion比AdamW慢44.4%，但仅比Muon慢3.2%。对于追求极致稳定的场景，这份开销或许物有所值。

九、Pion的额外优势：维持最小超球面能量

除了保持频谱稳定，Pion还有一个优雅的副产品：它能天然维持“最小超球面能量”配置。这是衡量高维空间中神经元分布均匀度的指标，分布越均匀，通常泛化能力越好。经典的随机初始化恰好产生最小能量配置，而Pion所做的正交变换被数学证明能保持这一配置。换言之，只要初始化得当，Pion能在整个训练过程中持续维持神经元分布的均匀性，而AdamW和Muon都会破坏这一性质。

说到底，Pion的核心哲学可以用一句话概括：只转动参数，不拉伸压缩。 这种看似简单的约束，引发了一系列连锁反应——训练更稳定、频谱不漂移、激活值可控，甚至在移除归一化层或训练超深网络等极端条件下也能存活，并在预训练、微调、强化学习三大主流场景中均展现出竞争力。

当然，Pion并非没有代价。其完整版的显存开销更高，将其扩展到更大参数量级的模型仍是未来工作。此外，关于μP兼容性的探索也仅是初步。

然而，一个能让模型在没有归一化层辅助下依然稳定训练的优化器，本身就提供了一个全新的视角：或许并非所有“稳定性”都必须由模型架构来提供，优化器本身也能承担更多责任。对于正在构思下一代模型训练方案的研究者与工程师而言，Pion无疑指明了一条值得深入探索的技术路径。

Q&A

Q1：Pion优化器和AdamW、Muon相比有什么本质区别？

A：AdamW和Muon通过直接“加减”数值来修改权重矩阵，而Pion只做“旋转”操作。它从左右两侧同时对权重矩阵施加正交变换，严格保持矩阵的奇异值频谱不变，从而确保了权重矩阵整体的“能量分布”在整个训练过程中始终与初始化时一致。这带来了更稳定的训练行为，尤其在移除归一化层或训练极深网络等极端情况下优势明显。

Q2：Pion的显存占用比AdamW高，实际使用时怎么处理这个问题？

A：完整版Pion因需在李代数空间维护额外的动量缓冲区，显存占用较高。但研究发现，去掉二阶动量（仅保留一阶动量）后，显存占用可降至低于AdamW的水平，且性能损失很小。此外，采用交替更新策略（奇偶步交替进行单侧旋转）可将每步计算量减少约一半，最终性能仅下降约0.23%。在资源受限的场景下，可通过这两项调整有效控制开销。

Q3：Pion在强化学习微调中为什么表现比监督微调更突出？

A：有观察表明，强化学习阶段的权重更新天然倾向于保持预训练时的频谱结构，这与Pion的设计原则高度契合，使得Pion成为该场景的“天然适配者”。此外，交替更新在强化学习中表现更优，推测原因是其每步只旋转一侧，引入了更多探索性变化，有助于模型在稀疏奖励信号下避免过早收敛到次优策略。