LLT: 用于PDE算子学习的局部线性Transformer

# LLT：用于偏微分方程算子学习的局部线性Transformer
来源：https://arxiv.org/html/2607.07718
Oded Ovadia通讯作者\. 邮箱：odedovadia@mail\.tau\.ac\.il
特拉维夫大学数学科学学院，以色列特拉维夫
Eli Turkel邮箱：turkel@tauex\.tau\.ac\.il
特拉维夫大学数学科学学院，以色列特拉维夫

###### 摘要

神经算子已成为学习PDE解映射和加速数值模拟的常用方法。基于Transformer的神经算子尤其引人关注，因为注意力机制可以学习计算域中的长程依赖关系。然而，标准注意力在应用于PDE时存在两个主要局限性：其计算复杂度随计算节点数量呈二次增长，且缺乏对局部交互的显式偏置。为解决这些问题，我们提出了用于PDE算子学习的局部线性Transformer（LLT）。该架构将线性全局注意力与局部空间混合相结合，并融入了坐标和几何信息。我们在多个PDE问题上评估了LLT，包括弹性、塑性、翼型流动、管道流动和达西流动。这些问题的参考数据涵盖了结构化网格和非结构化网格上的有限元、有限体积和有限差分离散格式。与先前研究中的其他神经算子和Transformer基线相比，LLT在这些问题上实现了具有竞争力或更低的相对L2L\_\{2\}误差。在匹配的结构化离散格式上，相对于Transolver，每个训练迭代的挂钟时间减少了1.8到2.5倍。我们还将该方法进行了扩展，并应用于一个三维汽车空气动力学数据集，每个样本包含32,186个非结构化网格点。这些结果共同表明，LLT为跨离散格式、网格类型和问题设置的PDE问题提供了一种准确且计算高效的算子。

关键词：偏微分方程，神经算子，Transformer，科学机器学习

## 1 引言

