PDEformer Towards a Foundation Model for One-Dimensional Partial Differential Equations

PDEformer Towards a Foundation Model for One-Dimensional Partial Differential Equations

简介

要求解的 PDE 涉及两部分信息：一是指定 PDE 数学形式的符号部分，另一部分是数字部分，包括 PDE 系数、初始值和边界值等。典型的神经算子如傅里叶神经算子（FNO）和 DeepONet 用于处理特定形式的偏微分方程，只需将数值信息作为网络输入。然而，为了构建可推广到不同偏微分方程的基础模型，符号信息必须无缝集成。

朝着这个方向的一些现有方法采用了语言模型，其中偏微分方程的数学表达式作为输入。这些方法可能难以完全捕获符号信息和数字信息之间的复杂交互。其他策略避免显式输入 PDE 形式，而是选择将其隐式编码在模型的数字输入中。此类隐式输入法可能不足以编码具有更大多样性和复杂性的偏微分方程类，并且可能必须配合另一个求解器来准备额外的解快照。

在本文中介绍 PDEformer。与以前的方法不同，本文建议将偏微分方程的符号形式表示为计算图，确保生成的图结构及其节点类型和特征向量封装求解偏微分方程所需的所有符号和数值信息。图 Transformer 和 INR 用于生成无网格预测解。在对具有一定多样性的偏微分方程进行预训练后，对基准数据集的评估表明，与基线专家模型相比，PDEformer 表现出更高的零样本预测精度，或者可以在有限数据的微调后实现这一目标。应用到各种下游任务的潜力主要通过偏微分方程系数恢复逆问题进行验证。尽管本文的实验目前仅限于一维偏微分方程，但相信这将成为构建基础偏微分方程模型的一个值得注意的里程碑。

方法

考虑定义在 $(t,x)\in [0,1]\times [-1,1]$ 上的一维时间相关偏微分方程，具有周期性边界条件，其一般形式如下：
$$\mathcal{F}(u,c_1,c_2,\dots )=0,u(0,x)=g(x)$$
其中 $c_1,c_2,\cdots \in \mathbb{R}$ 为实值系数，$g(x)$ 为初始条件。这里，假设算子 $\mathcal{F}$ 有一个符号表达式，可能涉及微分和代数运算。目标是构造解映射 $(\mathcal{F},g,c_1,c_2,\dots )\mapsto u$ 的代理，其本质上采用算子 $\mathcal{F}$ 的形式作为其输入。

下图中说明了整体网络架构。

下面的文本将介绍对所涉及元素的初步解释，进一步的详细信息将在附录中给出。

图构建

首先将 $\mathcal{F}$​​（即指定 PDE 形式的符号信息）表示为计算图。在这样的计算图中，节点可以代表未知场变量（记为UF）、标量系数（SC）、初始条件（IC）以及微分或代数运算，并且有向边可以是用于指定运算中涉及的操作数。这将构成具有异构节点和同构边的有向无环图（DAG）。

上图给出了计算图构建的说明。对 PDE 形式的计算图表示还有以下评论：

• 该计算图中仅涉及少量节点类型：UF（未知场变量）、SC（标量系数）、IC（初始条件）、dt、dx（分别相对于 t 和 x 的微分）、+ ( sum), × (乘积), − (负), $(\cdot {)}^2$ (平方), = 0 (等于零)。
• 请注意，“-”在计算图中代表求反而不是减法，使其成为一元运算而不是二元运算。这是因为该计算图中的节点是无序的，并且边是同构的。如果计算图中涉及二进制减法运算，无法区分减数和被减数。
• 出于同样的原因，不包括幂操作节点。然而，注意到可以表达具有正整数指数的幂。例如，$u^3=u\times u^2$，并且 $u^{11}=((u^2{)}^2{)}^2\times u^2\times u$，因为 $11=2^3+2^1+2^0$。
• 虽然实验中没有涉及，但也可以引入表示特殊函数（例如 sin、cos、exp、log）的节点类型，具体取决于所涉及的偏微分方程的形式。这也可以表达一般的幂运算，因为有 $a^b=\exp (b\times \log (a))$​。
• 类型为 UF 和 SC 的节点的入度为零，+ 和 × 的入度等于或大于 2，所有其余节点的入度恰好为 1 。

