Pion优化器:等谱流形上解决AdamW和Muon失稳
摘要
针对大模型训练中的失稳与超参数迁移问题,提出基于等谱流形的优化器Pion,通过正交变
当大语言模型参数规模突破百亿乃至千亿级别,研究焦点已从“如何更快收敛”转向两个更深层问题:一是如何在训练全程维持稳定性,二是如何实现从小模型到大模型的高效超参数迁移。 拆解其底层逻辑才能找到根本解法。
具体而言,前者体现为训练中各类数值与动力学失稳现象:注意力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代数中的斜对称矩阵生成,并通过矩阵指数映射回正交群。
这一更新带来关键特性:Pion不再直接“拉伸”权重矩阵,而是在特征空间中执行旋转。
由于左右两侧始终为正交变换,Pion严格保持权重矩阵的奇异值不变。训练过程中:
- spectral norm不会被无约束放大;
- Frobenius norm保持稳定;
- 权重的行空间与列空间持续演化,但整体尺度不会失控。
从几何视角看,传统优化器更新往往同时混合:
- 参数长度(magnitude)的变化;
- 参数方向(direction)的变化。
而Pion将更新完全转化为“旋转运动”。因此,Pion的更新范数不再对应参数缩放,而是直接刻画特征空间中的旋转强度。这意味着:模型不是在无约束地放大参数,而是在稳定地旋转特征空间。
基于上述规则,我们进一步探索了加速更新技术,详情参见论文。将最终得到的Pion算法总结为伪代码(图1):
图1:Pion优化器算法流程。
Pion与μP:谱保持优化器实现尺度迁移
μP(Maximal Update Parametrization)具体要求:
- 权重矩阵的spectral norm满足固定尺度规律;
- 参数更新量的spectral norm也满足对应尺度规律。
先前μP兼容优化器大多基于Muon路线,原因在于Muon更新天然易满足“更新谱条件”,prior work主要关注如何修正其权重尺度。而Pion恰好相反:由于Pion更新始终保持权重谱结构稳定,它天然满足μP对权重矩阵spectral norm的尺度规律。于是问题转变为:如何让Pion的更新幅度也符合μP的scaling law。
为此,我们进一步设计了μP兼容的Pion版本,对更新中的Lie代数因子实施谱归一化控制。在LLaMA架构与Qwen架构上,针对不同模型宽度验证学习率迁移性。如图2所示,Pion的最优学习率几乎可跨模型尺度直接迁移。
图2:Pion与μP。
Pion谱保持优化器实现稳定高效训练
我们从预训练与后训练两个维度评估Pion的训练稳定性。预训练环节除常规结构外,额外设置压力测试:完全去除归一化层以及超深层网络,检验Pion在极端条件下的稳定性。后训练环节采用SFT与RLVR两个标准流程测试。
稳定预训练
图3: 稳定性指标
在LLaMA-like 1.3B模型上进行预训练。除验证损失外,我们监控多项训练稳定性指标(如图3)。其中SwiGLU激活范数与最大attention logit已被广泛视为大规模预训练的关键稳定性指标。如图所示:AdamW的attention logit持续增长,激活范数快速放大;Muon虽显著抑制attention logit增长,但其激活值与down-projection相关范数在整个训练中持续上升;相比之下,Pion对所有监控指标保持近乎平坦且稳定的演化轨迹。这种截然不同的训练动态源于Pion有效的谱保持特性,如图4所示。
图4:谱的保持。
得益于稳定训练与更均匀的谱分布,Pion在zero-shot测试任务上取得更优泛化性(图5):
图5: Benchmark性能。
无归一化训练
为进一步压力测试Pion的训练稳定性,我们移除了一个60M LLaMA-like模型中的所有归一化层。归一化长期以来被视为控制激活尺度、稳定梯度反向传播的关键机制,缺少归一化后训练将极端不稳定,能更直接检验优化器本身的尺度控制能力。
图6: 无归一化训练曲线
实验结果非常显著(图6)。该设置下,AdamW与Muon虽在训练初期取得一定进展,但很快因梯度溢出而训练崩溃,最终产生NaN。而Pion在完整9.6B token训练过程中始终保持稳定,并成功收敛。这一结果表明:谱保持优化在一定程度上可替代架构层面的尺度控制机制,为模型训练提供来自优化器本身的稳定性来源。
超深层网络结构
为进一步压力测试Pion稳定性,我们在极端深度LLM结构上实验。深层网络通常放大优化稳定性问题,容易引发梯度消失与表示坍塌。实验中将60M LLaMA基线模型从8层逐步扩展至200层,在50B tokens的C4子集上训练。如图7所示,为清晰展示训练动态,用局部loss轨迹标准差均值衡量训练稳定性,阴影面积可视化波动程度。
图7: 深层网络训练
实验结果显示:AdamW出现最明显loss spike,整体稳定性最差;Muon在训练过程中仍存在持续波动累积;Pion在整个训练中保持最平滑的loss轨迹。对应标准差统计:AdamW:0.0931;Muon:0.0927;Pion:0.0892。这表明在极端深度设定下,Pion展现出最优训练稳定性,并在中期实现更快loss下降。
监督微调
监督微调(SFT)阶段长期存在一个问题:模型学习新任务时容易遗忘原有能力。
这本质上是stability-plasticity tradeoff:
- plasticity过强 → 学得快,但易灾难性遗忘;
- stability过强 → 保留旧能力,但新任务适应困难。
我们在Qwen2.5-1.5B与Llama3.2-3B两个基础模型上进行全参数微调实验,覆盖数学推理与代码生成两类典型任务。如图8所示,Pion在多个维度展现出更好的平衡能力。尤其在代码生成任务中,Pion同时取得最高ID(in-domain)与OOD(out-of-domain)表现;数学微调任务中,Pion在保持接近最优ID性能的同时,更有效地维持OOD泛化。
图8: Pion应用于监督微调。
表明:Pion不仅能更好学习目标任务,还能更稳定地保留原始模型能力。换句话说:Pion的谱保持更新,不仅稳定了训练过程,也稳定了模型知识本身。相比于传统优化器频繁改变参数尺度,Pion倾向在已有表示空间中进行“结构化旋转”,因此不易破坏预训练阶段形成的特征结构。
基于可验证奖励的强化学习
图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任务上的训练曲线。
上述结果表明:谱保持不仅适用于预训练,也可能是更适合RL的优化归纳偏置(inductive bias)。
结论:从“收敛优化器”走向“稳定优化器”
过去很长一段时间,优化器默认职责仅为尽快降低损失。但随着大模型训练规模不断扩张,“稳定性”本身正成为优化器最核心的能力之一。Pion提供了一条不同于传统路线的方向:它不依赖大量训练补丁,而是通过参数更新本身的几何约束,从源头抑制谱结构失控。因此Pion的意义或许不止是“一个更稳定的优化器”,更预示着:大模型优化器的下一阶段,也许不只是更快下降,而是更可控、更结构化、更长期稳定。
来源:互联网
本网站新闻资讯均来自公开渠道,力求准确但不保证绝对无误,内容观点仅代表作者本人,与本站无关。若涉及侵权,请联系我们处理。本站保留对声明的修改权,最终解释权归本站所有。