ICML2024 Spotlight | Transolver: 几何结构通用的快速PDE求解网络，高效完成车型、机翼仿真...

点击蓝字

关注我们

AI TIME欢迎每一位AI爱好者的加入！

本文介绍清华大学软件学院机器学习组ICML 2024 科学学习方向的最新工作：Transolver: A Fast Transformer Solver for PDEs on General Geometries

摘要：实际应用中，偏微分方程（PDE）通常被离散化到大规模非结构化网格上，造成了计算效率与物理建模上的巨大挑战。本文提出了Transolver模型，将6个标准数据集的误差大幅降低22%，同时在车型、机翼等大规模工业仿真任务上表现突出。

作者：吴海旭，罗华坤，王皓雯，王建民，龙明盛

链接：https://arxiv.org/pdf/2402.02366

代码：https://github.com/thuml/Transolver

1. 引言

求解偏微分方程（PDE）是科学与工程共有的基础性问题，对材料分析、工业设计、气象预报等重大需求至关重要。由于很难获得解析解，在实际应用中PDE通常被离散化到大规模网格上，再使用经典的数值方法进行求解，但是往往需要数小时甚至数天才能完成一个复杂几何结构的仿真模拟。近期，深度模型在PDE高效求解上展现了巨大潜力。得益于强大的非线性拟合能力，它可以从数据中学习从几何结构到物理场的映射，并以极快的速度（秒级）完成推理仿真。

特别地，作为深度学习领域重要的主干网络，Transformer也被应用于PDE求解中。虽然取得了一定效果，但是在面临大规模工业模拟场景时，Transformer仍然面临着计算效率与物理建模的内在短板，阻碍它成为理想的PDE求解器。具体包括：

• 海量网格带来的巨大计算开销：工业模拟需要处理上万甚至上百万的网格，然而Transformer面临二次复杂度难题，即使是最先进的Transformer模型GPT-4，能够处理的序列长度也只有几万；

• 几何结构背后复杂的物理关联：上图展示的车型仿真任务，需要模型准确捕捉数万离散网格点之间极其复杂、高阶的多物理交互，而单个网格点往往仅有坐标信息，难以支撑准确的关联建模。

因此，如何高效地捕捉复杂几何结构下潜在的物理关联是将Transformer变为实用PDE求解器的关键。

为解决上述挑战，我们提出了Transolver模型，具有以下特点：

• 抛弃冗余并且流于表面的网格，我们提出学习几何结构背后内在的物理状态，并在物理状态间应用注意力机制，天然具备线性复杂度和几何结构通用性；

• 在6个标准PDE数据集上平均比之前的SOTA误差降低22%，在大规模工业仿真场景中表现最优；

• 展现了优秀的计算效率，可扩展性（Scalability）以及分布外泛化能力（OOD Generalizability）。

2. 问题分析

针对Transformer在大规模PDE求解中的复杂度问题，之前的工作尝试通过线性注意力机制（Linear Attention）来实现加速，但是处理过长的序列可能会造成注意力分散的问题。近期英伟达提出的RULER基准证明：即使是最先进的大语言模型GPT-4，其实际有效长度也只有6.4万。

此外，仅仅依赖单一网格点的特征也难以得出可靠的物理关联。虽然计算机视觉中的分块（Patchify）机制可以一定程度上扩充单点的特征，但是只能应用在规则网格上，依然不足以支撑实际生产应用。

跳出当前神经PDE求解器的思维定式，我们从本质上思考这一问题，发现离散化网格只是对连续物理场的采样，求解PDE更重要的是挖掘其内在的物理交互过程。因此，我们提出学习内在的物理状态，在物理空间中求解PDE。

如下图所示，Transolver可以准确地捕捉PDE中多样化的物理状态，例如a. Darcy流中固液交互过程；b.弹性材料不同的挤压受力区域；c. 机翼周围的激波和尾流；d-e. 车辆在行驶过程中前后面以及上下空间。

3. Transolver

如前所述，Transolver提出在学习到的物理状态间应用注意力机制，称为Physics-Attention，包含以下三个步骤：（1）将处于相似物理状态的网格点特征汇聚为物理感知令牌（Physics-aware Tokens）（2）在学习到的物理空间中应用注意力机制（3）将交互之后的物理状态投影回离散网格。

3.1 物理感知令牌（Physics-aware Tokens）

