SS-eSOAP: Self-Scaled Adaptive Preconditioning for Physics-Informed Learning¶
会议: NeurIPS 2026
arXiv: 2608.29448
领域: 优化/理论
关键词: 物理信息神经网络、Kronecker 分解、自适应预条件、自标度拟牛顿法、特征基更新
一句话总结¶
本文提出 SS-eSOAP 优化器，通过将 Kronecker 结构预条件与沿位移方向的标量割线能量校正相结合，并引入基于非对角能量的自适应基底更新与方差状态下采样，成功克服了 PINN 高精度优化中的病态刚性难题。
研究背景与动机¶
物理信息神经网络（PINNs）在流体力学奇点分析与高精度科学机器学习中扮演着关键角色，例如在探索三维不可压缩流体方程的潜在有限时间爆破时，数值验证依赖将残差压制到极低量级。然而，PINN 的损失函数高度复合，同时结合了偏微分方程内部配置点残差、边界条件与初值约束项。不同损失项之间往往存在严重的梯度量级失衡、刚性谱分布以及极度病态的损失曲面，导致以 Adam 为代表的标准一阶优化算法在高精度收敛区间极易陷入停滞或出现平台期。
为了克服一阶优化的病态曲率瓶颈，二阶与拟牛顿方法被广泛引入，但它们在深度网络上面临难以两全的权衡。以 BFGS 和自标度拟牛顿法（SS-BFGS）为代表的稠密方法通过割线条件校正局部曲率谱，具有出色的局部收敛性质，但在全连接网络中需要存储稠密逆 Hessian，单层参数复杂度高达 \(\mathcal{O}(d_{\mathrm{in}}^2 d_{\mathrm{out}}^2)\)，在稍宽网络上立即内存溢出（OOM）。相反，K-FAC、Shampoo 以及在特征基下运行 Adam 的 SOAP 等 Kronecker 分解方法将预条件子解耦为小矩阵张量积，仅需 \(\mathcal{O}(d_{\mathrm{in}}^2 + d_{\mathrm{out}}^2)\) 空间。但现有 Kronecker 优化器大多采用固定的周期性基底重算策略，无法感知训练阶段动态变化的曲率漂移；更关键的是，在基底剧烈旋转时，直接重投影累积的二阶矩统计量会导致严重的数值方差突变与步长爆炸。
本文切入的角度是：在保持 Kronecker 分解极低显存开销的前提下，如何内生性地引入拟牛顿自标度曲率校准，并设计平滑的非平稳基底过渡机制。核心 idea：在 SOAP 算法框架中融合沿参数更新位移的 Kronecker 割线能量标量校正因子 \(\tau_k\)，配合基于非对角质量比的自适应特征基触发器与二阶矩方差下采样（Soft Reset），实现稳定且具备局部超线性下降能力的刚性物理学习优化。
方法详解¶
整体框架¶
SS-eSOAP 针对神经网络每层权重矩阵 \(W_t \in \mathbb{R}^{m \times n}\) 进行层级结构化预条件更新。整个优化循环分为四个有机阶段：首先累积梯度并更新左右 Kronecker 因子 \(L_t\) 与 \(R_t\)；其次依据因子偏离对角化的程度自适应触发特征基重分解，并在基变换时重投影一阶动量、对二阶矩方差执行软重置收缩；接着将梯度投影至当前特征基下计算 Adam 风格的一阶与二阶矩统计量；最后利用无需显式求逆的特征值缩放高效计算标量割线校正因子 \(\tau_t\)，并将预条件更新旋转回参数空间完成权重更新。
%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
A["梯度与 Kronecker 因子更新<br/>计算 Gt 并平滑累积 Lt 与 Rt"] --> B["自适应特征基触发与方差下采样<br/>检测非对角质量比与软重置 Vt"]
B --> C["特征基投影动量与自标度割线校正<br/>计算割线条修正因子 τt 与 eigenspace 更新"]
C --> D["参数更新与逆变换<br/>投影回原参数空间应用步长下降"]关键设计¶
1. Kronecker 适配的标量割线能量校正：无需矩阵求逆的曲率全局对齐
标准 SOAP 和 Shampoo 依赖未中心化的经验梯度协方差或 Fisher 统计量来提供相对缩放，但在损失平坦区或剧烈非凸峡谷中无法感知真实的曲率绝对尺度。为了使预条件矩阵与真实 Hessian 谱对齐，算法引入自标度拟牛顿因子的 Kronecker 变体。记参数位移为 \(S_t = W_{t-1} - W_{t-2}\)，梯度差为 \(Y_t = G_t - G_{t-1}\)，所期望的逆 Hessian 近似具有形式 \(\tau_t^{-1} (L_t^{-1} \otimes R_t^{-1})\)。标量因子定义为：
当沿位移方向的割线曲率小于 Kronecker 预条件能量时，\(\tau_t < 1\) 激活自标度收缩；反之则截断为 1。关键实现技巧在于避免显式计算大矩阵的逆 \(L_t^{-1}\) 和 \(R_t^{-1}\)。利用已有的正交特征基 \(L_t = Q_L \operatorname{diag}(D_L) Q_L^\top\) 和 \(R_t = Q_R \operatorname{diag}(D_R) Q_R^\top\)，将位移矩阵与梯度差旋转到特征空间 \(\tilde{S}_t = Q_L^\top S_t Q_R\) 与 \(\tilde{Y}_t = Q_L^\top Y_t Q_R\)。此时分母二次型与分子迹简化为纯元素级操作：
这一变换将曲学校正计算开销从 \(\mathcal{O}(n^3)\) 彻底降至 \(\mathcal{O}(mn)\)，不仅在数值上极度稳定，而且严谨地保证了沿当前更新方向精确拟合目标割线能量。
2. 净化非对角自适应基底触发器：按需启动高开销特征分解
传统方法以固定步长间隔周期性重新计算特征基，在曲率剧烈变化时基底严重滞后，而在平稳下降期则浪费大量算力。SS-eSOAP 采用基于未对角化质量的指标监控当前基底对于 \(L_t\) 和 \(R_t\) 的保真度：
在每隔 \(I_{\mathrm{check}}\)（默认每步）检测时，当最大偏离比 \(\rho_t = \max\{\rho(L_t, Q_L), \rho(R_t, Q_R)\} > \tau_{\mathrm{trigger}}\)（默认 0.2）时，才重新进行精确的特征分解更新 \(Q_L\) 与 \(Q_R\)。该机制将计算昂贵的特征分解与实际曲率几何演化解耦，实现了在误差失配严重时自适应刷新、平稳期零冗余分解的高效平衡。
3. 动量投影与方差状态下采样：消除基底旋转下的步长激增
在基底发生旋转从 \(Q_{\mathrm{old}}\) 切换到 \(Q_{\mathrm{new}}\) 时，标准方法往往对动量和二阶矩均执行相似性旋转变换 \(Q_{\mathrm{new}}^\top Q_{\mathrm{old}} V Q_{\mathrm{old}}^\top Q_{\mathrm{new}}\)。然而二阶矩 \(V_t\) 追踪的是坐标轴对齐的未中心化二阶矩，强行旋转破坏了坐标方差守恒假设，经常导致除以极小方差分母而发生灾难性的梯度步长尖峰。SS-eSOAP 采取差异化过渡策略：一阶动量保持空间一致性，通过正交变换精确重投影：
而二阶矩 \(V_{t-1}\) 则不作几何旋转，直接在坐标系下执行自适应幅值下采样（Soft Reset）：\(V_{t-1} \leftarrow \gamma_t V_{t-1}\)。其中缩放系数根据基底失配程度分段确定：若 \(\rho_t > 0.8\) 则 \(\gamma_t = 0.25\)；若 \(\rho_t > 0.5\) 则 \(\gamma_t = 0.5\)；否则取 \(\gamma_t = 0.75\)。由于 \(V\) 出现在 Adam 预条件分母中，\(\gamma_t < 1\) 暂时放大了即时更新阻尼和步长敏感性，既保留了历史方差的能量分布轮廓，又有效消除了基底跳变带来的更新不稳定尖峰。
损失函数 / 训练策略¶
在求解控制方程 \(\mathcal{F}[u](x,t) = 0\) 的 PINN 场景中，网络采用复合损失函数优化：
其中 \(\mathcal{L}_r\) 是内部配置点上的方程残差，\(\mathcal{L}_b\) 强加边界与初值条件，\(\mathcal{L}_d\) 代表观测数据误差。在每步参数更新中，特征基空间内的更新方向为 \(\widetilde{U}_t = \tau_t^{-1/2} \widehat{M}_t \oslash (\sqrt{\widehat{V}_t} + \epsilon)\)，恢复到参数空间后结合权重衰减 \(\lambda\) 进行下降：\(W_t \leftarrow (1 - \eta \lambda) W_{t-1} - \eta Q_L \widetilde{U}_t Q_R^\top\)。
实验关键数据¶
主实验¶
论文在包括 Wave、Burgers、Allen-Cahn、Boussinesq、Korteweg-de Vries (KdV)、Gray-Scott、Ginzburg-Landau 和 Lid-Driven Cavity 等八大标准偏微分方程基准上，将 SS-eSOAP 与 Adam、Muon、SOAP、Purifying Shampoo 进行全方位对比。
| 基准任务 (PDE Benchmark) | 评价指标 (Final PDE Residual) | SS-eSOAP (本文) | SOAP | Purifying Shampoo | Adam | 最优对比基线 |
|---|---|---|---|---|---|---|
| Wave (波动方程) | 最终 PDE 残差 | 4.12 × 10⁻⁵ | 1.45 × 10⁻³ | 1.18 × 10⁻³ | 2.15 × 10⁻³ | 1.18 × 10⁻³ (P. Shampoo) |
| 1D Burgers (伯格斯方程) | 最终 PDE 残差 | 5.54 × 10⁻¹⁰ | 2.15 × 10⁻⁹ | 1.12 × 10⁻⁹ | 4.21 × 10⁻⁶ | 1.12 × 10⁻⁹ (P. Shampoo) |
| Allen-Cahn (相变方程) | 最终 PDE 残差 | 1.85 × 10⁻⁹ | 8.45 × 10⁻⁷ | 3.15 × 10⁻⁹ | 2.10 × 10⁻⁷ | 3.15 × 10⁻⁹ (P. Shampoo) |
| 2D Boussinesq (流体方程) | 最终 PDE 残差 | 8.42 × 10⁻⁷ | 9.15 × 10⁻⁶ | 2.85 × 10⁻⁶ | 5.12 × 10⁻⁴ | 2.85 × 10⁻⁶ (P. Shampoo) |
| KdV (孤立子方程) | 最终 PDE 残差 | 3.52 × 10⁻¹⁰ | 8.12 × 10⁻¹⁰ | 1.75 × 10⁻⁷ | 1.85 × 10⁻⁸ | 8.12 × 10⁻¹⁰ (SOAP) |
| Lid-Driven Cavity (顶盖驱动方腔流) | 最终 PDE 残差 | 8.23 × 10⁻⁸ | 3.81 × 10⁻⁷ | 2.10 × 10⁻⁷ | 1.88 × 10⁻⁶ | 2.10 × 10⁻⁷ (P. Shampoo) |
| Gray-Scott (反应扩散方程) | 最终 PDE 残差 | 7.46 × 10⁻⁷ | 2.10 × 10⁻⁷ | 5.12 × 10⁻⁹ | 3.10 × 10⁻⁶ | 5.12 × 10⁻⁹ (P. Shampoo) |
| Ginzburg-Landau (复超导方程) | 最终 PDE 残差 | 1.75 × 10⁻¹¹ | 9.12 × 10⁻¹² | 2.45 × 10⁻¹¹ | 1.15 × 10⁻⁷ | 9.12 × 10⁻¹² (SOAP) |
在八项测试中，SS-eSOAP 在六项上斩获最低残差。尤其在 Wave 方程上相较最优基准降低了接近两个数量级（从 \(1.18 \times 10^{-3}\) 降至 \(4.12 \times 10^{-5}\)），在极度刚性的 2D Boussinesq 方程中将残差压至 \(8.42 \times 10^{-7}\)，而 Adam 和 Muon 甚至难以跨越 \(10^{-4}\) 的收敛停滞瓶颈。
消融实验¶
论文针对 2D Boussinesq 和 Allen-Cahn 开展了细致的机制消融，并系统评估了方差过渡机制与曲率自标度因子的作用：
| 实验配置 (Ablation Setting) | 2D Boussinesq 最终收敛状态 | 稳定性表现 (Loss Spikes) | 说明 |
|---|---|---|---|
| SS-eSOAP (Full Model: 软重置下采样 + 自标度 \(\tau_k\)) | 残差 8.42 × 10⁻⁷ (4.1 小时达 \(10^{-5}\)) | 极少尖峰，轨迹极其平稳 | 完整提出的自适应 Kronecker 拟牛顿框架 |
| 方差机制：标准旋转重投影 (Reprojection) | 陷入早期次优平台或震荡 | 每次基底更新伴随巨大 Loss 突增脉冲 | 几何旋转破坏二阶矩非各向同性方差分布 |
| 方差机制：完全硬重置 (Hard Reset, \(V \leftarrow 0\)) | 收敛至 \(6.54 \times 10^{-5}\) 精度变差 | 无尖峰但损失了所有累积曲率记忆 | 抹除历史统计导致有效学习率初期漂移 |
| 曲率机制：无缩放基线 (\(\tau = 1\), 标准 eSOAP) | 收敛缓慢，精度显著受限 | 平稳但缺乏深层步长缩放增益 | 无法捕捉局部非线性割线能阶 |
| 曲率机制：普通 Oren-Luenberger 缩放 (\(\tau_{\mathrm{OL}}\)) | 局部波动剧烈，难以稳定下降 | 步长调节不匹配张量几何导致不稳定 | 稠密向量标度未考虑 Kronecker 矩阵特征结构 |
此外，在 2D Boussinesq 达到 \(10^{-5}\) 残差目标的实际硬件开销对比中，Adam 耗时超过 14 小时仍未收敛（显存 8.2 GB），而 SS-eSOAP 仅需 4.1 小时即可达标（显存仅微增至 9.2 GB），印证了结构化二阶优化的实际落地效率。
关键发现¶
- 自标度割线因子的非平稳调节机制：实验统计显示，在绝大多数常规迭代步中 \(\tau_k = 1\)（未激活截断），而在遇到曲率陡峭或梯度快速变化的非平稳过渡步时，\(\tau_k < 1\) 迅速自适应缩小等效更新，正是这少数关键迭代步的精细规整避免了优化轨迹跌入发散区。
- 方差下采样 \(\gamma < 1\) 抹平更新奇点：对于二阶矩进行自适应缩放 \(\gamma \in \{0.25, 0.5, 0.75\}\) 在理论上提供了稳态方差距离上界，在实践中彻底消除了正交基旋转造成的除以极小方差引发的数值不稳定。
- 残差降低真实反映物理场精度提升：在 1D Burgers、2D Boussinesq 等四组典型方程的三随机种子实验中，SS-eSOAP 的相对 \(L^2\) 和 \(H^1\) 误差均处于全场最优（如 Burgers 相对 \(L^2\) 误差仅 \((8.45 \pm 0.34) \times 10^{-5}\)，大幅优于 Adam 的 \(4.12 \times 10^{-3}\)），证实低残差与精确恢复物理动力学之间具有极强的一致性。
亮点与洞察¶
- 巧用特征基对角化绕过高维矩阵逆运算：在计算 Kronecker 割线二次型时，将位移矩阵直接投射到左右因子的特征基空间，把求逆 \(L_t^{-1}\) 和 \(R_t^{-1}\) 转化为一维特征值的简单标量除法，成功以 \(\mathcal{O}(mn)\) 代价实现了自标度曲率修正。
- 打破"二阶矩必须随基变换正交重投影"的定式思维：深刻揭示了二阶矩的对角化本质是坐标轴方差统计量而非张量向量，证明基底大幅度旋转时重投影必然引入病态扰动，转而以方差状态下采样这一简明操作达成了更优雅的数值稳定过渡。
- 兼顾稠密拟牛顿收敛性与结构化张量紧凑性：成功克服了 SS-BFGS 因全网络展平参数矩阵导致内存爆炸的致命缺陷，使得自标度拟牛顿思想首次能够稳定且可扩展地应用于千万参数级的大型前沿物理网络架构。
局限与展望¶
- 缺乏全非凸问题的全局充分下降理论保证：论文的理论分析主要集中在稳态方差失配上界和单方向割线能量的最优拟合，尚未证明在复杂非凸非齐次目标下的全局收敛速率。
- 三次幂特征分解开销限制极端超宽层：自适应基底触发机制依然需要执行 \(\mathcal{O}(d^3)\) 的左右特征分解，在层宽极大（如 \(d > 4096\)）的超大模型中可能成为每步时钟周期的计算瓶颈。
- 收益具有 PDE 依赖性：在 Gray-Scott 和 Ginzburg-Landau 方程上，SS-eSOAP 略逊于 Purifying Shampoo 或纯 SOAP，表明割线能量自标度并非对所有物理动力学系统都产生正向增益。未来可探索与 PINNsFormer 等更加平滑的注意力架构融合，或进一步研究特征分解在线低秩近似。
相关工作与启发¶
- vs SOAP & Shampoo: 标准 Shampoo 和 SOAP 采用固定步频刷新特征基，并在基底改变时采用重投影或 Adam 接枝（Grafting）；SS-eSOAP 采用非对角质量自适应触发刷新，以方差下采样取代二阶矩重投影，并用内生性的割线能量校准消除了对外置 Adam 接枝的启发式依赖。
- vs SS-BFGS & SS-Broyden: 传统自标度拟牛顿方法对参数展平向量维护稠密逆 Hessian 近似，在网络宽度达到 \(W = 600\) 时即因 \(\mathcal{O}(W^4)\) 空间复杂度面临 OOM 崩溃；SS-eSOAP 利用 Kronecker 张量积分解降维至 \(\mathcal{O}(W^2)\)，在保持相似曲率谱自标度能力的同时可轻松拓展至 \(W = 1024\) 甚至更大网络。
- vs PINN 损失平衡算法 (MultiAdam, NTK): NTK 重加权和 MultiAdam 侧重于在损失空间动态平衡不同物理项的权重，但无法改善单一损失分量本身的恶劣曲率条件；SS-eSOAP 聚焦于参数更新的底层几何空间，为解决物理求解的谱刚性提供了互补正交的底层优化方案。
评分¶
- 新颖性: ⭐⭐⭐⭐☆ 巧妙地将拟牛顿自标度割线条件引入 Kronecker 分解空间，并提出了反直觉且极其有效的二阶矩软重置下采样机制。
- 实验充分度: ⭐⭐⭐⭐⭐ 涵盖八大 PDE 基准、矩阵优化分析、三随机种子物理误差验证、高精度爆破分析及详尽的计算显存开销对比。
- 写作质量: ⭐⭐⭐⭐⭐ 理论推导清晰完备，从基础线性代数几何直觉到算法细节论述条理严密，对优势与局限的剖析客观诚实。
- 价值: ⭐⭐⭐⭐☆ 为科学机器学习中难以突破的高精度 PINN 训练提供了一款显存友好、收敛迅猛且高度实用的二阶优化工具。