AdamW和Muon的失稳根源有解了,港中文新优化器Pion:在等谱流形上更新大模型


本文作者 施柯煊和李瀚轩是香港中文大学计算机科学与工程系博士生,邱泽钜是德国马克斯普朗克研究所博士生,温研东是西湖大学助理教授,Simon Buchholz是德国马克斯普朗克研究所研究员,刘威杨是香港中文大学计算机科学与工程系助理教授。

当大语言模型不断迈向百亿乃至千亿参数规模时,人们的核心关注已不再仅仅是 “如何更快收敛”,而是逐渐转向两个更加本质的问题: 其一,是如何在训练过程中维持稳定性;其二,是如何实现从小模型到大模型的有效超参数迁移

具体而言,前者主要表现为训练过程中的一系列数值与动力学失稳现象,例如注意力 logits 持续增大、激活输出范数逐步漂移、深层网络中损失尖峰频繁出现,甚至最终导致数值溢出与训练崩溃。后者则依赖于 μP(Maximal Update Parameterization)等尺度化理论,其核心思想是在宽度变化的条件下,通过对参数矩阵与更新尺度进行系统性约束,使不同规模模型之间的超参数具有可迁移性。

围绕这两个问题,当前主流的工程应对方式主要集中在 “训练稳定性补丁” 和 “尺度控制框架” 两个方向。一类方法通过梯度裁剪、激活截断、学习率衰减等手段,直接抑制训练过程中的数值爆炸;另一类方法则借助归一化或谱约束等机制,对参数或更新的尺度进行显式控制,从而缓解训练过程中的漂移问题。 这些方法通过抑制不稳定现象的外在表现,取得了一定的成效,但是并非从优化动力学层面约束其根源

近期,研究团队提出了一类基于 等谱流形(iso-spectral manifold)的优化器 Pion 。其核心思路是在参数矩阵对应的等谱流形上进行优化,从而在根源上规避动力学失稳。这一视角为长程训练中的尺度漂移与稳定性问题提供了一种从 “几何约束优化动力学” 出发的替代思路。

  • 论文标题: Pion: A Spectrum-Preserving Optimizer via Orthogonal Equivalence Transformation

  • 论文链接: https://arxiv.org/pdf/2605.12492

  • 代码链接: https://github.com/Sphere-AI-Lab/pion

加法更新的困境

追本溯源,无论是 μP 尺度的失效亦或是是训练频频失稳,传统优化器的加法更新范式是主要原因之一。 

过去几年,AdamW 与 Muon 等优化器的设计重点都集中在如何更快下降训练损失与提升训练效率, 但它们本质上都采用加法更新

。这种更新方式天然关注的是 “如何沿梯度快速下降”,但并不会主动约束权重矩阵的几何结构。

随着训练的不断进行,无约束的加法累积同时改变了参数的长度与方向,逐渐破坏了参数矩阵的谱几何。具体表现为奇异值谱范数被持续放大,不同特征方向的尺度越来越失衡,整体矩阵范数不断漂移。

这些底层的几何变化进一步放大了网络中的激活值,摧毁了 μP 的前向尺度前提。换句话说,许多训练失稳与参数化崩溃现象,并不仅仅来自梯度太大,而是因为参数矩阵本身的结构在长期更新过程中逐渐失控。

从谱保持重新理解稳定训练

近期工作 POET [1] 提出了一个新的思路:通过控制参数矩阵本身的谱结构实现稳定的训练。POET 的核心做法,是通过正交矩阵对权重进行双侧变换:

,其中 (R,P) 为正交矩阵。由于正交变换本质上只会旋转特征空间,而不会改变奇异值,因此这样的更新天然具有保谱性质(spectrum-preserving)。这意味着:

  • spectral norm 不会被无约束放大;

  • 参数矩阵整体范数更加稳定;

  • 特征空间可以持续演化,但矩阵尺度不会失控。

但与此同时,POET 仍然依赖重参数化训练框架,需要额外维护两个可训练正交矩阵,并固定原始权重矩阵。这不仅给训练系统兼容性与跨架构适配带来额外复杂度,也要求更加复杂的一阶动量设计。

Pion:不做重参数化,直接把 “保谱” 写进优化器

基于这一观察,我们进一步提出了 Pion(POET-induced Optimizer with No Reparameterization)。与 POET 不同,Pion 不再依赖显式重参数化,而是直接将 “谱保持更新” 写进优化器本身。

Pion 从一个非常简单但关键的观察出发。对于任意权重矩阵 

,都可以写成:

这里的两个单位矩阵,本质上可以被视为 “零旋转” 的正交变换。

Pion 的做法,并不是显式学习新的 (R,P),而是直接在正交群上更新这两个 “单位因子”,从而对权重矩阵施加左右两侧的正交变换。我们将

,Pion 的更新规则写为:

的梯度记为