如上图所示，我们观察到处于行驶的车来说，其正面（前挡风玻璃、车牌、车灯）、倾斜面（引擎盖）以及后面处于完全不同的物理状态（如迎风、尾流等），这启发我们将现有的网格点进行归类，得到一些“内在物理属性一致”的子集，进而汇聚出不同物理状态的特征。这也为求解PDE提供了一个更加“本质”的视角。

具体地，对于网格点集合记录了个点处的位置和物理量。我们首先将它通过线性嵌入层映射为深度特征，其中包含了几何结构和物理的信息。

为了将定位整个输入空间隐含的物理状态，我们提出了一种自底而上的思路：基于特征，将个网格点划分到预先设定的个潜在的、灵活形状的切片（Slices）中，处于相似物理状态的网格点被归类于同一切片。这一设计可以被形式化为：其中将每个特征从个通道映射为个权重，表示第个网格点归类到第个切片的权重，。而表示第个切片的特征，它按照学习到的权重对网格特征进行加权。注意，由于的参数共享，具备相似特征的网格点将天然导出相似的切片权重，这也意味着它们更有可能被划分到同一个切片中。

特别地，为了防止平凡的切片划分（例如均匀分配到个切片），我们在权重的学习过程使用函数以降低所得权重的熵，进而获得更加有信息量的切片分配方案。默认被设置为一个通道维度的线性层，可以天然适配复杂几何结构；对于结构化网格或者均匀网格，它也可以被设置为一个局部的卷积层以得到更好的表征，这样配置依然可以适配不同尺度的输入。

由于每个切片包含了内在物理状态一致的网格点，我们可以进一步将切片特征全局编码，从而得到几何无关的物理感知令牌：其中，。同时，由于切片是内在物理状态一致，每一个令牌都包含了某一特征物理状态的信息。

（1）为什么切片可以学习到内在一致物理状态？一方面，我们使用网格点的特征来学习切片划分权重，因此具有相似特征的网格点更有可能被归为同一个切片；另一方面，因为后续我们会在物理感知令牌之间计算注意力，如果切片间没有学习到足够可区分的物理信息，注意力机制将难以得到很好的效果。为了降低损失函数，模型的优化过程会驱使切片逐步学习到物理状态。

（2）划分切片不同于划分计算区域。在经典数值方法中，如有限元，通过会将网格划分为不同的计算区域，以提升模拟精度。但是这一过程往往需要大量的专业知识与人工设计，并且只能覆盖局部区域。不同于这一方法，我们提出的切片可以覆盖空间远离但是物理状态相近的网格点。例如，第2节展示的可视化结果，我们的设计可以将前挡风玻璃、车牌、车灯这些处于迎风状态的区域划分到一个切片中。

3.2 Physics-Attention

基于学习物理状态的思想，我们提出了Transolver模型。如上图所示，它通过物理注意力机制（Physics-Attention）来捕捉复杂的物理交互。

对于输入的网格特征，我们遵循上一节提出的方案，首先将它按照可学习权重划分为个切片，再继续编码得到物理感知令牌。

接下来，为了建模不同物理状态的交互，我们在物理感知令牌之间应用注意力机制：随后，我们将变换之后的物理令牌，按照之前的分配权重，转换回网格特征，即上述过程被总结为。其计算复杂度为，由于，Physics-Attention实际具备线性复杂度。

我们将Transformer中的标准注意力机制替换为Physics-Attention，即得到Transolver。

3.3 理论分析

为了对Transolver有更加本质的理解，我们将从理论角度分析Physics-Attention。

已有结论：前人将PDE求解形式化为一个迭代更新过程，并且证明了标准的注意力机制是对于全局积分的一个蒙特卡洛采样，因此可以被用于近似求解过程的单步迭代。

理论分析：不同于直接计算网格点关联的注意力机制，Physics-Attention在学习到的物理空间计算注意力，根据前人结论，我们可以直接推出其本质是物理空间的可学习积分。但是这依然无法证明Physics-Attention对于原始定义域PDE的求解能力，为此我们进一步证明了Physics-Attention等价于原始定义域的积分，主要证明过程如下：

上述证明过程也可以直接导出我们在Physics-Attention中的全部结构设计，这也从理论上证实了Transolver的PDE求解能力。

详细的理论分析及证明请见论文。

4. 实验

如下图所示，我们将Transolver在6个广泛使用的数据集中进行了测试，也在车辆、机翼设计等实际复杂任务上进行了实验，包含多样的几何结构，网格点数量从1k-32k。

