# Unisolver: PDE-Conditional Transformers Are
Universal PDE Solvers

Hang Zhou, Yuezhou Ma11footnotemark: 1, Haixu Wu🖂, Haowen Wang, Mingsheng
Long🖂
School of Software, BNRist, Tsinghua University, China
<EMAIL_ADDRESS><EMAIL_ADDRESS>Equal Contribution

###### 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
the PDEs reasonably well, they are mainly restricted to a specific set of
PDEs, e.g. a certain equation or a finite set of coefficients. This bottleneck
limits the generalizability of neural solvers, which is widely recognized as
its major advantage over numerical solvers. In this paper, we present the
Universal 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. Instead of simply 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 initial and
boundary conditions. Inspired by the mathematical structure of PDEs, we define
a complete set of PDE components and correspondingly 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 gains and endowing
favorable generalizability and scalability.

##  1 Introduction

Partial differential equations (PDEs) are essential for numerous scientific
and engineering problems [8, 2], such as meteorology, electromagnetism and
thermodynamics [41]. Since it is usually hard to obtain an analytic solution
for PDEs, numerical methods are widely explored [1]. However, these numerical
methods often require huge computation costs to generate a precise solution
for each PDE. Recently, deep models have empowered many breakthroughs in wide
areas [5, 22] and have also been applied to solve PDEs, i.e. neural PDE
solvers. Owing to the excellent capability in approximating nonlinear
mappings, deep models can learn to fit pre-collected data [19] or physics-
informed loss function [33] and generalize in a flash to new samples, offering
an efficient way for solving PDEs.

As shown in Figure 1, previous neural solvers involve two different paradigms:
physics-informed neural works (PINNs) [33] and neural operators [19, 44],
which train the deep models with formalized PDE loss function and pre-
calculated simulation data respectively. However, for the first paradigm,
while formalizing the PDE equations as objective functions can ensure a
relatively accurate solution, PINNs are hard to generalize to new scenarios
and we have to retrain the model for different tasks. As for neural operators,
they directly learn from the simulation data, which can generalize better for
diverse initial states or conditions than PINNs. Nevertheless, purely based on
the training data may be insufficient to guide PDE solving. For example, given
a fluid governed by Navier-Stokes equations, the typical task of neural
operators is to predict the future based on past observations [19], while
different viscosity coefficients and forcing terms will lead to distinct
solutions even under the same initial state. Thus, due to the omission of PDE
information, current neural operators are mainly trained and tested on a
specific set of PDEs. Notably, as neural solvers are expected to be efficient
surrogates of neural methods, generalization to various PDEs is essential for
a practical neural solver.

To tackle the generalization deficiency, several works are proposed by
incorporating the PDE information into deep models or training models with a
large-scale dataset. For example, message-passing neural PDE solver [3]
concatenates the PDE coefficients with inputs. PDEformer [45] formalizes the
PDE equation as a computation graph and employs the graph Transformer [46] to
aggregate PDE information. Although these methods have explored the potential
of training models with both data and PDE information, they fail to support a
complete set of PDE information, thereby limiting the generalizability of
models in some aspects. As for the other branches, such as DPOT [10], they
just scale up the training sets with diverse PDEs and expect the
generalizability emerges from large data and parameters. However, without the
complete PDE information, it is theoretically impossible to generate a correct
solution because what the models essentially perform is fitting an incomplete
data distribution, so it is hard to expect satisfactory generalizability.

Figure 1: Neural PDE solvers typically consist of two paradigms: physics
informed and data driven. Our proposed Unisolver combines data-driven methods
with physical insights from complete PDE components in a conditional modeling
framework, thereby boosting generalizability and scalability.

Going beyond prior methods, as shown in Figure 1, this paper presents
Unisolver as a Universal PDE solver. Concretely, Unisolver takes the
advantages from both data-driven and physics-informed paradigms and empowers
Transformer with favorable PDE generalizability by introducing complete
physics information as conditions. Instead of simply scaling up data and
parameters, we start from the theoretical analysis of PDE solving and conclude
with a complete set of PDE components. Further, drawing inspiration from the
mathematical structure of PDEs, we propose to classify PDE components into
domain-wise and point-wise categories according to their effect on the final
solution and aggregate them as two types of deep PDE conditions. Afterward, to
capture the special influence of different condition types on the hidden input
feature, we separate the Transformer representation into two subspaces and
integrate these deep conditions into Transformer in a decoupled way. We
conduct extensive experiments on our own generated dataset and two large-scale
benchmarks with various PDE conditions, where Unisolver achieves consistent
state-of-the-art with sharp relative gains. Overall, our contributions are
summarized as follows:

  * •

We introduce Unisolver as a conditional Transformer architecture utilizing
embedded PDE information completely, marking the first demonstration of the
potential of the canonical Transformer as a scalable backbone for solving
multiple PDEs with diverse conditions.

  * •

Motivated by the mathematical structure of PDEs, we classify universal PDE
conditions into domain-wise and point-wise categories, which further derives a
decoupled conditioning mechanism incorporating designs for introducing PDE
conditions to deep representations.

  * •

Unisolver achieves consistent state-of-the-art performances across three
large-scale benchmarks with impressive relative gains and presents favorable
generalizability and scalability.

##  2 Related Work

###  2.1 Neural PDE Solvers

Previous neural PDE solvers can be roughly categorized into the following two
paradigms. The first paradigm is physics-informed neural networks (PINNs)
[33], which optimize the deep model by formalizing the PDE equations as
objective functions. During training, the model outputs and gradients will
gradually satisfy the targeted PDE, thereby successfully instantiating the
solution as a deep model. However, PINNs are usually hard to generalize to
unseen PDEs, limiting their applications in broader practice [41]. Another
booming direction is neural operators, which learn from extensive data to
approximate functional dependence between input and output Banach spaces [25,
15]. Among various neural operators, FNO [19] and its variants [18, 32, 42]
are popular and well-established. FNO [19] effectively approximates the kernel
integral operator in the frequency domain through Fourier transformation.
Besides, LSM [43] generalizes the classical spectral method in latent space to
avoid the curse of dimensionality. Recently, given the impressive progress
achieved by Transformers [40], they have also been applied to solve PDEs.
However, due to the quadratic complexity of the attention mechanism, the main
challenge in applying Transformers is defining the tokens used to compute
attention in PDE-solving tasks. OFormer [16] and GNOT [11] treat each mesh
point as a token and utilize the linear Transformer to avert the complexity
problem. Factformer [17] axially factorizes the attention block to boost the
model efficiency. Recently, Transolver [44] proposes to learn the intrinsic
physical states as tokens behind the input meshes. Despite the success of
neural operators, they are only tested on the dataset that contains one
certain PDE. The effectiveness of these methods under large and diverse
datasets has not been fully explored.

###  2.2 Generalizable PDE Solvers

In addition to model architectures, the generalizability of neural solvers,
the major advantage over numerical solvers, has also been explored. The
research mainly lies in the following two directions.

#### Incorporating PDE information

To guide the PDE-solving process, the PDE information has been introduced to
deep models. For example, PINO [20] imposes explicit equation constraints at a
higher resolution to assist the learning of neural operators. CAPE [36] embeds
PDE coefficients to adapt neural solvers to unseen PDE parameters. Prose [21]
and PITT [23] tokenize PDEs and embed the mathematical expressions to make the
transformer backbone aware of the underlying physics. PDEformer [45]
represents the symbolic form of equations as a graph and the numeric
components as nodes to optimize the processing of complex interactions between
symbolic and numeric information. Still, all of these methods, while
incorporating equation information, do not leverage the mathematical structure
of PDEs for complete and categorized embedding or integrating the meanings of
equation symbols within the context of natural language. In contrast,
Unisolver applies powerful Large Language Models (LLMs) [38] to semantically
embed the equation symbolic information and categorize the complete equation
conditions based on mathematical insights, thereby better at modeling
intricate physical correlations.

#### Large-scale pre-training

As a vital cornerstone of deep learning [4, 12], recent research has also
started to explore the effectiveness of large-scale training in solving PDEs.
In [35], authors examined the scaling capabilities and transfer learning
behaviors of FNO on three time-independent equation families. MPP [27]
proposes an auto-regressive strategy to pre-train on a broad fluid mechanics-
oriented benchmark. DPOT [10] enhances MPP with the denoising method and pre-
trains a Fourier transformer on massive PDE data comprised of 12 datasets.
PDEformer [45] focuses on a 1D time-dependent PDE family and pre-trains a
graph transformer on 3M samples under various equation conditions. However,
most of the existing methods fall short in effectively and completely
integrating PDE conditions. This will be well addressed by Unisolver in a
natural and insightful way.

##  3 Unisolver

To tackle the incapability in generalization behind neural PDE solvers, we
deeply dive into the PDE-solving process and present Unisolver to model the
intricate interactions between initial observations and complete equation
components, leading to a novel PDE-conditional Transformer model.

#### Problem setup

Let domain $\mathcal{D}\subset\mathbb{R}^{d}$ be a bounded continous set and
$\mathcal{M}=\\{x_{1},\dots,x_{n}\\}$ be a $n$-point discretization of
$\mathcal{D}$, recording the coordinate information of each point. Assume that
we have observations of initial conditions $\mathbf{X}$ as input and target
quantities $\mathbf{Y}$ as output on mesh $\mathcal{M}$ with governed PDE
components $\mathbf{C}_{\text{PDE}}$ (e.g. equation symbols, coefficients,
etc) for each observation pair. The PDE-solving task is to approximate the
input-PDE-output mapping
$G:(\mathbf{X},\mathcal{M},\mathbf{C}_{\text{PDE}})\rightarrow\mathbf{Y}$. For
instance, in fluid prediction, we need to predict the future vorticity field
based on the initial field, observation grid, as well as PDE equations,
viscosity coefficients, and force terms.

###  3.1 Complete PDE Components

