PDE-Transformer: Efficient and Versatile Transformers for Physics Simulations

Authors: Benjamin Holzschuh, Qiang Liu, Georg Kohl, Nils Thuerey

For P3D: see https://github.com/tum-pbs/P3D

PDE-Transformer is a state-of-the-art neural architecture for physics simulations, specifically designed for partial differential equations (PDEs) on regular grids. This work will be presented at ICML 2025.

Key Highlights

• Production Ready: Available as a pip package for easy installation and experimentation.
• State-of-the-Art: Outperforms existing methods across 16 different types of PDEs and three challenging downstream tasks involving complex dynamics.
• Transfer Learning: Improved performance when adapting pre-trained models to new physics problems with limited training data.
• Open Source: Full implementation with pre-trained models and comprehensive documentation.

Quick Installation

# Install from PyPI
pip install pdetransformer

# Or install from source
git clone https://github.com/tum-pbs/pde-transformer.git
cd pde-transformer
pip install -e .

Model Description

PDE-Transformer is designed to efficiently process and predict the evolution of physical systems described by partial differential equations (PDEs). It can handle multiple types of PDEs, different resolutions, domain extents, boundary conditions, and includes deep conditioning mechanisms for PDE- and task-specific information.

Key features:

• Multi-scale architecture with token down- and upsampling for efficient modeling.
• Shifted window attention for improved scaling to high-resolution data.
• Mixed Channel (MC) and Separate Channel (SC) representations for handling multiple physical quantities.
• Flexible conditioning mechanism for PDE parameters, boundary conditions, and simulation metadata.
• Pre-training and fine-tuning capabilities for transfer learning across different physics domains.

Training Objectives

The model supports both supervised and diffusion training:

• Supervised Training: Direct MSE loss for deterministic, unique solutions. Fast training and inference.
• Flow Matching: For probabilistic modeling and uncertainty quantification.

Supported PDE Types

PDE-Transformer has been trained and evaluated on 16 different types of PDEs including:

• Linear PDEs: Diffusion
• Nonlinear PDEs: Burgers, Korteweg-de-Vries, Kuramoto-Sivashinsky
• Reaction-Diffusion: Fisher-KPP, Swift-Hohenberg, Gray-Scott
• Fluid Dynamics: Navier-Stokes (Decaying Turbulence, Kolmogorov Flow)

Quick Start

from pdetransformer.core.mixed_channels import PDETransformer
import torch

# Load pre-trained model
model = PDETransformer.from_pretrained('thuerey-group/pde-transformer', subfolder='mc-s').cuda()

# For physics simulation
x = torch.randn((1,2,256,256), dtype=torch.float32).cuda()
predictions = model(x)

See the following notebook for an example of how to initialize the model from pretrained weights and use it for inference:

Documentation

For detailed documentation, visit tum-pbs.github.io/pde-transformer.

Citation

If you use PDE-Transformer in your research, please cite:

@article{holzschuh2025pde,
  title={PDE-Transformer: Efficient and Versatile Transformers for Physics Simulations},
  author={Holzschuh, Benjamin and Liu, Qiang and Kohl, Georg and Thuerey, Nils},
  booktitle={Forty-second International Conference on Machine Learning, {ICML} 2025, Vancouver, Canada, July 13-19, 2025},
  year={2025}
}

License

This project is licensed under the MIT License. See the LICENSE file for details.

Note: This is a research project from the Technical University of Munich (TUM) Physics-based Simulation Group. For questions and support, please refer to the GitHub repository or contact the authors.