Pion: A Spectrum-Preserving Optimizer via Orthogonal Equivalence Transformation¶
会议: NeurIPS 2026
arXiv: 2605.12492
OpenReview: https://openreview.net/forum?id=xFbqlwBtwM
代码: https://spherelab.ai/pion
领域: 优化/理论
关键词: 谱保持优化器, 正交等价变换, 训练稳定性, 李代数动量, 大语言模型预训练
一句话总结¶
针对大语言模型训练中权重谱漂移和激活爆炸的痛点，本文提出基于正交等价变换的优化器 Pion，通过左右正交李群旋转完全保持权重矩阵的奇异值谱与超球面能量，结合 RMS 尺度一致性缩放与李代数动量，实现了无归一化层与极端深度下的稳定高效训练。
研究背景与动机¶
随着大语言模型（LLM）参数规模与训练上下文的持续扩展，优化算法的数值稳定性成为制约大模型训练的瓶颈之一。最大更新参数化（\(\mu\)P）表明，控制权重矩阵与更新量的谱范数对于防止激活爆炸、维持不同宽度下的特征学习常数尺度至关重要。近年来以 Muon 为代表的优化器通过对更新量进行正交化实现了谱范数下的最速下降，在预训练中展现出相比 AdamW 的收敛优势；但 Muon 本质上仍属于加性更新，其权重矩阵本身的奇异值谱在优化轨迹中仍会不可避免地发生漂移，不得不引入后验归一化或复杂的谱回退算子。
这种不稳定的根源在于传统加性更新（如 AdamW 和 Muon）将权重幅度缩放与方向旋转紧密耦合，使得深层 Transformer 极易出现注意力 logit 持续攀升、前馈网络（FFN）激活爆炸或中间层表征退化。虽然此前 POET 提出了将权重矩阵显式重参数化为两个正交矩阵与固定基底矩阵的乘积 \(\bm{W} = \bm{R}\bm{W}_0\bm{P}\)，从结构上锁定了权重谱，但引入辅助正交参数后破坏了原本清晰的优化动态，带来了严重的训练尖峰（loss spikes）与梯度动量失配难题。
本文的切入角度是打破显式重参数化的约束，将正交等价变换直接内化为作用在模型参数自身的李群连续优化机制。核心 idea：将权重矩阵更新表述为恒等旋转因子的李群映射，利用反对称投影导数驱动左右正交变换 \(\bm{W}_{t+1} = \exp(-\eta \bm{G}^{\text{out}}_t) \bm{W}_t \exp(-\eta \bm{G}^{\text{in}}_t)\)，在直接优化参数的同时严格保全矩阵奇异值谱与超球面极小能量配置。
方法详解¶
整体框架¶
Pion 的优化机制立足于等谱流形（iso-spectral manifold）。对于任意给定的权重矩阵 \(\bm{W} \in \mathbb{R}^{d_{\text{out}} \times d_{\text{in}}}\)，在第 \(t\) 步时均可恒等写为 \(\bm{W}_t = \bm{I}_{d_{\text{out}}} \bm{W}_t \bm{I}_{d_{\text{in}}}\)。Pion 不引入任何外部辅助参数，而是直接在正交李群上更新左右两侧的恒等算子，产生耦合的左右正交等价变换，在几何上仅旋转权重矩阵的行与列子空间，严格保持其全部奇异值。
整个优化流程首先从反向传播得到的反向梯度矩阵构建左右两侧的切空间 Lie 代数反对称梯度；随后在李代数内部累积一阶动量与二阶自适应矩统计量；接着根据权重矩阵维度实施一阶近似的 RMS 尺度一致性缩放，控制各层的旋转角速度；最后采用二阶截断的矩阵指数泰勒展开将更新映射回正交群，并支持交替更新策略以大幅削减计算量。
%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
A["反向传播梯度 G_t 与权重 W_t"] --> B["李代数反对称投影<br/>构建 G_in 与 G_out"]
B --> C["李代数动量累加<br/>一阶方向 m 与二阶方差 v"]
C --> D["自适应步长与方向 A_in, A_out"]
D --> E["RMS 尺度一致性缩放<br/>计算标量因子 alpha_t"]
E --> F["二阶指数近似与更新<br/>单双侧正交旋转作用于 W_t"]
F --> G["保持奇异值谱的新权重 W_t+1"]关键设计¶
1. 李代数反对称投影与谱保持机制：彻底解耦旋转方向与幅度漂移
标准欧氏梯度下降将权重方向旋转与范数缩放混杂在一起，容易引发激活值尺度爆炸。Pion 针对两端恒等变换因子求导，并通过投影到李代数 \(\mathfrak{so}(d)\) 强制满足反对称性： $\(\bm{G}^{\text{in}}_{t} = \bm{W}_{t}^{\top} \bm{G}_{t} - (\bm{W}_{t}^{\top} \bm{G}_{t})^{\top}, \quad \bm{G}^{\text{out}}_{t} = \bm{G}_{t} \bm{W}_{t}^{\top} - (\bm{G}_{t} \bm{W}_{t}^{\top})^{\top}\)$ 利用矩阵指数映射 \(\bm{R}_t = \exp(-\eta \bm{G}^{\text{out}}_t)\) 与 \(\bm{P}_t = \exp(-\eta \bm{G}^{\text{in}}_t)\)，若 \(\bm{W}_t = \bm{U}_t \bm{\Sigma}_0 \bm{V}_t^{\top}\)，则更新后 \(\bm{W}_{t+1} = (\bm{R}_t \bm{U}_t) \bm{\Sigma}_0 (\bm{P}_t^{\top} \bm{V}_t)^{\top}\)，其奇异值谱与初始奇异值 \(\bm{\Sigma}_0\) 完全恒等。在几何上，\(\bm{G}^{\text{in}}_t\) 和 \(\bm{G}^{\text{out}}_t\) 在两两正交的 2D 不变子空间中产生独立的平面旋转，其 Frobenius 范数分别表征了平均旋转强度，谱范数控制最大旋转角，使得更新增量 \(\|\Delta \bm{W}\|_F\) 纯粹对应于角度位移，消除了谱发散隐患。
2. RMS 尺度一致性缩放：跨参数尺寸的更新幅度自适应对齐
朴素李群更新存在两重尺度失配问题：同一层内同尺寸矩阵的更新模长方差巨大，且输入端与输出端的平均旋转角常年失衡，导致模型在较大初始学习率下快速发散。Pion 引入标量自适应因子 \(\alpha_t\) 对旋转步长进行再调控，保证更新步长与矩阵尺寸成比例（\(\mathrm{RMS}(\Delta \bm{W}/\eta) \approx \text{const}\)）。为避免计算精确指数映射带来的双倍 FLOPs 开销，利用一阶泰勒展开 \(\Delta \bm{W}/\eta \approx - \bm{G}^{\text{out}}_t \bm{W}_t - \bm{W}_t \bm{G}^{\text{in}}_t\)，推导闭式缩放系数： $\(\alpha_{t} \approx \frac{c \sqrt{d_{\text{out}} d_{\text{in}}}}{\|\bm{A}^{\text{out}}_t \bm{W}_{t-1} + \bm{W}_{t-1} \bm{A}^{\text{in}}_t\|_F + \epsilon}\)$ 该设计使不同形状与参数量的层获得与其空间维度相符的旋转速率，成功使 Pion 能够在 \(1\times 10^{-3}\) 等大初始学习率下高速稳定收敛。
3. 李代数内在动量跟踪与二阶矩自适应：曲率空间下的准确梯度累积
由于等谱流形具有非零黎曼曲率，常规在周围欧氏空间（ambient space）中直接做梯度滑动平均会遭遇不同切空间之间的方向失准。Pion 提出在李代数内部维护一阶动量 \(\bm{m}^{\text{in}}_t, \bm{m}^{\text{out}}_t\) 与二阶自适应统计量 \(\bm{v}^{\text{in}}_t, \bm{v}^{\text{out}}_t\)： $\(\bm{m}_t^{\text{in}} = \beta_1 \bm{m}_{t-1}^{\text{in}} + (1-\beta_1) \bm{G}_t^{\text{in}}, \quad \bm{v}_t^{\text{in}} = \beta_2 \bm{v}_{t-1}^{\text{in}} + (1-\beta_2) (\bm{G}_t^{\text{in}} \odot \bm{G}_t^{\text{in}})\)$ 输入输出方向向量相应化为 \(\bm{A}^{\text{in}}_t = - \bm{m}^{\text{in}}_t / (\sqrt{\bm{v}^{\text{in}}_t} + \epsilon)\) 与 \(\bm{A}^{\text{out}}_t = - \bm{m}^{\text{out}}_t / (\sqrt{\bm{v}^{\text{out}}_t} + \epsilon)\)。由于每一步更新均以内蕴于单位元处的固定切空间表达，该机制在保持李代数反对称结构的同时，提供了无偏且几何一致的历史梯度积分，收敛速度显著优于环境空间动量。
4. 二阶指数截断近似与交替更新：计算与显存开销的实用化折衷
精确矩阵指数计算极其昂贵，常规 Cayley 变换亦涉及矩阵求逆。Pion 证明，由于每一步迭代都以当前点为基准从李群单位元发起微小旋转，步长很小使得截断误差不会跨迭代级联放大，因此二阶幂级数展开 \(\mathcal{E}_2(\bm{A}, \alpha) = \bm{I} + \eta \alpha \bm{A} + \frac{1}{2}(\eta \alpha \bm{A})^2\) 足以维持严格谱保持与高精度收敛。此外，Pion 设计了奇偶步交替更新机制（Alternate Update）：奇数步仅更新输入侧正交变换，偶数步仅更新输出侧。在保持两端动量连续积累的前提下，交替更新将更新侧的计算开销直接降低近 50%，在仅产生 0.23% 的极小最终验证损失差距下大幅提升吞吐效率。
损失函数 / 训练策略¶
Pion 的优化目标完全兼容通用的经验风险最小化 \(\min_{\bm{W}} f(\bm{W})\)。在预训练和后训练中，Pion 采用余弦退火学习率调度（最低退火至初始学习率的 0.01 倍），默认动量系数设为 \(\beta_1=0.9, \beta_2=0.95\)，RMS 缩放系数常数设为 \(c=0.2\)。对于多头注意力机制，Pion 推荐分头（per-head）或独立对 Query、Key、Value 与 FFN Projection 矩阵进行分块正交变换；在显存极度敏感时，可直接省略二阶矩状态 \(\bm{v}\)（Pion w/o 2nd Moment），显存占用可降低 24.3%，甚至低于 AdamW 与 Muon。
实验关键数据¶
主实验¶
在 LLaMA-1.3B 预训练（54B tokens, C4 数据集, 约 400K 步迭代）及 downstream evaluation benchmark 评测中，Pion、Muon 和 AdamW 的综合表现对比如下：
| 模型 / 数据集 | 指标 | Pion | Muon | AdamW | 相对优势 |
|---|---|---|---|---|---|
| LLaMA-1.3B (Pretrain) | Val Loss (C4) | 2.7350 | 2.7225 | 2.7700 | 显著优于 AdamW，与 Muon 相当 |
| ARC-Challenge | Acc (%) | 26.79 | 25.34 | 25.94 | +0.85% vs AdamW, +1.45% vs Muon |
| ARC-Easy | Acc (%) | 49.41 | 47.94 | 45.96 | +3.45% vs AdamW, +1.47% vs Muon |
| BoolQ | Acc (%) | 57.58 | 51.56 | 46.30 | +11.28% vs AdamW, +6.02% vs Muon |
| HellaSwag | Acc (%) | 47.34 | 46.70 | 45.10 | +2.24% vs AdamW, +0.64% vs Muon |
| PIQA | Acc (%) | 71.27 | 72.20 | 71.27 | 持平 AdamW |
| SciQ | Acc (%) | 73.40 | 71.60 | 70.80 | +2.60% vs AdamW, +1.80% vs Muon |
| TriviaQA | Acc (%) | 2.17 | 1.64 | 1.06 | +1.11% vs AdamW, +0.53% vs Muon |
| WinoGrande | Acc (%) | 53.59 | 53.75 | 51.46 | +2.13% vs AdamW |
| 8 项基准平均 (Avg) | 平均准确率 (%) | 47.69 | 46.34 | 44.74 | 相比 AdamW +2.95%，相比 Muon +1.35% |
在强化学习可验证奖励（RLVR, 使用 GRPO 算法微调）的数学推理任务中，Pion 同样展现了优异的泛化与收敛性能：
| 基础模型 | 评测任务 (avg@K) | Pion | Muon | AdamW | Base 模型 |
|---|---|---|---|---|---|
| Qwen3-1.7B | AIME24 (avg@32) | 25.42 | 20.42 | 22.71 | 4.06 |
| Qwen3-1.7B | AIME25 (avg@32) | 21.98 | 19.27 | 20.94 | 10.10 |
| Qwen3-1.7B | AMC23 (avg@8) | 59.94 | 54.22 | 58.43 | 30.27 |
| Qwen3-1.7B | 5 项数学基准平均 (Avg) | 36.12 | 32.08 | 34.82 | 16.87 |
| DeepSeek-R1-Distill-1.5B | AIME24 (avg@32) | 30.00 | 29.06 | 25.42 | 20.52 |
| DeepSeek-R1-Distill-1.5B | 5 项数学基准平均 (Avg) | 38.32 | 37.32 | 35.97 | 30.20 |
消融实验¶
在 60M LLaMA 模型（9.6B tokens, C4 数据集）上对更新策略、动量设计和矩阵指数近似的核心消融结果如下：
| 配置 | 验证损失 (Val Loss) | 训练稳定性说明 |
|---|---|---|
| 完整配置 (Lie+Lie 动量 + RMS 尺度缩放 + 双边更新) | 3.3575 | 收敛最快，参数空间尺度严格对齐 |
| 交替更新 (Alternate Update) | 3.3654 | 损失仅增加 0.23%，每步更新侧计算量减少约 50% |
| 去除 RMS 缩放 (Naive Update) | 发散 / 差 | 大学习率下直接发散，仅在极微小学习率下勉强收敛 |
| 纯双边归一化 (Bilateral Norm w/o RMS) | 3.3820 | 无法解决跨矩阵更新方差失衡，无法稳定适应大学习率 |
| 仅输运环境空间一阶动量 (Transported Ambient + Ambient) | 3.3760 | 比 Lie+Lie 动量稍慢，但无需额外李代数持久显存 |
| 矩阵指数一阶近似 (1st-order Series) | 3.4210 | 奇异值发生剧烈漂移，破坏谱保持几何性质 |
| 矩阵指数二阶近似 (2nd-order Series) | 3.3575 | 谱保持极好，几乎完全重合高阶 Cayley 变换且无需矩阵求逆 |
| Cayley 变换近似 (Cayley Transform) | 3.3580 | 表现相当但需要 \(O(d^3)\) 求解线性方程组，耗时更高 |
在实际物理资源与极限场景测试下： 1. 极端无归一化层测试（Normalization-free）：移除 Transformer 中所有 RMSNorm/LayerNorm，AdamW 与 Muon 均迅速遭遇梯度溢出并报 NaN，而 Pion 成功平稳跑完 9.6B tokens 并顺利收敛。 2. 极端超深层网络（200 层 DeepNet）：在 200 层深度下，AdamW 出现剧烈 Loss 震荡（标准差 0.0931），中间层雅可比矩阵范数骤降（发生表征退化）；Muon 范数持续衰减（标准差 0.0927）；Pion 全程最平稳（标准差 0.0892），且沿整个 200 层维持均匀且高的雅可比范数。 3. 显存与速度实测（8×H100 NVLink）：Pion 完整版显存 59,839 MB（比 AdamW 高 16.0%），步耗时 0.5679s（仅比 Muon 的 0.5505s 慢 3.2%）；Pion 省略二阶矩版显存降至 45,289 MB（比 AdamW 低 12.2%，比 Muon 低 4.2%）。
关键发现¶
- 谱保持约束具备固有的尺度归一化效应：即使在模型结构中剥离所有显式归一化层，Pion 凭借固定的奇异值谱与李群旋转机制，天然锁定了信号传输增益，阻止了内部激活值的无边界放大。
- 李代数动量优于环境空间动量：由于流形曲率的存在，在流形内蕴的李代数切空间中累积一阶动量与二阶方差，其梯度方向准确性明显高于在欧氏环境空间平移或平行输运的累积。
- 深度扩展与特征表征保真度：极深模型训练中的梯度消失和中间层坍缩主要源于加性更新带来的幅度缩放畸变。Pion 通过保全矩阵的正交等价不变性，促使深层网络每一层都能充分贡献特征变换能力。
亮点与洞察¶
- 从重参数化退回原权重的逆向推导：前作 POET 依赖参数分解 \(\bm{R}\bm{W}_0\bm{P}\) 引入了沉重的优化器复杂性与训练震荡；Pion 敏锐地利用 \(\bm{W} = \bm{I}\bm{W}\bm{I}\) 的虚设因子，将等谱约束直接推演为作用在模型实际权重上的李群正交流动，化繁为简。
- \(\mu\)P 最大更新参数化的天然兼容性：现有的 Muon 变体因满足前向条件而需要后验调整谱范数，而 Pion 天然继承前向谱条件 \(\|\bm{W}\|_2 = \Theta(\sqrt{d_{\text{out}}/d_{\text{in}}})\)，仅需归一化李代数梯度即能实现跨模型尺寸的学习率零样本迁移。
- 可迁移的轻量级工程实现：二阶矩阵泰勒展开搭配奇偶步交替更新（Alternate Update），配合分头投影与无二阶矩版本，打破了黎曼流形优化“慢而耗显存”的固有偏见，使李群优化器能够在百亿级 Token 预训练中直接落地。
局限与展望¶
- 长宽比极高矩阵的显存开销：当权重矩阵高度长方形（如词表嵌入层 \(d_{\text{out}} \gg d_{\text{in}}\)）时，维护 \(d_{\text{out}} \times d_{\text{out}}\) 的李代数张量会占用较多显存，通常需要退回为 AdamW 或采用低秩/分块处理。
- 更大规模模型的验证留白：目前实验主要覆盖至 1.3B-3B 规模模型及 54B Token 训练，超大规模参数（如 70B+）及数万亿 Token 级别下的全量扩展性仍待进一步工程化验证。
- 动态权重谱调整能力的丧失：完全锁定初始奇异值谱虽然带来了极高的稳定性，但也限制了权重矩阵有效秩的动态伸缩；若初始化谱分布不够理想，可能对特定任务的表达能力带来上限约束。
相关工作与启发¶
- vs AdamW: AdamW 采用坐标轴对齐的对角自适应学习率与加性更新，无法感知权重矩阵的整体谱分布，在长周期训练中容易导致特征范数持续膨胀和注意力 Logit 发散；Pion 通过正交等价变换将谱范数严格固化。
- vs Muon / NorMuon: Muon 对更新量进行牛顿-舒尔茨正交化（满足更新侧谱条件），但无法保证权重矩阵自身的谱稳定性，需要外加谱回缩；Pion 则从根本上构建在等谱流形之上，天然满足前向谱条件。
- vs POET: POET 显式将矩阵解耦为左右正交矩阵并冻结中间权重，训练中频繁遭遇损失突刺（Loss spikes）；Pion 摒弃重参数化，直接作用于原参数，兼具 POET 的几何归纳偏置与标准优化器的稳定平滑轨迹。
评分¶
- 新颖性: ⭐⭐⭐⭐⭐ 巧妙跳过显式参数分解，直接在原权重上用恒等算子李代数推导等谱正交优化流，理论构思极其优美。
- 实验充分度: ⭐⭐⭐⭐⭐ 涵盖 LLM 预训练、消融、无归一化极限测试、200 层超深网络、RLVR 数学推理及真实 H100 显存吞吐评测，论据非常扎实。
- 写作质量: ⭐⭐⭐⭐⭐ 几何直观阐述清晰，公式推导严谨自洽，李代数动量与尺度缩放的设计动机层层递进。
- 价值: ⭐⭐⭐⭐⭐ 为大模型训练提供了除传统加性更新（AdamW/Muon）之外全新且切实可行的非欧几何优化器范式。