# PDE-Transformer: Efficient and Versatile Transformers for Physics Simulations

arXiv: https://arxiv.org/abs/2505.24717 · HTML: https://mlsanity.com/paper/2505.24717/

## Key facts

- Cited by 32 papers in the map (data through 14 Oct 2026).
- Builds on 65 of its 92 references that are in the map.
- First posted on arXiv 30 May 2025.

## Builds on (most cited first, 30 of 65)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2006.11239/) (2020)
- [U-Net: Convolutional Networks for Biomedical Image Segmentation](https://mlsanity.com/paper/1505.04597/) (2015)
- [High-Resolution Image Synthesis with Latent Diffusion Models](https://mlsanity.com/paper/2112.10752/) (2021)
- [Swin Transformer: Hierarchical Vision Transformer using Shifted Windows](https://mlsanity.com/paper/2103.14030/) (2021)
- [A Simple Framework for Contrastive Learning of Visual Representations](https://mlsanity.com/paper/2002.05709/) (2020)
- [Diffusion Models Beat GANs on Image Synthesis](https://mlsanity.com/paper/2105.05233/) (2021)
- [Masked Autoencoders Are Scalable Vision Learners](https://mlsanity.com/paper/2111.06377/) (2021)
- [Scalable Diffusion Models with Transformers](https://mlsanity.com/paper/2212.09748/) (2022)
- [Flow Matching for Generative Modeling](https://mlsanity.com/paper/2210.02747/) (2022)
- [Scaling Rectified Flow Transformers for High-Resolution Image Synthesis](https://mlsanity.com/paper/2403.03206/) (2024)
- [Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flow](https://mlsanity.com/paper/2209.03003/) (2022)
- [FiLM: Visual Reasoning with a General Conditioning Layer](https://mlsanity.com/paper/1709.07871/) (2017)
- [DeepONet: Learning nonlinear operators for identifying differential equations based on the universal approximation theorem of operators](https://mlsanity.com/paper/1910.03193/) (2019)
- [Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour](https://mlsanity.com/paper/1706.02677/) (2017)
- [Root Mean Square Layer Normalization](https://mlsanity.com/paper/1910.07467/) (2019)
- [Swin Transformer V2: Scaling Up Capacity and Resolution](https://mlsanity.com/paper/2111.09883/) (2021)
- [Florence: A New Foundation Model for Computer Vision](https://mlsanity.com/paper/2111.11432/) (2021)
- [Neural Operator: Graph Kernel Network for Partial Differential Equations](https://mlsanity.com/paper/2003.03485/) (2020)
- [Scaling Vision Transformers to 22 Billion Parameters](https://mlsanity.com/paper/2302.05442/) (2023)
- [Physics-Informed Neural Operator for Learning Partial Differential Equations](https://mlsanity.com/paper/2111.03794/) (2021)
- [All are Worth Words: A ViT Backbone for Diffusion Models](https://mlsanity.com/paper/2209.12152/) (2022)
- [Machine learning accelerated computational fluid dynamics](https://mlsanity.com/paper/2102.01010/) (2021)
- [Simple diffusion: End-to-end diffusion for high resolution images](https://mlsanity.com/paper/2301.11093/) (2023)
- [Axial Attention in Multidimensional Transformers](https://mlsanity.com/paper/1912.12180/) (2019)
- [Choose a Transformer: Fourier or Galerkin](https://mlsanity.com/paper/2105.14995/) (2021)
- [Transolver: A Fast Transformer Solver for PDEs on General Geometries](https://mlsanity.com/paper/2402.02366/) (2024)
- [Transformer for Partial Differential Equations' Operator Learning](https://mlsanity.com/paper/2205.13671/) (2022)
- [Foundation Models for Time Series Analysis: A Tutorial and Survey](https://mlsanity.com/paper/2403.14735/) (2024)

## Built on by (most cited first, 30 of 32)

- [Physics vs Distributions: Pareto Optimal Flow Matching with Physics Constraints](https://mlsanity.com/paper/2506.08604/) (2025)
- [MORPH: PDE Foundation Models with Arbitrary Data Modality](https://mlsanity.com/paper/2509.21670/) (2025)
- [Neural Emulator Superiority: When Machine Learning for PDEs Surpasses its Training Data](https://mlsanity.com/paper/2510.23111/) (2025)
- [(Sparse) Attention to the Details: Preserving Spectral Fidelity in ML-based Weather Forecasting Models](https://mlsanity.com/paper/2604.16429/) (2026)
- [Emergent Transfer of a Physics Foundation Model from Simulation to Laboratory Turbulence](https://mlsanity.com/paper/2606.01470/) (2026)
- [On the Generalization in Topology Optimization via Sensitivity-Conditioned Bernoulli Flow Matching](https://mlsanity.com/paper/2606.02179/) (2026)
- [SuperWing: a comprehensive transonic wing dataset for data-driven aerodynamic design](https://mlsanity.com/paper/2512.14397/) (2025)
- [Out-of-distribution generalization of deep-learning surrogates for 2D PDE-generated dynamics in the small-data regime](https://mlsanity.com/paper/2601.08404/) (2026)
- [Latent Generative Solvers for Generalizable Long-Term Physics Simulation](https://mlsanity.com/paper/2602.11229/) (2026)
- [ClothTransformer: Unified Latent-Space Transformers for Scalable Cloth Simulation](https://mlsanity.com/paper/2605.27852/) (2026)
- [Efficient Generative Transformer Operators For Million-Point PDEs](https://mlsanity.com/paper/2512.04974/) (2025)
- [A Mechanistic Analysis of Transformers for Dynamical Systems](https://mlsanity.com/paper/2512.21113/) (2025)
- [Iterative Refinement Neural Operators are Learned Fixed-Point Solvers: A Principled Approach to Spectral Bias Mitigation](https://mlsanity.com/paper/2605.24041/) (2026)
- [Learning Where to Simulate: Generative Active Sampling for Online PDE Surrogate Training](https://mlsanity.com/paper/2606.09949/) (2026)
- [TRIE: An Evaluation Framework for Stochastic PDE Surrogates](https://mlsanity.com/paper/2607.00196/) (2026)
- [StocBench: A Benchmark for Generative Modeling of Stochastic Dynamics](https://mlsanity.com/paper/2608.22309/) (2026)
- [Genomic Next-Token Predictors are In-Context Learners](https://mlsanity.com/paper/2511.12797/) (2025)
- [Neural-HSS: Hierarchical Semi-Separable Neural PDE Solver](https://mlsanity.com/paper/2602.18248/) (2026)
- [Factorized Neural Implicit DMD for Parametric Dynamics](https://mlsanity.com/paper/2603.10995/) (2026)
- [Extending Fourier Neural Operators for Modeling Parameterized and Coupled PDEs](https://mlsanity.com/paper/2607.23466/) (2026)
- [Unsupervised Adaptation of PDE Foundation Models](https://mlsanity.com/paper/2608.07053/) (2026)
- [Sim+Real: Joint Simulation-Experiment Training Improves Balanced Prediction in Physical Systems](https://mlsanity.com/paper/2610.01974/) (2026)
- [Hierarchical Inference and Closure Learning via Adaptive Surrogates for ODEs and PDEs](https://mlsanity.com/paper/2603.03922/) (2026)
- [Towards a Foundation-Model Paradigm for Aerodynamic Prediction in Three-dimensional Design](https://mlsanity.com/paper/2604.18062/) (2026)
- [Inpainting physics: self-supervised learning for context-driven fluid simulation](https://mlsanity.com/paper/2605.08832/) (2026)
- [DiffUNet^2: Bidirectional Conditional Diffusion for Probabilistic Scientific Spatiotemporal Modeling](https://mlsanity.com/paper/2606.03926/) (2026)
- [RD-JEPA: Predictive latent pretraining for few-trajectory transfer across reaction--diffusion equations](https://mlsanity.com/paper/2609.29403/) (2026)
- [Scale-Split Neural Operator for Memory- and Data-Efficient 3D Turbulence Prediction](https://mlsanity.com/paper/2609.38977/) (2026)
- [Geometry Meets Physics: Data-Efficient Pre-Training for Unstructured Neural PDE Solvers](https://mlsanity.com/paper/2610.03363/) (2026)
- [EC-EarthFlow: Probabilistic emulation of daily transient global climate model simulations with flow matching](https://mlsanity.com/paper/2610.09715/) (2026)

## Abstract (arXiv)

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.