用消息传递求解偏微分方程，ML大牛Max Welling等用全神经求解器做到了更强、更快

2022 年 2 月 20 日 机器之心

机器之心报道

编辑：杜伟

对于求解偏微分方程来说，阿姆斯特丹大学、高通 AI 研究院的研究者最近推出的 MP-PDE 求解器又提供了一个选择。

在科学领域，常年的工作已经面向各种物理现象生成了极其详细的数学模型。很多这些模型通过微分方程（Olver, 2014）的形式进行自然地表达，大多数时候表现为时间偏微分方程（partial differential equation, PDE）。求解这些微分方程对于解决天气预报、天文数字模拟、分子建模、喷气式发动机设计等所有数学学科中的问题至关重要。大多数重要方程的求解难以分析，因此不得不反溯至数值近似方法。想要以最小的计算开销获得有界误差的精确解需要手动求解器（handcrafted solver），通常根据手头的方程量身定制。

设计一个「好的」PDE 求解器绝非易事。完美的求解器应该满足大量的条件。首先是用户需求，比如速度快、使用最少的计算开销、提供不确定性估计、跨 PDF 泛化以及易于使用；然后是问题的结构需求，比如空间分辨率和时间尺度、域采样正则性、域拓扑和几何、边界条件、维数和解空间平滑度；接着是实现需求，比如在长时间 rollout 时保持稳定性和不变形。正是由于上述大量的多样化需求，数值法（ numerical method）是一个 splitter 领域，而不是一个 lumper 领域，旨在为每个子问题构建手动手动求解器。

近日，阿姆斯特丹大学、高通 AI 研究院的三位研究者在论文《Message Passing Neural PDE Solvers》中提出使用端到端神经求解器来从数值上求解 PDE。

论文地址：https://arxiv.org/pdf/2202.03376.pdf

具体而言，这篇论文主要做出了以下贡献：

提出一个基于神经消息传递（message passing, MP）的端到端全神经 PDE 求解器，其灵活性能够满足典型 PDE 问题的所有结构需求。这一设计的灵感来源于一些经典求解器（有限差分、有限体积和 WENO 格式）可以作为消息传递的特例；
提出时间捆绑（temporal bundling）和前推（pushforward）技巧，以在训练自回归模型中鼓励零稳定性（zerostability）；
在给定类中实现跨多个 PDE 的泛化。在测试期间，新的 PDE 稀疏可以成为求解器的输入。

方法

研究者基于最近该领域令人兴奋的工作进展来学习 PDE 求解器。这些神经 PDE 求解器的背后离不开这一快速发展且有影响力的研究领域。用于时间 PDE 的神经 PDE 求解器可以分为两大类，分别为自回归方法和神经算子方法，具体如下图 1a 所示。

研究者通过两部分详细描述了他们的方法，即训练框架和架构。其中训练框架解决自回归求解器中的分布位移问题，该问题会导致不稳定性；网络架构是一个消息传递神经网络。

训练框架

自回归求解器将解 u^k 映射到因果后续（causally consequent）解 u^k+1。一种直接的训练方法是单步训练。如果 p_0(u^0 ) 在训练集中是初始条件的分布，则

是迭代为 k 时的真值分布。研究者最小化如下公式（6）

下图 2 为不同的训练策略。图左为单步训练，只能预测接下来一步的解；图中为展开（unrolled）训练，可以预测接下来 N 步的解；图右为对抗性训练，可以预测接下来 N 步的解，但只能在最后一步反向传播。

架构

在网络架构选择上，研究者遵循 Battaglia et al. (2018) 和 Sanchez-Gonzalez et al. (2020) 提出的编码器 - 处理器 - 解码器（Encode-Processor-Decode）框架，并做了调整。他们并不是首个将 GNN 用作 PDE 求解器的，但自己的方法具有一些显著特征。下图 3 为本文 MP-PDE 求解器的概览：

具体而言，编码器用来计算节点嵌入。

处理器计算学得消息传递的第 M 步，中间图表示为

。具体更新如下公式（8）和（9）

最后来说解码器。在消息传递后，研究者使用了一个浅层 1D 卷积网络，并在空间位置上共享权重，以在网格点 x_i 处输出 K 接下来的时间步预测。对于每个节点 i，处理器输出向量 f^M_i。他们将该向量视为时间连续的信号，并随时间推移将它馈入到 CNN。

实验

研究者在不同难度的任务上展示了 MP-PDE 求解器的有效性。其中，在 1D 方程中，研究者探究了 MP-PDE 泛化到给定族中未见过方程的能力，周期性、狄利克雷（Dirichlet）边界条件和诺伊曼（Neumann）边界条件下的边界处理能力，以及建模冲击波（shock wave）的能力。然后，他们又展示了 MP-PDE 有能力求解 2D 方程。

此外，研究者还针对前推技巧和变体进行了消融实验，以验证实用性。作为基线，他们比较了几种不同的标准经典 PDE 求解器，即 FDM、伪谱方法和 WENO5 求解器。不仅如此，研究者还与 SOTA 神经算子方法——傅里叶神经算子（Fourier Neural Operator, FNO）进行了比较。

在实验中，研究者考虑了三种场景，分别如下：

E1 伯格斯（Burgers）方程，没有用于冲击建模的扩散θ_PDE = (1, 0, 0)；
E2 伯格斯方程，有可用扩散θ_PDE = (1, η, 0)，其中 0 ≤ η ≤ 0.2；
E3：θ_PDE = (α, β, γ) 的混合场景，其中 0.0 ≤ α ≤ 3.0、0.0 ≤ β ≤ 0.4 和 0.0 ≤ γ ≤ 1.0。

具体而言，他们观察 E1 方程上的求解器生存时间，定义为「解偏离真值之前的时间」。该求解器展开到 n_t = 1000 时间步，其中 T = 16 s。下图 4 底部展示了一个示例，研究者观察到大约 8 秒后发散增加。该现象在下图 5a 中得到了验证，他们发现了生存率与时间步的关系。

在第二个实验中，研究者比较了前推技巧的效用。他们观察到，前推技巧加上时间捆绑可以提升自回归任务中的 FNO 效果。在下图 5b，研究者绘制了使用和未使用前推技巧训练的模型的生存率。

下表 2 比较了 MP-PDE 求解器与 SOTA 数值伪谱求解器。结果可知，MP-PDE 求解器在伪谱求解器中断工作的低分辨率条件下获得了准确的结果。有趣的是，MP-PDE 求解器可以在不同的边界条件上泛化，并且如果边界条件通过θ_PDE 特征注入到方程中，泛化更加明显。

最后，研究者测试了 MP-PDE 到更多空间维度上的可扩展性，尤其是在 2D 实验中。他们使用来自开源流模拟工具包 PHIFLOW1 中的数据。具体而言，研究者观察了基于纳维 - 斯托克斯方程（Navier-Stokes equation），并将烟雾流模拟成 32 × 32 网格，在每个时间步后添加更多烟雾。结果显示，MP-PDE 求解器能够准确地捕获给定时间阶段内的烟雾流入，表明它可以扩展到更高维度。