SS-ESOAP：针对物理信息学习的自缩放自适应预条件方法

# SS-eSOAP：面向物理信息学习的自缩放自适应预条件
来源：https://arxiv.org/html/2608.29448
Mads Toftrup*单位：* 奥胡斯大学
Sebastian Loeschcke*单位：* 哥本哈根大学
Yixuan Wang*单位：* 加州理工学院
Anima Anandkumar*单位：* 加州理工学院

###### 摘要
物理信息神经网络（PINNs）常面临目标函数病态条件问题，这限制了高精度训练。稠密拟牛顿方法能改善局部条件性，但需要昂贵的优化器状态；而Kronecker分解方法（如SOAP）能扩展到更大网络，但依赖周期性基底更新。我们引入**SS-eSOAP**，该方法在SOAP式预条件基础上，增加了一个适应Kronecker几何结构的标量割线能量校正，以及一个跟随方差状态降尺度的自适应基底更新。我们刻画了标量校正引起的定向割线匹配特性，并给出了基底变化时方差状态失配的上界。在八个PDE基准测试中，**SS-eSOAP**在六个（包括Burgers和Boussinesq方程）上取得了最低的最终残差，而SOAP系基线方法在Gray-Scott和Ginzburg-Landau方程上表现更好。在Boussinesq方程上，**SS-eSOAP**在4.1小时内以9.2GB峰值显存达到了10⁻⁵的残差，而Adam在14小时内未达到该目标。在四个代表性PDE上的三种子L²和H¹误差表明，较低的残差与改进的解精度相关联。这些结果将**SS-eSOAP**定位为处理刚性、高精度物理信息训练的可扩展选项，而非现有优化器的普遍替代品。

## 1 引言
高精度优化在科学机器学习中至关重要。一个例子是寻找三维Euler和Navier-Stokes方程中的有限时间奇点。PINNs已被用于重建无界域中的候选自相似剖面[35, 34, 33]。此类候选解在计算机辅助分析之前需要低残差。PINN目标函数在此精度区间常呈病态条件，标准一阶训练可能在此处停滞[31, 17, 25, 29, 36]。更广泛地说，PINN优化所面临的挑战与标准监督学习不同。PINN损失结合了PDE残差与边界和初始条件项，可能导致梯度不平衡、数值刚性和病态损失景观[31, 17, 25]。因此，标准一阶优化器（如Adam）在高精度区间常收敛缓慢或停滞，这促使了针对PINN训练的曲率感知方法[27]。二阶方法通过Hessian近似更新利用曲率信息，在许多非凸问题中能产生比一阶方法更高效的下降方向。拟牛顿方法如BFGS[5]和L-BFGS[18]通过梯度差构建这些近似，避免显式计算二阶导数，同时保持良好的局部收敛性。自缩放变体（如SS-BFGS）对Hessian近似进行缩放，以更好地匹配局部曲率谱[23, 2, 22, 27]。然而，这些方法在展开的向量上操作，对于单个d_out × d_in矩阵参数，需要O(d_in² d_out²)的存储，使得精确的自缩放拟牛顿更新对深度网络不切实际。

Kronecker分解预条件器提供了一种可扩展的替代方案。如K-FAC[20]、Shampoo[8]和SOAP[28]等方法通过较小矩阵的张量积近似曲率或梯度协方差，无需维护稠密Hessian近似即可实现结构化预条件。这些方法不强制执行拟牛顿割线条件，而是使用二阶统计量（如经验Fisher信息或未中心化梯度协方差）来缩放梯度。然而，现有的Kronecker分解方法在非平稳区间仍面临限制。标准Shampoo和SOAP通常以固定间隔更新预条件器特征空间，这可能与PINN训练过程中曲率统计量的演变速度不匹配。此外，累积的二阶矩统计量在基底更新后常被重新投影，当基底发生显著旋转时，这可能是不稳定的[7]。

在本工作中，我们提出了**SS-eSOAP**，一种用于高精度PINN训练的逐层Kronecker分解优化器。该方法在SOAP式预条件上增加了两个操作。首先，一个标量因子使Kronecker度量与沿最新参数位移观测到的割线能量匹配。其次，一个EShampoo式的非对角准则触发基底更新[7]。在每次触发的更新中，**SS-eSOAP**重新投影动量状态并降尺度坐标方向二阶矩，而不是重新投影该二阶矩。对于一个d_out × d_in矩阵参数，**SS-eSOAP**除了Adam式状态外，还存储O(d_in² + d_out²)的Kronecker因子。标量割线计算在所需特征空间量可用后，成本为O(d_in d_out)，或计入所有矩阵条目时为O(d_in d_out)。基底检查和特征分解保留了三次层宽成本，在附录F中分析。

在八个PDE基准测试中，**SS-eSOAP**在六个上记录了最低的最终残差。在Boussinesq方程上，它在4.1小时内达到了10⁻⁵，而Adam在14小时内运行未达到相同目标。这些增益依赖于具体问题。Purifying Shampoo在Gray-Scott方程上表现更好，而SOAP在Ginzburg-Landau方程上表现更好。因此，我们研究**SS-eSOAP**作为刚性、高精度区间的专用优化器。我们的贡献包括：
- • 我们为Kronecker分解度量推导了定向割线能量校正。
- • 我们将自适应基底触发器与动量重新投影和为突变基底变化设计的方差状态转换相结合。
- • 我们阐述了两个理论结果的范围：一个单向割线匹配结果和一个稳态方差失配上界。
- • 我们与一阶、结构化、PINN特定以及自缩放拟牛顿基线进行了比较，包括运行时间、内存、多种子和物理误差测量。
- • 我们识别了自缩放有益和有害的区间。

