Geometry Aware Operator Transformer as an Efficient and Accurate Neural Surrogate for PDEs on Arbitrary Domains

TL;DR

GAOT introduces Geometry Aware Operator Transformer, a neural surrogate for PDEs on arbitrary domains by integrating MAGNO encoders/decoders, geometry embeddings, and a Vision Transformer processor to deliver high accuracy and efficiency. It handles inputs from irregular point clouds and outputs at any query point, with a time-dependent extension via flexible time-stepping strategies. Extensive benchmarks, including large-scale 3D industrial datasets, show state-of-the-art accuracy and superior scalability compared to baselines, with substantial speedups over classical solvers. The work positions GAOT as a potential backbone for PDE foundation models, enabling efficient UQ, inverse problems, and PDE-constrained optimization across diverse geometries.

Abstract

The very challenging task of learning solution operators of PDEs on arbitrary domains accurately and efficiently is of vital importance to engineering and industrial simulations. Despite the existence of many operator learning algorithms to approximate such PDEs, we find that accurate models are not necessarily computationally efficient and vice versa. We address this issue by proposing a geometry aware operator transformer (GAOT) for learning PDEs on arbitrary domains. GAOT combines novel multiscale attentional graph neural operator encoders and decoders, together with geometry embeddings and (vision) transformer processors to accurately map information about the domain and the inputs into a robust approximation of the PDE solution. Multiple innovations in the implementation of GAOT also ensure computational efficiency and scalability. We demonstrate this significant gain in both accuracy and efficiency of GAOT over several baselines on a large number of learning tasks from a diverse set of PDEs, including achieving state of the art performance on three large scale three-dimensional industrial CFD datasets.

Figures (46)

• Figure 1: Normalized performance of GAOT and baselines across eight axes, covering accuracy (Acc.), robustness (Robust), throughput (Tput), scalability (Scal.) on time-dependent (TD) and time-independent (TI) tasks.
• Figure 2: Schematic of the GAOT with an equispaced latent token grid. The encoder uses a multiscale attentional graph neural operator (MAGNO) to aggregate the input data into geometry-aware tokens. A vision transformer (ViT) block with residual connections processes tokens, enabling global exchange of information. A MAGNO decoder identifies the nearest tokens around a given query point to decode the final field.
• Figure 3: Training throughput (samples/s) with increasing input grid size (a) and model size (b) for proposed GAOT, GINO, RIGNO and Transolver. (c) Transfer learning performance of GAOT on unseen bluff body shapes (See SM Sec. \ref{['app:tl']} for dataset details). FT (fine-tuning) adapts a pretrained GAOT model from Table \ref{['tab:main-overall results']}, while TFS denotes training from scratch. FT consistently outperforms TFS across varying numbers of task-specific training samples.
• Figure 4: Comparison of predicted and ground-truth (GT) results for the pressure and wall shear stress in the x-direction (WSS-x) on the DrivAerNet++ test sample N_S_WWS_WM_172.
• Figure 5: Comparison of predicted and ground-truth (GT) results for the pressure on the test sample of DrivAerML and NASA-CRM.