然后，为了包含数字信息，为每个图节点赋予 ${\mathbb{R}}^{d_f}$ 中的特征向量。对于标量系数 $c$，该值重复 $d_f$ 次以形成相应 SC 节点的特征向量。就初始条件 $g(x)$ 而言，假设它是在具有 $n_x$ 个点的等距网格上给出的。受 ViT 的启发，将这些网格值划分为 $N=n_x/d_f$ 块，产生 $N$ 个向量 ${\mathbf{g}}_1,\dots ,{\mathbf{g}}_N\in {\mathbb{R}}^{d_f}$ 这些将用作 $N$ 个新引入的“补丁”节点的特征向量，其类型分别表示为 ${\mathfrak{p}}_1,{\mathfrak{p}}_2,\dots ,{\mathfrak{p}}_N$ 。将这些补丁节点与相应的 IC 节点连接起来。所有剩余节点的特征向量被设置为零。

• 就辅助节点而言，让每个补丁节点 $p_i$ 从相应的 IC 节点接收一条边，并且每个潜在调制节点 $m_{\ell }$ 向 UF 发出一条边。采用这种边缘方向约定是为了提高最终 DAG 的连接性，因为将屏蔽图 Transformer 模块中断开的节点对之间的注意力

在实验中，初始条件 $g(x)$ 在 $n_x=256$ 个点的空间网格上离散化，并将值划分为 $N=16$ 个长度 $d_f=16$ 的块。与 $N=4$ 或 $N=1$ 的情况相比，这表现出更好的求解精度。对于 INR 解码器的外部条件，观察到每个隐藏层使用不同的潜在向量 ${\mu }^1,\dots ,{\mu }^L\in {\mathbb{R}}^{d_e}$ 与使用共享的 $\mu \in {\mathbb{R}}^{d_e}$​ 性能提高。因此辅助节点的引入被认为是有意义的。

此外，引入了 $L$ 个类型为 ${\mathrm{m}}_1,{\mathrm{m}}_2,\dots ,{\mathrm{m}}_L$​ 的附加节点，并将它们连接到相应的 UF 节点。这些节点将用于解码预测的解决方案。

编码图数据

封装在图数据中的符号和数字信息被集成到潜在代码 $\mathbf{\mu }=[{\mu }^1,\dots ,{\mu }^L{]}^{\mathrm{T}}\in {\mathbb{R}}^{L\times d_e}$ 中，这是通过图 Transformer 实现的，图 Transformer 是一类基于 Transformer 的现代图神经网络，具有令人印象深刻的表示能力。Graphormer 的改编版本被用作实验中的特定架构，而更多潜在的替代方案可以在 Min 等人中找到。对于 $\ell =1,\dots ,L$，令 ${\mu }^{\ell }\in {\mathbb{R}}^{d_e}$ 为分配给该图Transformer的输出层中类型为 $m_{\ell }$ 的节点的嵌入向量。

解码 PDE 解

采用 INR，将坐标 $(t,x)$ 作为输入，并根据 $\mathbf{\mu }$ 生成无网格预测 $\hat{u}(t,x)$。具有此类外部条件的各种 INR 架构已在神经算子、数据压缩和生成模型中采用。在实验中，由于其效率，使用具有 $L$ 个隐藏层的 Poly-INR 的改编版本，并且第 $\ell$ 个隐藏层的调制是基于 ${\mu }^{\ell }$​ 生成的。

结果

训练设置

上表概述了实验设置，包括模型超参数和配置。为了全面了解实验中使用的基线模型，提供了所有模型的概述：

• DeepONet：DeepONet 采用独特的架构，具有两个子网络：分支网络和主干网络。分支网络处理固定数量的传感器观测值（在我们的例子中，初始条件为 256 个点），而主干网络处理用于推理的坐标输入，类似于 PDEformer 的输入机制。两个网络的输出相结合以产生解决方案值。每个子网络由具有 256 个隐藏神经元的六层 MLP 组成，并利用 ReLU 激活函数。值得注意的是，DeepONet 的无网格特性允许使用分散的数据点进行训练，使我们能够在 DeepONet 的训练和 PDEformer 的微调过程中为每个数据样本每次迭代从 256 × 256 网格中采样 8192 个点。
• FNO：傅立叶神经算子 (FNO) 根据网格相关但与分辨率无关的原理进行操作。它首先通过逐点全连接层将规则网格数据转换为多通道隐藏特征，然后通过几个傅里叶层进行处理，最后映射到解网格。在傅里叶层中，FNO 保留最低的 12 个傅里叶模式。在我们的实验中，使用FNO2D模型，将初始条件（256个空间点）扩展到形成256×256输入网格，允许同时全场输出。
• U-Net：U-Net采用基于CNN的编码器-解码器框架，其特点是其4层下采样和上采样卷积，通过中间残差连接桥接。与 FNO2D 类似，输入和输出尺寸均设置为 256 × 256。与无网格 DeepONet 或 PDEformer 不同，FNO 和 U-Net 需要以规则网格组织的训练数据。
• PDEformer：基于 Transformer 的 Graphhormer 配置有 9 层、512 维嵌入空间和 32 个注意力头。 Poly-INR 部分采用具有 256 个神经元的 L = 8 个隐藏层，每个隐藏层都使用单独的缩放和移位超网络进行动态调制，每个隐藏层由具有独立参数的 3 层 MLP 组成。

