[Paper Review] GNOT: A General Neural Operator Transformer for Operator Learning

Zhongkai Hao, Zhengyi Wang+7 more|arXiv (Cornell University)|Feb 28, 2023

Model Reduction and Neural NetworksPhysics and Astronomy36 citations

TL;DR

GNOT introduces a scalable Transformer-based neural operator with heterogeneous normalized attention and a geometric gating mechanism to handle irregular meshes, multiple input functions, and multi-scale problems, achieving strong improvements across diverse PDE-related datasets.

ABSTRACT

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}.

Motivation & Objective

• Motivate the need for efficient operator learning for PDEs across irregular meshes, multiple inputs, and multi-scale domains.
• Propose a flexible Transformer-based architecture (GNOT) to learn operators under challenging practical conditions.
• Develop mechanisms (heterogeneous normalized attention and geometric MoE gating) to enable scalability and accuracy on diverse datasets.

Proposed method

• Introduce a General Neural Operator Transformer (GNOT) with a heterogeneous normalized cross-attention block for arbitrary input embeddings.
• Use a cross-attention followed by self-attention block as the core transformer module.
• Implement a geometric gating mechanism based on mixture-of-experts to perform soft domain decomposition and improve multi-scale learning.
• Encode diverse inputs (boundary shapes, domain-distributed functions, parameter vectors) via separate MLP encoders to obtain conditional embeddings.
• Adopt a linear-attention variant to achieve O((N+sum N_l)n_e^2) complexity, enabling handling of large point sequences.
• Train end-to-end with MSE loss on discretized input and solution meshes over Omega.

Experimental results

Research questions

• RQ1 Can GNOT accurately learn PDE solution operators on irregular meshes and across multiple input types?
• RQ2 Does the heterogeneous normalized attention enable better integration of diverse inputs compared to existing neural operators?
• RQ3 Is the geometric gating (MoE) mechanism effective for multi-scale problems and how does it affect scalability?
• RQ4 How does GNOT perform relative to baselines (MIONet, FNO, Geo-FNO, GK-Transformer, OFormer) across varied domains?

Key findings

• GNOT achieves significant improvements over baselines on several datasets, reducing prediction error by roughly 40-50% on tasks like Elasticity, Inductor2d, and Heatsink.
• On NS2d, GNOT reduces error from 13.7% to 4.42% with more data; on Heat, error drops from 4.13% to 2.58%.
• GNOT scales with data and model size, with larger embedding capacity yielding better performance and a roughly linear data-efficiency trend in the reported scaling experiments.
• A cross-attention followed by self-attention block (cross + self) consistently performs best in ablations across NACA, Elasticity, and NS2d-c datasets.
• Using 3 experts in the MoE-based geometric gating is advantageous for multi-subdomain problems (e.g., Heat); more than 8 experts can hurt performance.