Pion：保持谱稳定的大语言模型训练优化器

研究者提出Pion优化器，通过正交等价变换而非传统加法更新来优化权重矩阵，在训练过程中保持奇异值不变，为大语言模型的预训练和微调提供了稳定且有竞争力的替代方案。

大语言模型（LLM）的训练是一个复杂而精密的工程，涉及数据、模型架构、计算资源和优化算法等多个方面。其中，优化器的选择对训练稳定性、收敛速度和最终模型质量有着决定性影响。从经典的SGD到广泛使用的Adam，再到近年来涌现的各种变体，优化器的发展始终是深度学习研究的核心议题之一。

当前主流的优化器，如Adam、AdamW以及最近受到关注的Muon，都采用**加法更新（additive update）**的方式调整模型参数。具体来说，它们计算一个更新量（通常基于梯度的一阶和二阶矩估计），然后将其加到当前参数上。这种方式简单直观，在实践中表现良好，但也存在一些潜在问题，特别是在大规模语言模型训练中。

在深入讨论Pion之前，我们需要理解一个关键概念：谱（spectrum），即矩阵奇异值的集合。对于神经网络中的权重矩阵而言，奇异值分布反映了矩阵的几何特性和表达能力。

为什么谱稳定性重要？

训练稳定性：权重矩阵的谱范数（最大奇异值）如果发生剧烈变化，可能导致梯度爆炸或消失，影响训练的稳定性。

泛化能力：研究表明，保持适当的谱特性有助于模型获得更好的泛化能力，避免过拟合。

优化几何：从优化的几何视角看，权重矩阵的奇异值分布影响着损失曲面的形状，进而影响优化的难易程度。

传统的加法更新优化器在更新权重时，往往会改变权重矩阵的奇异值分布。虽然这种改变通常是朝着有益的方向，但它引入了额外的复杂性，使得训练过程更难预测和控制。

Pion优化器的核心思想是通过正交等价变换而非加法更新来优化权重矩阵。具体来说，Pion在每次更新时，对权重矩阵应用左正交变换和右正交变换，而不是简单地加上一个更新矩阵。

正交等价变换的形式化表达为：

W_new = Q_left @ W @ Q_right

其中，Q_left和Q_right是正交矩阵（满足Q^T @ Q = I），W是原始权重矩阵。

关键性质：正交变换保持奇异值不变。也就是说，无论进行怎样的正交等价变换，权重矩阵的奇异值集合保持不变，只有奇异向量会发生旋转。

这意味着Pion在优化过程中严格保持权重矩阵的谱特性，仅通过调整矩阵的几何方向（即奇异向量的方向）来寻找更优的参数配置。

为了更直观地理解Pion的创新，我们可以将其与Adam进行对比：

Adam的更新：W_new = W + ΔW

• 更新量是任意的，可能显著改变奇异值
• 需要学习率等超参数精细调节
• 可能破坏原有的谱结构

Pion的更新：W_new = Q_left @ W @ Q_right

• 保持奇异值完全不变
• 通过正交矩阵的几何旋转进行优化
• 天然保持谱稳定性

论文详细推导了Pion的更新规则，展示了如何从优化目标出发，得到具体的正交变换参数。核心步骤包括：

1. 定义优化目标

Pion的目标是在保持谱不变的前提下，找到能够降低损失函数的正交变换。这可以形式化为一个带约束的优化问题。

2. 计算自然梯度

利用正交群的几何结构，计算在正交变换空间中的自然梯度方向。这与传统欧几里得空间中的梯度不同，需要考虑流形约束。

3. 投影到正交群

将计算得到的更新投影到正交群上，确保更新后的变换矩阵仍然是正交的。这通常通过Cayley变换或指数映射实现。

4. 应用到权重矩阵

将得到的左、右正交变换应用到当前权重矩阵，完成一次更新。

论文系统地考察了Pion设计中的各种选择，包括：

正近似的精度：计算正交变换时数值近似的精度选择，在计算效率和更新质量之间权衡。

学习率调度：虽然Pion保持谱不变，但学习率仍然影响收敛速度，需要适当调度。

动量的使用：是否以及如何在正交变换空间中引入动量机制，以加速收敛。

初始化策略：权重矩阵的初始谱分布对训练的影响，以及如何设计合适的初始化。

这些分析为实际应用Pion提供了实用的指导。