Pion：基于正交等价变换的谱保持优化器 Pion: A Spectrum-Preserving Optimizer via Orthogonal Equivalence Transformation

Pion 优化器的完整流程（Algorithm 1, Lie Algebra 版本）包括以下步骤：（1）**初始化**：初始权重 $W_0 \in \mathbb{R}^{d_{out} \times d_{in}}$，四个动量变量 $m_{in}^0, m_{out}^0, v_{in}^0, v_{out}^0$ 全部初始化为零矩阵。（2）**计算李代数梯度**：在第 $t$ 步，给定当前梯度 $G_t = \nabla_W f(W_{t-1})$，计算输入侧和输出侧的 skew-symmetric 梯度 $G_{in} = W_{t-1}^\top G_t - G_t^\top W_{t-1}$ 和 $G_{out} = G_t W_{t-1}^\top - W_{t-1} G_t^\top$。这两个矩阵分别属于李代数 $\mathfrak{so}(d_{in})$ 和 $\mathfrak{so}(d_{out})$，其 F 范数刻画两侧的平均旋转角，谱范数刻画最大旋转角。（3）**更新李代数动量**：$m_{in} \leftarrow \beta_1 m_{in} + (1-\beta_1) G_{in}$，$m_{out} \leftarrow \beta_1 m_{out} + (1-\beta_1) G_{out}$，$v_{in} \leftarrow \beta_2 v_{in} + (1-\beta_2)(G_{in} \odot G_{in})$，$v_{out} \leftarrow \beta_2 v_{out} + (1-\beta_2)(G_{out} \odot G_{out})$。这里 $v$ 模仿 AdamW 的二阶动量，但作用于李代数元素上。（4）**自适应缩放**：计算 $A_{in} = -m_{in} / (\sqrt{v_{in}} + \epsilon)$，$A_{out} = -m_{out} / (\sqrt{v_{out}} + \epsilon)$，与 AdamW 类似但作用于李代数元素。（5）**RMS 控制的尺度一致性**：计算缩放系数 $\alpha_t = c \sqrt{d_{out} d_{in}} / (\|A_{out} W_{t-1} + W_{t-1} A_{in}\|_F + \epsilon)$，保证不同大小的权重矩阵接收与之规模成比例的更新量。RMS 常数 $c$ 控制总体学习率缩放。（6）**二阶指数近似**：定义 $E_2(A, \alpha) = I + \eta \alpha A + \frac{1}{2}(\eta \alpha A)^2$ 作为 $\exp(\eta \alpha A)$ 的二阶截断泰勒近似。（7）**双侧或交替更新**：双侧模式 $W_t = E_2(A_{out}, \alpha_t) W_{t-1} E_2(A_{in}, \alpha_t)$，每步同时更新两侧；交替模式则按 $\psi \in \{0, 1\}$ 标志每隔一步仅更新一侧。（8）**谱保持性**：由于 $E_2(A, \alpha)$ 对 skew-symmetric $A$ 是严格正交的（误差 $O(\|\eta \alpha A\|^3)$），且 Pion 的指数近似从恒等矩阵出发（无误差累积），$W_t$ 的奇异值始终与 $W_0$ 高度一致。

Pion 的技术新颖性体现在以下方面。第一，**严格保持权重谱的优化器设计**：在所有现有 LLM 优化器中，Pion 是第一个不依赖重参数化或外部约束、而是从李群优化出发直接保持权重谱的工作。Muon 系列、AdamW 系列均不能保证权重谱不变，POET 虽能保持但需要重参数化。第二，**收敛性定理 2.2 的证明**：作者证明了在 $L$-smooth 和下界条件下，Pion 的随机版本以 $O(1/\sqrt{T})$ 速率收敛到稳定点，且 $T+1$ 步内 $\min_{0 \le t \le T}(\|G_{in}\|_F^2 + \|G_{out}\|_F^2) \le (C_1 + C_2)/\sqrt{T+1}$，其中 $C_1$ 依赖初始最优性 gap，$C_2$ 依赖 $L, \gamma, \sigma^2, C$。这是首个谱保持优化器的收敛保证。第三，**几何解析的清晰度**：Pion 显式揭示了 $\|\Delta W\|_F$ 是「总旋转强度」、$\|G_{in}\|_F/\sqrt{d_{in}}$ 是输入侧平均旋转角、$\|G_{in}\|_2$ 是最大旋转角——这种几何透明度是 AdamW 完全不具备的。第四，**最小能量保持性**：Pion 证明零均值高斯初始化对应的最低超球能量配置被严格保持，这一点在 Proposition 2.1 中通过奇异值不变性间接得到。第五，**工程设计选择的系统性探索**：作者对比了 transported ambient-space 与 Lie-algebra 动量（Figure 4）、bilateral 与 alternate 更新（Figure 5）、一阶到四阶矩阵指数近似与 Cayley 变换（Figure 6），最终给出有原则的设计选择——Lie+Lie 动量、alternate 更新、二阶指数近似。第六，**稳定性极限测试**：在移除所有归一化层（AdamW/Muon 都会 NaN）和 200 层超深网络（验证深度可扩展性）这两个极端设置中，Pion 是唯一不崩溃的方法，这是非常强的工程价值证据。