To enable complete modeling of the PDE, we attempt to incorporate all the
underlying components that affect final solutions into neural solvers. Here,
we elucidate the complete PDE components by considering the classical
_vibrating string equation_ with fixed endpoints as a motivating example:

$\displaystyle\partial_{tt}u-a^{2}\partial_{xx}u=f(x,t),$ $\displaystyle(x,t)$
$\displaystyle\in(0,L)\times(0,T),$ (1a) $\displaystyle u(0,t)=0,\;u(L,t)=0,$
$\displaystyle t$ $\displaystyle\in(0,T],$ (1b) $\displaystyle
u(x,0)=\phi(x),\;\partial_{t}u(x,0)=\psi(x),$ $\displaystyle x$
$\displaystyle\in[0,L].$ (1c)

The equation formulation precisely models the physical laws governing string
vibrations. In this PDE, the coefficient $a$ is determined by physical
properties, such as tension and linear density, while $f$ represents the
external force that drive the vibrations of the string. The domain geometry
spans the range $[0,L]\times[0,T]$, which affects the mesh discretization in
numerical solving, especially for irregular geometries. In addition, Equation
(1b) sets Dirichlet boundary conditions at the endpoints, maintaining zero
displacement throughout the vibration period. Equation (1c) specifies the
string’s position $\phi(x)$ and velocity $\psi(x)$ at time zero. These
complete equation components theoretically determine a unique solution. In the
following theorem, we dive into the PDE-solving process and explore how the
equation components affect the final solution.

######  Theorem 3.1 (The connections between PDE solutions and complete PDE
components).

The analytical solution of equation (1a) with boundary conditions (1b) and
initial conditions (1c) is

$u(x,t)=\frac{1}{2}(\underbrace{\Phi(x+\text{{\color[rgb]{0.2627,0.4196,0.6510}a}}t)+\Phi(x-\text{{\color[rgb]{0.2627,0.4196,0.6510}a}}t)}_{\text{{\color[rgb]{0.6784,0.0275,0.0510}Initial
position}}})+\frac{1}{2\text{{\color[rgb]{0.2627,0.4196,0.6510}a}}}\underbrace{\int^{x+\text{{\color[rgb]{0.2627,0.4196,0.6510}a}}t}_{x-\text{{\color[rgb]{0.2627,0.4196,0.6510}a}}t}\Psi(\xi)}_{\text{{\color[rgb]{0.6784,0.0275,0.0510}Initial
velocity}}}\mathrm{d}\xi+\frac{1}{2\text{{\color[rgb]{0.2627,0.4196,0.6510}a}}}\underbrace{\int_{0}^{t}\mathrm{d}\tau\int_{x-\text{{\color[rgb]{0.2627,0.4196,0.6510}a}}(t-\tau)}^{x+\text{{\color[rgb]{0.2627,0.4196,0.6510}a}}(t-\tau)}}_{\text{{\color[rgb]{0.2627,0.4196,0.6510}Geometry}}}\underbrace{F(\xi,\tau)}_{\text{{\color[rgb]{0.2627,0.4196,0.6510}Force}}}\mathrm{d}\xi,$
(2)

where $\Phi(x)$, $\Psi(x)$ and $F(x,t)$ are odd, periodic functions with
period 2L defined on the upper half plane, extended from $\phi(x)$, $\psi(x)$
and $f(x,t)$. The boundary conditions will be explicit by extending the
equation to the upper half plane and solving it by operator splitting and
characteristic lines.

###### Proof.

See Appendix B for the explicit solution and the complete proof. ∎

From the analytical solution of the simple motivating example, we pinpoint
that the PDE is solved under the complex interactions between a series of
equation components. Specifically, the coefficient $a$ exerts a consistent
influence over the domain, while the impact of the external force $f$ is
imposed point-wisely. This distinction inspires us to classify the PDE
components into two categories, _domain-wise_ and _point-wise_ , which better
capture the intricate interactions.

Table 1: Concrete categorization of complete PDE components: domain-wise and point-wise. Domain-wise components | Point-wise components
---|---
Equation formulation | External force
Equation coefficient | Domain geometry
Boundary condition type | Boundary value function

As presented in Table 1, the equation formulation is a domain-wise component
due to its consistency across all locations. Domain geometry is a point-wise
component since it is represented by a binary mask and each point’s inclusion
is determined individually. Boundary conditions are slightly more complex to
handle. Commonly used boundary conditions are divided into two main types:
periodic boundary conditions and Robin (non-periodic) conditions. We classify
the boundary condition type as a domain-wise component and the specific
boundary value functions with zero values at non-boundary points as point-wise
components.

###  3.2 Universal Components Embedding

Figure 2: Overview of Unisolver. We universally embed all PDE components into
deep conditions and employ a conditional Transformer to aggregate deep
conditions in the decoupled subspace.

As described in Section 3.1, PDE solutions can be obtained by intricate
interactions between initial conditions and multiple equation components
classified into two categories. In previous works, these components are
coarsely and incompletely included as conditions by using them to modulate the
input observations. In this paper, we will elaborate on how Unisolver finely
and completely embeds all equation components (Table 1) into deep PDE
conditions based on our insights from the analysis.

#### Equation formulation

Since the mathematical symbols convey rich mathematical information, we
utilize a Large Language Model (LLM) for symbolic embedding. Specifically, we
employ the recently released LLaMA 3 8B model
111https://ai.meta.com/blog/meta-llama-3/ to embed the equation formulation,
leveraging its symbolic understanding ability from pre-training on 15 TB
language tokens. The input to the LLM is the LaTeX code of the equation
formulation. For example, the string vibration equation is prompted as

Prompt: "\partial_{tt} u - a^2\partial_{xx} u = f(x,t)"

Then we take the output of the last Transformer block of the LLM and compute
the mean across the sequence dimension, resulting in a 4096-dimensional
embedding for each equation. Since we embed the complete equation components
separately, we represent the other components, e.g. coefficients and external
force, as symbols in the prompt rather than using their corresponding values.
After LLM embedding, the hidden formulation is encoded by a two-layer MLP to
obtain deep conditions.

#### Other components

Other domain-wise components encompass coefficients, which are real-valued
vectors, and boundary types, which are analogous to class labels. Thus we
embed them using two MLP layers with in-between SiLU activation function [7].
Moreover, point-wise components include external force, binary geometry mask,
and boundary value functions, which are physical fields observed on mesh
$\mathcal{M}$. We employ the same patchify method used for the input
observations, transforming them into token sequences and linearly mapping each
token to a deep representation.

#### Deep condition consolidation

After universal components embedding, deep conditions within the same category
are aggregated together to consolidate their influence. This strategy is
advantageous because it prevents excessive separation of component conditions
that could weaken the model’s expressive capabilities, and thus will enhance
representation learning for PDE solving via pre-training. Additionally, it
allows for easy adaptation to new conditions in downstream tasks by only
requiring finetuning a zero-initialized MLP.

###  3.3 PDE-Conditional Transformer

We propose a conditional Transformer to adaptively fuse these deep conditions
embedded from the equation components, empowered by the projections of input
observations in decoupled subspaces.

#### Subspace decoupling

We split the hidden inputs along the feature dimension guided by a
hyperparameter, the domain-wise condition ratio $\alpha$, which is set to 0.5
by default. This ratio determines the subspace for domain-wise conditions with
dimension $d_{\text{domain}}=\lfloor\alpha\cdot d_{\text{feature}}\rfloor$,
while the remaining subspace of dimension
$d_{\text{point}}=d_{\text{feature}}-d_{\text{domain}}$ is used by point-wise
conditions. When combined with multi-head attention [40], subspace decoupling
is equivalent to assigning some heads to learn the impact of domain-wise
conditions while others focusing on point-wise conditions. This leads to
improved representation learning for both categories, and minimized
interference across the subspace.

#### Deep condition aggregation

We utilize MLPs to individually project domain-wise conditions and point-wise
conditions into the corresponding subspace. After projection, domain-wise
conditions are repeated along the sequence dimension to match the length of
hidden tokens and ensure consistent physical guiding along the entire
sequence. The transformed conditions convey both token-wise and feature-wise
information, which are then integrated adaptively by aggregation functions.

As shown in Figure 2, we aggregate conditions either before or after the
attention and feedforward modules. Inspired by recent conditional Transformer
advances like DiT [29] and other conditional normalization approaches [28,
30], we take the aggregation paradigm to finely capture the intricate
correlations between hidden inputs of initial observations and deep equation
conditions. Specifically, we scale and shift the hidden inputs based on the
conditions. After passing through the Transformer modules, we use the
conditions to softly select whether this information should be retained.

#### Overall design

We embed input of initial observations $\mathbf{X}$ using a patchification
layer similar to ViT [6] to obtain the input embeddings $\mathbf{X}^{0}$, and
embed the complete PDE equation components $\mathbf{C}_{\text{PDE}}$ following
Section 3.2 to obtain deep conditions $\mathbf{C}_{\text{domain}}\text{ and
}\mathbf{C}_{\text{point}}$. Suppose there are $N$ layers in the PDE-
conditional Transformer architecture, the $n$-th layer of Unisolver can be
formalized as follows:

$\begin{split}\mathbf{I}_{\text{op}}&=\operatorname{Concat}\big{(}\operatorname{MLP}(\mathbf{C}_{\text{domain}}).\texttt{repeat},\operatorname{MLP}(\mathbf{C}_{\text{point}})\big{)},\\\
{\mathbf{X}}^{n-\frac{1}{2}}&=\mathbf{I}_{\text{select}}\odot\operatorname{Self-
Attention}\big{(}\mathbf{I}_{\text{scale}}\odot\operatorname{{LayerNorm}}({\mathbf{X}}^{n-1})+\mathbf{I}_{\text{shift}}\big{)}+{\mathbf{X}}^{n-1},\\\
{\mathbf{I}}^{\prime}_{\text{op}}&=\operatorname{Concat}\big{(}\operatorname{MLP}^{\prime}(\mathbf{C}_{\text{domain}}).\texttt{repeat},\operatorname{MLP}^{\prime}(\mathbf{C}_{\text{point}})\big{)},\\\
{\mathbf{X}}^{n}&={\mathbf{I}}^{\prime}_{\text{select}}\odot\operatorname{FeedForward}\big{(}{\mathbf{I}}^{\prime}_{\text{scale}}\odot\operatorname{LayerNorm}({\mathbf{X}}^{n-\frac{1}{2}})+{\mathbf{I}}^{\prime}_{\text{shift}}\big{)}+{\mathbf{X}}^{n-\frac{1}{2}},\end{split}$
(3)