近年来，机器学习（ML）方法在科学计算中的应用迅速增长，已出现许多成功的数学问题建模方法（Raissi等人，2019 (https://arxiv.org/html/2607.07718#bib.bib43)；Lu等人，2021 (https://arxiv.org/html/2607.07718#bib.bib2)；Li等人，2021 (https://arxiv.org/html/2607.07718#bib.bib12)；Long等人，2018 (https://arxiv.org/html/2607.07718#bib.bib44)；Xu等人，2019 (https://arxiv.org/html/2607.07718#bib.bib45)；Takamoto等人，2022 (https://arxiv.org/html/2607.07718#bib.bib46)；Gupta and Brandstetter，2022 (https://arxiv.org/html/2607.07718#bib.bib47)）。这些方法已在许多不同领域展现出潜力，包括计算力学（Cai等人，2021 (https://arxiv.org/html/2607.07718#bib.bib49)；Zhang等人，2022 (https://arxiv.org/html/2607.07718#bib.bib50)）、波传播（Ovadia等人，2021 (https://arxiv.org/html/2607.07718#bib.bib48)，2024b (https://arxiv.org/html/2607.07718#bib.bib32)）、材料科学（Dingreville等人，2020 (https://arxiv.org/html/2607.07718#bib.bib51)；Oommen等人，2022 (https://arxiv.org/html/2607.07718#bib.bib52)）、流体动力学（Sharma等人，2023 (https://arxiv.org/html/2607.07718#bib.bib55)；Zhao等人，2024 (https://arxiv.org/html/2607.07718#bib.bib56)）以及湍流（Wu等人，2018 (https://arxiv.org/html/2607.07718#bib.bib53)；Wang等人，2020 (https://arxiv.org/html/2607.07718#bib.bib54)）。

在许多应用中，同一类PDE需要针对变化的系数、几何形状、初始条件或边界条件反复求解。经典求解器在准确性方面仍是参考方法，但当需要数千次相关求解时，它们可能代价高昂。这促使了神经算子的发展，它们可以直接从数据中学习解映射。神经算子通过学习输入函数与输出解场之间的映射来提供这样一种框架（Kovachki等人，2023 (https://arxiv.org/html/2607.07718#bib.bib1)；Azizzadenesheli等人，2024 (https://arxiv.org/html/2607.07718#bib.bib10)）。

特别地，Transformer架构是算子学习的一个有前途的候选，因为其注意力机制可以学习计算域中远距离点之间的信息（Vaswani等人，2017 (https://arxiv.org/html/2607.07718#bib.bib25)；Li等人，2023a (https://arxiv.org/html/2607.07718#bib.bib19)；Hao等人，2023 (https://arxiv.org/html/2607.07718#bib.bib6)）。在PDE问题中，边界条件、材料界面和椭圆耦合可能会引入跨域的非局部依赖性（Kovachki等人，2023 (https://arxiv.org/html/2607.07718#bib.bib1)；Azizzadenesheli等人，2024 (https://arxiv.org/html/2607.07718#bib.bib10)）。然而，Transformer在此设置中也面临重要困难。首先，自注意力的计算复杂度为O(N2)\mathcal{O}(N^{2})（N为空间点数），这在高分辨率网格上代价高昂（Vaswani等人，2017 (https://arxiv.org/html/2607.07718#bib.bib25)；Katharopoulos等人，2020 (https://arxiv.org/html/2607.07718#bib.bib26)；Choromanski等人，2021 (https://arxiv.org/html/2607.07718#bib.bib27)；Wu等人，2024 (https://arxiv.org/html/2607.07718#bib.bib23)）。其次，标准注意力没有内置对局部交互的偏置。许多PDE解包含强烈的局部结构，经典的数值方法通常通过模板或单元来利用这种结构（Hughes，2012 (https://arxiv.org/html/2607.07718#bib.bib39)；Godunov and Bohachevsky，1959 (https://arxiv.org/html/2607.07718#bib.bib40)；Eymard等人，2000 (https://arxiv.org/html/2607.07718#bib.bib41)）。这种局部性对于双曲型PDE尤其重要，其中信息沿特征线传播，且短时间间隔内的解仅取决于域中有限区域（LeVeque，2002 (https://arxiv.org/html/2607.07718#bib.bib42)）。因此，纯粹的全局注意力机制可能会将计算花费在局部更新中微弱或无关的交互上。

这些需求引出了特定的架构设计。基于Transformer的算子应在采样的域内进行通信，同时强调解场中存在的局部结构。它还必须随节点数量良好扩展，以便在现实网格上运行。最后，它应同时适用于结构化网格和非结构化节点集，因为算子学习基准中的参考数据可能来自有限元、有限体积和有限差分离散格式。这促使我们设计一个模型，将高效的域范围通信与显式的局部混合路径相结合，并使用线性注意力来使注意力成本接近点数的线性增长（Katharopoulos等人，2020 (https://arxiv.org/html/2607.07718#bib.bib26)；Choromanski等人，2021 (https://arxiv.org/html/2607.07718#bib.bib27)）。

我们提出LLT（LocalLinearTransformer）用于监督式PDE算子学习。该模型接收在有限点集上采样的输入场，并预测相同点处的相应输出场。它使用核化线性注意力进行全局通信，并为空间邻域提供局部混合路径。它还包括坐标编码、到参考网格的距离编码以及跳跃连接的解码器。这些组件共同解决了上述Transformer的两个主要局限性，即降低了注意力成本并增加了对局部交互的显式偏置。

我们在五个广泛使用的PDE问题上评估了LLT，涵盖弹性、塑性、翼型、管道流动和达西流动（Li等人，2023c (https://arxiv.org/html/2607.07718#bib.bib4)；Wu等人，2024 (https://arxiv.org/html/2607.07718#bib.bib23)；Luo等人，2025 (https://arxiv.org/html/2607.07718#bib.bib24)）。这些问题的参考数据由有限元、有限体积和有限差分求解器生成（第4节 (https://arxiv.org/html/2607.07718#S4)）。我们还将该模型应用于三维汽车空气动力学数据集，以测试其在大型非结构化网格上的扩展性。计算研究提供了五个基准问题上的相对L2L\_\{2\}误差，以及在结构化离散格式上单独的匹配时序和内存比较。这些结果用于评估局部-全局注意力设计的准确性和计算成本。我们发现，与其他领先的神经算子方法相比，LLT在这些问题上取得了准确的结果。计算成本分析进一步表明，在相同大小的匹配结构化网格上，LLT扩展良好，并且能够保持相对快速和高效。

参见图注图1：LLT架构。NN个点的输入坐标和场值通过基于距离、傅里叶和卷积的分支通道并行编码，连接后由多层感知机（MLP）投影为每个点的固定宽度特征向量。然后，LL个Transformer编码器块的堆栈精炼此表示；每一层对其输入进行归一化（RMSNorm），应用跨所有点的全局信息交换步骤以及局部空间混合步骤，并通过残差路径将结果加回输入，随后进行第二次归一化和逐点非线性变换（SwiGLU）。解码器将处理后的输出与初始嵌入的直接连接合并，并通过最终的MLP映射以产生预测的解场。
## 2 相关工作

#### 神经算子

神经算子学习函数空间之间的映射，例如偏微分方程的解算子（Azizzadenesheli等人，2024 (https://arxiv.org/html/2607.07718#bib.bib10)）。给定输入函数（如系数、源项或边界及初始条件），神经算子预测相应的解场（Kovachki等人，2023 (https://arxiv.org/html/2607.07718#bib.bib1)）。DeepONet是这种方法的一个早期例子，它使用分支-主干架构来编码输入函数和空间查询位置（Lu等人，2021 (https://arxiv.org/html/2607.07718#bib.bib2)）。

另一类神经算子基于谱表示。傅里叶神经算子（FNOs）使用谱卷积在规则网格上进行高效的全局混合，是PDE建模的标准基线（Li等人，2021 (https://arxiv.org/html/2607.07718#bib.bib12)）。相关工作已将算子学习扩展到更复杂的离散格式和几何形状，包括用于不规则网格的基于图的构造（Li等人，2020 (https://arxiv.org/html/2607.07718#bib.bib3)）以及基于变形的方法，这些方法在应用谱算子之前将一般域映射到潜在规则网格（Li等人，2023c (https://arxiv.org/html/2607.07718#bib.bib4)）。混合公式还可以增加物理结构，例如通过在训练期间强制执行PDE残差（Li等人，2024 (https://arxiv.org/html/2607.07718#bib.bib5)）。

#### 用于PDE的Transformer

Transformer通过注意力机制建模长程交互，并且比绑定到固定网格的架构更直接地处理不规则采样（Vaswani等人，2017 (https://arxiv.org/html/2607.07718#bib.bib25); Hao等人，2023 (https://arxiv.org/html/2607.07718#bib.bib6)）。例如，OFormer使用观测输入样本与查询点之间的自注意力和交叉注意力来公式化算子学习，减少了对网格特定归纳偏置的依赖，同时在标准PDE测试问题上取得了强劲性能（Li等人，2023a (https://arxiv.org/html/2607.07718#bib.bib19)）。后续工作开发了用于多尺度结构和异构离散格式的注意力变体（Hao等人，2023 (https://arxiv.org/html/2607.07718#bib.bib6)）。

其他工作专注于混合方法，将Transformer与其他组件结合。例如，ViTO将U-Net编码器与视觉Transformer块配对，以捕获局部和全局结构（Ovadia等人，2024a (https://arxiv.org/html/2607.07718#bib.bib7); Taccari等人，2023 (https://arxiv.org/html/2607.07718#bib.bib9)）。TC-UNet将此设计扩展到时间相关算子，通过使用Transformer注意力将U-Net骨干网络基于时间进行条件化，从而在测试时无需时间离散化即可实现时间连续推理（Ovadia等人，2025 (https://arxiv.org/html/2607.07718#bib.bib8)）。另一种方法是将神经算子与经典数值方法相结合，以在大尺度和小尺度上实现收敛（Zhang等人，2024 (https://arxiv.org/html/2607.07718#bib.bib21); Ovadia等人，2024b (https://arxiv.org/html/2607.07718#bib.bib32)）。其他工作则专注于将物理信息方面集成到Transformer架构和训练过程中（Zhao等人，2023 (https://arxiv.org/html/2607.07718#bib.bib35); Lorsung等人，2024 (https://arxiv.org/html/2607.07718#bib.bib36)）。

基于Transformer的PDE模型的一个关键挑战是全局注意力的二次成本（Vaswani等人，2017 (https://arxiv.org/html/2607.07718#bib.bib25); Katharopoulos等人，2020 (https://arxiv.org/html/2607.07718#bib.bib26); Choromanski等人，2021 (https://arxiv.org/html/2607.07718#bib.bib27); Wu等人，2024 (https://arxiv.org/html/2607.07718#bib.bib23)）。这在空间分辨率高时尤其麻烦，因此最近的方法强调可扩展的注意力机制和结构化标记化（Li等人，2023b (https://arxiv.org/html/2607.07718#bib.bib20); Wu等人，2024 (https://arxiv.org/html/2607.07718#bib.bib23); Luo等人，2025 (https://arxiv.org/html/2607.07718#bib.bib24)）。Transolver引入了物理注意力（Physics-Attention），将网格点分组为可学习的切片，并在切片级别执行注意力，从而在线性计算复杂度下保持对复杂几何形状的准确性（Wu等人，2024 (https://arxiv.org/html/2607.07718#bib.bib23)）。Transolver++通过改进并行性和使用更高效的注意力将此方法扩展到更大的输入（Luo等人，2025 (https://arxiv.org/html/2607.07718#bib.bib24)）。

#### 线性注意力

线性注意力方法通过用特征图替换softmax核来降低注意力成本，允许键值聚合仅计算一次并为每个查询重用。这一思想出现在线性Transformer和相关的核化注意力方法中（Katharopoulos等人，2020 (https://arxiv.org/html/2607.07718#bib.bib26); Choromanski等人，2021 (https://arxiv.org/html/2607.07718#bib.bib27)）。多项工作已成功将线性注意力应用于PDE建模（Hao等人，2023 (https://arxiv.org/html/2607.07718#bib.bib6); Tran等人，2021 (https://arxiv.org/html/2607.07718#bib.bib57)）。同时，其他研究表明，某些形式的线性注意力可能会降低性能，并阻止模型学习有用的交互（Wu等人，2022 (https://arxiv.org/html/2607.07718#bib.bib58)）。

#### 空间和几何编码

当神经网络表示空间场时，坐标编码是常用的。傅里叶特征映射可以帮助多层感知机表示坐标中的高频变化（Mildenhall等人，2020 (https://arxiv.org/html/2607.07718#bib.bib28); Tancik等人，2020 (https://arxiv.org/html/2607.07718#bib.bib29)），相关的周期性表示已被用于建模空间信号及其导数（Sitzmann等人，2020 (https://arxiv.org/html/2607.07718#bib.bib59)）。在算子学习中，坐标信息也很重要，因为输入和输出场是在物理位置（而非抽象标记索引）上采样的。因此，多种神经算子架构包含坐标或几何信息以补充