[论文解读] Separable Physics-Informed Neural Networks

Jun Woo Cho, Seungtae Nam+4 more|arXiv (Cornell University)|Jun 28, 2023

Model Reduction and Neural Networks被引用 27

一句话总结

SPINN 使用轴可分离的子网络和前向模式自动微分来大幅降低 PINN 的计算量，从而在多维偏微分方程中实现非常大的取点集合，并提供更快、更精确的解。

ABSTRACT

Physics-informed neural networks (PINNs) have recently emerged as promising data-driven PDE solvers showing encouraging results on various PDEs. However, there is a fundamental limitation of training PINNs to solve multi-dimensional PDEs and approximate highly complex solution functions. The number of training points (collocation points) required on these challenging PDEs grows substantially, but it is severely limited due to the expensive computational costs and heavy memory overhead. To overcome this issue, we propose a network architecture and training algorithm for PINNs. The proposed method, separable PINN (SPINN), operates on a per-axis basis to significantly reduce the number of network propagations in multi-dimensional PDEs unlike point-wise processing in conventional PINNs. We also propose using forward-mode automatic differentiation to reduce the computational cost of computing PDE residuals, enabling a large number of collocation points (>10^7) on a single commodity GPU. The experimental results show drastically reduced computational costs (62x in wall-clock time, 1,394x in FLOPs given the same number of collocation points) in multi-dimensional PDEs while achieving better accuracy. Furthermore, we present that SPINN can solve a chaotic (2+1)-d Navier-Stokes equation significantly faster than the best-performing prior method (9 minutes vs 10 hours in a single GPU), maintaining accuracy. Finally, we showcase that SPINN can accurately obtain the solution of a highly nonlinear and multi-dimensional PDE, a (3+1)-d Navier-Stokes equation. For visualized results and code, please see https://jwcho5576.github.io/spinn.github.io/.

研究动机与目标

• 动机并解决在多维和复杂 PDE 上训练 PINN 的计算瓶颈。
• 提出一种可分离网络架构，降低计算 PDE 残差所需的前向传播次数。
• 利用前向模式自动微分，使在通用 GPU 上能够处理大规模取点集合。
• 在扩散、亥姆霍兹、克莱因-戈登以及 Navier–斯托克斯方程的 3D 和 4D 问题上，展示精度和速度的提升。

提出的方法

• 引入 SPINN，将输入分成 d 个一维轴并送入 d 个主体网。
• 每个主体网 f^(θ_i): R -> R^r 输出轴 i 的一组特征向量。
• 预测解记为 û(x1,...,xd) = sum_{j=1}^r prod_{i=1}^d f_j^(θ_i)(x_i)（一个低秩张量表示）。
• 利用前向模式 AD 高效计算导数和 PDE 残差，实现梯度的 O(Nd) 雅可比-向量乘积评估。
• 训练利用通过坐标轴样本的笛卡尔积形成的可分解（网格状）取点，使在输入点数适中的情况下实现密集评估。
• 提供一个基于 L^2 的普适近似论证及在多种 PDE 上的经验验证。

实验结果

研究问题

• RQ1 与传统 PINN 相比，SPINN 是否能够在多维 PDE 中降低计算成本和内存使用，同时保持或提高精度？
• RQ2 轴向分离结合前向模式 AD 是否能够在标准 GPU 上处理非常大的取点集合（>10^7）？
• RQ3 SPINN 是否能够比先前的 PINN 方法更快地求解具有挑战性的 PDE（例如混沌的 Navier–Stokes），同时保持精度？

主要发现

• SPINN 实现了显著的时钟时间下降（最高可达 62×）和 FLOP 下降（最高可达 1,394×），相较于具有相同取点的基线 PINN。
• 对分离输入使用前向模式 AD 降低高阶导数的计算成本，并在通用 GPU 上实现非常大的取点集合。
• SPINN 能够显著更快地求解一个混沌的 (2+1)-d Navier–Stokes 方程（在单个 GPU 上从 10 小时降至 9 分钟），无需因果性归纳偏置。
• SPINN 在 3+1-d Navier–Stokes 以及其他 PDE（扩散、亥姆霍兹、克莱因-戈登）上展示了在效率和精度方面优于传统PINNs 的准确解。
• 该方法在 L^2 中支持普遍近似性质，并通过合并操作给出解的低秩张量解释。