where $\forall\text{op}\in\\{\text{scale, shift, select}\\}$,
$n\in\\{1,\ldots,N\\}$, and ${\mathbf{X}}^{n}$ is the output of the $n$-th
layer. As the PDE components have a crucial impact on the scale of the output,
we apply conditional scale and shift operations before a linear projection on
$\mathbf{X}^{N}$ to obtain the final output $\mathbf{Y}$, as shown in Figure
2.

While our model takes similar aggregation mechanisms as DiT, it includes
several PDE-inspired differences. First, we do not use a diffusion model
because PDE solving demands accuracy more than diversity. Second, our
conditions are drawn from equation components with physical insights, unlike
the timesteps and text descriptions typically used in generative models [13].
Third, our conditions are not confined to scalar values; instead, we decouple
the feature space and leverage point-wise conditions special to PDEs to learn
interactions between initial observations and PDE components.

##  4 Experiments

We conduct extensive experiments to evaluate Unisolver on a challenging, self-
generated Heterogeneous 2D Navier-Stokes Equations dataset and two large-scale
benchmarks proposed by PDEformer [45] and DPOT [10], covering a broad range of
PDEs and diverse generalization scenarios.

#### Benchmarks

As summarized in Table 2, our experiments involve three large-scale datasets
and numerous equation conditions in both 1D and 2D spaces. Specifically, we
generate the HeterNS dataset as an extension for the NS dataset of FNO [19],
including multiple coefficients and forces. PDEformer proposes a large-scale
dataset with 3M samples of structured 1D PDEs and evaluates the pre-trained
model on PDEBench [37]. DPOT collects 12 datasets from FNO [19], PDEBench
[37], PDEArena [9] and CFDBench [26] for diversity. More details can be found
in Appendix C.

Table 2: Summary of benchmarks. $\checkmark$ indicates the PDE component will change among different samples, while $\times$ refers to unchanged ones. Unisolver can handle all kinds of changed components. Benchmarks | #Dim | #Resolution | # Samples | #Size | Symbols | Coefficient | Force | Geometry | Boundary
---|---|---|---|---|---|---|---|---|---
HeterNS | 2D+Time | (64,64,10) | 15k | 4.6 GB | $\times$ | $\checkmark$ | $\checkmark$ | $\times$ | $\times$
PDEformer [45] | 1D+Time | (256,100) | 3M | 300 GB | $\checkmark$ | $\checkmark$ | $\checkmark$ | $\times$ | $\checkmark$
DPOT [10] | 2D+Time | (128,128,10) | 74.1k | 384 GB | $\checkmark$ | $\checkmark$ | $\checkmark$ | $\checkmark$ | $\checkmark$

#### Baselines

We compare Unisolver with three advanced baselines on the HeterNS dataset to
demonstrate its generalizability under varied conditions, including classical
ViT [6], FNO [19] and current state-of-the-art neural operator Factformer
[17]. Specifically, to ensure a fair comparison, we also include the
additional PDE components used in Unisolver to these baselines by
concatenating them with inputs. We also compare Unisolver with other
generalizable solvers: PDEformer [45] and DPOT [10] on pre-training loss and
zero-shot and fine-tuning settings.

#### Implementations

For fairness, in the HeterNS dataset, all the methods are trained with
relative L2 loss for 300 epochs, using the ADAM [14] optimizer with an initial
learning rate of $0.0005$ and a cosine annealing learning rate scheduler [24].
The batch size is set to 60. In the large-scale PDEformer and DPOT benchmarks,
we follow their training strategies to compare unbiasedly. We adopt relative
L2 as the evaluation metric. See Appendix D for thorough implementation
details.

Table 3: Performance comparison on HeterNS with different viscosity coefficients and fixed force frequency coefficient $\omega=2$. Relative L2 is recorded. Promtotion refers to the relative improvement over the second-best method. "ID" means in-distribution case and "OOD" means the out-of-distribution. HeterNS |  Viscosity $\nu$ | 8e-6 | 1e-5 | 3e-5 | 5e-5 | 8e-5 | 1e-4 | 3e-4 | 5e-4 | 8e-4 | 1e-3 | 2e-3
---|---|---|---|---|---|---|---|---|---|---|---|---
Params | OOD | ID | OOD | ID | OOD | ID | OOD | ID | OOD | ID | OOD
FNO | 4.7M | 0.0702 | 0.0669 | 0.0373 | 0.0225 | 0.0141 | 0.0114 | 0.0088 | 0.0031 | 0.0084 | 0.0011 | 0.2057
Factformer | 9.4M | 0.0489 | 0.0438 | 0.0489 | 0.0128 | 0.0297 | 0.0064 | 0.1386 | 0.0018 | 0.0631 | 0.0010 | 0.3207
ViT | 4.8M | 0.0458 | 0.0432 | 0.0353 | 0.0206 | 0.0119 | 0.0098 | 0.0100 | 0.0031 | 0.0174 | 0.0015 | 0.1878
Unisolver | 4.1M | 0.0336 | 0.0321 | 0.0178 | 0.0094 | 0.0064 | 0.0051 | 0.0066 | 0.0015 | 0.0096 | 0.0008 | 0.1504
Promotion | - | 26.6% | 25.7% | 49.6% | 26.6% | 46.2% | 20.3% | 34.0% | 16.7% | - | 20.0% | 19.9%
Table 4: Comparison (relative L2) on HeterNS with varied force and fixed viscosity $\nu=10^{-5}$. HeterNS |  Force $\omega$ | 0.5 | 1 | 1.5 | 2 | 2.5 | 3 | 3.5
---|---|---|---|---|---|---|---|---
Params | OOD | ID | OOD | ID | OOD | ID | OOD
FNO | 4.7M | 1.110 | 0.0640 | 0.1742 | 0.0661 | 0.1449 | 0.1623 | 0.2974
Factformer | 9.4M | 0.9998 | 0.0326 | 0.1110 | 0.0438 | 0.1243 | 0.0803 | 0.2257
ViT | 4.8M | 0.7900 | 0.0348 | 0.1412 | 0.0432 | 0.1240 | 0.1000 | 0.2080
Unisolver | 4.1M | 0.0980 | 0.0244 | 0.0770 | 0.0321 | 0.0720 | 0.0720 | 0.1740
Promotion | - | 87.6% | 25.2% | 30.6% | 25.7% | 41.9% | 10.3% | 16.4%

###  4.1 Heterogeneous 2D Navier-Stokes Equation (HeterNS)

#### Setups

To test the model performance on diverse varied coefficients, we generate
HeterNS, which is extended from the widely used 2D NS dataset [19]. The
original 2D NS dataset only contains the fixed viscosity coefficient and
external force for each task, which cannot be used to evaluate the
effectiveness of the model in handling varied equation components. In
contrast, HeterNS involves five different viscosity coefficients and three
different external force within the family of trigonometric functions,
resulting in 15 different equation component pairs and 15k samples.
Specifically, HeterNS considers viscosity coefficient
$\nu\in\\{10^{-5},5\times 10^{-5},10^{-4},5\times 10^{-5},10^{-3}\\}$ and
force frequency coefficient $\omega\in\\{1,2,3\\}$. In this benchmark, we
evaluate the in-distribution fitting ability and out-of-distribution
generalization performance on different viscosity $\nu$ and force $\omega$.

#### Results

As shown in Tables 3-4, Unisolver achieves the best performance on 10 of 11
tasks, covering both ID and OOD settings. Notably, for external force
generalization, Unisolver beats other methods more significantly in OOD
conditions (average 44.1%) than ID conditions (average 20.4%), demonstrating
the effectiveness of our design in capturing generalizable physics relations
between external force and model inputs. Even though we explicitly concatenate
the condition information with the model inputs, without special designs for
PDE conditions, most advanced neural operators perform unremarkably on
HeterNS. Specifically, all neural operators fail to solve the case of
$\omega=0.5$ in Table 4, highlighting the benefits of our conditional
Transformer design in learning generalizable features. We also include
experiments in Appendix E.1, where both $\nu$ and $\omega$ are OOD. Unisolver
still achieves considerable improvement (average 41.3%) on this challenging
double OOD setting.

###  4.2 1D Time-dependent PDEs

#### Setups

Proposed by PDEformer [45], this benchmark contains 3 million high-quality 1D
time-dependent PDE samples with varied equation symbols, coefficients, forces,
and boundary conditions for pre-training, with periodic and Robin boundary
conditions each comprising half of the dataset. The input for the pre-training
task consists of all the PDE components, and the output is the full space-time
field. Similar to PDEformer, we utilize an adapted version of Poly-INR [34] to
decode the PDE solution from the features extracted by Unisolver. After pre-
training, the pre-trained model is evaluated using Burgers equation and
Advection equation from PDEBench[37], where the PDE coefficients fall within
the distribution of the pre-training data. Concretely, we recorded the results
for zero-shot on Burgers and Advection, and finetuning on the these two
dataset after 100 epochs.

