# Scalable Transformer for PDE Surrogate Modeling

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

## Key facts

- Cited by 128 papers in the map (data through 14 Oct 2026).
- Builds on 85 of its 123 references that are in the map.
- First posted on arXiv 27 May 2023.

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

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [High-Resolution Image Synthesis with Latent Diffusion Models](https://mlsanity.com/paper/2112.10752/) (2021)
- [End-to-End Object Detection with Transformers](https://mlsanity.com/paper/2005.12872/) (2020)
- [RoFormer: Enhanced Transformer with Rotary Position Embedding](https://mlsanity.com/paper/2104.09864/) (2021)
- [Neural Machine Translation by Jointly Learning to Align and Translate](https://mlsanity.com/paper/1409.0473/) (2014)
- [Longformer: The Long-Document Transformer](https://mlsanity.com/paper/2004.05150/) (2020)
- [Fourier Neural Operator for Parametric Partial Differential Equations](https://mlsanity.com/paper/2010.08895/) (2020)
- [Fourier Features Let Networks Learn High Frequency Functions in Low Dimensional Domains](https://mlsanity.com/paper/2006.10739/) (2020)
- [Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention](https://mlsanity.com/paper/2006.16236/) (2020)
- [Big Bird: Transformers for Longer Sequences](https://mlsanity.com/paper/2007.14062/) (2020)
- [Reformer: The Efficient Transformer](https://mlsanity.com/paper/2001.04451/) (2020)
- [Rethinking Attention with Performers](https://mlsanity.com/paper/2009.14794/) (2020)
- [Generating Long Sequences with Sparse Transformers](https://mlsanity.com/paper/1904.10509/) (2019)
- [Linformer: Self-Attention with Linear Complexity](https://mlsanity.com/paper/2006.04768/) (2020)
- [Instance Normalization: The Missing Ingredient for Fast Stylization](https://mlsanity.com/paper/1607.08022/) (2016)
- [Multi-Scale Context Aggregation by Dilated Convolutions](https://mlsanity.com/paper/1511.07122/) (2015)
- [Effective Approaches to Attention-based Neural Machine Translation](https://mlsanity.com/paper/1508.04025/) (2015)
- [Learning to Simulate Complex Physics with Graph Networks](https://mlsanity.com/paper/2002.09405/) (2020)
- [Learning Mesh-Based Simulation with Graph Networks](https://mlsanity.com/paper/2010.03409/) (2020)
- [FourCastNet: A Global Data-driven High-resolution Weather Model using Adaptive Fourier Neural Operators](https://mlsanity.com/paper/2202.11214/) (2022)
- [Neural Operator: Graph Kernel Network for Partial Differential Equations](https://mlsanity.com/paper/2003.03485/) (2020)
- [Neural Turing Machines](https://mlsanity.com/paper/1410.5401/) (2014)
- [DeepXDE: A deep learning library for solving differential equations](https://mlsanity.com/paper/1907.04502/) (2019)
- [Solving high-dimensional partial differential equations using deep learning](https://mlsanity.com/paper/1707.02568/) (2017)
- [Learning the solution operator of parametric partial differential equations with physics-informed DeepOnets](https://mlsanity.com/paper/2103.10974/) (2021)
- [Physics-Informed Neural Operator for Learning Partial Differential Equations](https://mlsanity.com/paper/2111.03794/) (2021)

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

- [Transolver: A Fast Transformer Solver for PDEs on General Geometries](https://mlsanity.com/paper/2402.02366/) (2024)
- [Improved Operator Learning by Orthogonal Attention](https://mlsanity.com/paper/2310.12487/) (2023)
- [Pretraining Codomain Attention Neural Operators for Solving Multiphysics PDEs](https://mlsanity.com/paper/2403.12553/) (2024)
- [AB-UPT: Scaling Neural CFD Surrogates for High-Fidelity Automotive Aerodynamics Simulations via Anchored-Branched Universal Physics Transformers](https://mlsanity.com/paper/2502.09692/) (2025)
- [Unisolver: PDE-Conditional Transformers Towards Universal Neural PDE Solvers](https://mlsanity.com/paper/2405.17527/) (2024)
- [PDE-Transformer: Efficient and Versatile Transformers for Physics Simulations](https://mlsanity.com/paper/2505.24717/) (2025)
- [On the Benefits of Memory for Modeling Time-Dependent PDEs](https://mlsanity.com/paper/2409.02313/) (2024)
- [Recent Advances on Machine Learning for Computational Fluid Dynamics: A Survey](https://mlsanity.com/paper/2408.12171/) (2024)
- [Generative Latent Neural PDE Solver using Flow Matching](https://mlsanity.com/paper/2503.22600/) (2025)
- [Text2PDE: Latent Diffusion Models for Accessible Physics Simulation](https://mlsanity.com/paper/2410.01153/) (2024)
- [Erwin: A Tree-based Hierarchical Transformer for Large-scale Physical Systems](https://mlsanity.com/paper/2502.17019/) (2025)
- [Neural Chaos: A Spectral Stochastic Neural Operator](https://mlsanity.com/paper/2502.11835/) (2025)
- [Transolver-3: Scaling Up Transformer Solvers to Industrial-Scale Geometries](https://mlsanity.com/paper/2602.04940/) (2026)
- [Multi-scale Time-stepping of Partial Differential Equations with Transformers](https://mlsanity.com/paper/2311.02225/) (2023)
- [Strategies for Pretraining Neural Operators](https://mlsanity.com/paper/2406.08473/) (2024)
- [P2C2Net: PDE-Preserved Coarse Correction Network for efficient prediction of spatiotemporal dynamics](https://mlsanity.com/paper/2411.00040/) (2024)
- [Vectorized Conditional Neural Fields: A Framework for Solving Time-dependent Parametric Partial Differential Equations](https://mlsanity.com/paper/2406.03919/) (2024)
- [Machine learning for modelling unstructured grid data in computational physics: a review](https://mlsanity.com/paper/2502.09346/) (2025)
- [Principled Approaches for Extending Neural Architectures to Function Spaces for Operator Learning](https://mlsanity.com/paper/2506.10973/) (2025)
- [Hierarchical Neural Operator Transformer with Learnable Frequency-aware Loss Prior for Arbitrary-scale Super-resolution](https://mlsanity.com/paper/2405.12202/) (2024)
- [Spectral-Refiner: Accurate Fine-Tuning of Spatiotemporal Fourier Neural Operator for Turbulent Flows](https://mlsanity.com/paper/2405.17211/) (2024)
- [Holistic Physics Solver: Learning PDEs in a Unified Spectral-Physical Space](https://mlsanity.com/paper/2410.11382/) (2024)
- [Latent Mamba Operator for Partial Differential Equations](https://mlsanity.com/paper/2505.19105/) (2025)
- [GeoPT: Scaling Physics Simulation via Lifted Geometric Pre-Training](https://mlsanity.com/paper/2602.20399/) (2026)
- [OmniArch: Building Foundation Model For Scientific Computing](https://mlsanity.com/paper/2402.16014/) (2024)
- [FUSE: Fast Unified Simulation and Estimation for PDEs](https://mlsanity.com/paper/2405.14558/) (2024)
- [Predicting Change, Not States: An Alternate Framework for Neural PDE Surrogates](https://mlsanity.com/paper/2412.13074/) (2024)
- [TANTE: Time-Adaptive Operator Learning via Neural Taylor Expansion](https://mlsanity.com/paper/2502.08574/) (2025)
- [AMR-Transformer: Enabling Efficient Long-range Interaction for Complex Neural Fluid Simulation](https://mlsanity.com/paper/2503.10257/) (2025)
- [FEDONet : Fourier-Embedded DeepONet for Spectrally Accurate Operator Learning](https://mlsanity.com/paper/2509.12344/) (2025)

## Abstract (arXiv)

Transformer has shown state-of-the-art performance on various applications and has recently emerged as a promising tool for surrogate modeling of partial differential equations (PDEs). Despite the introduction of linear-complexity attention, applying Transformer to problems with a large number of grid points can be numerically unstable and computationally expensive. In this work, we propose Factorized Transformer (FactFormer), which is based on an axial factorized kernel integral. Concretely, we introduce a learnable projection operator that decomposes the input function into multiple sub-functions with one-dimensional domain. These sub-functions are then evaluated and used to compute the instance-based kernel with an axial factorized scheme. We showcase that the proposed model is able to simulate 2D Kolmogorov flow on a 256× 256 grid and 3D smoke buoyancy on a 64×64×64 grid with good accuracy and efficiency. The proposed factorized scheme can serve as a computationally efficient low-rank surrogate for the full attention scheme when dealing with multi-dimensional problems.