[论文解读] Message Passing Neural PDE Solvers

J. Brandstetter, Daniel E. Worrall+1 more|arXiv (Cornell University)|Feb 7, 2022

Model Reduction and Neural Networks被引用 61

一句话总结

介绍 MP-PDE，一种基于图神经网络的完全神经、自回归 PDE 求解器，能够在跨 PDE、域以及边界条件上泛化，并具备如 pushforward trick 和 temporal bundling 这样的稳定性技巧。

ABSTRACT

The numerical solution of partial differential equations (PDEs) is difficult, having led to a century of research so far. Recently, there have been pushes to build neural--numerical hybrid solvers, which piggy-backs the modern trend towards fully end-to-end learned systems. Most works so far can only generalize over a subset of properties to which a generic solver would be faced, including: resolution, topology, geometry, boundary conditions, domain discretization regularity, dimensionality, etc. In this work, we build a solver, satisfying these properties, where all the components are based on neural message passing, replacing all heuristically designed components in the computation graph with backprop-optimized neural function approximators. We show that neural message passing solvers representationally contain some classical methods, such as finite differences, finite volumes, and WENO schemes. In order to encourage stability in training autoregressive models, we put forward a method that is based on the principle of zero-stability, posing stability as a domain adaptation problem. We validate our method on various fluid-like flow problems, demonstrating fast, stable, and accurate performance across different domain topologies, equation parameters, discretizations, etc., in 1D and 2D.

研究动机与目标

• 推动端到端的神经 PDE 求解器在分辨率、拓扑、几何、边界条件、离散化和维度方面的泛化能力。
• 提出一个神经消息传递架构，使经典求解器成为特殊情形（FDM/FVM/WENO）。
• 通过 pushforward trick 与 temporal bundling 解决自回归训练不稳定性，以改善零稳定性。
• 使 MPDE 在同一类别内对多种 PDE 泛化，并支持在推断阶段对未知 PDE 系数进行测试。

提出的方法

• 把计算网格建模为一个图，并在 Encode-Process-Decode 框架中使用图神经网络。
• 将 PDE 求解器表示为一系列学习的消息传递更新，模仿经典导数/通量计算。
• 使用方程嵌入向量将 PDE 参数和边界条件注入网络。
• 应用解码器联合输出多个未来时间步，强制保持一致性以实现稳定性。
• 使用包含一步项与稳定性（pushforward）项的联合损失进行训练，以减轻分布偏移。
• 采用 temporal bundling 以并行预测若干时间步，降低 rollout 误差和计算量。

实验结果

研究问题

• RQ1 神经消息传递求解器是否能在同一类内不同 PDE、以及不同域几何和边界条件之间实现泛化？
• RQ2 面向稳定性的训练策略（pushforward 与 temporal bundling）是否提高长时间 rolls 的准确度并减缓自回归误差？
• RQ3 MP-PDE 求解器在多大程度上能够复现经典求解器行为（FDM/FVM/WENO），并对非规则网格和高维（1D/2D）具泛化性？
• RQ4 对 PDE 系数（θPDE）进行条件化在求解未知参数配置时对泛化能力有何影响？
• RQ5 在多物理场类问题（如 Burgers、热方程、KdV 和 Navier–Stokes 启发的设置）上，与最先进的神经算子与经典求解器相比，MP-PDE 的性能如何？

主要发现

• MP-PDE 求解器在分辨率、域拓扑、几何、边界条件和维度（1D 和 2D）上具有泛化能力。
• 神经消息传递可以在其架构中表示并再现经典求解器，如有限差分、有限体积和 WENO。
• 以稳定性为导向的训练（pushforward trick）和 temporal bundling 能显著提高 rollout 稳定性并降低自回归预测中的误差累积。
• 在同一族的未见 PDE 上，MP-PDE 优于 WENO5 和自回归 FNO-RNN，并从 PDE-parameter conditioning（θPDE）中获益。
• 在不规则网格和非周期边界条件下，MP-PDE 可以超越伪谱方法，并在提供 θPDE 特征时实现对边界条件的泛化。
• 对 Navier–Stokes 启发的流动（PhiFlow 数据）的初步 2D 实验表明具备扩展到更高维度的可扩展性。