Table 5: Performance (relative L2) comparison of pre-training and finetuning on 1D time-dependent PDEs. PDEformer-L means the large version of PDEformer. FT-100 means finetuning 100 epochs. 1D Time-dependent PDEs | Tasks | Pretrain | Burgers | Advection
---|---|---|---|---
Params | Periodic | Robin | $\nu=0.1$ | $\nu=0.01$ | $\nu=0.001$ | $\beta=0.1$
PDEformer-L | 22M | 0.0211 | 0.0238 | 0.00744 | 0.0144 | 0.0393 | 0.0178
Unisolver | 19M | 0.0107 | 0.0108 | 0.00513 | 0.00995 | 0.0299 | 0.0138
Promotion | - | 49.3% | 54.6% | 31.0% | 30.9% | 23.9% | 22.5%
PDEformer-L (FT-100) | 22M | - | - | 0.00364 | 0.0112 | 0.0331 | 0.00975
Unisolver (FT-100) | 19M | - | - | 0.00105 | 0.00474 | 0.0170 | 0.00420
Promotion | - | - | - | 71.2% | 57.7% | 48.6% | 59.6%

#### Results

Table 5 presents that the pre-training performance of Unisolver is
significantly better than that of PDEformer-L on both periodic and Robin
boundary conditions, which indicates that our design of incorporating PDE
conditions is more effective than the computational graph utilized by
PDEformer in modeling complex distributions. Additionally, UniSolver achieves
better performance in zero-shot evaluations on downstream tasks, with an
average improvement of 27.1% over PDEformer-L. Furthermore, fine-tuning on
downstream tasks yields even more pronounced improvements for Unisolver
(average 59.3%), highlighting its larger model capacity despite relatively
fewer parameters.

###  4.3 2D Mixed PDEs

#### Setups

This benchmark is gathered by DPOT [10], which involves 12 datasets from four
different sources, covering a wide range of PDEs. For conciseness, we adopt
the “source-PDE” to record different tasks in Table 6, e.g. FNO-NS. The second
row lists the primary conditions considered for each dataset. For example,
1e-5 represents the viscosity coefficient $\nu=10^{-5}$. Unlike the well-
balanced data in PDEformer, these datasets exhibit significant imbalance among
different conditions. To address this issue, we adopted the balanced data
sampling method used in DPOT, but it still presents thorny challenges in
tackling intricate and multifarious PDE samples.

Table 6: Performance comparison across 12 pre-training datasets. Relative L2 ($\times 10^{-2}$) is recorded. 2D |  Equations |  FNO-NS-$\nu$ |  PDEBench-CNS-(M,$\zeta$) |  DR |  SWE |  PDEArena-NS |  CFDBench-NS |  Mean
---|---|---|---|---|---|---|---|---
Mixed PDEs |  Params |  1e-5 |  1e-4 |  1e-3 |  (1, 0.1) |  (1, 0.01) |  (0.1, 0.1) |  (0.1, 0.01) | - | - | - |  Force |  Geometry | -
DPOT-S |  30M |  5.53 |  4.42 |  1.31 |  1.53 |  3.37 |  1.19 |  1.87 |  3.79 |  0.66 |  9.91 |  31.6 |  0.70 |  5.50
Unisolver |  33M |  4.17 |  3.36 |  0.61 |  1.23 |  2.89 |  1.01 |  1.59 |  4.39 |  0.45 |  6.87 |  27.4 |  0.54 |  4.54
Promotion | - |  24.6% |  24.0% |  53.4% |  19.6% |  14.2% |  15.1% |  15.0% | - |  31.8% |  30.7% |  13.3% |  22.9% |  17.5%
DPOT-S-FT |  30M |  4.49 |  3.42 |  0.68 |  1.52 |  2.11 |  1.50 |  1.51 |  1.71 |  0.22 |  8.92 |  29.0 |  0.44 |  4.63
Unisolver-FT |  33M |  3.82 |  2.79 |  0.31 |  0.95 |  1.99 |  1.01 |  1.34 |  1.37 |  0.20 |  6.67 |  26.8 |  0.47 |  3.98
Promotion | - |  14.9% |  18.4% |  54.4% |  37.5% |  5.69% |  32.6% |  11.3% |  19.8% |  9.1% |  25.2% |  7.6% | - |  14.0%

#### Results

We pre-trained Unisolver with 33M parameters to compare with DPOT-S, the small
version of DPOT. As shown in Table 6, our method outperforms DPOT in 11 out of
12 datasets except for a small dataset DR whose relative L2 is less than 5%,
verifying the effectiveness of our design in modeling complex conditions.
After 200 epochs of fine-tuning on each dataset, Unisolver achieves a
reduction in error of more than 12% compared with zero-shot results,
demonstrating its ability to extract useful knowledge from complex pre-
training datasets. These results further highlight the potential of Unisolver
for serving as a foundation PDE solver.

###  4.4 Model Analysis

Table 7: Ablations on different LLM choices, removing our designs (W/o), and replacing domain-wise or point-wise conditions from our design to input concat (Concat-D and Concat-P). L2RE | Different LLMs | W/o | Condition Modeling | Ours
---|---|---|---|---
Dataset |  T5 [31] |  LLaMa-2 [39] | LLM Embedding | Subspaces | Concat-D | Concat-P
Pretrain-periodic | 0.0329 | 0.0329 | 0.0358 | 0.0331 | 0.0433 | 0.0327 | 0.0330
Pretrain-robin | 0.0223 | 0.0224 | 0.0232 | 0.0243 | 0.0427 | 0.0306 | 0.0223
Burgers $\nu=0.01$ | 0.0664 | 0.0647 | 0.0692 | 0.0675 | 0.0950 | 0.0769 | 0.0659

#### Ablations

We conduct detailed ablations on 1D time-dependent PDEs of 50k samples, with
50% periodic and 50% Robin boundaries. From Table 7, we have the following
observations. Firstly, without equation symbols, model performance on all
tasks will drop consistently, demonstrating the necessity of introducing
symbolic information. However, the choice of LLM has minimal impact on
performance, probably due to the limited ablation data and the prompt design.
Secondly, removing subspace decoupling impairs the performance on robin
boundaries more significantly for the insufficient modeling of point-wise
conditions influenced by domain-wise conditions. Thirdly, replacing the
conditional modeling with directly concatenating conditions to the initial
inputs severely harms relation learning, further highlighting the advantages
of our design.

#### Case study

To provide an intuitive comparison, we provide showcases on the HeterNS
benchmark in Figure3. First, all the presented trajectories are generated from
the same initial condition but exhibit distinct final fields, underscoring the
importance of PDE components. Further, it is observed that Unisolver
significantly outperforms FNO under complex conditions, such as smaller $\nu$
and larger $\omega$, particularly in OOD conditions. More showcases can be
found in Appendix A.

Figure 3: Showcases on error maps of FNO and Unisolver on HeterNS benchmark.
All the cases have the same initial condition but different viscosity
coefficients and force-frequency coefficients, which are presented in the
first row. The left shows ID conditions, while the right depicts OOD
conditions.

#### Scalability

Scalability is vital for building a universal PDE solver. Figure 4 illustrates
our exploration of Unisolver’s scalability, where we progressively increase
the training data by 60 times and the model parameters by 21 times. Unisolver
exactly displays the scaling law, achieving better performance with increased
data and parameters, posing the potential for a universal PDE solver.

Figure 4: Data scalability (60x) and model scalability (21x) on the PDEformer
benchmark. Relative L2 error for pre-training (periodic and Robin) and Burgers
equation are plotted on a log-log scale.

##  5 Conclusions

To break the generalization bottleneck, this paper presents Unisolver as a
PDE-conditional Transformer, which stems from the theoretical analysis of the
PDE-solving process. Concretely, Unisolver proposes a complete set of PDE
components and embeds these components into domain-wise and point-wise deep
conditions separately for deep conditions. By integrating these conditions
with Transformers through a decoupled mechanism, Unisolver can handle
universal PDE components and achieve consistent state-of-the-art results
across three challenging, large-scale benchmarks. Extensive analyses are
provided to verify the performance, generalizability and scalability of our
model.