在 PDEformer 的预训练阶段，我们采用归一化均方根误差（nRMSE）损失函数，因为它可以有效提高训练效率。实施学习率计划，在预定时期逐步降低学习率，以提高训练过程的稳定性。此外，在训练开始时利用预热期，通过将学习率从零逐渐增加到初始预定值来降低早期训练失败的风险。

预训练阶段

生成一个包含 50 万个样本的数据集，按方程类型、系数和初始条件进行区分。具体来说，所解决的偏微分方程遵循以下形式：$u_t+f_0(u)+f_1(u{)}_x-\nu u_{xx}=0,(t,x)\in [0,1]\times [-1,1]$，具有周期性边界和初始条件 $u(0,x)=g(x),x\in [-1,1]$，其中 $f_i(u)={\sum }_{k=0}^3{c}_{ik}{u}^k$ 对于 $i=0,1$。使用随机生成的 $c_{ik},\nu$ 以及 $g(x)$ ，利用Dedalus 包求解相应的偏微分方程来创建预训练数据集。预训练涉及 90% 数据的 1,000 个 epoch，保留剩余 10% 用于测试。 PDEformer 在训练数据集上实现了 0.0104 的相对 $L^2$ 误差，在测试数据集上实现了 0.0128 的相对 $L^2$​​​ 误差。

所有实验均使用MindSpore5 2.0进行，涉及1, 000个epoch的预训练在8个NPU上大约需要79小时（如果考虑内部测试评估则需要84小时）。下图说明了PDEformer的预训练过程。

下图说明了预训练的 PDEformer 对测试数据集的预测，强调了其在学习不同 PDE 表示方面的高精度和熟练程度。

正问题

预训练的 PDEformer 在处理各种方程方面具有高度通用性。它在前向问题上的性能是使用 PDEBench 中的参数偏微分方程进行评估的，包括 Burgers’、Advection 和 1D Reaction-Diffusion PDE。使用针对各个偏微分方程定制的神经算子模型进行比较分析。在这种情况下，神经算子接收初始条件并预测整个解域。值得注意的是，之前的所有方法都是在不同的数据集上单独训练和测试的，而 PDEformer 模型在预训练后的所有测试数据集上显示了零样本推理。此外，PDEformer-FT 模型涉及对相应 PDEBench 数据集的额外微调过程。

在上表中，预训练的 PDEformer 模型通过在所有分布内测试中获得合理的精度来展示零样本熟练程度。值得注意的是，对于 $\nu =0.1$ 和 $0.01$​ 的 Burgers 方程，零样本 PDEformer 优于经过充分训练的 FNO 模型。这种卓越的性能源于预训练的 PDEformer 在预训练期间接触到不同的 PDE，模型从中学习了一个可推广的规律，从而优于专门针对单个 PDE 训练的模型。经过微调的 PDEformer 在所有测试中始终表现出色，进一步凸显了方法在解决各种偏微分方程时的稳健性和多功能性。

为了比较微调 PDEformer 与训练传统神经算子的效率，对具有有限数量的 100 个训练样本的 Advection 方程（$\beta =1$​，OoD）数据集进行了比较分析。如下图所示，PDEformer 只需大约 100 次迭代即可快速收敛。相反，从头开始训练的 FNO 模型需要数千次迭代才能收敛，并导致较高的测试误差。这种对比凸显了 PDEformer 快速而准确的适应性，标志着相对于神经算子的显着进步。

