Unisolver: PDE-Conditional Transformers Towards Universal Neural PDE Solvers

Hang Zhou, Yuezhou Ma, Haixu Wu, Haowen Wang, Mingsheng Long

TL;DR: This paper presents the Universal neural PDE solver (Unisolver) capable of solving a wide scope of PDEs by leveraging a Transformer pre-trained on diverse data and conditioned on diverse PDEs.

Abstract: Deep models have recently emerged as promising tools 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 their generalization to diverse PDEs, preventing them from being practical surrogate models of numerical solvers. In this paper, we present Unisolver, a novel Transformer model trained on diverse data and conditioned on diverse PDEs, aiming towards a universal neural PDE solver capable of solving a wide scope of PDEs. Instead of purely scaling up data and parameters, Unisolver stems from the theoretical analysis of the PDE-solving process. Inspired by the mathematical structure of PDEs that a PDE solution is fundamentally governed by a series of PDE components such as equation symbols and boundary conditions, we define a complete set of PDE components and flexibly embed them as domain-wise and point-wise deep conditions for Transformer PDE solvers. Integrating physical insights with recent Transformer advances, Unisolver achieves consistent state-of-the-art on three challenging large-scale benchmarks, showing impressive performance and generalizability. Code is available at [https://github.com/thuml/Unisolver](https://github.com/thuml/Unisolver).

Lay Summary: Solving physical equations that describe natural phenomena—called partial differential equations (PDEs)—is a critical task in science and engineering. Classical numerical methods can be slow and require re-computing for each specific task. Recently, deep models have been explored as efficient surrogates for numerical solvers. However, previous methods may struggle to handle diverse types of PDEs. In this work, we propose Unisolver, a powerful AI model based on Transformer architecture capable of solving a wide scope of PDEs. Instead of just training the model on more data, we carefully design the model to incorporate different components of PDEs, like the equations, coefficients and boundary conditions, just like what a numerical solver would receive. Unisolver learns from diverse PDE samples and achieves excellent results on challenging tests, making meaningful progress towards a universal AI tool for solving physical equations efficiently.

Primary Area: Deep Learning->Algorithms

Keywords: Neural PDE Solver, Deep Learning

Submission Number: 3382