## References

  * [1] William F Ames.  Numerical methods for partial differential equations.  Academic press, 2014.
  * [2] Vladimir Igorevich Arnol’d.  Mathematical methods of classical mechanics.  Springer Science & Business Media, 2013.
  * [3] Johannes Brandstetter, Daniel Worrall, and Max Welling.  Message passing neural pde solvers.  In ICLR, 2022.
  * [4] Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al.  Language models are few-shot learners.  In NeurIPS, 2020.
  * [5] Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova.  Bert: Pre-training of deep bidirectional transformers for language understanding.  In NAACL, 2019.
  * [6] Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al.  An image is worth 16x16 words: Transformers for image recognition at scale.  In ICLR, 2020.
  * [7] Stefan Elfwing, Eiji Uchibe, and Kenji Doya.  Sigmoid-weighted linear units for neural network function approximation in reinforcement learning.  Neural networks, 2018.
  * [8] Lawrence C Evans.  Partial differential equations.  American Mathematical Society, 2022.
  * [9] Jayesh K Gupta and Johannes Brandstetter.  Towards multi-spatiotemporal-scale generalized pde modeling.  TMLR, 2023.
  * [10] Zhongkai Hao, Chang Su, Songming Liu, Julius Berner, Chengyang Ying, Hang Su, Anima Anandkumar, Jian Song, and Jun Zhu.  Dpot: Auto-regressive denoising operator transformer for large-scale pde pre-training.  In ICML, 2024.
  * [11] Zhongkai Hao, Chengyang Ying, Zhengyi Wang, Hang Su, Yinpeng Dong, Songming Liu, Ze Cheng, Jun Zhu, and Jian Song.  Gnot: A general neural operator transformer for operator learning.  In ICML, 2023.
  * [12] Kaiming He, Xinlei Chen, Saining Xie, Yanghao Li, Piotr Dollár, and Ross Girshick.  Masked autoencoders are scalable vision learners.  In CVPR, 2022.
  * [13] Tero Karras, Miika Aittala, Timo Aila, and Samuli Laine.  Elucidating the design space of diffusion-based generative models.  NeurIPS, 2022.
  * [14] Diederik P. Kingma and Jimmy Ba.  Adam: A method for stochastic optimization.  In ICLR, 2015.
  * [15] Nikola Kovachki, Zongyi Li, Burigede Liu, Kamyar Azizzadenesheli, Kaushik Bhattacharya, Andrew Stuart, and Anima Anandkumar.  Neural operator: Learning maps between function spaces with applications to pdes.  JMLR, 2023.
  * [16] Zijie Li, Kazem Meidani, and Amir Barati Farimani.  Transformer for partial differential equations’ operator learning.  TMLR, 2023.
  * [17] Zijie Li, Dule Shu, and Amir Barati Farimani.  Scalable transformer for pde surrogate modeling.  In NeurIPS, 2024.
  * [18] Zongyi Li, Daniel Zhengyu Huang, Burigede Liu, and Anima Anandkumar.  Fourier neural operator with learned deformations for pdes on general geometries.  JMLR, 2023.
  * [19] Zongyi Li, Nikola Borislavov Kovachki, Kamyar Azizzadenesheli, Burigede liu, Kaushik Bhattacharya, Andrew Stuart, and Anima Anandkumar.  Fourier neural operator for parametric partial differential equations.  In ICLR, 2021.
  * [20] Zongyi Li, Hongkai Zheng, Nikola Kovachki, David Jin, Haoxuan Chen, Burigede Liu, Kamyar Azizzadenesheli, and Anima Anandkumar.  Physics-informed neural operator for learning partial differential equations.  J. Data Sci., 2021.
  * [21] Yuxuan Liu, Zecheng Zhang, and Hayden Schaeffer.  Prose: Predicting operators and symbolic expressions using multimodal transformers.  arXiv preprint arXiv:2309.16816, 2023.
  * [22] Ze Liu, Yutong Lin, Yue Cao, Han Hu, Yixuan Wei, Zheng Zhang, Stephen Lin, and Baining Guo.  Swin transformer: Hierarchical vision transformer using shifted windows.  In ICCV, 2021.
  * [23] Cooper Lorsung, Zijie Li, and Amir Barati Farimani.  Physics informed token transformer for solving partial differential equations.  Mach. Learn.: Sci. Technol, 2024.
  * [24] Ilya Loshchilov and Frank Hutter.  Sgdr: Stochastic gradient descent with warm restarts.  In ICLR, 2016.
  * [25] Lu Lu, Pengzhan Jin, Guofei Pang, Zhongqiang Zhang, and George Em Karniadakis.  Learning nonlinear operators via deeponet based on the universal approximation theorem of operators.  Nat. Mach. Intell, 2021.
  * [26] Yining Luo, Yingfa Chen, and Zhen Zhang.  Cfdbench: A comprehensive benchmark for machine learning methods in fluid dynamics.  arXiv preprint arXiv:2310.05963, 2023.
  * [27] Michael McCabe, Bruno Régaldo-Saint Blancard, Liam Holden Parker, Ruben Ohana, Miles Cranmer, Alberto Bietti, Michael Eickenberg, Siavash Golkar, Geraud Krawezik, Francois Lanusse, et al.  Multiple physics pretraining for physical surrogate models.  In NeurIPS AI for Science Workshop, 2023.
  * [28] Taesung Park, Ming-Yu Liu, Ting-Chun Wang, and Jun-Yan Zhu.  Semantic image synthesis with spatially-adaptive normalization.  In CVPR, 2019.
  * [29] William Peebles and Saining Xie.  Scalable diffusion models with transformers.  In ICCV, 2023.
  * [30] Ethan Perez, Florian Strub, Harm De Vries, Vincent Dumoulin, and Aaron Courville.  Film: Visual reasoning with a general conditioning layer.  In AAAI, 2018.
  * [31] Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J Liu.  Exploring the limits of transfer learning with a unified text-to-text transformer.  2020\.
  * [32] Md Ashiqur Rahman, Zachary E Ross, and Kamyar Azizzadenesheli.  U-no: U-shaped neural operators.  TMLR, 2023.
  * [33] Maziar Raissi, Paris Perdikaris, and George Em Karniadakis.  Physics-informed neural networks: A deep learning framework for solving forward and inverse problems involving nonlinear partial differential equations.  J. Comput. Phys., 2019.
  * [34] Rajhans Singh, Ankita Shukla, and Pavan Turaga.  Polynomial implicit neural representations for large diverse datasets.  In CVPR, 2023.
  * [35] Shashank Subramanian, Peter Harrington, Kurt Keutzer, Wahid Bhimji, Dmitriy Morozov, Michael W Mahoney, and Amir Gholami.  Towards foundation models for scientific machine learning: Characterizing scaling and transfer behavior.  In NeurIPS, 2023.
  * [36] Makoto Takamoto, Francesco Alesiani, and Mathias Niepert.  Learning neural pde solvers with parameter-guided channel attention.  In ICML, 2023.
  * [37] Makoto Takamoto, Timothy Praditia, Raphael Leiteritz, Daniel MacKinlay, Francesco Alesiani, Dirk Pflüger, and Mathias Niepert.  Pdebench: An extensive benchmark for scientific machine learning.  NeurIPS, 2022.
  * [38] Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al.  Llama: Open and efficient foundation language models.  arXiv preprint arXiv:2302.13971, 2023.
  * [39] Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, et al.  Llama 2: Open foundation and fine-tuned chat models.  arXiv preprint arXiv:2307.09288, 2023.
  * [40] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Ł ukasz Kaiser, and Illia Polosukhin.  Attention is all you need.  In NeurIPS, 2017.
  * [41] Hanchen Wang, Tianfan Fu, Yuanqi Du, Wenhao Gao, Kexin Huang, Ziming Liu, Payal Chandak, Shengchao Liu, Peter Van Katwyk, Andreea Deac, et al.  Scientific discovery in the age of artificial intelligence.  Nature, 2023.
  * [42] Gege Wen, Zongyi Li, Kamyar Azizzadenesheli, Anima Anandkumar, and Sally M Benson.  U-fno—an enhanced fourier neural operator-based deep-learning model for multiphase flow.  Adv. Water Resour., 2022.
  * [43] Haixu Wu, Tengge Hu, Huakun Luo, Jianmin Wang, and Mingsheng Long.  Solving high-dimensional pdes with latent spectral models.  In ICML, 2023.
  * [44] Haixu Wu, Huakun Luo, Haowen Wang, Jianmin Wang, and Mingsheng Long.  Transolver: A fast transformer solver for pdes on general geometries.  In ICML, 2024.
  * [45] Zhanhong Ye, Xiang Huang, Leheng Chen, Hongsheng Liu, Zidong Wang, and Bin Dong.  Pdeformer: Towards a foundation model for one-dimensional partial differential equations.  In ICLR AI4Differential Equations In Science Workshop, 2024.
  * [46] Chengxuan Ying, Tianle Cai, Shengjie Luo, Shuxin Zheng, Guolin Ke, Di He, Yanming Shen, and Tie-Yan Liu.  Do transformers really perform badly for graph representation?  NeurIPS, 2021.

##  Appendix A More Showcases

We provide more showcases here as a supplement to the numerical results of the
main text.

Figure 5: Showcases comparison for in-distribution (ID) data with three
baselines on the HeterNS dataset. See Table 3 and 4 for numerical comparison.
All data has the same initial condition but different viscosity coefficients
and force frequency coefficients, as shown in the first row. Error maps are
provided. A darker color indicates a lower error.

Figure 6: Showcases comparison for out-of-distribution (OOD) data with three
baselines on the HeterNS dataset with the same initial conditions but
different viscosity coefficients and force frequency coefficients. See Table 3
and 4 for numerical comparison. Error maps are provided. A darker color
indicates a lower error.

Figure 7: Showcases comparison with PDEformer on the pre-training and
downstream tasks of 1D time-dependent PDEs. See Table 5 for numerical
comparison. The pre-training dataset contains both periodic and robin boundary
conditions. The numbers in the Burgers columns refer to the diffusion
coefficient $\nu$ and the number in the Advection column indicates the
constant advection speed $\beta$. Error maps are provided. A darker color
indicates a lower error.

Figure 8: Showcases of Unisolver on 2d mixed PDE benchmark that covers a wide
range of PDE components. See Table 6 for numerical comparison with DPOT. Both
prediction results and error maps are provided. As shown in the CFDBench-NS
columns, Unisolver presents an impressive ability to handle different geometry
conditions.

##  Appendix B Analytical Solution for the String Vibration Equation

Here we derive the analytical solution for the PDE (1a) with boundary
conditions (1b) and initial conditions (1c).

We extend $\varphi(x)$, $\psi(x)$ and $f(x,t)$ as odd, periodic functions of
$x$ with period $2L$, defining $\Phi(x)$, $\Psi(x)$ and $F(x,t)$ as follows:

$\displaystyle\Phi(x)=-\Phi(-x),\quad\Phi(x)=\Phi(x+2L),$
$\displaystyle-\infty<x<\infty,$ (4) $\displaystyle\Phi(x)=\varphi(x),$
$\displaystyle 0\leq x\leq L,$
$\displaystyle\Psi(x)=-\Psi(-x),\quad\Psi(x)=\Psi(x+2L),$
$\displaystyle-\infty<x<\infty,$ $\displaystyle\Psi(x)=\psi(x),$
$\displaystyle 0\leq x\leq L.$ $\displaystyle F(x,t)=-F(-x,t),\quad
F(x,t)=F(x+2L,t),$ $\displaystyle-\infty<x<\infty,0\leq t\leq T$
$\displaystyle F(x,t)=f(x,t),$ $\displaystyle 0\leq x\leq L,0\leq t\leq T.$

Thus, we obtain a wave equation defined on the upper half-plane, with initial
conditions $\Phi(x)$ and $\Psi(x)$.