在对前向问题的研究中，我们展开了详细的研究，以在有限的数据下评估模型的学习效率。具体来说，我们将微调训练数据集大小从 9k 减少到 100 和 1k 样本。

如上图所示，经过微调的 PDEformer 模型非常出色，优于所有其他方法。此外，零样本 PDEformer 建立了值得称赞的高基准，无需任何微调即可展现强大的性能。特别值得注意的是，在 OoD 条件下，例如在平流 ($\beta =1$​) 和反应扩散场景中，经过微调的 PDEformer 迅速获得了优异的结果。这凸显了该模型适应不熟悉场景的小样本学习能力。

反问题

除了正向问题之外，还可以利用预训练的 PDEformer 来基于一个有噪声的观察到的解决方案实例来解决 PDE 系数恢复逆问题。对于每个 PDE，将 PDE 系数的当前估计输入到预训练的 PDEformer 中以获得预测解，并最小化相对于观测值的相对 $L^2$​ 误差以获得恢复的系数。由于该优化问题表现出许多局部极小值，因此采用了粒子群优化算法。

上图说明了涉及前文中解释的测试集的 40 个偏微分方程的结果。每个方程要恢复的系数数量各不相同（范围从 1 到 7）。在没有噪声的情况下，恢复的系数与真实值一致，分散点主要沿着 $y=x$​​ 线分布。一些异常值的存在可以归因于这个逆问题的固有不适定性质。即使在高噪声水平下，大部分偏微分方程系数也能得到有效恢复。

就反问题而言，每个网格点的加性噪声值是从 $U([-r\Vert u{\Vert }_{L^{\infty }},r\Vert u{\Vert }_{L^{\infty }}])$ 中随机采样的，其中 $u$ 是没有噪声的真解，$r$ 是噪音水平。

上表显示了 40 个随机样本中三个 PDE 的恢复系数，相应的噪声观测值和 PDEformer 预测如下图所示。请注意，PDEformer 的输入是恢复的 PDE 系数而不是真实值。结果表明，即使在观测有噪声的情况下，PDEformer 在正向和逆向问题中也具有良好的准确性。

推理时间

上表展示了一系列模型（包括 DeepONet、FNO、U-Net 和 PDEformer）的参数数量、每个样本的推理时间和预测精度的比较。我们还包括两个传统数值求解器的结果。前者是基于一阶迎风有限差分 (FD) 方案，利用 SciPy Python 包提供的solve_ivp 函数，后者是 Dedalus，用于生成真实解数据的基于谱方法的求解器。使用 256 × 256 时空网格上的一维平流方程（$\beta =1.0$​）作为测试用例进行评估，神经网络模型在单个 NPU 上测试，传统求解器在 CPU 上执行。神经网络模型在相应的数据集上进行了充分的训练，并且在测试过程中将批量大小标准化为 10。我们对所有样本中每个模型的总时间消耗进行平均，以显示每个样本的推理时间。由于FD求解器的精度较低，其求解过程中空间网格分辨率被细化为16×256=4096。

虽然比较显示 Dedalus 的推理时间明显更长，但必须承认计算平台和模型本身性质的固有差异。这种并置虽然并不严格公平，但旨在说明机器学习方法在解决偏微分方程方面的潜在效率。

其它细节

数据集

预训练数据集

数据集由以下形式的偏微分方程的解组成：
$$\begin{gathered} u_t+f_0(u)+f_1(u{)}_x-\nu u_{xx}=0,(t,x)\in [0,1]\times [-1,1], \\ u(0,x)=g(x),x\in [-1,1], \end{gathered}$$
其中 $f_i(u)={\sum }_{k=0}^3{c}_{ik}{u}^k$ 其中 $i=0,1$。每个系数 $c_{ik}$ 以 0.5 的概率设置为零，否则从 $U([-3,3])$ 中随机抽取。粘度 $\nu$ 满足$\log \nu \sim U([\log 10^{-3},\log 1])$。对于线性通量的情况，即当 $c_{12}=c_{13}=0$ 时，以 0.5 的概率设置 $\nu =0$。请注意，偏微分方程的计算图中将排除系数为零的项。随机初始条件 $g(x)$​ 的生成方式与 PDEBench 数据集相同。

