面向高精度物理信息神经网络的曲率感知优化

《COMPUTER METHODS IN APPLIED MECHANICS AND ENGINEERING》：Curvature-aware optimization for high-accuracy physics-informed neural networks

　　

摘要：高效且鲁棒的优化对于神经网络至关重要，它使科学机器学习模型能够快速收敛至极高的精度——忠实地捕捉由微分方程支配的复杂物理行为。本工作提出了先进的优化策略，以加速物理信息神经网络（PINNs）在求解具有挑战性的偏微分方程（PDEs）和常微分方程（ODEs）时的收敛。特别地，本工作聚焦于曲率感知方法，即利用二阶信息或其近似替代来使更新方向自适应于损失景观局部几何结构的优化方法。具体而言，我们提供了自然梯度（NG）优化器、自缩放BFGS和Broyden优化器的有效实现，并展示了其在Helmholtz方程、Stokes流动、无粘Burgers方程、高速流动Euler方程以及在药代动力学和药效动力学中出现的刚性ODE等问题上的性能。除优化器开发外，我们还提出了新的基于PINN的求解无粘Burgers方程和Euler方程的方法，并将所得解与高阶数值方法进行对比，以提供严格且公正的评估。最后，我们解决了这些拟牛顿优化器在批量训练中的扩展问题，从而为大尺度数据驱动问题提供了高效且可扩展的解决方案。

引言

优化神经网络是现代机器学习的核心挑战，也是理论研究领域的活跃方向[1]、[2]、[3]、[4]。尽管神经网络优化问题具有高度非凸的结构，但其在实践中往往比最坏情况下的非凸问题更易处理，这表明其损失景观中存在可利用的结构。大多数训练算法是梯度方法的变体，其中梯度通过反向传播高效计算，参数则通过精心调优的步长进行迭代更新。相关的优化困难可大致分解为三个方面：确保收敛至驻点、实现快速收敛以及达到高质量的全局解。这些挑战表现为局部问题（如梯度消失或爆炸和缓慢收敛）以及全局问题（如低质量的局部极小值、平坦高台和复杂的损失景观）。理解这些现象并开发有原则的算法和理论工具来应对它们，是使神经网络成为可靠且强大的优化模型的关键。

由于训练本质上归结为选择使损失最小化的参数，优化充当了将数据和模型结构转化为学习行为的机制。一个有用的类比是在浓雾中穿越崎岖地形：损失曲面可能有山谷、悬崖和平坦盆地，算法必须仅利用有限的局部信息来决定如何移动。梯度提供了最速下降的方向，但缺乏对曲率的感知，这可能在平坦区域导致进展缓慢，或在急剧弯曲的区域导致不稳定。

这一局限性促使了方法层次结构的自然形成。形式上，如梯度下降等一阶方法仅使用梯度评估来更新参数。它们简单且可扩展，但可能在狭窄的山谷中左右摇摆，并且需要仔细调整步长以避免发散。相比之下，二阶方法通过Hessian矩阵、Gauss-Newton矩阵或Fisher信息矩阵等矩阵来纳入曲率信息[5]、[6]。通过捕捉局部几何结构，这些方法旨在产生更好缩放的搜索方向和步长，通常能够实现更快且更可靠的收敛。

物理直觉与地形类比相呼应。起初，你可能只知道哪个方向是上坡或下坡——这对应于梯度下降，它遵循局部坡度（梯度）但忽略了景观如何弯曲。因此，优化器可能缓慢地穿越高台或在陡峭的峡谷上来回震荡。现在想象你发现了一张详细的地图，它不仅揭示了坡度，还显示了你脚下地面的弯曲程度。这种曲率信息编码在Hessian、Gauss-Newton和Fisher信息矩阵等二阶对象中。利用这种结构使更新能够适应损失曲面的局部形状，这是一阶优化背后的核心理念。这些研究共同表明，关键瓶颈不在于曲率是否有帮助，而在于如何在计算和内存限制下利用它。这促使了那些以可扩展且鲁棒的方式近似二阶信息的方法的发展，旨在弥合一阶训练的效率与预条件梯度更新之间在迭代速度上的差距。特别地，曲率感知的二阶优化方法在PINNs之外也已展现出强大的前景，包括在大规模视觉和语言建模任务中[7]、[8]，这进一步激励了在科学机器学习中寻找可扩展且鲁棒的方式利用曲率信息。从这一角度来看，本工作将这些想法引入科学机器学习这一特别具有挑战性的设置中，并特别聚焦于PINNs。