$\displaystyle\partial_{tt}U-a^{2}\partial_{xx}U=F,\quad(x,t)\in\mathbb{R}\times(0,T),$
(5) $\displaystyle U(x,0)=\Phi(x),\quad\partial_{t}U(x,0)=\Psi(x),\quad
x\in\mathbb{R}.$

######  Lemma B.1.

The solution of the simplified equation (9) is
$U(x,t)=\frac{1}{2a}\int_{x-at}^{x+at}\Psi(\xi)\mathrm{d}\xi$.

$\displaystyle\partial_{tt}U-a^{2}\partial_{xx}U=0,\quad(x,t)\in\mathbb{R}\times(0,T),$
(6) $\displaystyle U(x,0)=0,\quad\partial_{t}U(x,0)=\Psi(x),\quad
x\in\mathbb{R}.$

###### Proof.

The above expression for $U$ can be derived using the method of operator
splitting and characteristic lines.

Note that

$0=\partial_{tt}U-c^{2}\partial_{xx}U=(\partial_{t}-a\partial_{x})(\partial_{t}+a\partial_{x})U.$

Let

$V=(\partial_{t}+a\partial_{x})U$

Then

$\partial_{t}V-a\partial_{x}V=0,\quad V(x,0)=\Psi(x).$

Using the method of characteristic lines for first-order equations, we obtain

$V(x,t)=V(x+at,0)=\Psi(x+at).$

Thus, $U$ satisfies

$(\partial_{t}+a\partial_{x})U=V=\Psi(x+at).$

For this first-order equation, along the characteristic line passing through
the point $(X,T^{\prime})$, defined by $x=at+X-aT^{\prime}$, we have

$\frac{\mathrm{d}u(X-aT^{\prime}+at,t)}{\mathrm{d}t}=\Psi(X-aT^{\prime}+2at).$

By integrating on both sides, we get

$U(X-aT^{\prime}+at,t)=\int_{0}^{t}\Psi(X-aT^{\prime}+2as)\,\mathrm{d}s.$

Therefore,

$U(x,t)=\int_{0}^{t}\Psi(x-at+2as)\,\mathrm{d}s=\frac{1}{2a}\int_{x-at}^{x+at}\Psi(\xi)\,\mathrm{d}\xi.$

It is straightforward to verify it’s indeed a solution by substituting it into
the equation, given that $\Psi\in C^{1}(\mathbb{R})$. ∎

Inspired by the principle of linear superposition for linear problems, the
solution $U=U(t,x)$ for equation (5) can be decomposed as

$U=U_{1}+U_{2}+U_{3},$

where $U_{1},U_{2},$ and $U_{3}$ are the solutions to the following three
equations, respectively,

$\displaystyle\begin{cases}\partial_{tt}U-a^{2}\partial_{xx}U=0,\\\
U(x,0)=\Phi(x),\\\ \partial_{t}U(x,0)=0,\end{cases}$
$\displaystyle\begin{cases}\partial_{tt}U-a^{2}\partial_{xx}U=0,\\\
U(x,0)=0,\\\ \partial_{t}U(x,0)=\Psi(x),\end{cases}$
$\displaystyle\begin{cases}\partial_{tt}U-a^{2}\partial_{xx}U=F(x,t),\\\
U(x,0)=0,\\\ \partial_{t}U(x,0)=0.\end{cases}$

######  Lemma B.2.

The three solutions have the following simple relationships. Denoting that
$U_{2}=S_{\Psi}$, we have

$U_{1}=\frac{\mathrm{d}S_{\Phi}}{\mathrm{d}t},\quad
U_{3}=\int_{0}^{t}S_{F}(\tau,x)(t-\tau,x)d\tau.$

###### Proof.

These relations can be verified by substitution. ∎

Thus, the explicit expression of $U$ is,

$\displaystyle U(x,t)$ $\displaystyle=(U_{1}+U_{2}+U_{3})(x,t)$
$\displaystyle=\frac{1}{2}(\Phi(x+at)+\Phi(x-at))+\frac{1}{2a}\int^{x+at}_{x-at}\Psi(\xi)\mathrm{d}\xi+\frac{1}{2a}\int_{0}^{t}\mathrm{d}\tau\int_{x-a(t-\tau)}^{x+a(t-\tau)}F(\xi,\tau)\mathrm{d}\xi$

Noting that due to the construction of $\Phi(x)$ and $\Psi(x)$, the solution
$U(x,t)$ of equation 5 is an odd function of $x$ with a period of $2L$.
Moreover, it satisfies $U(kL,t)=0,\forall k\in\mathbb{Z},t\in\mathbb{R}$,
ensuring that the boundary conditions of the original string vibrating
equation (1a) are met. Therefore, we obtain the solution of equation (1a):

$u(x,t)=U(x,t),\quad\forall(x,t)\in[0,L]\times[0,T]$ (7)

##  Appendix C Benchmarks

###  C.1 HeterNS

Similar to FNO [19], we consider the 2-d Navier-Stokes equation in vorticity
formulation for the viscous, incompressible fluid on a unit torus:

$\displaystyle\partial_{t}w(x,t)+u(x,t)\cdot\nabla w(x,t)$
$\displaystyle=\nu\Delta w(x,t)+f(x),$ $\displaystyle
x\in(0,1)^{2},\,t\in(0,T]$ (8a) $\displaystyle\nabla\cdot u(x,t)$
$\displaystyle=0,$ $\displaystyle x\in(0,1)^{2},\,t\in[0,T]$ (8b)
$\displaystyle w(x,0)$ $\displaystyle=w_{0}(x),$ $\displaystyle x\in(0,1)^{2}$
(8c)

We experiment with viscosity $\nu\in[8\times 10^{-6},2\times 10^{-3}]$ and
external force in the form
$f(x)=0.1(\sin(\omega\pi(x_{1}+x_{2}))+\cos(\omega\pi(x_{1}+x_{2})))$.
Specifically, we pre-train models on a dataset with $\nu\in\\{1\times
10^{-5},5\times 10^{-5},1\times 10^{-4},5\times 10^{-4},1\times 10^{-3}\\}$
and $\omega\in\\{1,2,3\\}$, and test the models on both in-domain and out-of-
domain equation components. Pretraining dataset can be found here
222https://drive.google.com/drive/folders/1te5IyQHTznu_Kw7v3zDHg0i_KCHysPKw.

###  C.2 1D Time-dependent PDEs

As described in PDEformer [45], this benchmark contains 3 million high
equality 1D time-dependent PDEs with various equation components for pre-
training and then evaluates the pre-trained model using two downstream tasks
from PDEBench, e.g. advection equation and Burgers equation.

#### Pre-training

The pre-training dataset is generated following the PDE family:

$\displaystyle\partial_{t}u+f_{0}(u)+s(x)+\partial_{x}(f_{1}(u)-\kappa(x)\partial_{x}u)=0$
$\displaystyle,\;(x,t)\in[-1,1]\times[0,1]$ (9a) $\displaystyle u(0,x)=g(x)$
$\displaystyle,\;x\in[-1,1]$ (9b)

where $f_{i}(u)=c_{i1}u+c_{i2}u^{2}+c_{i3}u^{3},i=0,1$. Each coefficient
$c_{ik}$ is set to zero with probability 0.5, and uniformly sampled from
$[-3,3]$ otherwise.

Notably, terms with a zero coefficient will be excluded in the equation
symbolic expression embedded by a LLM. The initial condition $g(x)$ is
randomly generated within the family of trigonometric functions, which is same
as downstream tasks.

This dataset includes both periodic and non-periodic boundary conditions, with
1.5 million samples each. For the non-periodic cases, the boundary condition
type at each endpoint are randomly selected from Dirichlet, Neumann, and Robin
types. The Dirichlet condition specifies the solution value at the boundary,
the Neumann condition specifies the derivative value at the boundary, and the
Robin condition is a linear combination of the two, where Dirichlet and
Neumann conditions can be seen as corner cases of the Robin condition.

In conclusion, the domain-wise components of the pre-training dataset include
equation symbolic expression, boundary condition types, and coefficients of
$f_{i}$ and point-wise components include the force fields $s(x)$ and
$\kappa(x)$ and boundary value functions. The input consists of the equation’s
initial condition with a resolution of $256$ and above classified components,
and the output is solution $u(x,t)$ with a resolution of $256\times 100$.

#### Downstream tasks

We employed the following two 1D PDE datasets from PDEBench as downstream
evaluation tasks. All tasks adhere to periodic boundary conditions and the
same initial condition family. The resolution of these samples is $1024\times
201$. For each dataset, we utilize the first 9000 samples for finetuning and
the rest 1000 samples for testing. We downsample the spatial resolution of
these datasets to 256 and maintain the temporal resolution unchanged. The
downstream PDEs consist of the Advection equation and the Burgers equation.

#### (1) Advection equation

The advection equation models pure advection behavior without non-linearity,
which can be formalized as:

$\displaystyle\partial_{t}u(t,x)+\beta\partial_{x}u(x,t)$
$\displaystyle=0,\;x\in(0,1),t\in(0,2]$ (10a) $\displaystyle u(0,x)$
$\displaystyle=u_{0}(x),\;x\in(0,1),$ (10b)

where constant advection speed $\beta$ and equation symbols are domain-wise
components in this dataset. The exact solution of the equation is
$u(t,x)=u_{0}(x-\beta t)$. Only the periodic boundary condition is considered.
The initial condition uses a super-position of sinusoidal waves as,
$u_{0}(x)=\sum_{k_{i}=k_{1},...,k_{N}}A_{i}\sin(k_{i}x+\phi_{i})$, where
$k_{i}=2\pi{n_{i}}/L_{x}$ are wave numbers and ${n_{i}}\in\mathbb{N}$ are
selected randomly in $[1,n_{max}]$.

#### (2) Burgers equation

This fundamental equation in fluid mechanics models the non-linear behavior
and diffusion process of fluid dynamics as

