[论文解读] GNOT: A General Neural Operator Transformer for Operator Learning
GNOT 引入了一个可扩展的基于 Transformer 的神经算子,具备异质归一化注意力和几何门控机制,能够处理不规则网格、多个输入函数和多尺度问题,在各类与 PDE 相关的数据集上取得显著提升。
Learning partial differential equations' (PDEs) solution operators is an essential problem in machine learning. However, there are several challenges for learning operators in practical applications like the irregular mesh, multiple input functions, and complexity of the PDEs' solution. To address these challenges, we propose a general neural operator transformer (GNOT), a scalable and effective transformer-based framework for learning operators. By designing a novel heterogeneous normalized attention layer, our model is highly flexible to handle multiple input functions and irregular meshes. Besides, we introduce a geometric gating mechanism which could be viewed as a soft domain decomposition to solve the multi-scale problems. The large model capacity of the transformer architecture grants our model the possibility to scale to large datasets and practical problems. We conduct extensive experiments on multiple challenging datasets from different domains and achieve a remarkable improvement compared with alternative methods. Our code and data are publicly available at \url{https://github.com/thu-ml/GNOT}.
研究动机与目标
- 在不规则网格、多个输入和多尺度域中,为 PDE 提高效率的算子学习的需求。
- 提出一个灵活的基于 Transformer 的架构(GNOT),在具有挑战的实际条件下学习算子。
- 开发机制(异质归一化注意力和几何 MoE 门控),以实现对多样数据集的可扩展性和准确性。
提出的方法
- 引入一个通用神经算子 Transformer (GNOT),具有用于任意输入嵌入的异质归一化跨注意力块。
- 以跨注意力后接自注意力块作为核心 Transformer 模块。
- 实现基于混合专家的几何门控机制,执行软域划分并改善多尺度学习。
- 通过单独的 MLP 编码器对不同输入(边界形状、域分布函数、参数向量)进行编码,以获得条件嵌入。
- 采用线性注意力变体,实现 O((N+sum N_l)n_e^2) 的复杂度,支持处理大规模点序列。
- 用 Omega 上离散化的输入和解的 MSE 损失进行端到端训练。
实验结果
研究问题
- RQ1GNOT 是否能够在不规则网格和多种输入类型下准确学习 PDE 解算子?
- RQ2与现有神经算子相比,异质归一化注意力是否能更好地整合多样输入?
- RQ3几何门控(MoE)机制在多尺度问题上是否有效,以及对可扩展性有何影响?
- RQ4在不同域上,与基线(MIONet、FNO、Geo-FNO、GK-Transformer、OFormer)相比,GNOT 的表现如何?
主要发现
| 数据集 | 类型 | MIONet | FNO(-interp) | GK-Transformer | Geo-FNO | OFormer | Ours | 挑战子集 |
|---|---|---|---|---|---|---|---|---|
| Darcy2d | A | - | - | 5.45e-2 | 1.09e-2 | 1.09e-2 | 1.24e-2 | 1.05e-2 |
| NS2d (part) | - | – | – | 1.56e-1 | 1.40e-1 | 1.56e-1 | 1.71e-1 | 1.38e-1 |
| NS2d (full) | - | – | – | 8.20e-2 | 7.92e-2 | 8.20e-2 | 6.46e-2 | 4.43e-2 |
| Elasticity | A | - | 9.65e-2 | 5.08e-2 | 2.01e-2 | 2.20e-2 | 1.83e-2 | 8.65e-3 |
| NS2d-c (u) | A,C | $u$ | 2.74e-2 | 6.56e-2 | 1.52e-2 | 1.41e-2 | 2.33e-2 | 6.73e-3 |
| NS2d-c (v) | A,C | $v$ | 5.51e-2 | 1.15e-1 | 3.15e-2 | 2.98e-2 | 4.83e-2 | 1.55e-2 |
| NS2d-c (p) | A,C | $p$ | 2.74e-2 | 1.11e-2 | 1.59e-2 | 1.62e-2 | 2.43e-2 | 7.41e-3 |
| NACA | A,C | - | 1.32e-1 | 4.21e-2 | 1.61e-2 | 1.38e-2 | 1.83e-2 | 7.57e-3 |
| Inductor2d (Az) | A,C | $A_{z}$ | 3.10e-2 | – | 2.56e-1 | – | 2.23e-2 | 1.21e-2 |
| Heat (part) | A,B,C | part | 1.74e-1 | – | – | – | 4.13e-2 | |
| Heat (full) | A,B,C | full | 1.45e-1 | – | – | – | 2.56e-2 | |
| Heatsink (T) | A,B,C | $T$ | 4.67e-1 | – | – | – | 2.53e-1 | |
| Heatsink (u) | A,B,C | $u$ | 3.52e-1 | – | – | – | 1.42e-1 | |
| Heatsink (v) | A,B,C | $v$ | 3.23e-1 | – | – | – | 1.81e-1 | |
| Heatsink (w) | A,B,C | $w$ | 3.71e-1 | – | – | – | 1.88e-1 |
- GNOT 在多个数据集上显著优于基线,在如弹性、Inductor2d 和 Heatsink 等任务上将预测误差降低约 40-50%。
- 在 NS2d 上,通过更多数据,GNOT 将误差从 13.7% 降至 4.42%;在 Heat 上,误差从 4.13% 降至 2.58%。
- GNOT 的数据和模型大小呈线性增长关系,嵌入容量更大时性能更好,在报道的扩展实验中呈现大致线性的数据效率趋势。
- 在 NACA、Elasticity、NS2d-c 数据集的消融中,跨注意力后接自注意力块(cross + self)始终表现最好。
- 在 MoE 基于几何门控中使用 3 个专家对多子域问题(如 Heat)有利;超过 8 个专家可能会降低性能。
更好的研究,从现在开始
从论文设计到论文写作,大幅缩短您的研究时间。
无需绑定信用卡
本解读由 AI 生成,并经人工编辑审核。