Source-linked AI summary
GNOT: A General Neural Operator Transformer for Operator Learning
Zhongkai Hao, Zhengyi Wang, Hang Su, Chengyang Ying, Yinpeng Dong, Songming Liu, Ze Cheng, Jian Song, Jun Zhu
TL;DR
Practical PDE operator learning must handle irregular meshes, multiple input functions, and multi-scale solution complexity while supporting repeated simulation. GNOT addresses these challenges with heterogeneous normalized attention and coordinate-based geometric gating in a scalable transformer framework. Across challenging scientific and engineering datasets, it reports substantial improvements over competing baselines, including about 50% lower prediction error on several practical datasets.
Problem
Practical operator learning remains challenging because real-world problems involve irregular meshes, multiple input functions, and multi-scale behavior.
Method
GNOT uses heterogeneous normalized attention to encode arbitrary input functions and geometric mixture-of-experts gating based on query coordinates as soft domain decomposition.
Results
About 50% lower prediction error than baselines is reported on practical Elasticty, Inductor2d, and Heatsink datasets.
Takeaways & Limitations
GNOT provides a general transformer framework for operator learning across multiple inputs, irregular meshes, multi-scale problems, and large datasets.
Takeaways & Limitations
Multi-scale problems remain difficult, with Heatsink velocity-distribution prediction errors exceeding 10% in failure cases.
Abstract
from arXiv · showhide
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}.
1. Introduction
PDE simulation is important but computationally expensive, motivating neural operators that accelerate repeated solution prediction. Practical operator learning remains challenging because of irregular meshes, multiple input functions, and multi-scale problems; GNOT addresses these challenges with heterogeneous normalized attention and geometric gating, and reports substantial improvements over baselines.
- Motivation: Finite element simulations can take seconds to days, whereas pre-trained neural operators predict PDE solutions with a single forward computation.Neural operators learn mappings from input functions to PDE solutions and generalize to unseen inputs.
- Challenges: Irregular meshes challenge methods such as FNO and U-Net, which are limited to uniform regular grids.Practical geometries, such as airfoils, can be highly complex.
- Challenges: Multiple input functions require models to flexibly process different input types, including boundary shapes, global parameter vectors, and source functions.Existing approaches lack a general encoder framework for this setting.
- GNOT: GNOT combines heterogeneous normalized attention for arbitrary input functions with geometric gating based on mixture-of-experts for multi-scale problems.The attention block provides a general encoding interface with linear sequence-length complexity, while coordinate-based gating acts as soft domain decomposition.
- Results: About 50% lower prediction error is reported than baselines on practical datasets including Elasticty, Inductor2d, and Heatsink.The experiments span fluids, elastic mechanics, electromagnetism, and thermology.
2. Related Work
Neural operator research includes branch-trunk, tensor-product, Fourier, and attention-based architectures. Efficient attention is important because operator-learning sequences can contain thousands to millions of elements, while existing transformer approaches still struggle with multiple inputs and scales.
- Neural operators: DeepONet separates input-function and query-point processing with branch and trunk networks, while MIONet extends this architecture to multiple input functions.DeepONet is described as capable of approximating nonlinear operators with a sufficiently large network.
- Attention-based operators: FNO uses Fourier-based processing, whereas attention mechanisms offer learnable kernels, large capacity, flexible input lengths, and permutation equivariance.Galerkin Transformer interprets attention as an integral transform with a learnable kernel.
- Attention-based operators: Transformer architectures are flexible for irregular meshes but remain unsatisfactory for challenging operators with multiple inputs and scales.OFormer and HT-Net apply transformer-based designs to operator learning and multi-scale problems.
- Efficient transformers: Original attention has quadratic complexity in sequence length, which is problematic when operator-learning sequences range from thousands to millions.Efficient attention operations are therefore necessary for operator learning.
3. Proposed Method
GNOT learns PDE solution operators from heterogeneous inputs on irregular meshes using an attention-based architecture with linear-cost attention and geometry-aware expert gating. Its design encodes varied inputs, processes query features through cross- and self-attention, and uses multiple experts for multi-scale problems.
- 3.1. Problem Formulation: GNOT formulates operator learning as mapping multiple input functions, boundary shapes, and system parameters to PDE solution functions.Inputs may include functions in H and parameter vectors in R^p, while outputs are solution functions over the domain.
- 3.3. General Input Encoding: The model encodes query points and heterogeneous input functions with MLPs, using separate MLP parameters for each input function.It supports parameter vectors, boundary-point sets, domain-distributed functions, and additional mesh information.
- 3.4. Heterogeneous Normalized Attention Block: Its basic attention block cascades heterogeneous normalized cross-attention with normalized self-attention to update query-point features.Cross-attention combines query features with multiple conditional embeddings, while heterogeneous projections compute keys and values from different inputs.
- 3.4. Heterogeneous Normalized Attention Block: Linear attention reduces the sequence-length cost from quadratic self-attention and bilinear cross-attention to O((M + N)n_e^2).Vanilla self-attention costs O(N^2n_e) and cross-attention costs O(NMn_e), which is problematic for sequences containing thousands to millions of points.
- 3.5. Geometric Gating Mechanism: Geometric gating computes coordinate-dependent weights for multiple expert FFNs, providing a soft domain-decomposition mechanism for multi-scale problems.The gate network uses query-point coordinates to produce expert-mixture scores, and each layer averages K subnetworks.
4. Experiments
The experiments evaluate GNOT across challenging operator-learning datasets and show strong performance, scalability with more data and layers, and component-specific design effects. Multi-scale problems remain difficult for all evaluated models.
- Experimental Setup: Experiments cover datasets from fluids, elastic mechanics, electromagnetism, heat conduction, and multi-physics problems.The datasets include irregular meshes, multiple input functions, and multi-scale settings.
- Main Results: GNOT performs significantly better than baselines on nearly all tasks, including irregular-mesh and multi-scale datasets such as NACA, NS2d-c, and Inductor2d.It also handles multiple-input datasets such as Heat and Heatsink.
- Main Results: GNOT reduces prediction error by about 40% ∼50% on some tasks compared with baselines.The evaluation uses mean l2 relative error, with lower scores indicating better performance.
- Main Results: GK-Transformer performs slightly better on Darcy2d, a simple uniform-grid dataset.This is the reported exception to GNOT’s broader task-level advantage.
- Scaling Experiments: 13.7% to 4.42%: on NS2d, GNOT reduces error by more than three times as data increases; on Heat, error falls from 4.13% to 2.58%.The authors interpret these results as evidence of greater scalability with accessible data.
- Limitations: All models perform worse on multi-scale Heatsink, where velocity-field prediction errors exceed 10% in reported failure cases.The authors suggest incorporating physical prior as a possible improvement.
- Scaling Experiments: Prediction error generally decreases with more GNOT layers, but gains become small beyond four layers on Elasticity.The authors identify four layers as an efficient choice because additional layers increase computational cost.
- Ablation Experiments: The cross+self attention block performs best across NACA, Elasticity, and NS2d-c, supporting self-attention after cross-attention.On Heat, three experts perform best, while using at least eight experts deteriorates performance; attention-head count has little impact.
5. Conclusion
GNOT combines heterogeneous normalized attention with geometric gating to address practical operator-learning challenges. Experiments across science and engineering datasets show effectiveness relative to baselines and suggest a direction for large-scale neural surrogates.
- GNOT combines heterogeneous normalized attention and geometric gating to address challenges in practical operator learning.The paper presents these as its two new components.
A. Details and visualization of datasets
The supplementary datasets cover fluids, electromagnetism, heat conduction, and coupled heat–fluid systems, with varied geometries, inputs, and predicted fields. Several datasets use generated samples with separate training and testing sets.
- NS2d-c: NS2d-c is a steady-state two-dimensional Navier–Stokes problem on a rectangle with four circular regions removed.Its outputs are velocity fields u and v plus pressure p, given the input mesh.
- NS2d-c: NS2d-c contains 1100 samples with different circle positions, using 1000 for training and 100 for testing.
- Inductor2d: Inductor2d is a two-dimensional steady-state inductor governed by Maxwell’s equation, with varying geometry, coil current, and material parameters.The geometry and solution fields are visualized in Figures 6 and 7.
- Heat: Heat is a two-dimensional steady-state heat-equation problem whose rectangular domain is divided into three parts by two splines.Its inputs include top-boundary temperature and spline parameters; the dataset has 1100 small-dataset samples and 5500 full-dataset samples.
- Heatsink: Heatsink is a three-dimensional steady-state multiphysics problem coupling heat and fluid behavior through convection and conduction.Its inputs include geometric parameters and inlet velocity, and the goal is to predict the velocity field.
B. Hyperparameters and details for models.
The supplementary material specifies model-training choices and baseline configurations, including architectures, search ranges, optimizers, and runtime-related implementation details. Figures and tables document dataset visualizations, hyperparameters, and runtime comparisons.
- MIONet: MIONet uses four-layer, width-256 MLP branch and trunk networks, with multiple branch networks for multiple input functions.With one branch network, MIONet degenerates to DeepONet.
- Supplementary tables: The supplementary material includes a table of main-experiment hyperparameters and a table comparing runtimes across methods.
- FNO and Geo-FNO: FNO uses four layers with modes from 12, 16, or 32 and widths from 16, 32, or 64; irregular-grid data are interpolated onto selected uniform resolutions.Geo-FNO degenerates to vanilla FNO on Darcy2d and NS2d.
- Transformer baselines and GNOT: GK-Transformer, OFormer, and GNOT vary heads across 1, 4, 8, or 16 and layers across 2 through 6, with embedding and FFN sizes from 64 to 256.The models use AdamW with a one-cycle learning-decay strategy.
C. Other Supplementary Results
Training GNOT and other transformer-based methods is slower than training FNO, identifying runtime as a practical drawback of the transformer approach.
- Transformer-based methods train more slowly than FNO.The paper identifies this as a drawback shared by the transformer methods in its runtime comparison.
D. Broader Impact
GNOT targets neural-operator applications across several scientific and engineering fields. However, transformer-based neural networks lack theoretical guarantees and interoperability, creating risks in risk-sensitive deployments.
- Neural operators have potential applications in physics, quantum mechanics, heat engineering, fluid dynamics, and aerospace.
- Transformer-based neural networks lack theoretical guarantees and interoperability.
- In risk-sensitive areas, deploying unexplainable models can make accident investigation more difficult.
- The paper suggests developing more explainable and robust methods with stronger guarantees or corner-case protection for such deployments.