$\displaystyle\partial_{t}u(t,x)+\partial_{x}({u(t,x)}^{2}/2)$
$\displaystyle=\nu/\pi\partial_{xx}u(t,x),\;x\in(0,1),t\in(0,2],$ (11a)
$\displaystyle u(0,x)$ $\displaystyle=u_{0}(x),\;x\in(0,1),$ (11b)

where diffusion coefficient $\nu$ and equation formulation are the domain-wise
components. To align with the solution domain with pretraining dataset, we
rescale the spatial-temporal coordinates $(x,t)$ to
$(x^{\prime},t^{\prime})\in[-1,1]\times[0,1]$. The PDEs after the coordinate
transformation are given by:

  * •

Burgers’ equation:
${\partial_{t^{\prime}}u}+{\partial_{x^{\prime}}}(2u^{2})-\frac{8\nu}{\pi}{\partial_{x^{\prime}x^{\prime}}u}=0,\quad\text{where
}t^{\prime}=\frac{t}{2},x^{\prime}=2x-1.$

  * •

Advection equation: ${\partial_{t^{\prime}}u}+{\partial_{x^{\prime}}}(4\beta
u)=0,\quad\text{where }t^{\prime}=\frac{t}{2},x^{\prime}=2x-1.$

###  C.3 2D Mixed PDEs

This dataset is presented by DPOT [10], which consists of the following 12
datasets from 4 benchmarks. These datasets are aggregated to pre-train a
foundational neural PDE solver.

#### FNO-$\nu$

This benchmark considers the 2D Navier-Stokes equation for a viscous,
incompressible fluid in vorticity form on the unit torus. The task is to
estimate the vorticity field of the future ten timesteps on a regular
$64\times 64$ grid based on the initial ten timesteps observations of the
vorticity field. The only varying PDE component in this dataset is the
viscosity coefficient, which varies in $\\{1\times 10^{-5},1\times
10^{-4},1\times 10^{-3}\\}$. We utilize 1000, 9800, and 1000 instances for
each viscosity value to pre-train or finetune our model, and use the rest 200
instances to test its performance.

#### PDEBench

The following three datasets are derived from PDEBench, encompassing the
compressible Navier-Stokes equation, the diffusion-reaction equation, and the
shallow-water equation. All datasets considered in PDEBench adhere to periodic
boundary conditions.

(1) The compressible Navier-Stokes equation models compressible fluid
dynamics, such as shock wave formation and propagation. In this dataset, two
domain-wise components are considered, namely Mach number and shear viscosity.
For these two components, this dataset provides four combinations, which are
$(1,0.1),(1,0.01),(0.1,0.1),(0.1,0.01)$. Each combination provides 9000
instances for training and 200 instances for testing. The instances in this
dataset take the spatial resolution of $128\times 128$, and the task is to
predict the future 11 timesteps of vorticity, pressure and density given the
initial 10 timesteps of observations.

(2) The shallow-water equations, derived from the general Navier-Stokes
equations, model free-surface flow problems,

$\displaystyle\partial_{t}h+\nabla\cdot(h\boldsymbol{u})$ $\displaystyle=0,$
(12a)
$\displaystyle\partial_{t}(h\boldsymbol{u})+\nabla\cdot\left(\cfrac{1}{2}h\boldsymbol{u}^{2}+\cfrac{1}{2}g_{r}h^{2}\right)$
$\displaystyle=-g_{r}h\nabla b$ (12b)

where $h$ describes the water depth, $b$ describes a spatially varying
bathymetry, $g_{r}$ describes the gravitational acceleration, and
$\nabla\cdot(h\boldsymbol{u})$ can be interpreted as the directional momentum.

The spatial resolution of this benchmark is $128\times 128$, and the task is
to predict the future 91 timesteps of water depth given the first 10 timesteps
of observations.

(3) The 2D Diffusion-Reaction Equation involves two non-linearly coupled
variables, namely the activator $u=u(t,x,y)$ and the inhibitor $v=v(t,x,y)$.
It is primarily applicable for modeling biological pattern formation.

$\displaystyle\partial_{t}u$
$\displaystyle=D_{u}\partial_{xx}u+D_{u}\partial_{yy}u+R_{u}$ (13a)
$\displaystyle\partial_{t}v$
$\displaystyle=D_{v}\partial_{xx}v+D_{v}\partial_{yy}v+R_{v}$ (13b)

where $D_{u}=1\times 10^{-3}$ and $D_{v}=5\times 10^{-3}$ are the diffusion
coefficient for the activator and inhibitor, respectively, and
$R_{u}=R_{u}(u,v)$ and $R_{v}=R_{v}(u,v)$ are the corresponding reaction
functions for the activator and inhibitor, which are defined by the Fitzhugh-
Nagumo equation as,

$\displaystyle R_{u}(u,v)$ $\displaystyle=u-u^{3}-k-v,$ (14a) $\displaystyle
R_{v}(u,v)$ $\displaystyle=u-v,$ (14b)

where $k=5\times 10^{-3}$. The initial condition is generated as standard
normal random noise $u(0,x,y)\sim N(0,1.0)$ for $x\in(-1,1)$ and $y\in(-1,1)$.
The dataset is discretized into $N_{x}=128$, $N_{y}=128$, and $N_{t}=101$. The
task is to predict the future 91 timesteps of $u$ and $v$ given the initial 10
timesteps of observations.

#### PDEArena-NS1/2

The velocity function formulation of the incompressible Navier-Stokes
equations is very common in real-life applications.

$\displaystyle\partial_{t}\boldsymbol{v}$
$\displaystyle=-\boldsymbol{v}\cdot\nabla\boldsymbol{v}+\mu\nabla^{2}\boldsymbol{v}-\nabla
p+\boldsymbol{f},$ (15a) $\displaystyle\nabla\cdot\boldsymbol{v}$
$\displaystyle=0$ (15b)

where $v\cdot\nabla v$ is the convection which means the rate of change of $v$
along $v$, $\mu\nabla^{2}v$ is the viscosity, i.e. the diffusion or net
movement of $v$, $\nabla p$ is the internal pressure, and $f$ is the external
buoyancy force. The inclusion of the incompressibility constraint $\nabla\cdot
u=0$ ensures mass conservation in the Navier-Stokes equations. The spatial
resolution is $128\times 128$. This benchmark contains 2 datasets, one with
fixed external force and another with varying externel force. Given the
initial 10 timesteps of observations, the task on the fixed-force dataset is
to predict the future 4 timesteps of velocity, while on the varying-force
dataset, the number of timesteps to predict is 46.

#### CFDBench

We consider three important and representative fluid problems that can
comprehensively evaluate methods’ capabilities to generalize to unseen
operating conditions. They are (1) the flow in the lid-driven cavity, (2) the
flow into the circular tube, (3) the flow around the cylinder.

$\displaystyle\partial_{t}(\rho\boldsymbol{u})+\nabla\cdot(\rho\boldsymbol{u}^{2})$
$\displaystyle=-\nabla
p+\nabla\cdot\mu(\nabla\boldsymbol{u}+\nabla\boldsymbol{u}^{T}),$ (16a)
$\displaystyle\nabla\cdot(\rho\boldsymbol{u})$ $\displaystyle=0$ (16b)

where $\rho$ is the constant density, $\mu$ is the dynamic viscosity,
$\boldsymbol{u}=(u,v)^{T}$ is the velocity field, and $p$ is the pressure. For
each problem, flows are generated with different operating parameters, which
is the combination of the three kinds of components: (1) the Boundary
conditions, (2) the fluid physical properties including the density and the
viscosity, and (3) the geometry of the field. The boundary conditions refer to
the inlet velocity or the movement velocity according to different cases. Each
kind of operating parameter corresponds to one subset. In each subset, the
corresponding operating conditions are varied while other parameters remain
constant. The task is to predict the future 10 timesteps of velocity given the
initial 10 timesteps of observations.

##  Appendix D Implementation Details

###  D.1 Metrics

#### Relative L2 for physics fields

We can calculate the relative L2 error of model prediction $\hat{u}$ and
ground truth $u$ as follows:

$\frac{\|u-\hat{u}\|_{L^{2}}}{\|u\|_{L^{2}}}.$ (17)

where $\|u-\hat{u}\|_{L^{2}}$ is the $L^{2}$-distance between the predicted
solution $\hat{u}$ and the ground-truth solution $u$, and $\|u\|_{L^{2}}$ is
the $L^{2}$-norm of the ground-truth solution.

#### Relative Promotion

Given the error of our model $\epsilon_{\text{ours}}$ and the error of the
second best model $\epsilon_{\text{second-best model}}$, we can calculate the
relative promotion as follows:

$\text{Relative \
Promotion}=1-\frac{\epsilon_{\mathrm{ours}}}{\epsilon_{\text{second-best
model}}}.$ (18)

###  D.2 Implementations for Each Benchmark

#### HeterNS

As for the HeterNS dataset, the model hyperparameters of the baselines and
Unisolver are summarized in Table 8. The count of trainable parameters of each
model and the training configurations are already shown in Section 4. For
baselines, the viscosity coefficient and the external force are concatenated
to the model input along the channel dimension to ensure fair comparison.

#### 1D Time-dependent PDEs

We conduct experiments using the checkpoints and dataset provided by the
official repository of
PDEformer333https://gitee.com/mindspore/mindscience/tree/master/MindFlow/applications/pdeformer1d/.
We randomly select one subset with periodic boundaries and another subset with
Robin boundaries from the pre-training datasets. Both of the subset contain
10,000 samples and are used to calculate the pre-training relative L2 error.

For finetuning on the downstream tasks, we utilize the finetuning script
provided in the repository and set the finetuning epochs to 100 for fair
comparison. The pre-training and finetuing configurations for our model and
the finetuing configurations for PDEformer are shown in Table 10.

As shown in the model scalability experiments in Section 4.4, we progressively
increase the Unisolver parameter from 3M to 63M, thus resulting in 4 different
model configurations. We present the configurations of these models in Table
9.