其中,两侧更新都由 Lie algebra 中的斜对称矩阵生成,并通过矩阵指数映射回正交群。

这一更新方式带来了一个非常重要的性质: Pion 不再直接 “拉伸” 权重矩阵,而是在特征空间中对其进行旋转

由于左右两侧始终是正交变换,Pion 会严格保持权重矩阵的奇异值不变。换句话说,训练过程中:

  • spectral norm 不会被无约束放大;

  • Frobenius norm 保持稳定;

  • 权重的行空间与列空间持续演化,但整体尺度不会失控。

从几何视角来看,传统优化器的更新往往同时混合了:

  • 参数长度(magnitude)的变化;

  • 参数方向(direction)的变化。

而 Pion 则将更新完全转化为 “旋转运动”。因此,Pion 的更新范数不再对应参数缩放,而直接刻画了特征空间中的旋转强度。这意味着: 模型不是在无约束地放大参数,而是在稳定地旋转特征空间。

基于上述规则,我们进一步对加速更新的技术进行了探索,详情见论文链接。我们将最后的得到的 Pion 算法总结成伪代码(图 1):

图 1:Pion 优化器算法流程。图 1:Pion 优化器算法流程。

Pion 与 μP:谱保持优化器实现尺度迁移

μP(Maximal Update Parametrization)具体来说要求:

  • 权重矩阵的 spectral norm 满足固定尺度规律;

  • 参数更新量的 spectral norm 也满足对应尺度规律。

过去已有的 μP-compatible 优化器,大多建立在 Muon 路线之上。原因在于:Muon 的更新天然容易满足 “更新谱条件”,因此 prior work 主要关注如何进一步修正它的权重尺度。而 Pion 刚好相反。由于 Pion 的更新始终保持权重谱结构稳定,因此: 它天然满足 μP 对 权重矩阵的 spectral norm 的尺度规律。于是问题变成:如何让 Pion 的更新幅度也满足 μP 的 scaling law。

为此,我们进一步设计了 μP-compatible 的 Pion 版本,对更新中的 Lie algebra 因子进行谱归一化控制。我们分别在 LLaMA 架构和 Qwen 架构,对不同模型宽度下验证 learning rate transferability。如图 2 所示, Pion 的最优学习率几乎可以跨模型尺度直接迁移

图 2:Pion 与 μP。图 2:Pion 与 μP。

Pion 谱保持优化器实现稳定高效训练

我们从预训练、后训练两个方面观察 Pion 的训练的稳定性。对于预训练,除了常规的结构的训练,我们额外增加了压力测试:完全去除归一化层以及超深层网络,来测试 Pion 在极端情况下的稳定性。 对于后训练,我们采用 SFT 和 RLVR 两个常规 Pipeline 进行测试。

稳定预训练

图 3: 稳定性指标图 3: 稳定性指标

我们在 LLaMA-like 1.3B model 上进行预训练。除了验证损失之外,我们进一步监控了多项训练稳定性指标(如图 3 所示)。其中,SwiGLU 激活范数以及最大 attention logit,已经被广泛认为是大规模预训练中的关键稳定性指标。如图所示: AdamW 的 attention logit 持续增长,同时激活范数迅速放大; Muon 虽然显著抑制了 attention logit 的增长,但其激活值与 down-projection 相关范数仍在整个训练过程中持续上升; 相比之下,Pion 对所有监控指标都保持了近乎平坦且稳定的演化轨迹。 这种截然不同的训练动态,来自于 Pion 有效的谱保持特性,如图 4 所示。

图 4:谱的保持。图 4:谱的保持。

得益于稳定训练以及更均匀的谱,Pion 在 zero-shot 的测试任务上,取得了更好的泛化性(图 5):

图 5: Benchmark 性能。图 5: Benchmark 性能。

Normalization-free 训练. 

为了进一步对 Pion 的训练稳定性进行压力测试,我们移除了一个 60M LLaMA-like 模型中的所有 normalization 层。之所以采用这一设置,是因为 normalization 长期以来都被认为是控制激活尺度、稳定梯度反向传播的关键机制。因此,在缺少 normalization 的情况下,训练会变得极其不稳定,也能够更直接地检验:优化器本身,是否具备足够的尺度控制能力。

图 6: Normlaization-free 训练曲线图 6: Normlaization-free 训练曲线

实验结果非常明显,如图 6 所示。在这一设置下,AdamW 与 Muon 虽然能够在训练初期取得一定进展,但很快便由于梯度溢出而训练崩溃,并最终产生 NaN。相比之下,Pion 在完整的 9.6B token 训练过程中始终保持稳定,并最终成功收敛。这一结果表明: 谱保持优化在一定程度上能够替代架构层面的尺度控制机制,为模型训练提供一种来自优化器本身的稳定性来源

超深层网络结构

