Skip to content

PDE-Transformer Documentation

Welcome to the documentation for PDE-Transformer, a state-of-the-art neural architecture for physics simulations, specifically designed for partial differential equations (PDEs) on regular grids.

Overview

PDE-Transformer is designed to efficiently process and predict the evolution of physical systems described by partial differential equations (PDEs). Our model provides:

  • Production Ready: Available as a pip package for easy installation and experimentation
  • State-of-the-Art Performance: Outperforms existing methods across a wide range different types of PDEs
  • Transfer Learning: Improved performance when adapting pre-trained models to new physics problems
  • Open Source: Full implementation with pre-trained models and documentation

Key Features

Architecture

  • Multi-scale transformer 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
  • Flexible conditioning mechanism for PDE parameters and metadata

Different Pretraining Datasets

  • 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)

Training Objectives

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

Code

The implementation is available on GitHub: tum-pbs/pde-transformer

# Install via pip
pip install pdetransformer

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

For detailed documentation, see our Documentation.

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

Acknowledgments

This work was supported by the ERC Consolidator Grant SpaTe (CoG-2019-863850). The authors gratefully acknowledge the scientific support and resources of the AI service infrastructure LRZ AI Systems provided by the Leibniz Supercomputing Centre (LRZ) of the Bavarian Academy of Sciences and Humanities (BAdW), funded by Bayerisches Staatsministerium fur Wissenschaft und Kunst (StMWK).