#### 2D Mixed PDEs

The training hyperparameter and model configurations are presented in Table
11. As this dataset includes multiple distinct PDEs, each accompanied by its
unique equation components, we introduce a binary masking channel, where an
"1" indicates that the current component is a valid condition and a "0"
indicates an invalid component.

Table 8: Model hyperparameters of Unisolver and all baselines on the HeterNS benchmark. Hyperparameter | Value | Description
---|---|---
FNO
modes | 12 | The truncation number of Fourier modes
channels | 64 | The number of channels in the hidden layers
depths | 4 | The number of Fourier Layers in the neural network
Factformer
dim | 128 | hidden dimension of the transformer
n_head | 12 | num of attention heads
dim_head | 64 | hidden dimension of each attention heads
depths | 8 | The number of Transformer Blocks in the neural network
ViT
Attention dim | 256 | The hidden dimension of the transformer attention layer
MLP dim | 256 | The hidden dimension of the transformer FFN layer
patch_size | 4 | The height and width of the ViT patches
n_head | 8 | num of attention heads
dim_head | 32 | The hidden dimension of each attention heads
depths | 12 | The number of Transformer Blocks in the neural network
Unisolver
Attention dim | 256 | The hidden dimension of the transformer attention layer
MLP dim | 256 | The hidden dimension of the transformer FFN layer
patch_size | 4 | The height and width of the Unisolver patches
n_head | 8 | num of attention heads
dim_head | 32 | The hidden dimension of each attention heads
depths | 8 | The number of Transformer Blocks in the neural network
Table 9: Model configurations of Unisolver with different sizes. Parameter Count | Attention dim | MLP dim | Layers | Heads
---|---|---|---|---
3M | 256 | 256 | 6 | 4
10M | 384 | 384 | 8 | 8
19M | 512 | 512 | 8 | 8
63M | 768 | 768 | 12 | 12
Table 10: Pre-training and finetuning configurations on the 1D time-dependent PDE benchmark. Parameter | Value | Description
---|---|---
Unisolver Pre-training
batch_size | 1024 | Total batchsize used in one iteration
learning_rate | 6e-4 | The initial learning rate for the optimizer
epochs | 500 | The total number of training epochs
loss_type | Relative-l2 | Use relative L2-Norm for pretraining
optimizer | Adam | The optimization algorithm
lr_scheduler | Cosine Annealing | The learning rate scheduler
Unisolver Finetuning
batch_size | 256 | Total batchsize used in one iteration
learning_rate | 1e-5 | The initial learning rate for the optimizer
epochs | 100 | The total number of training epochs
loss_type | Relative-l2 | Use relative L2-Norm for finetuning
optimizer | Adam | The optimization algorithm
lr_scheduler | Cosine Annealing | The learning rate scheduler
PDEformer Finetuning
batch_size | 80 | Total batchsize used in one iteration
learning_rate | 5e-6 | The initial learning rate for the optimizer
epochs | 100 | The total number of training epochs
loss_type | Relative-l2 | Use relative L2-Norm for finetuning
optimizer | Adam | The optimization algorithm
lr_scheduler | Cosine Annealing | The learning rate scheduler
warmup_epochs | 10 | Epochs to linearly increase the learning rate
Table 11: Training configurations on the 2D mixed PDE benchmark. Parameter | Value | Description
---|---|---
Unisolver Training Configurations
batch_size | 320 | Total batchsize used in one iteration
learning_rate | 1e-3 | The initial learning rate for the optimizer
epochs | 1000 | The total number of training epochs
loss_type | Relative-l2 | Use relative L2-Norm for pretraining
optimizer | AdamW | The optimization algorithm
lr_scheduler | OneCycle | The learning rate scheduler
warmup_epochs | 200 | Epochs to linearly increase the learning rate
Unisolver Model Configurations
Attention dim | 768 | The hidden dimension of the transformer attention layer
MLP dim | 768 | The hidden dimension of the transformer FFN layer
patch_size | 8 | The height and width of the ViT patches
n_head | 8 | num of attention heads
dim_head | 96 | The hidden dimension of each attention heads
depths | 6 | The number of Transformer Blocks in the neural network

##  Appendix E Additional Experiments

###  E.1 OOD Viscosity and OOD External Force

In addition to Table 3 and 4, we further compare the generalizability of
Unisolver with other baselines on the both OOD viscosity coefficient and OOD
external force of the HeterNS benchmark. Specifically, we generate nine
different component pairs $(\nu,\omega)$, where $\nu$ and $\omega$ are both
out of the pre-training distribution. The full results are recorded in Table
12.

Table 12: Performance comparison on OOD viscosity $\nu$ and OOD force frequency coefficients $\omega$. The pairs in the first row are in the form of ($\nu,\omega$). For clarity, the best result is in bold. L2RE |  (2e-5, 0.7) |  (2e-5, 1.7) |  (2e-5, 2.7) |  (2e-5, 3.7) |  (2e-5, 4.7) |  (4e-5, 0.8) |  (4e-5, 1.4) |  (4e-5, 2.3) |  (4e-5, 6)
---|---|---|---|---|---|---|---|---|---
FNO | 0.1862 | 0.0640 | 0.1176 | 0.2404 | 0.4226 | 0.0873 | 0.1516 | 0.0655 | 1.3102
Factformer | 0.2727 | 0.0735 | 0.1262 | 0.3238 | 0.2868 | 0.0589 | 0.1230 | 0.0495 | 0.3147
ViT | 0.1961 | 0.0690 | 0.1075 | 0.2057 | 0.2226 | 0.0488 | 0.1305 | 0.0772 | 0.2276
Unisolver | 0.0781 | 0.0378 | 0.0471 | 0.1421 | 0.1364 | 0.0399 | 0.0433 | 0.0374 | 0.1431
Promotion | 58.06% | 40.94% | 56.19% | 30.92% | 38.72% | 18.24% | 66.82% | 24.44% | 37.13%

###  E.2 Hyperparamter Sensitivity

As shown in Table 13, we test the hyperparameter sensitivity of our model by
changing the domain-wise condition ratio $\alpha$ and keeping the other
hyperparameters fixed. We pre-train models with domain-wise condition ratios
of 0.25, 0.5 and 0.75 on the 1D time-dependent PDE dataset, which contains 50k
samples, all adhering to periodic boundary conditions, and report the relative
L2 on the pre-training dataset and downstream Burgers dataset with
$\nu=0.001$. In this case, the domain-wise ratio of 0.5 leads to the best
results, with the ratio of 0.25 being the worst, which suggests that an
appropriate ratio should be chosen in practice.

Table 13: Hyperparameter sensitivity of domain-wise condition ratio. L2RE | Domain-wise condition ratio
---|---
Dataset | 0.25 | 0.50 | 0.75
PDEformer-50k | 0.0331 | 0.0298 | 0.0314
Burgers $\nu=0.001$ | 0.1127 | 0.1040 | 0.1070

###  E.3 Discussion of Ablation

Why Concat-P is better than ours in pretrain-periodic dataset? Since the
periodic boundary does not have the boundary value functions, the model can
concentrate on the dominated domain-wise conditions for concatenating the
point-wise conditions with model input directly. It is similar to set the
domain-wise condition ratio to 1, which is beneficial for learning the complex
relations of domain-wise conditions.

###  E.4 Full Scalability

As a supplement to Figure 4 in the main text, we also conduct experiments on
different downstream tasks and record the concrete data in Table 14 for
clarity.

Table 14: Scalability results on pre-training dataset and downstream tasks, as depicted in Figure 4. L2RE | Data Scalability (Samples) | Model Scalability (Parameters)
---|---|---
Scale | 50k | 100k | 200k | 3M | 3M | 10M | 19M | 63M
Pretraining Dataset / Circ | 0.0282 | 0.0241 | 0.0198 | 0.0112 | 0.0400 | 0.0267 | 0.0241 | 0.0183
Pretraining Dataset / Robin | 0.0181 | 0.0162 | 0.0141 | 0.0099 | 0.0283 | 0.0185 | 0.0162 | 0.0128
Burgers $\nu=0.1$ | 0.0161 | 0.0116 | 0.0081 | 0.0051 | 0.0143 | 0.0134 | 0.0116 | 0.0091
Burgers $\nu=0.01$ | 0.0649 | 0.0412 | 0.0260 | 0.0144 | 0.0552 | 0.0421 | 0.0412 | 0.0351
Burgers $\nu=0.001$ | 0.1399 | 0.1003 | 0.0689 | 0.0299 | 0.1188 | 0.0976 | 0.1003 | 0.0889

###  E.5 Detailed Compute Resources

Our models were trained on servers with A100-40G GPUs. Here we present the
compute resources needed to reproduce our experiment results in Table 15.

Table 15: Compute resources for each of our benchmarks. For the HeterNS dataset, we report the compute resources needed for all the baselines and our model. For the rest two datasets, we report the compute resources needed for pre-training Unisolver only. Benchmarks | HeterNS | 1D Time-dependent PDEs | 2D Mixed PDEs
---|---|---|---
Models | FNO | Factformer | ViT | Unisolver | Unisolver | Unisolver
A100 GPU hours |  12 |  100 |  24 |  24 | 3000 | 800

##  Appendix F Limitations and Future Work

This paper presents Unisolver to solve PDEs under universal PDE components,
which achieves impressive performance supported by extensive analyses and
visualizations. However, our method is currently limited to grid data due to
the patchifying process during the embedding of point-wise components.
Actually, this limitation is shared in all the generalizable PDE solvers, such
as PDEformer and DPOT. One foundation reason is the lack of suitable and
large-scale irregular mesh datasets, which will require extremely high
computation costs for generation. Since we mainly focus on the foundation
model design in this paper, we would like to leave the mesh dataset as a
future work. And the capability to handle irregular meshes of Unisolver can be
achieved by replacing the canonical Transformer with the latest geometry-
general PDE solver: Transolver [44].