[论文解读] Hutchinson Trace Estimation for High-Dimensional and High-Order Physics-Informed Neural Networks

Zheyuan Hu, Zekun Shi+2 more|arXiv (Cornell University)|Dec 22, 2023

Neural Networks and Applications参考文献 57被引用 7

一句话总结

该论文提出Hutchinson Trace Estimation（HTE）以加速高维及高阶偏微分方程（PDEs）的物理信息神经网络（PINNs）。通过用Hessian-Vector Products（HVPs）替代完整的Hessian矩阵计算，HTE降低了内存占用，并实现了高效的泰勒模式自动微分，从而在方差较低的情况下实现无偏、收敛的训练，尤其在高方差场景下优于随机维度梯度下降（SDGD）。

ABSTRACT

Physics-Informed Neural Networks (PINNs) have proven effective in solving partial differential equations (PDEs), especially when some data are available by seamlessly blending data and physics. However, extending PINNs to high-dimensional and even high-order PDEs encounters significant challenges due to the computational cost associated with automatic differentiation in the residual loss. Herein, we address the limitations of PINNs in handling high-dimensional and high-order PDEs by introducing Hutchinson Trace Estimation (HTE). Starting with the second-order high-dimensional PDEs ubiquitous in scientific computing, HTE transforms the calculation of the entire Hessian matrix into a Hessian vector product (HVP). This approach alleviates the computational bottleneck via Taylor-mode automatic differentiation and significantly reduces memory consumption from the Hessian matrix to HVP. We further showcase HTE's convergence to the original PINN loss and its unbiased behavior under specific conditions. Comparisons with Stochastic Dimension Gradient Descent (SDGD) highlight the distinct advantages of HTE, particularly in scenarios with significant variance among dimensions. We further extend HTE to higher-order and higher-dimensional PDEs, specifically addressing the biharmonic equation. By employing tensor-vector products (TVP), HTE efficiently computes the colossal tensor associated with the fourth-order high-dimensional biharmonic equation, saving memory and enabling rapid computation. The effectiveness of HTE is illustrated through experimental setups, demonstrating comparable convergence rates with SDGD under memory and speed constraints. Additionally, HTE proves valuable in accelerating the Gradient-Enhanced PINN (gPINN) version as well as the Biharmonic equation. Overall, HTE opens up a new capability in scientific machine learning for tackling high-order and high-dimensional PDEs.

研究动机与目标

• 解决在高维及高阶PDEs上训练PINNs时面临的高计算成本和内存负担问题。
• 克服现有方法（如SDGD）在维度间存在显著差异性时方差过高的局限性。
• 通过引入张量-向量积（TVPs）管理巨大张量，将PINN训练扩展至四阶PDEs（如双调和方程）。
• 在特定条件下，确保基于HTE的损失函数在理论上的收敛性与无偏性，使其逼近原始PINN损失。
• 通过高效的迹估计，实现在梯度增强PINNs（gPINN）和高阶PDE求解器中的实际训练加速。

提出的方法

• 将Hutchinson迹估计（HTE）应用于PINN损失函数中Hessian矩阵的迹近似，替代完整的Hessian矩阵计算。
• 将Hessian矩阵计算转化为Hessian-Vector Product（HVP），将内存需求从O(d²)降低至O(d)，并实现在JAX中的高效泰勒模式自动微分。
• 使用Rademacher分布的随机向量计算无偏迹估计器：$ \frac{1}{V}\sum_{i=1}^{V} \boldsymbol{v}_i^T A \boldsymbol{v}_i \approx \mathrm{Tr}(A) $。
• 通过将HVP推广为张量-向量积（TVP），将HTE扩展至高阶PDEs，实现对双调和方程等四阶PDE项的高效计算。
• 将HTE集成至PINN框架中，确保损失函数在适当条件下保持无偏性并收敛至原始PINN损失。
• 通过理论方差分析与高维及高阶PDEs上的实证评估，将HTE与SDGD进行对比。

实验结果

研究问题

• RQ1 HTE能否在高维PINNs中提供一种可扩展的替代方案，以减少完整Hessian矩阵计算带来的内存与计算开销？
• RQ2 在具有显著维度间差异性的高维PDEs中，HTE的方差与SDGD相比如何？
• RQ3 能否通过张量-向量积（TVPs）将HTE扩展至四阶PDEs（如双调和方程）？
• RQ4 在标准假设下，基于HTE的PINN损失是否无偏且收敛于原始PINN损失？
• RQ5 HTE是否能在梯度增强PINNs（gPINN）和高阶PDEs中实现更快且更内存高效的训练？

主要发现

• 通过用标量HVP替代矩阵运算，HTE将内存消耗从完整Hessian的O(d²)降低至O(d)。
• 当迹估计器应用于残差损失的Hessian时，基于HTE的损失函数保持无偏，并收敛至原始PINN损失。
• 在维度间方差较高的场景中，HTE表现出低于SDGD的方差，使其在复杂PDEs中更具稳定性。
• 对于双调和方程，HTE结合TVPs通过将O(d⁴)张量转化为一系列张量-向量运算，实现了高效计算。
• 实证结果表明，HTE在收敛速度上与SDGD相当，同时显著降低内存使用，并在高维设置下实现更快训练。
• HTE通过高效估计梯度增强项，实现了gPINN中的训练加速，且不损失模型表达能力。