数值解是使用名为 Dedalus v3的开源 Python 包获得的，这是一个基于谱方法的灵活求解器。为了生成数据样本，使用具有 256 个网格点的统一空间网格。求解器以 $\delta t_{\mathrm{solver}}=4\times 10^{-4}$ 的时间步长进行，并以时间步长 $\delta t_{\mathrm{data}}=0.01$ 记录解快照，为每个数据样本生成总共 101 个时间值。当Dedalus无法求解偏微分方程时，或者当解的 $L^{\infty }$ 范数超过10时，相应的数据样本将被丢弃，并且不包含在最终数据集中。

由于偏微分方程具有周期性边界条件，并且在均匀网格点上离散化，因此我们在预训练阶段通过沿 x 轴的随机平移引入数据增强。对于每个数据实例，从101×256网格中随机采样总共8192个时空坐标点。这些采样点作为解决方案解码器INR的输入，我们将模型预测与真实数值进行比较以计算损失。

PDEBENCH数据集

在本小节中，我们概述了从 PDEBench 派生的三个一维 PDE 数据集，我们在实验分析中使用了这些数据集。每个数据集根据特定的 PDE 类型和系数配置定制，包含 10k 个实例。出于训练目的，我们使用每个数据集中的 9k 个样本，保留剩余的 1k 个样本用于测试。值得注意的是，PDEBench 中的每个数据集都遵守周期性边界条件。

• Burgers’ 方程： ${\partial }_{t}u+{\partial }_x(u^2)=\frac{\nu }{\pi }{\partial }_{xx}u$ 其中 $(t,x)\in [0,2]\times [-1,1]$, ν ∈ { 0.1 , 0.01 , 0.001 } \nu \in \{0.1,0.01,0.001\} ν∈{0.1,0.01,0.001} 该方程是流体力学的基本偏微分方程。
• 平流方程：${\partial }_{t}u+\beta {\partial }_{x}u=0$ 其中 $(t,x)\in [0,2]\times [0,1]$， β ∈ { 0.1 , 1 } β \in \{0.1, 1\} β∈{0.1,1}。该方程模拟了量 $u$​ 的传输，且不改变其形式。
• 反应扩散方程： ${\partial }_{t}u=\nu {\partial }_{xx}u+\rho u(1-u)$ 其中 $(t,x)\in [0,1]\times [0,1]$，其中我们只考虑 $\nu =1,\rho =1$​。方程表示化学反应和扩散动力学相结合的过程。

每个数据集的初始条件由 $u_0(x)={\sum }_{k_i=k_1,...,k_N}{A}_i\sin (k_{i}x+{\varphi }_i)$ 给出，频率数 $k_i=\frac{2\pi n_i}{L_x}$，其中 $n_i$ 是在预训练中在预定范围随机选择的整数。$L_x$ 是空间域的长度，幅度 $A_i$ 是 $[0,1]$ 内的随机数，相位 ${\varphi }_i$ 从区间 $(0,2\pi )$​ 中随机选择。随后以 10% 的概率应用具有随机签名的绝对值函数，以及通过乘以窗口函数来限制随机子区间。对于反应-扩散方程，初始条件的范围重新调整为单位区间 [0, 1]。

为了利用预训练的 PDEformer 模型进行预测，我们将时空坐标重新缩放到范围 $(t\prime ,x\prime )\in [0,1]\times [-1,1]$，并将所得的 PDE 作为输入PDEformer 具有以下形式：

• Burgers’ 方程： ${\partial }_{t^{\prime }}u+{\partial }_{x^{\prime }}(2u^2)-\frac{2\nu }{\pi }{\partial }_{x^{\prime }{x}^{\prime }}u=0,$ 其中 $t^{\prime }=t/2,x^{\prime }=x$。
• 平流方程：${\partial }_{t^{\prime }}u+{\partial }_{x^{\prime }}(4\beta u)=0$ 其中 $t^{\prime }=t/2,x^{\prime }=2x-1$。
• 反应扩散方程： ${\partial }_{t^{\prime }}u-4\nu {\partial }_{x^{\prime }{x}^{\prime }}u+(-\rho )u+\rho u^2=0$ 其中 $t^{\prime }=t,x^{\prime }=2x-1$​。

为了确保基线模型之间的公平比较，我们将所有 PDEBench 样本的分辨率标准化为 256 × 256。更具体地说，原始 PDEBench 数据集的空间分辨率为 1024，下采样为 256。原始记录的时间步数Burgers 和 Advection 数据集的值为 201，一维反应扩散数据集的值为 101，并利用线性插值获得 256 的时间分辨率。值得注意的是，PDEformer 进行无网格预测，使我们能够将预训练数据集的时间分辨率设置为 101，将 PDEBench 数据集的时间分辨率设置为 256。对于FNO和U-Net，初始值重复256次，形成分辨率为256×256的二维数据，然后作为网络输入。