为了更好地理解曲率信息在PINN训练中的作用，以下对二阶方法的函数解释是有益的。具体而言，PINN的训练可以解释为最小化如下形式的非线性最小二乘目标：L(θ) = ?‖r(θ)‖2，其中r(θ)表示与控制方程、边界条件及数据约束相关的残差集合。在当前迭代点θ_k附近对残差进行线性化可得：r(θ) ≈ r(θ_k) + J_k(θ - θ_k)，其中J_k是残差对参数的雅可比矩阵。在此设置下，经典梯度下降更新取以下形式：θ_{k+1} = θ_k - η_k J_k^T r(θ_k)，这可解释为参数空间中的最速下降。相比之下，Gauss-Newton等预条件方法引入了近似Hessian：H_k ≈ J_k^T J_k，从而产生如下形式的更新：θ_{k+1} = θ_k - η_k (J_k^T J_k)^{-1} J_k^T r(θ_k)，由此纳入曲率信息。

这种解释为二阶方法在PINNs中性能提升提供了统一的解释。特别地，PINNs中通常观察到的较差训练行为可归因于神经切线核（NTK）的极端病态[9]，这导致收敛缓慢以及向低频分量的有偏学习（谱偏差）。通过纳入曲率信息，Gauss-Newton类方法有效地正则化和重标定NTK谱，改善了条件数并使误差在各模式间更均匀地降低。所提出的Gram-Gauss-Newton近似因此可被视为这一理想预条件的可行替代方案，在过参数化区间中平衡计算成本与条件数改善。

在本工作中，我们对优化器性能进行了广泛的实证研究，包括SSBFGS、SSBroyden[10]、NG[11]、预条件特征基中的Adam的Shampoo（SOAP）[8]和AdamW[12]。我们的主要关注点是PINNs中出现的具有挑战性的训练机制，这是一个最初被引入用于求解微分方程正问题和逆问题的框架，通过将控制物理直接嵌入训练目标[13]。自其引入以来，PINNs在科学机器学习领域吸引了大量关注，因为它们提供了一种无网格且灵活的方法来将物理约束纳入神经网络训练。与此同时，越来越多的工作表明，PINN的训练往往比标准监督学习更为微妙，原因在于梯度流病理、不平衡的损失项、谱偏差以及优化失效模式，这些对于刚性、多尺度、高振荡或激波主导的问题尤为严重[14]、[15]、[16]、[17]、[18]、[19]。近期研究还强调，除了架构和采样选择之外，优化器本身可能是决定PINNs是否能收敛至低误差解以及是否能在实践中达到极高精度的关键因素[20]、[21]、[22]、[23]、[24]、[25]、[26]。受这些观察启发，我们在一组具有挑战性的问题上对一大类二阶优化器进行了基准测试，这些问题涵盖了一个刚性ODE系统（PK-PD）和三类PDE：椭圆问题（2D和3D Helmholtz）、抛物问题（2D粘性Burgers）和双曲问题（无粘Burgers和1D Euler方程）。这些测试用例代表了不同的训练困难来源，包括刚性动力学、高频振荡解、非线性扩散行为和激波诱导的不连续性。由于我们早期的综合研究[20]已表明双精度在此类问题上优于单精度，除特别说明外，我们将范围限定为双精度计算。单精度和双精度性能的详细对比见[20]。这一观察也与Xu等人[27]的近期发现一致，他们表明算术精度不足本身即可引发PINN训练中的失效模式，且从FP32切换到FP64可定性地改变优化动力学。

