PDE-Transformer: Efficient and Versatile Transformers for Physics Simulations

TL;DR

PDE-Transformer is a new transformer architecture tailored for physics simulations that outperforms existing models in accuracy and scalability across multiple PDE types, enabling more effective large-scale scientific modeling.

Contribution

The paper introduces PDE-Transformer, a scalable and versatile transformer architecture specifically designed for surrogate modeling of physics simulations on regular grids.

Findings

Outperforms state-of-the-art transformers on PDE datasets

Maintains consistent information density across multiple PDE types

Achieves improved downstream task performance with pre-training

Abstract

We introduce PDE-Transformer, an improved transformer-based architecture for surrogate modeling of physics simulations on regular grids. We combine recent architectural improvements of diffusion transformers with adjustments specific for large-scale simulations to yield a more scalable and versatile general-purpose transformer architecture, which can be used as the backbone for building large-scale foundation models in physical sciences. We demonstrate that our proposed architecture outperforms state-of-the-art transformer architectures for computer vision on a large dataset of 16 different types of PDEs. We propose to embed different physical channels individually as spatio-temporal tokens, which interact via channel-wise self-attention. This helps to maintain a consistent information density of tokens when learning multiple types of PDEs simultaneously. We demonstrate that our pre-trained models achieve improved performance on several challenging downstream tasks compared to training from scratch and also beat other foundation model architectures for physics simulations.