Training Transformers for Mesh-Based Simulations

TL;DR

This paper introduces a novel Graph Transformer architecture for mesh-based physics simulations, achieving superior scalability and performance on large 3D CFD datasets compared to existing GNN methods.

Contribution

The paper presents a new Graph Transformer design with adjacency-based attention and augmentation techniques, enabling efficient large-scale mesh simulations with improved accuracy.

Findings

Models scale to meshes with 300k nodes and 3 million edges.

Smallest model matches MeshGraphNet performance, is 7x faster and 6x smaller.

Largest model outperforms previous state-of-the-art by 38.8%.

Abstract

Simulating physics using Graph Neural Networks (GNNs) is predominantly driven by message-passing architectures, which face challenges in scaling and efficiency, particularly in handling large, complex meshes. These architectures have inspired numerous enhancements, including multigrid approaches and $K$-hop aggregation (using neighbours of distance $K$), yet they often introduce significant complexity and suffer from limited in-depth investigations. In response to these challenges, we propose a novel Graph Transformer architecture that leverages the adjacency matrix as an attention mask. The proposed approach incorporates innovative augmentations, including Dilated Sliding Windows and Global Attention, to extend receptive fields without sacrificing computational efficiency. Through extensive experimentation, we evaluate model size, adjacency matrix augmentations, positional encoding and $K$-hop configurations using challenging 3D computational fluid dynamics (CFD) datasets. We also train over 60 models to find a scaling law between training FLOPs and parameters. The introduced models demonstrate remarkable scalability, performing on meshes with up to 300k nodes and 3 million edges. Notably, the smallest model achieves parity with MeshGraphNet while being $7\times$ faster and $6\times$ smaller. The largest model surpasses the previous state-of-the-art by $38.8$\% on average and outperforms MeshGraphNet by $52$\% on the all-rollout RMSE, while having a similar training speed. Code and datasets are available at https://github.com/DonsetPG/graph-physics.