## 2 背景与相关工作
### 2.1 物理信息神经网络（PINNs）
物理信息神经网络（PINNs）[24, 14]通过使用神经网络表示解场，并通过可微残差损失强制执行控制方程，来解决PDE约束学习问题。形式上，给定一个定义在域Ω上的PDE：ℱ[u](x, t) = 0，PINN通过神经网络u_θ(x, t)逼近解u(x, t)。自动微分评估ℱ[u_θ]，并通过最小化复合目标函数ℒ(θ) = λ_r ℒ_r(θ) + λ_b ℒ_b(θ) + λ_d ℒ_d(θ)来训练网络，其中ℒ_r惩罚配点处的PDE残差，ℒ_b强制边界或初始条件，ℒ_d在可用时纳入观测数据。

### 2.2 物理信息学习优化
PINN目标函数结合了残差、边界、初始条件和数据项。其梯度在尺度和方向上可能差异显著[31, 32, 17]。损失平衡方法通过学习率退火、基于NTK的重加权或参数归一化来解决此问题[31, 32, 37]。其他工作则针对完整目标函数的几何结构。示例包括NysNewton-CG[25]、PDE预条件[19]、能量和Gauss-Newton自然梯度[21, 11]、ANaGRAM[26]、对偶自然梯度[12]以及用于PINNs的K-FAC[6]。近期工作还研究了SOAP下的梯度对齐[29]、自缩放BFGS和Broyden方法[16, 10]。**SS-eSOAP**针对此设计空间的不同点。它不重新加权PINN损失分量，也不求解全局参数空间或残差空间系统。它将逐层Kronecker统计与定向割线校正相结合。当其状态和线搜索成本可控时，稠密自缩放方法仍是更强的参考。

### 2.3 结构化矩阵预条件
Shampoo使用Kronecker因子近似全矩阵AdaGrad[8]。对于矩阵参数W_t ∈ R^{m×n}和梯度G_t，它维护：
L_t = β₂ L_{t-1} + (1 - β₂) G_t G_t^⊤,
R_t = β₂ R_{t-1} + (1 - β₂) G_t^⊤ G_t。
SOAP在这些因子的特征基上应用Adam[28]。Muon则正交化矩阵更新，与零衰减Shampoo极限相关[13, 4]。SOAP通常按固定计划刷新其基底。Purifying Shampoo分离了特征值和特征基误差，并为热启动基底更新开发了自适应停止准则[7]。**SS-eSOAP**采用了这种非对角基底诊断。其新元素是Kronecker适配的割线校正和触发基底变化后使用的二阶矩转换。Muon[13]是一种最近的优化算法，旨在以与一阶方法相当的成本近似大规模神经网络训练的二阶缩放。该方法使用Newton-Schulz迭代正交化梯度为G' = UV^⊤，其中G = UΣV^⊤是线性层矩阵梯度的SVD。此外，它可以被视为β₂=0时Shampoo的特例，因为(GG^⊤)^{-1/4} G (G^⊤G)^{-1/4} = UV^T [4]。

### 2.4 SOAP、Purifying Shampoo和自适应特征值校正的局限性
标准SOAP实现在固定间隔更新特征基(Q_L, Q_R)。这种静态调度无法区分训练的不同阶段或具有不同动态的层。此外，当基底更新时，标准方法会重新投影累积的二阶矩估计V_k。我们认为，通过相似变换V_k^{new} = (Q_L^{new})^T Q_L V_k Q_R^T Q_R^{new}重新投影V_k，假设了旋转下方差的保持，当Hessian主方向显著移动时，这无效，导致不稳定步长。对Shampoo算法的近期分析表明，其实际成功严重依赖于两个启发式：陈旧预条件（以固定、不频繁的间隔更新特征基）和学习率嫁接（强制更新量级与Adam匹配以缓解谱误差）。[7]通过将预条件器更新分解为其特征向量（基底）和特征值（曲率）分量，形式化了这些问题，称为“净化”框架。

### 2.5 拟牛顿迭代算法中的自缩放
一类广义的拟牛顿迭代算法可以归入自缩放Broyden公式下[2]。如果我们定义辅助变量 s_k = Θ_{k+1} - Θ_k，y_k = ∇J(Θ_{k+1}) - ∇J(Θ_k)，v_k = √{y_k · H_k y_k} [s_k/(y_k · s_k) - H_k y_k/(y_k · H_k y_k)]，那么，每次迭代的逆Hessian矩阵的下一次近似可以按[3, 1]计算：
H_{k+1} = (1/τ_k) [H_k - (H_k y_k ⊗ H_k y_k)/(y_k · H_k y_k) + φ_k v_k ⊗ v_k] + (s_k ⊗ s_k)/(y_k · s_k)，
其中我们定义 τ_k = min{1, (y_k · s_k)/(α_k s_k · H_k^{-1} s_k)}。（2）
这里，⊗表示两个向量的张量积，τ_k, φ_k分别是缩放和更新参数，它们在迭代之间经常变化。当τ_k=1和φ_k=1时，得到标准BFGS算法。为了有效计算缩放参数τ_k，[27]设置 H_k^{-1} s_k = -α_k ∇J(Θ_k)，这样τ_k对H_k^{-1}的显式依赖就消失了。这是可行的，因为步长α_k和梯度∇J(Θ_k)是已知的。