大量工作试图通过改变学习问题本身的公式化来应对PINNs的优化困难和精度局限。保守型PINNs在离散子域上强制施加守恒律[28]；变分法和基于能量的公式化用弱形式或能量原理替代强形式残差[29]；混合公式化引入辅助变量以降低导数阶数并改善异质介质处理；激波捕捉型PINN变体修改残差以考虑不连续性和熵相容性[30]、[31]、[32]。与此同时，为力学中的实时预测开发了神经网络代理模型，特别是当需要大量查询仿真时[33]。我们的关注点不同且互补。我们研究了二阶和拟牛顿优化器在应用于多个PDE区间的PINNs时的表现，并识别出何时优化是主导瓶颈以及何时必须改变公式化本身。在光滑的椭圆和抛物问题中，改进优化可直接带来高精度的解。然而，对于双曲守恒律，PINN残差还必须编码守恒性和激波结构；否则，优化问题与熵解的对齐仍然不佳。因此，本工作的贡献不仅是对优化器的基准测试，也是对优化、公式化和计算成本在力学导向的PINN求解器中如何交互的研究。

除基准测试外，本工作还引入了用于求解具有不连续解的双曲PDE的新型基于PINN的方法学，并展示了其在无粘Burgers方程和可压缩Euler方程上的收敛性。此外，我们提出了SSBFGS和SSBroyden算法的新型批量训练变体，以提高大规模训练设置中的效率和可扩展性。因此，本论文的贡献不仅限于对现有优化器的实证比较。相反，它呈现了对PINNs优化的更广泛研究，结合了跨椭圆、抛物和双曲PDE以及一个刚性PK-PD ODE系统的系统实验，以及新的双曲PINN公式化、损失景观与函数空间分析、基于屋顶线的性能研究和自缩放拟牛顿方法的可扩展批量训练策略。特别地，虽然Helmholtz问题包含了在更简单设置中的优化器直接对比，但优化器性能的综合研究已在我们的早期工作[20]中完成。因此，本研究超越了该设置，转向更具挑战性的高频和三维Helmholtz问题，在这些问题上诸如域扩展和可分离物理信息神经网络（SPINN）[34]表示等额外要素变得重要。这一进展在Burgers和Euler问题中继续推进，结果表明仅选择更强的优化器是不够的，因为精确解还需要激波捕捉和与守恒一致的公式化。上述函数视角激励了针对PINNs设计预条件优化策略的评估。本工作的主要贡献总结如下：

• 我们对用于PINNs的曲率感知优化方法进行了全面研究，包括自缩放BFGS、自缩放Broyden、自然梯度、SPRING和SOAP，涵盖椭圆、抛物和双曲PDE以及刚性ODE。该研究提供了对收敛速度、鲁棒性、可扩展性和可达精度的系统评估，并在NG方法中引入了Jacobi缩放以提高稳定性，同时提出了一种新的用于PINNs拟牛顿优化的批量训练策略。

• 我们引入了用于求解高速流动的双曲守恒律的新型PINNs方法，应用于Burgers和Euler方程。

• 我们开发了一个新框架，用于通过函数空间和参数空间中的损失景观来分析PDE学习动力学。

• 我们进行了详细的屋顶线分析，以表征拟牛顿优化器在PINNs中的计算扩展和效率。

• 我们提出了一种方法来缓解高波数Helmholtz问题中的谱偏差。

• 我们展示了PINNs在药代动力学和药效动力学中出现的刚性ODE系统上的适用性。

• 尽管本研究聚焦于正问题，我们概述了所提方法如何自然地扩展至逆问题，包括参数估计和数据同化。

第2节回顾了本工作所考虑的优化方法。第3节介绍了参数空间和函数空间的可视化工具。第4节呈现了计算基准测试，包括Helmholtz、Stokes、Burgers、Euler和PK-PD问题。第7节总结了主要发现。