为了进一步对 Pion 的稳定性进行压力测试,我们在极端深度的 LLM 结构上开展实验。 深层网络通常被认为是优化稳定性的 “放大镜”,容易引发严重的训练问题,例如梯度消失以及表示坍塌等现象。在该实验中,我们将一个 60M 规模的 LLaMA 基线模型从 8 层逐步扩展至 200 层,并在 50B tokens 的 C4 子集上进行训练。如图 7 所示,为了更清晰地展示训练动态,我们通过局部 loss 轨迹的标准差均值来衡量训练稳定性,其对应的阴影面积可视化了波动程度。

图 7: 深层网络训练图 7: 深层网络训练

实验结果显示:AdamW 出现最明显的 loss spike,整体稳定性最差;Muon 在训练过程中仍然存在持续的波动累积;Pion 则在整个训练过程中保持最平滑的 loss 轨迹。对应的标准差统计结果则为:AdamW:0.0931; Muon:0.0927; Pion:0.0892。这表明在极端深度设置下,Pion 表现出最优的训练稳定性,同时仍然能够在中期阶段实现更快的 loss 下降。

监督微调

在监督微调(SFT)阶段,一个长期存在的问题是: 模型在学习新任务时,很容易遗忘原有能力

这本质上是一个 stability-plasticity tradeoff:

  • plasticity 太强 → 学得快,但容易 catastrophic forgetting;

  • stability 太强 → 保留旧能力,但新任务适应困难。

我们在 Qwen2.5-1.5B 和 Llama3.2-3B 两个基础模型上进行了全参数微调实验,并覆盖:数学推理,代码生成两类典型任务。如图 8 所示,Pion 在多个维度都展现出了更好的平衡能力。尤其是在代码生成任务中,Pion 同时取得了最高的 ID(in-domain)与 OOD(out-of-domain)表现;数学微调任务中,Pion 也在保持接近最优 ID 性能的同时,更有效地维持了 OOD 泛化。

图 8: Pion 应用于监督微调。图 8: Pion 应用于监督微调。

表明:Pion 不仅能更好学习目标任务,同时还能更稳定地保留原始模型能力。换句话说: Pion 的谱保持更新,不仅稳定了训练过程,也稳定了模型知识本身 。相比于传统优化器频繁改变参数尺度,Pion 更倾向于在已有表示空间中进行 “结构化旋转”,因此不容易破坏预训练阶段已经形成的特征结构。

Reinforcement Learning with Verifiable Reward

图 9: Pion 应用于 RLVR。图 9: Pion 应用于 RLVR。

我们进一步测试 Pion 在 RLVR(Reinforcement Learning with Verifiable Reward)中的表现。RL 往往是大模型训练中最不稳定的阶段之一,具体来说,reward variance 大;optimization noise 强;容易出现模式崩塌与训练震荡。最近的一些研究发现,一个非常有意思的现象是:RL 阶段的参数更新,其实往往天然倾向于保留预训练权重的谱结构。这意味着:强化学习可能本身就更偏向 “结构保持型更新”。而这与 Pion 的更新几何几乎天然一致。因此,我们进一步在:Qwen3-1.7B, DeepSeek-R1-Distill-Qwen-1.5B 上进行了 RLVR 实验,并采用 GRPO 训练框架进行数学推理强化学习。如图 9 所示,Pion 在所有 RL 设置下都取得了最佳平均表现。与此同时,验证集 accuracy 曲线(图 10)也显示:Pion 收敛更快;训练更稳定;后期性能波动更小。

图 10: Pion 在 RLVR 任务上的训练曲线。图 10: Pion 在 RLVR 任务上的训练曲线。

上述结果表明: 谱保持不仅适用于预训练,也可能是一种更适合 RL 的优化归纳偏置(inductive bias)

结论:从 “收敛优化器” 到 “稳定优化器”

过去很长时间里,人们默认优化器的职责只有一个:尽快降低 loss。 但随着大模型训练规模不断扩大,“稳定性” 本身,正在变成优化器最核心的能力之一。而 Pion 提供了一个不同于传统路线的方向:它不依赖大量训练补丁,而是通过参数更新本身的几何约束,从源头抑制谱结构失控。因此,Pion 的意义可能不仅仅是 “一个更稳定的优化器”,更预示着: 大模型优化器的下一阶段,也许不只是更快下降,而是更可控、更结构化、更长期稳定

]article_adlist–>

            <!-- 非定向300*250按钮    17/09  wenjing  begin -->
            <!-- 非定向300*250按钮  end -->
        </div>



Source link

未经允许不得转载:紫竹林-程序员中文网 » AdamW和Muon的失稳根源有解了,港中文新优化器Pion:在等谱流形上更新大模型
关于我们 免责申明 意见反馈 隐私政策
程序员中文网:公益在线网站,帮助学习者快速成长!
关注微信 技术交流
推荐文章
每天精选资源文章推送
推荐文章
随时随地碎片化学习
推荐文章
发现有趣的