同时我们将Transolver与超过20个基线模型进行了细致对比，包含神经算子网络：FNO (2021)，U-NO (2023)，LSM (2023)等，基于Transformer的PDE求解网络：GNOT (2023), FactFormer (2023), ONO (2024)等，经典的几何深度学习模型：PointNet (2017), MeshGraphNet (2021)等。特别对比工业设计任务，我们也对比了GINO (2023)和3D-GeoCA (2024)。

4.1 标准数据集

如上表所示，Transolver取得了一致的最优结果，与各个数据集上之前的SOTA方法相比，误差平均降低了22%。特别对于具有复杂几何结构的Elasticity和Plasticity，Transolver效果提升尤为明显。同时也可以很好地处理定义在标准网格上的Navier-Stokes和Darcy。

4.2 车型、机翼设计

设计任务需要估计不同车行驶、飞机飞行过程中的物理场、拉力与升力系数，具体为：

• 车型设计（Shape-Net Car）：估计不同车型在20m/s速度行驶下的表面压强以及周围空间风速。

• 机翼设计（AirfRANS）：基于美国NACA提供的翼型进行仿真，数据集包含不同的攻角、雷诺数、翼型。

如上表所示，Transolver不仅在物理场估计上表现优秀，还可以准确估计拉力与升力系数，预测与真实的spearman相关系数超过了99%，证明了它在工业设计中的实用价值。

如下图所示，Transolver在车结构的曲折连接处效果尤其突出，例如前挡玻璃与引擎盖连接处（涉及风与压强等多物理场交互），验证了模型对于复杂几何结构的建模能力。

4.3 效率分析

与其他基于Transformer的求解器相比，Transolver可以取的很好的“效果-效率”平衡，在处理大规模网格时效率优势明显。

5. 作为PDE基础模型的潜力

5.1 可扩展性（Scalability）

可扩展性是Transformer作为基础模型的重要能力。我们验证了Transolver在不同尺度下的效果稳定性以及对于数据与参数良好的可扩展性。通过增加训练数据与模型参数，Transolver可以得到持续的优化。

此外，下图在更加广泛的任务上证明了Transolver的参数泛化性：当我们将Transolver从8层（主实验设置）增加到40层时，可以获得远超当前基准的效果。

5.2 分布外泛化能力（OOD Generalizability）

分布外泛化能力是构建基础模型的关键。为了验证模型的这一能力，我们将AirfRANS上训练好的模型在完全分布外的雷诺数和攻角进行测试，发现Transolver依然可以保持最优的效果，同时对于升力系数的估计依然保持了接近99%的相关系数。

5.3 处理多样化任务（Versatile）

Transolver的物理建模能力也可以辅助拉格朗日流体模拟（追踪流体中粒子随时间的流动，需要应对随时间变化的几何结构）。如上所示，经过Transolver的辅助，GNN可以准确地预测出400步之后的液体飞溅。

6. 总结

针对复杂几何结构上PDE求解的计算效率以及物理建模难题，我们革新了以往直接在网格点上建模的方案，基于学习物理状态的思想提出快速PDE求解网络Transolver。不仅大幅度降低了标准数据集上模拟误差，也在车型、机翼仿真等工业生产设计场景表现突出，具有广阔应用前景。

特别地，Transolver具有良好的可扩展性、分布外泛化性，并且可以辅助多样任务，展现出了作为PDE求解基础模型的潜力。在未来，我们也将进一步探索Transolver在大规模数据上的训练以及泛化。

欢迎感兴趣的朋友阅读我们的论文(https://arxiv.org/abs/2402.02366)或者访问GitHub项目页面(https://github.com/thuml/Transolver)查看更多细节内容。

往期精彩文章推荐

记得关注我们呀！每天都有新知识！

 关于AI TIME 

AI TIME源起于2019年，旨在发扬科学思辨精神，邀请各界人士对人工智能理论、算法和场景应用的本质问题进行探索，加强思想碰撞，链接全球AI学者、行业专家和爱好者，希望以辩论的形式，探讨人工智能和人类未来之间的矛盾，探索人工智能领域的未来。

迄今为止，AI TIME已经邀请了1800多位海内外讲者，举办了逾600场活动，超700万人次观看。

我知道你

在看

提出观点，表达想法，欢迎

留言

点击 阅读原文 查看讲者直播回放！