章节片段：所考虑优化器的回顾

本节简要描述了本研究中每种优化器的基本原理及实现。

我们重点关注核心更新规则、曲率近似以及计算需求，突出与大规模训练和物理信息神经网络（PINN）应用最为相关的实际应用差异。图1提供了PINN训练流程的概述，并指出了优化在流程中的介入位置。让我们考虑偏微分方程（PDE）问题 $Au=f$ 在 $\Omega$ 中，$u=g$ 在 $\Gamma$ 上，其中 $A$ 是一个PDE算子。

**PDE景观可视化**

本节从两个互补的视角可视化优化动力学。首先，我们检查参数空间中损失景观的二维切片，以说明不同优化器遇到的局部几何结构。其次，我们将参数空间的更新推进一步，映射到函数空间，在那里可以直接与当前预测误差进行比较。这提供了一种更具可解释性的视角，展示每种优化器如何改变所学习的PDE解。

**计算实验**

本节介绍了一组具有挑战性的PINN基准问题，用于评估优化器在刚性、振荡性、多尺度以及以激波为主导的损失景观下的鲁棒性。该基准测试套件包括：二维和三维的不同波数的亥姆霍兹方程、无粘Burgers方程、二维粘性Burgers方程、可压缩Euler方程、二维Stokes方程以及一个药代动力学–药效动力学（PK–PD）系统。

**基于屋顶线模型的优化器性能分析**

屋顶线模型 [75], [76] 是一种性能界限，它将内核的可实现浮点吞吐量刻画为其算术强度的函数。它提供了一个直观的可视化框架，用于判断内核是受计算吞吐量还是内存带宽的限制，并量化实测性能与硬件峰值性能之间的差距。

**可复现性**

为了促进可复现性并方便更广泛的采用，我们概述了实现背后的设计原则，并提供了公开可用的代码库。所有实现均基于JAX框架开发，确保了构建优化器变体的效率和灵活性。每个优化器实例由以下几个关键组件定义：

1. 相对容差（rtol）
2. 绝对容差（atol）
3. 收敛准则的范数选择（例如，梯度的 $L_1$、$L_2$ 或 $L_\infty$ 范数）
4. 行搜索策略的选择

**总结**

本研究在广泛的具有挑战性的科学机器学习问题上，探讨了PINN的不同优化方法。计算实验围绕三类主要PDE——椭圆型、抛物型和双曲型——以及一个刚性PK–PD常微分方程系统进行组织。在所有情况下，我们比较了自适应缩放拟牛顿方法、自然梯度与高斯–牛顿方法以及结构化自适应优化器，重点关注收敛速度、鲁棒性、精度等方面。

**CRediT作者贡献声明**

Anas Jnini：撰写–初稿，可视化，验证，软件，方法论，形式化分析。Elham Kiyani：撰写–初稿，可视化，验证，软件，方法论，调查，形式化分析，数据整理，概念化。Khemraj Shukla：撰写–初稿，可视化，验证，软件，方法论，调查，形式化分析，数据整理，概念化。Jorge F. Urbán：撰写–审阅与编辑，可视化，验证。

**利益冲突声明**

作者声明他们没有已知的可能影响本研究报告工作的竞争财务利益或个人关系。

**致谢**

本研究主要得到AIM for Composites的支持，这是一个由美国能源部（DOE）、科学办公室、基础能源科学（BES）资助的能源前沿研究中心，项目资助编号为DE-SC0023389（计算研究、数据分析）。KS和GEK的工作部分得到了美国能源部（DOE）、科学办公室、高级科学计算研究（ASCR）计划下的高级计算科学发现（SciDAC）项目支持。

Anas Jnini | Elham Kiyani | Khemraj Shukla | Jorge F. Urbán | Nazanin Ahmadi Daryakenari | Johannes Müller | Marius Zeinhofer | George Em Karniadakis