We’d like to highlight the code release of our scalable & efficient PDE Transformer (P3D) at https://akanota.github.io/p3d/ , great work Benjamin with support from Florian & Georg. Please try out the pretrained models, and let us know how it works! Highlights are, e.g., stable inference of 1024^3 rollouts on a single GPU with 90GB.
Key Contributions:
- Hybrid CNN-Transformer: P3D combines convolutions for fast local features with windowed self-attention for deep representation learning in a hierarchical U-shape structure.
- Scalable Pretraining: Pretrain on small 128³ crops, then scale to the full 1024³ domain — reducing memory and compute while maintaining accuracy.
- Global Context Network: A sequence-to-sequence model links bottleneck layers for efficient global information processing, with region tokens for direct decoder feedback.
- Probabilistic Generation: Train P3D as a diffusion model to produce probabilistic samples of turbulent channel flows, accurately capturing flow statistics across Reynolds numbers.
Links:
- Main page: https://akanota.github.io/p3d/
- Paper: https://openreview.net/forum?id=8UdCE5nhFl
- Code: https://github.com/tum-pbs/P3D
- Weights & Data: https://huggingface.co/thuerey-group/p3d
