PDE-Transformer: Efficient and Versatile Transformers for Physics Simulations

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.Our source code is available at https://github.com/tum-pbs/pde-transformer.

The PDE-Transformer is a new AI model designed to predict how physical systems behave, like how heat spreads or fluids flow. It is built on a network architecture called a transformer, which is also used in fields such as computer vision and natural language processing. We have made some key changes to make this transformer better suited for complex scientific simulations. It can handle larger simulations and is more flexible, making it a good base for creating powerful AI models for various scientific problems. Our tests show that the PDE-Transformer is better than other leading models at predicting the behaviour of 16 different types of physics problems. We also found a better way to represent different physical properties, like density or velocity, allowing the model to learn multiple types of physics at the same time more effectively. When we train our PDE-Transformer on a lot of data first, it performs better on new, challenging tasks compared to models trained from scratch.