Unisolver: PDE-Conditional Transformers Towards Universal Neural PDE Solvers

TL;DR

Unisolver introduces a Transformer-based neural PDE solver trained on diverse PDE data and conditioned on PDE components, achieving broad generalization and state-of-the-art results across multiple benchmarks.

Contribution

It presents a novel PDE-conditional Transformer model that leverages PDE components for universal solving capability, extending beyond traditional specialized neural PDE solvers.

Findings

Achieves state-of-the-art performance on large-scale PDE benchmarks.

Demonstrates strong generalization across diverse PDE types.

Integrates physical insights with Transformer architecture effectively.

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.