网络架构及训练设置

图 Transformer 架构

实验中使用的特定图 Transformer 架构基于 Graphormer（Ying 等人，2021），并进行了一些调整以适应我们的设置。详情如下。

初始嵌入向量

在图中 Transformer 中，节点 i 的初始嵌入向量为：
$$h_i^{(0)}=x_{\text{type}(i)}+\text{Feat-Enc}(f_i)+z_{{\text{deg}}^{-}(i)}^{-}+z_{{\text{deg}}^{+}(i)}^{+}$$
其中 $x,z^{-},z^{+}\in {\mathbb{R}}^{d_e}$ 分别是由节点类型 $\mathrm{type}(i)$、入度 ${\mathrm{deg}}^{-}(i)$ 和出度 ${\mathrm{deg}}^{+}(i)$ 指定的可学习嵌入向量。为了对原始 Graphomer 中未涉及的节点特征向量 $f_i\in {\mathbb{R}}^{16}$​ 进行编码（Ying 等人，2021），我们利用特征编码器 Feat-Enc，它是一个三层多层感知器（MLP），每个隐藏层有 ReLU 激活和 256 个神经元。

注意偏差

表示 $\varphi (i,j)$ 为从节点 $i$ 到节点 $j$ 的最短路径长度。如果这样的路径不存在，或者长度大于14，我们将设置 $\varphi (i,j)=14$。对于图 Transformer 中涉及的每个注意力头，对应于节点对 $(i,j)$ 给出为：
$$B_{ij}=b_{\varphi (i,j)}^{+}+b_{\varphi (j,i)}^{-}+d_{ij}.$$
这里，$b_{\varphi (i,j)}^{+}$ 和 $b_{\varphi (i,j)}^{-}$ 是分别由 $\varphi (i,j)$ 和 $\varphi (j,i)$ 索引的可学习标量，并在所有层之间共享。引入附加项 $d_{ij}$（未出现在原始 Graphomer 中）是为了屏蔽断开连接的节点对之间的注意力。更具体地说，当节点 $i$ 和节点 $j$ 在图中是相连的，即存在从 $i$ 到 $j$ 或从 $j$ 到 $i$ 的路径，我们取 $d_{ij}=0$，否则设置 $d_{ij}=-\infty$​。我们在实验中观察到，通过这种额外的掩蔽操作可以提高整体预测精度。此外，由于我们的图具有同质边，因此我们没有引入原始 Graphomer 中出现的边编码术语。

图转换器层

图 Transformer 层的结构与原始 Graphormer 相同，为了方便读者，我们将其放在这里。每一层都采用以下形式：
$$\begin{gathered} {\bar{h}}^{(l)}=\mathrm{Attn}(\mathrm{LN}(h^{(l-1)}))+h^{(l-1)} \\ {h}^{(l)}=\mathrm{FFN}(\mathrm{LN}({\bar{h}}^{(l)}))+{\bar{h}}^{(l)} \end{gathered}$$
其中 FFN 表示具有单个隐藏层和 GeLU 激活函数的位置前馈网络，LN 表示层归一化。在自注意力块 Attn 方面，我们将遵循原始 Graphormer 论文中的约定，并且为了简单起见仅呈现单头情况。令 $H=[h_1^{\prime },\cdots ,h_n^{\prime }{]}^{\mathrm{T}}\in {\mathbb{R}}^{n\times d_e}$ 表示涉及 $n$ 个图节点的自注意力模块的输入，自注意力计算为：
$$\begin{gathered} Q=H{W}_Q,K=H{W}_K,V=H{W}_V, \\ A=\frac{Q{K}^{\mathrm{T}}}{\sqrt{d_e}}+B,\mathrm{Attn}(H)=\mathrm{softmax}(A)V \end{gathered}$$
其中 $W_Q,W_K,W_V\in {\mathbb{R}}^{d_e\times d_e}$ 是投影矩阵，$B$​ 是前文中给出的注意力偏差。多头注意力的扩展是标准且直接的。

进一步的实施细节