基于论文细节可以独立分析 Pion 的几个弱点并给出改进方向。**弱点 1：复杂度过高，难以应用于所有场景**。Pion 的双边更新主导开销为 $O(d_{out}^2 d_{in} + d_{out} d_{in}^2 + d_{out}^3 + d_{in}^3)$，其中立方项在 embedding/lm_head（vocab 维度 128k 以上）上会爆炸。虽然论文提到交替更新可减半，但在这些特殊层上仍可能成为瓶颈。**改进方向**：对 embedding/lm_head 等特殊形状层混合使用 Pion 和 AdamW（类似 Muon 在实践中对 1D 参数的处理）；或对这些层使用低秩 Pion 变体。**弱点 2：RMS 缩放系数 $c$ 是新引入的超参数**。论文并未给出 $c$ 的默认值，也未提供细致的消融。RMS 缩放是 Pion 在大学习率下不崩溃的关键，但 $c$ 的不当选择可能导致训练不稳定。**改进方向**：提供 $c$ 的自适应选择策略，例如基于 $\|W\|_F/\sqrt{d_{out} d_{in}}$ 的相对学习率。**弱点 3：缺少对归一化层影响的系统对比**。Pion 在无归一化层下表现亮眼，但论文未直接对比「有归一化 Pion」与「无归一化 Pion+其他优化器」的公平对比。**改进方向**：增加一个消融实验，固定有归一化，对比 Pion vs Muon；以及固定无归一化，对比不同优化器的训练稳定性。**弱点 4：实证规模有限**。1.3B 模型的对比虽然有意义，但远未达到 LLM 训练稳定性的真正挑战规模（7B+）。论文对「稳定性」的许多论断（如 Figure 8 的平坦曲线）需要更大规模验证。**改进方向**：在 7B 或 13B 规模上完整复现 AdamW/Muon/Pion 对比。**弱点 5：代码与超参数搜索细节不完整**。论文未提供 $c$、$\beta_1$、$\beta_2$ 等超参数的具体搜索范围和最优值，也不清楚 Pion 对这些超参数的敏感度。**改进方向**：在附录中提供完整的超参数 sweep 结果。

作者在论文中提出了以下未来方向：（1）**更大规模验证**：将 Pion 应用于 7B+ 模型预训练，验证其稳定性优势是否在更大规模上保持；（2）**更广泛的 µP 兼容变体**：探索 $G_t$-正交化等更精细的 µP 兼容 Pion 变体；（3）**RLVR 深度探索**：作者观察到 RLVR 更新主要保持预训练谱结构，Pion 与此天然契合，但 Pion 在 RLVR 上的稳定性优势（如 KL 散度、奖励曲线）有待更细致研究。基于论文结果还可以延伸出几个方向：（4）**Pion 与稀疏化/量化的兼容性**：Pion 的旋转更新形式是否能在 FP8/INT8 量化训练中保持数值稳定性是一个值得探索的实际问题；（5）**Pion 用于非 LLM 任务**：在视觉 Transformer、扩散模型、多模态模型上测试 Pion 的通用性；（6）**Pion 与学习率调度器的协同**：Pion 的 RMS 缩放已部分实现自适应学习率，但与 warmup、cosine decay 等调度的交互尚未研究；（7）**Pion 的概率解释**：从信息几何或黎曼优化的角度，将 Pion 解释为沿自然梯度方向更新，可能得到更深入的理论理解；（8）**结合 Q-GDL 或 Shampoo 等二阶信息**：Pion 只用了一阶/二阶动量，是否能融合 Q-GDL 的二阶曲率信息进一步加速。

Pion 论文的复现评估总体上中等偏上。**积极方面**：（1）论文作者来自 CUHK/MPI/Westlake 等知名机构，并在 SphereLab Technical Report 系列下发布，方法论严谨；（2）Algorithm 1 给出了完整的 Pion 优化器伪代码，包括所有超参数符号（$\eta, \beta_1, \beta_2, c, \epsilon$, alternating flag）；（3）附录 D 列出了三个主要实验场景的详细配置（数据、模型、训练时长、批大小、token 数）；（4）C4 数据集和 LLaMA 架构都是公开标准，便于复现；（5）评测使用 LM Evaluation Harness 等成熟框架。**挑战方面**：（1）论文未声明代码是否开源，截至目前也未提供独立代码库链接；（2）$\beta_1, \beta_2, c, \epsilon$ 等超参数的具体搜索空间和最优值未充分披露，需要复现者自行调参；（3）1.3B 模型预训练 54B tokens 需要相当可观的算力（按标准估算，单卡 A100/H100 需数周到数月），RLVR 训练涉及 VeRL 框架和 GRPO 实现也较为复杂；（4）作者使用 VeRL 框架训练 RLVR，POLARIS 评测，复现需要熟悉这些工具链；（5）附录 C 的计算复杂度分析明确 Pion 在 LLM 预训练上的额外开销「相对较小」，但具体到 RLVR 等小 batch 场景的开销比例需要实测。**整体复现难度**：中高。研究者可以基于论文较为忠实地复现核心方法，但完整复现所有 5 个实验（预训练/无归一化/超深/SFT/RLVR）需要数月工作和相当算力。