对于求解偏微分方程来说,阿姆斯特丹大学、高通 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 求解器能够准确地捕获给定时间阶段内的烟雾流入,表明它可以扩展到更高维度。