实验中的图 Transformer 包含 9 层，嵌入维度 $d_e=512$ 和 32 个自注意力头。 FFN模块的隐藏层的宽度等于$d_e$。此外，为了简化实现，我们没有在原始Graphormer中包含特殊节点 [VNode] 。

INR架构

在隐式神经表示 (INR) 领域，数据样本被解释为基于坐标的函数，其中每个函数接受坐标 $(t,x)$ 作为输入，并在该特定坐标点生成近似函数值 $\hat{u}(t,x)$ 。在实验中，我们使用了 Poly-INR 的改编版本（Singh 等人，2023），与我们设置中的其他候选者相比，它表现出更好的预测准确性和训练稳定性。受 COIN++ (Dupont et al., 2022) 的启发，我们还采用了 $L$ 个超网络，其中第 $\ell$ 个超网络将 ${\mu }^{\ell }\in {\mathbb{R}}^{d_e}$ 作为输入，并为Poly-INR 的第 $\ell$​ 个隐藏层生成尺度和移位调制。

我们的 INR 解码器的复杂架构如上图所示，数学框架详述如下。我们将 $h_0=1$ 视为所有条目都等于 1 的向量。对于 $\ell =1,2,\dots ，L$​，我们计算
$$\begin{array}{rl}& g_{\ell }=W_{\ell }^{\mathrm{in}}\left[\begin{array}{c}t\\ x\end{array}\right]+b_{\ell }^{\mathrm{in}},s_{\ell }^{\mathrm{scale}}={\mathrm{MLP}}_{\ell }^{\mathrm{scale}}({\mu }^{\ell }),s_{\ell }^{\mathrm{shift}}={\mathrm{MLP}}_{\ell }^{\mathrm{shift}}({\mu }^{\ell }),\\ & q_{\ell }=s_{\ell }^{\mathrm{scale}}\odot \left(W_{\ell }^{\mathrm{h}}\left(h_{\ell -1}\odot g_{\ell }\right)+b_{\ell }^{\mathrm{h}}\right)+s_{\ell }^{\mathrm{shift}},h_{\ell }=\sigma \left(q_{\ell }\right),\end{array}$$
网络输出为 $\hat{u}(t,x)=W^{\mathrm{Last}}{h}_L+b^{\mathrm{Last}}$。这里，激活函数 $\sigma (\cdot )$ 是在负输入范围内斜率为 0.2 的 leaky-ReLU 操作，然后在区间 [−256, 256] 中进行裁剪操作以提高训练稳定性。超网对应于 ${\mathrm{MLP}}_{\ell }^{\mathrm{scalc}}$ 和 ${\mathrm{MLP}}_{\ell }^{\mathrm{shift}}$ 。请注意，在原始 Poly-INR 中，超网络用于生成 $W_{\ell }^{\mathrm{in}}$ 和 $b_{\ell }^{\mathrm{in}}$。与我们生成 $s_{\ell }^{\mathrm{scale}}$ 和 $s_{\ell }^{\mathrm{shift}}$​ 的做法相比，该方法表现出更好的精度，但降低了训练效率，因此在我们的实验中没有采用。

总结

将方程的符号形式通过计算图的形式嵌入到神经网络中，利用图 Transformer 和 INR 生成预测解。

PINN主要针对一个确定的PDE问题进行求解，DeepONet则是对一类算子问题进行求解，本文似乎想将任意PDE都纳入到同一个网络之中。但在具体预训练是使用的依然是预定义形式的PDE，当然，这里预定义的形式会比算子学习中针对的一类算子问题更加复杂，可以包含Burgers’、Advection 和 1D Reaction-Diffusion 等PDE的形式。但毕竟训练数据集总是有限的，或许可以讨论一下如何随机生成PDE？文章中虽然讨论了OoD的情况，但主要是针对SC和g(x)，个人认为，方程形式的OoD或许更有趣些。看结果的话，确实会比其他算子学习更好一些，但毕竟模型参数大太多了，我也不好说这是把方程嵌入的好处，还是模型参数大的好处。如果可以随机采样PDE，而不是按固定形式随机采样参数的话，感觉PDEformer的形式有成为PDE方向大模型的潜力啊。

相关链接：

• 原文：[2402.12652] PDEformer: Towards a Foundation Model for One-Dimensional Partial Differential Equations (arxiv.org)
• 原文代码：mindscience: MindScience is scientific computing kits for various industries based on the converged MindSpore framework. - Gitee.com