Unisolver: PDE-Conditional Transformers Are Universal PDE Solvers (2405.17527v3)
Abstract: Deep models have recently emerged as a promising tool to solve partial differential equations (PDEs), known as neural PDE solvers. While neural solvers trained from either simulation data or physics-informed loss can solve PDEs reasonably well, they are mainly restricted to a few instances of PDEs, e.g. a certain equation with a limited set of coefficients. This limits the generalization of neural solvers to diverse PDEs, impeding them from being practical surrogate models for numerical solvers. In this paper, we present the Universal PDE Solver (Unisolver) capable of solving a wide scope of PDEs by training a novel Transformer model on diverse data and conditioned on diverse PDEs. Instead of purely scaling up data and parameters, Unisolver stems from the theoretical analysis of the PDE-solving process. Our key finding is that a PDE solution is fundamentally under the control of a series of PDE components, e.g. equation symbols, coefficients, and boundary conditions. Inspired by the mathematical structure of PDEs, we define a complete set of PDE components and flexibly embed them as domain-wise (e.g. equation symbols) and point-wise (e.g. boundaries) conditions for Transformer PDE solvers. Integrating physical insights with recent Transformer advances, Unisolver achieves consistent state-of-the-art results on three challenging large-scale benchmarks, showing impressive performance gains and favorable PDE generalizability.
- William F Ames. Numerical methods for partial differential equations. Academic press, 2014.
- Vladimir Igorevich Arnol’d. Mathematical methods of classical mechanics. Springer Science & Business Media, 2013.
- Message passing neural pde solvers. In ICLR, 2022.
- Language models are few-shot learners. In NeurIPS, 2020.
- Bert: Pre-training of deep bidirectional transformers for language understanding. In NAACL, 2019.
- An image is worth 16x16 words: Transformers for image recognition at scale. In ICLR, 2020.
- Sigmoid-weighted linear units for neural network function approximation in reinforcement learning. Neural networks, 2018.
- Lawrence C Evans. Partial differential equations. American Mathematical Society, 2022.
- Towards multi-spatiotemporal-scale generalized pde modeling. TMLR, 2023.
- Dpot: Auto-regressive denoising operator transformer for large-scale pde pre-training. In ICML, 2024.
- Gnot: A general neural operator transformer for operator learning. In ICML, 2023.
- Masked autoencoders are scalable vision learners. In CVPR, 2022.
- Elucidating the design space of diffusion-based generative models. NeurIPS, 2022.
- Adam: A method for stochastic optimization. In ICLR, 2015.
- Neural operator: Learning maps between function spaces with applications to pdes. JMLR, 2023.
- Transformer for partial differential equations’ operator learning. TMLR, 2023.
- Scalable transformer for pde surrogate modeling. In NeurIPS, 2024.
- Fourier neural operator with learned deformations for pdes on general geometries. JMLR, 2023.
- Fourier neural operator for parametric partial differential equations. In ICLR, 2021.
- Physics-informed neural operator for learning partial differential equations. J. Data Sci., 2021.
- Prose: Predicting operators and symbolic expressions using multimodal transformers. arXiv preprint arXiv:2309.16816, 2023.
- Swin transformer: Hierarchical vision transformer using shifted windows. In ICCV, 2021.
- Physics informed token transformer for solving partial differential equations. Mach. Learn.: Sci. Technol, 2024.
- Sgdr: Stochastic gradient descent with warm restarts. In ICLR, 2016.
- Learning nonlinear operators via deeponet based on the universal approximation theorem of operators. Nat. Mach. Intell, 2021.
- Cfdbench: A comprehensive benchmark for machine learning methods in fluid dynamics. arXiv preprint arXiv:2310.05963, 2023.
- Multiple physics pretraining for physical surrogate models. In NeurIPS AI for Science Workshop, 2023.
- Semantic image synthesis with spatially-adaptive normalization. In CVPR, 2019.
- Scalable diffusion models with transformers. In ICCV, 2023.
- Film: Visual reasoning with a general conditioning layer. In AAAI, 2018.
- Exploring the limits of transfer learning with a unified text-to-text transformer. 2020.
- U-no: U-shaped neural operators. TMLR, 2023.
- Physics-informed neural networks: A deep learning framework for solving forward and inverse problems involving nonlinear partial differential equations. J. Comput. Phys., 2019.
- Polynomial implicit neural representations for large diverse datasets. In CVPR, 2023.
- Towards foundation models for scientific machine learning: Characterizing scaling and transfer behavior. In NeurIPS, 2023.
- Learning neural pde solvers with parameter-guided channel attention. In ICML, 2023.
- Pdebench: An extensive benchmark for scientific machine learning. NeurIPS, 2022.
- Llama: Open and efficient foundation language models. arXiv preprint arXiv:2302.13971, 2023.
- Llama 2: Open foundation and fine-tuned chat models. arXiv preprint arXiv:2307.09288, 2023.
- Attention is all you need. In NeurIPS, 2017.
- Scientific discovery in the age of artificial intelligence. Nature, 2023.
- U-fno—an enhanced fourier neural operator-based deep-learning model for multiphase flow. Adv. Water Resour., 2022.
- Solving high-dimensional pdes with latent spectral models. In ICML, 2023.
- Transolver: A fast transformer solver for pdes on general geometries. In ICML, 2024.
- Pdeformer: Towards a foundation model for one-dimensional partial differential equations. In ICLR AI4Differential Equations In Science Workshop, 2024.
- Do transformers really perform badly for graph representation? NeurIPS, 2021.