北京大学科研团队发布PDEformer-2——基于昇思MindSpore的二维偏微分方程基础模型

11月4-6日，汇集AI for Science领域“产、学、研、用”方向单位，旨在共同思考AI在科学研究中的应用与未来的“2024科学智能峰会”在北京召开。

会上，北京大学董彬教授发布了其团队联合华为 AI4SCI Lab 基于昇思MindSpore AI框架开发的PDEformer-2模型。PDEformer-2是一种可以直接处理任意 PDE 形式的端到端解预测模型，同时适用于含时与不含时的方程。模型使用约 40TB 的数据集进行预训练，能够对具有不同边界条件、求解区域、变量个数的二维方程直接推理，快速获得任意时空位置的预测解。此外，PDEformer-2 作为正问题解算子的可微分代理模型，也可以用于求解各类反问题，包括基于有噪声的时空散点观测进行全波反演以恢复方程中的波速场等。这为模型支持包括流体、电磁等领域的众多物理现象及工程应用的研究打下良好基础。

背景

偏微分方程（PDE）与众多物理现象及工程应用紧密相连，涵盖机翼设计、电磁场模拟、应力分析等广泛领域。在这些实际应用中，PDE 的求解往往需要反复进行。尽管传统的 PDE 求解方法往往具有较高的求解精度，但它们在计算资源和时间上的消耗较大，且难以适应各种 PDE 形式的统一求解。近年来，随着深度学习的发展，领域中提出了许多种基于深度学习方法的 PDE 求解器，其中的许多方法相比传统求解器有明显的速度优势，因而受到了很多关注。

2024 年上半年，华为 AI4Sci Lab 携手北京大学董彬教授，基于昇思 MindSpore 框架，开创性地构建了一个针对 PDE 数值求解的通用模型 PDEformer-1，旨在使用一个统一的模型对不同的一维 PDE 进行数值求解。与 FNO、DeepONet 等专用的神经算子模型不同，这一通用模型无需针对不同方程进行定制化的架构设计与训练，从而大幅降低模型的部署成本，提升了求解效率。

在这一出色结果的基础上，团队再接再厉，开发了针对二维 PDE 的通用求解基础模型 PDEformer-2。模型使用大规模、多样化的数据集进行预训练，获得了强大的求解能力，能够直接求解具有不同边界条件、区域形状、未知变量个数的二维方程。对于预训练分布外的场景，PDEformer-2 也展现出了强大的小样本学习能力，能利用有限的微调数据迅速适应新的方程求解任务。此外，PDEformer-2 作为正问题算子的可微分代理模型，还能用于各类反问题的求解，基于有噪声的时空散点观测估计方程中的常数系数、源项场和波方程的波速场。

技术路径

我们考虑定义在 (t,r)∈[0,1]×Ω 上的二维偏微分方程（PDE），其一般形式为

如图所示，PDEformer-2 先将方程形式表示为一个计算图，利用标量、函数编码器将方程涉及的数值信息嵌入到计算图的节点特征当中。接下来，PDEformer-2 通过 graph Transformer 编码这一计算图，并使用隐式神经表征（INR）对所得的隐向量进行解码，获得 PDE 的各个解分量在特定时空坐标下的预测值。

对于二维方程中可能出现的复杂区域形状和边界位置，PDEformer-2 将其表示为符号距离函数（SDF），并使用函数编码器将这部分信息嵌入到计算图当中。下图所示的例子展示了使用计算图表示方形区域上 Dirichlet 边界条件的方式：

实验结果

0****1

正问题

团队为 PDEformer-2 的预训练准备了大规模、多样化的数据集，其中包括时间一阶的反应对流扩散（DCR）方程与浅水波（SWE）方程，时间二阶的波动方程及其多分量（MC）版本，以及不含时的稳态弹性力学方程。这些数据分别使用开源软件包 Dedalus 和 FEniCSx 生成，涵盖了不同方程定义域、边界条件和解分量个数的样本。数据集的规模与预训练后的模型精度（nRMSE）如下表所示：

下图展示了模型在部分测试数据集上的预测结果，包括方形区域的波方程（周期边界）：

圆盘区域的反应对流扩散方程（Neumann 边界）：

以及方形区域的稳态弹性波方程（二分量，混合边界）：

充分的预训练不仅允许模型对预训练分布内的方程进行高效的推理，也为微调适应分布外的方程打下了良好的基础。团队使用 PDEBench 基准数据集提供的浅水波数据对模型进行了微调，其中所使用的分片常数初值在模型的预训练过程中并未出现。尽管微调过程仅仅使用了一个样本，模型还是迅速地适应了这一数据集的物理场景，在测试集上达到了仅 3.60% 的相对误差。其中一个测试样本的预测效果如下图所示：

0****2

反问题

为了验证PDEformer-2在下游任务上的应用潜力，我们将其作为正问题算子的代理模型应用到反问题中，并观察其结果。

(1) 反问题形式1：标量系数反演

该反问题是在已知方程形式和带噪声的初值条件下，根据在固定空间位置上部分时间点的带噪声观测，来恢复方程中的未知系数。

例如，已知方程形式为

(2) 反问题形式2：源项反演

该反问题是在已知方程形式和带噪声的初值条件下，根据在固定位置、时间的观测（带噪声），来恢复方程中的未知源项场。

在这个任务中，我们所采用的方程形式分布仍与预训练数据集相同，而需要恢复的对象是其中对应于方程源项的系数场 s(x)。与标量系数反演问题类似，对于每个PDE，我们假设能够获取25个不同的初始条件下的解，可观测到的样本点局限于128个固定的空间位置，平均每个空间位置能获取20个观测样本。这些点的观测值上会添加相对幅度为1%的噪声，而作为 PDEformer-2 输入的初值上添加1%的噪声。相应的优化问题使用梯度下降算法求解。

下图展示了 PDEformer-2 在反问题中的应用结果，从左到右依次为反演的源项的真实值、反演结果和两者的误差，能够看到在不同的方程中，PDEformer-2为正算子的反演结果与真实值较为接近。

总结与展望

PDEformer-2 模型基于 PDE 计算图表示的技术路线。在大量而多样的二维 PDE 数据集上进行预训练之后，模型展现出了较高的通用能力，能够针对不同 PDE 形式、边界条件、区域形状、时间依赖情况和解分量个数的方程进行直接推理，快速获得无网格的预测解，且可通过快速微调适应新的方程。在处理反问题这样的下游任务时，PDEformer-2 也表现出良好的适应性，能够在多种问题上给出较为准确的解。

基于 PDE 的物理仿真及与其相关的反问题求解、优化设计、控制等任务对许多工业应用而言十分关键。未来，团队计划对接航空、汽车、通信、能源等领域，根据实际工业需求推进这一基础模型的后续开发，进一步提高其解决现实问题的潜力与能力。同时，充分发挥 PDEformer 模型的通用性、灵活性优势，推动其在多种真实场景中落地应用，最终赋能千行百业。

团队成员：叶展宏（北大），刘子宁（北大），吴秉阳（北大），蒋弘杰（北大），吴闻道（北大），陈乐恒（北大），Gitee: zhang-minyan，Gitee: huangxiang360729，Gitee: hsliu_ustc，董彬（北大）