Research map / Architectures and efficiency / Attention and transformers
PDE-Transformer: Efficient and Versatile Transformers for Physics Simulations
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Introduces PDE-Transformer, a multi-scale diffusion-transformer-style architecture with shifted windows and channel-wise attention for pre-training and fine-tuning surrogate models on 2D PDE data.
Problem
Building general-purpose, scalable transformer backbones for surrogate modeling of physics simulations on regular 2D grids that can handle multiple PDE types with varying numbers of physical channels at high resolution.
Approach
Adapt diffusion transformer (DiT) design with PixelShuffle token down/upsampling, shifted-window attention with log-spaced relative positions, and a separate-channel (SC) embedding where each physical channel is tokenized independently and channels interact via axial channel-wise attention, keeping token information density constant.
Key results
- PDE-S reaches nRMSE_1 of 0.044 and nRMSE_10 of 0.36 on the pre-training set, beating DiT-S (0.78 at step 10) and UDiT-S (0.39) with a 7h 42m training time versus 18h 30m for UDiT-S.
- Pre-trained PDE-S achieves the lowest average rollout nRMSE on all three downstream tasks, e.g., 0.455 on active matter and 0.104 on RBC, with about 42% average improvement over the second best model (FNO).
- Against Poseidon's scOT-B, PDE-B reduces accumulated nRMSE by 75%, while scOT-B becomes unstable on long rollouts.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| pre-train PDE-S on generic PDEs before downstream fine-tuning | PDE-S · Active Matter, RBC, and Shear flow | average rollout nRMSE: 0.494 -> 0.455; 0.155 -> 0.104; 0.199 -> 0.125 | Helped |
| replace training from scratch with generic-PDE pre-training | PDE-S · Active Matter, first 20 rollout steps | nRMSE: 0.494 -> 0.455 | Helped |
| replace training from scratch with generic-PDE pre-training | PDE-S · RBC, first 20 rollout steps | nRMSE: 0.155 -> 0.104 | Helped |
| replace training from scratch with generic-PDE pre-training | PDE-S · Shear flow, first 20 rollout steps | nRMSE: 0.199 -> 0.125 | Helped |
| increase token embedding dimension from d=96 to d=192 (ablation) | PDE-Transformer · Pre-training PDE test set, supervised | nRMSE: 0.045 -> 0.038 | Helped |
| increase token embedding dimension from d=192 to d=384 (ablation) | PDE-Transformer · Pre-training PDE test set, supervised | nRMSE: 0.038 -> 0.035 | Helped |
| use separate-channel rather than mixed-channel representation (ablation) | PDE-S trained from scratch · Active Matter | nRMSE: 0.493 -> 0.494 | No clear effect |
| use separate-channel rather than mixed-channel representation (ablation) | pre-trained PDE-S · Active Matter, first 20 rollout steps | nRMSE: 0.479 -> 0.455 | Helped |
| use separate-channel rather than mixed-channel representation (ablation) | PDE-S trained from scratch · RBC | nRMSE: 0.147 -> 0.155 | Hurt |
| use separate-channel rather than mixed-channel representation (ablation) | pre-trained PDE-S · RBC, first 20 rollout steps | nRMSE: 0.130 -> 0.104 | Helped |
| use separate-channel rather than mixed-channel representation (ablation) | PDE-S trained from scratch · Shear flow | nRMSE: 0.178 -> 0.199 | Hurt |
| use separate-channel rather than mixed-channel representation (ablation) | pre-trained PDE-S · Shear flow, first 20 rollout steps | nRMSE: 0.163 -> 0.125 | Helped |
Surprising findings
- Diffusion-based sampling is slightly less accurate than supervised training on single-sample nRMSE, even though it captures the posterior distribution.
- DiT-S shows persistent loss spikes it does not recover from, requiring checkpoint selection by validation loss.
Limitations
- Limited to 2D regular grids.
- Focuses on autoregressive prediction; noisy/partial data, data assimilation and inverse problems left to future work.
Experimental setup
Datasets: APEBench-based Exponax-simulated pre-training datasets (16 PDEs) The Well: active matter The Well: Rayleigh-Benard convection The Well: shear flow
Benchmarks: 16-PDE pre-training dataset The Well active matter The Well RBC The Well shear flow
Metrics: nRMSE GFlops GPU memory training time
Baselines: DiT-S UDiT-S scOT-S FactFormer UNet FNO OFormer Galerkin Transformer Poseidon scOT-B
Builds on: Diffusion Transformer (DiT) Swin Transformer UDiT adaLN-Zero APEBench Exponax
Closest papers by meaning
- Transolver: A Fast Transformer Solver for PDEs on General Geometries (ICML 2024 · similarity 0.89)
- Universal Physics Transformers: A Framework For Efficiently Scaling Neural Operators (NeurIPS 2024 · similarity 0.89)
- SEA: State-Exchange Attention for High-Fidelity Physics Based Transformers (NeurIPS 2024 · similarity 0.88)
- Multiple Physics Pretraining for Spatiotemporal Surrogate Models (NeurIPS 2024 · similarity 0.88)
- Text2PDE: Latent Diffusion Models for Accessible Physics Simulation (ICLR 2025 · similarity 0.88)
- CALM-PDE: Continuous and Adaptive Convolutions for Latent Space Modeling of Time-dependent PDEs (NeurIPS 2025 · similarity 0.88)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.