Unisolver: PDE-Conditional Transformers Towards Universal Neural PDE Solvers
Authors: Hang Zhou
Yuezhou Ma
Haixu Wu
Haowen Wang
Mingsheng Long
Venue: ICML 2025 poster (2025)
Keywords: Neural PDE Solver
Deep Learning
Primary area: deep_learning->algorithms
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).

Unisolver:
PDE-Conditional Transformers Towards Universal Neural PDE Solvers
Hang Zhou *1
Yuezhou Ma"
Haixu Wu
Haowen Wang'
Mingsheng Long'
Abstract
Amoo'd, d, 2013), such as meteorology, electromagnetism and thermodynamics (Wang et al., 2023). Since it is usually hard to obtain an analytic solution for PDE, numerical methods are widely explored (Ames, 2014). However, these numerical methods often require huge computation costs to generate a precise solution for each PDE. Recently, deep learning models have facilitated significant advancements across wide range of domains (Devlin al., 2019; Liu et al., 2021; Jumper et al., 2021) and have been applied to solving PDEs, i.e. neural PDE solvers (Karniadakis et al., 2021). Owing to their excellent capability to approximate nonlinear mappings, deep learning models can learn to fit pre-collected data (Li et al., 2021a) or physics-iformed loss function (Raissi et al., 2019) and generalize a flash new samples, providing efficient approach to solving PDEs.
Deep models have recently emerged promis- ing tools to solve partial differential equations (PDEs), known as neural PDE solvers. While neu- ral 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 practi- cal 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 universal neu- ral 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 series of PDE components such as equation symbols and boundary conditions, we define a complete set of PDE components and flexibly embed them domain-wise and point-wise deep conditions for Transformer PDE solvers. Integrating physical insights with recent Transformer advances, Uni- solver achieves consistent -fate-of-the-art on three challenging large-scale benchmarks, showing im- pressive performance and generalizability. Code is available at http://thwwoooooo//o/ooooeetea
As shown in Figure I, previous neural solvers can be broadly categorized into two paradigms: physics-infomma networks (PINNs) (Raissi et al., 2019) and neural operators (Li et al., 2021a). The former trains deep models using a formalized PDE loss function, while the latter solely relies on pre-collected data. However, for PINNs, while formu- lating the PDE equations objective functions ensures relatively accurate solutions, they struggle to generalize to new scenarios, necessitating retraining for each new task. Neural operators, on the other hand, directly learn from data and generalize better to diverse initial states and PDEs than PINNs. Nevertheless, purely based on training data may be insufficient to guide PDE solving. For example, in the case of a fluid governed by renowned Navier-Stokes equations, the typical task of neural operators is predict future states based on past observations (Li et al., 2021a), while different viscosity coefficients and forcing terms will lead to distinct solutions even when the initial states stay the same. Thus, due to the omission of PDE information, current neural oper- ators are mainly trained and tested on a limited set of PDEs. Notably, as neural solvers expected be efficient surro- gate models of classical numerical solvers, generalization to various PDEs is essential for practical neural solver.
neural
1. Introduction
Partial differential equations (PDEs) are essential for nu- merous scientific and engineering problems (Evans, 2022:
'Equal contribution 'Scoool of Software, BN- Rist, Tsinghua University, China. Hang Zhou<<hou- Correspondence to: Haixu Wu <wuhx23 mails. tsinghua.ccu.cns. Mingsheng Long <mingsheng @ tsinghua. ssinghua.edu.cn>.
To tackle the generalization deficiency, several works have been proposed by incorporating the PDE information into deep models or training models with large-scale datasets. For example, pesssage-passing neural PDE solver (Brand- stetter et al., 2022) concatenates PDE coefficients with the
Proceedings of the ndnd International Conference Machine Learning, Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Unisolver: PDE-Conditional Transformers Towards Universal Neural PDE Solvers
Physics informed
Data driven
Coordinates
Time
Inputs
Time
Inputs
Time
Spoce
PINNs
Space
Unisolver
Spece
Neural Operators
(x.
u(x. 9
Complete PDE Components
Equation Symbols
Boundary Type
PDE Coefficients
Domain Geometry
Force Term
Boundary Value
Figure 1. Neural PDE solvers typically consist of two paradigms: physies informed and data driven. Unisolver combines data-driven methods with physical insights from PDE components conditional modeling framework, boosting generalizability and scalability.
inputs. PDEformer (Ye et al., 2024) formalizes the PDE as a computation graph and employs the graph Transformer (Ying al., 2021) aggregate PDE information. Although these methods explore the potential of training models with both data and PDE information, they do not consider complete set of PDE information, limiting their generaliz- ability in some aspects. As for the other branches, such as DPOT (Hao et al., 2024), they purely scale up train- ing sets with diverse PDEs and expect the generalizability emerges from large data and parameters. Although models can implicitly extract PDE information from observations, the extraction process is inherently complex and resembles the challenges associated with solving inverse problems (Karniadakis et al., 2021). Therefore, these models often end up fitting an insufficient or vague representation of the underlying observation distribution, which may ultimately hamper their generalizability to broader PDE solving.
We introduce Unisolver as a conditional Transformer architecture utilizing the embedded PDE information completely. marking the first demonstration of the po- tential of the canonical Transformer scalable back- bone for solving multitudinous PDEs universally.
Motivated by the mathematical structure of PDEs, define the concept of complete PDE components, clas- sify them into domain -iise and point-wise categories, and derive a decoupled conditioning mechanism for introducing physics information into PDE solving.
Unisolver achieves consistent state-of-the-arr perfor- mances across three challenging large-scale bench- marks with impressive relative gains and presents fa- vorable generalizability and scalability.
Going beyond prior methods, as shown in Figure 1, this pa- per introduces Unisolver-2 PDE-conditional Transformer progressing towards Universal Neural PDE solver. Uni- solver takes advantages from both data-driven and physics- informed paradigms and empowers Transformer with fa- vorable generalizability by introducing complete physics information conditions. Instead of simply scaling up data and parameters, we motivated by the theoretical analysis of PDE solving and propose complete set of PDE compo- nents. Drawing inspiration from the mathematical structure of PDEs, we classify PDE components into domain-wi 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 distinet influence of different condition types on input representations, we sep- arate the hidden space into subspaces and integrate these deep PDE conditions into hidden representations in decoupled way. We conduct extensive experiments our own generated dataset and two large-scale benchmarks with various PDE components, where Unisolver achieves consistent tfateoof-the-art with sharp relative gains. Overall, our contributions are summarized as follows:
2. Related Work
2.1. Neural PDE Solvers
Previous neural PDE solvers can be roughly categorized into two paradigms (Wu et al., 2024). The first is physics- informed neural networks (PINNs) (Raissi et al., 2019), which optimize deep models by formalizing PDEs as objec- tive functions. During training, model outputs and gradients gradually satisfy the targeted PDE, thereby instantiating the solution deep model. However, PINNs usually hard to generalize to unseen PDEs, limiting their broader applications (Wang et al., 2023). Another booming direc- tion is neural operators, which learn from extensive data to approximate functional dependence between input and out- put Banach spaces (Lu al., 2021: Kovachki et al., 2023). Among various neural operators, FNO (Li et al., 2021a) and its variants (Li et al., 2023a; Rahman et al., 2023; Wen et al., 2022) are popular and well-established. FNO effectively approximates the kernel integral operator in the frequency domain. Recently, given the impressive progress achieved by Transformers (Vaswani et al., 2017). they have also been applied to solve PDEs. Existing methods treat inputs as
2

---

Unisolver: PDE-Conditional Transformers Towards Universal Neural PDE Solvers
[a] Joint Training Diverse PDEs
(b) Inference unseen PDEs
Loss Function:
Periodic Boundary
DHHuiianReacfinn
Relative L2 of (Y,
Equatian
Y
Unseen Initial State
Déichlet
Shollow-Water
Boundory
Equation
Unseen Boundary Value
Robin
Navie-Stokes Equation
Unisolver
Domain Geometry
Extemal Force
Boundary
Varied
Varied Externl
Varied Boundary
Voried PDE Coefficients
Varied Boundary Type
Varied
Unseen Boundary
Unssen
Initial
Domain
Equation
PDE Coefficients
Equation
State
Geometry
force
Volue
Symbols
(X,
M.
Cree)
Type
Symbols
Figure 2. Overview of universal neural PDE solving, taking the 2D mixed PDEs in Sec 4.3 as example. Our model is jointly trained on diverse PDEs with varied initial conditions and governing PDE components, aiming for direct generalization to unseen PDEs in downstream tasks. The "Robin boundary" in gray is valid boundary type despite not included the example dataset.
sequences of tokens and adopt the attention mechanism to approximate integrals. GNOT (Hao et al., 2023) treats each mesh point as token and utilizes the linear Transformer to address complexity issues. FactFormer (Li al., 2023b) ax- ially factorizes the attention block to boost efficiency. Tran- solver (Wu et al., 2024) learns intrinsic physical states tokens behind input meshes, introducing a physics-attention mechanism. Despite the success of neural operators, they are only tested on the dataset with limited PDEs. Their effectiveness for various PDEs has not been fully explored.
training in solving PDEs. Subramanian et al. examine the scaling capabilities and transfer learning behaviors of FNO on three time- -ene-indeeendent PDE families. MPP (McCabe et al., 2023) proposes an auto-regressive strategy train on broad fluid mechanics-oriene benchmark. DPOT (Hao et al., 2024) enhances MPP with denoising method and trains a Fourier Transformer on massive PDE data com- prised of 12 datasets. PDEformer (Ye et al., 2024) focuses on I ID time-dependent PDE family and pre-trains graph transformer on 3M samples under various equation con- ditions. ICON (Yang et al., 2023) trains single neural operator capable of performing in-context learning across a wide range of differential equations. However, most existing methods fall short of effectively and completely integrating PDE information, which Unisolver will address well.
2.2. Generalizable PDE Solvers
In addition to model architectures, the generalizability of neural solvers has also been explored. The research mainly lies in the following two directions.
Incorporating PDE information To guide the PDE- solving process, PDE information has been explored in many deep models. For example, PINO (Li al., 202,1b) im-
3. Unisolver
To tackle the incapability in generalization behind neural PDE solvers, dive into the PDE-solving process and present Unisolver to model intricate interactions between ini- tial observations and complete equation components, lead- ing to a novel PDE-conditional Transformer model.
poses explicit equation constraints at higher resolution to aid neural operator learning. CAPE (Takamoto et al., 2023) directly embeds PDE coefficients to adapt neural solvers to unseen equation coefficients. PROSE (Liu al., 2023) and PITT (Lorsung al., 2024) tokenize PDEs and embed mathematical expressions, informing the transformer of the underlying physics. PDEformer (Ye et al., 2024) represents symbolic equations as graphs and numeric components as nodes to optimize interactions between them. However, all of these methods, while incorporating equation informa- tion, do not leverage the mathematical structure of PDEs for complete and categorized embedding or integrating the prior information of equation symbols within the context of natural language. In contrast, Unisolver leverages the capa- bilities of large language models (LLMs) (Touvron et al., 2023) to semantically embed the equation symbolic infor- mation and categorize the complete equation components based on mathematical insights for better PDE solving.
Problem setup To achieve ideal generalizability. we fo- cus on the task of universal neural PDE solving shown in Figure 2. Let D C IRd be bounded continuous set and M =
xn) an n-point discretization of record- ing coordinates of each point. For each observation pair, assume have initial condition observations X input and target quantities Y as output the mesh M,, with governing PDE components CPDE (e.g. PDE symbols, coef- ficients) that may vary for each observation. The universal neural PDE solving task is to approximate the input-PDE- output map G: (X, M, CPDE) + Y across a diverse train- ing dataset and generalize in a flash to unseen PDEs.
3.1. Complete PDE Components
Large-scale training As a vital cornerstone deep leam- ing (Brown et al., 2020; He et al., 2022), recent research has also started to explore the effectiveness of large-scale
To enable complete modeling of the PDE, we attempt to incorporate the complete PDE components, iee. all the under-
3

---

Unisolver: PDE-Conditional Transformers Towards Universal Neural PDE Solvers
PDE:Conditional Transformer
x
Initial State
Solving Results
Scale Shift
Scale & Shift
Patch Embed
Layer Norm
Self- Attention
Layer Norm
Feed Forward
De- patch
Select
Select
Point-wise
Domain-wise Deep Conditions
Deep Conditions
Universal
Components
PatchEmb
PatchEmb
PatchEmb
LLM
MLP
MLP
Embedding
Domain
External
Boundary
Equation Symbols
PDE Coefficients
Boundary Type
Diickler
Geometry
Force
Value
Figure 3. Overview of Unisolver. We universally embed all PDE components, such as domain geometries, equation symbols and boundary types into deep conditions and employ conditional Transformer aggregate deep conditions in the decoupled subspace.
lying components that affect solutions, into neural solvers.
as a domain-wiss component due to its consistency across all locations. Domain geometry is categorized as point- wise component since is usually recorded as a binary mask and each point's inclusion is determined individually, Boundary conditions are more complicated due their di- verse forms, e.g. periodic and Robin boundary conditions. As a result, we use two components represent boundary conditions precisely: the boundary condition type, treated as Nomain-wise component, and the boundary value function, considered as a point-wise component.
motivating example Here, we clarify the concept in the context of deep learning by considering the classical vi- brating string equation with fixed endpoints as a motivating example, which can be solved explicitly, as shown in (Evans, 2022). The analytical solution is provided in Appendix F.
Duu d'0mau = [(r,)), (c, () E L) x (0, u(0, t) t( = 0, te (0,
(la)
(lb)
u(x, 0) = ((x), 0,u(z, = (x2), x L..
(1c)
3.2. Universal Components Embedding
As described in Section 3.1, PDE solutions are obtained by intricate interactions between initial conditions and com- plete equation components which can be grouped into two categories. In previous works (Brandstetter et al., 2022; Takamoto et al., 2023), these equation components are coarsely and incompletely included as conditions to modu- late the input observations. In this paper, we will elaborate how Unisolver finely and completely embeds all consid- ered PDE components (Table 1) into deep PDE conditions based our insights from the mathematical analysis.
From the analytical solution of the motivating example, we pinpoint that the PDE is solved through complex interactions between a series of equation components, as detailed in Table These components are referred to as the complete PDE components and exhibit key shared characteristics. Specifically, the coefficient exerts the same influence over the entire domain, while the impact of the force f is imposed wiinl- This distinction inspires us to classify these components into two categories, domain-wise and point- wise, which better capture the intricate interactions.
Equation formulation Since the mathematical symbols convey rich mathematical information, utilize Large Language Model (LLM) for symbolic embedding. Specifi- cally, we adopt the LLaMA-3 8B model to embed the equa- tion formulation. We attempt to leverage its understanding of prior mathematical information, which was learned from pre-training on 15 TB of language tokens, as well its flex- ible encoding of unstructured PDE information. Technically, the input to the L.M is the LaTTeX code of the equation. For example, the Eq. (la) is prompted
Table 1. Categorization of complete PDE components.
Groups
Components
Examples
Input
Initial condition
Eq. (1c)
Equation symbols
Eq. (la)
Dimain-wise components
Equation coefficient
a
Boundary condition type
Robin
External force
f(n,
Point-wise components
Domain geometry
(0, L) x [)
Boundary value function
Eq. (lb)
Prompt: "u_(tt} a u
fxx t)
Moreover, we explain the classification of the other compo- nents shown in Table 1. The equation formulation is defined
'nttps://a1 .et.i.etta...m meeta.com/tt com/bl og wog/metaaa1a1 ama--3/

---

Unisolver: PDE-Conditionll Transformers Towards Universal Neural PDE Solvers
Then we take the output of the last Transformer block of the LLM and average the representations along the sequence dimension, resulting in 4096-dimensional embedding for each PDE. In the LLM embedding stage, we use mathemat- ical symbols of the remaining equation components (e.g., coefficients and force terms) instead of their actual values in the prompt. For instance, we use the symbol "a" in the prompt rather than its concrete value to make the LLM focus on the key physics meaning of the PDEs. The embedding of concrete values for the other components is detailed in the next paragraph. After the LLM embedding stage, the hidden representations of PDE symbols are encoded by an MLP to align channel dimensions and obtain deep conditions.
responding subspaces. After projection, domain-wise condi- tions are repeated along the sequence dimension match the token sequence length and ensure consistent physical guiding throughout the sequence. The transformed condi- tions convey both domain-wise and point-wise information and are integrated adaptively by aggregation functions.
As shown in Figure 3, we aggregate conditions either before or after the attention and feedforward modules within Trans- former. Inspired by recent conditional Transformers like DiT (Peebles & Xie, 2023) and other conditional normaliza- tion approaches (Park et al., 2019; Perez et al., 2018), 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 representations of inputs based on the equa- tion conditions. After passing through the Transformer modules, we use the equation conditions to softly select whether this information should be retained.
Other components As shown in Table 1, other compo- nents can be categorized as domain-wise and point-wise based on their effect on the final solution. Correspondingly, we adopt different embedding methods for these two types. For domain- -wis components, including coefficients repre- sented as real-valued vectors and boundary types akin to class labels, we embed them using two linear layers with intermediate SiLU activation function (Elfwing et al., 2018). For point-wise components like external force, binary ge- ometry mask, and boundary value functions are essentially physical fields observed on mesh M. We apply the same patchify embedding method used for input observations, transforming them into deep representation sequences.
Overall design Summarizing the above designs, we pro- pose the Unisolver (Figure 3). Given input X, is projected to embeddings X using patchify layer (Dosovitskiy et al., 2020). The complete PDE equation components CYDE embedded into deep conditions Cdomain and Cpoint following Section 3.2. Suppose there are N layers, the n-th layer of Unisolver can be formalized as:
MLP. (CpeiaC X n- = Lakeet @ SA (Ieale C LN((X-1))++ + Initt) + x
I. = Concat [MLP, MML.,
Deep condition consolidation As shown in Figure 3, after universal components embedding, deep conditions within the same category are added together to consolidate their impact. This strategy prevents excessive separation of deep PDE conditions that could weaken the model's 's expressive capabilities, and thus will enhance representation learning for diverse PDE solving via joint training.
(2)
i. = Concat (MLP. SHILP.(Carnaaa
MLP. (Cpoin)).
X Lutes C PF (Locae O LN(X"-1)+ Inin) -=
where n E (1,..., NJ. is the output of the n-th layer. Mcanwhile, SA denotes self- attention, LN denotes layer normalization, and FF denotes feedforward layers in the Transformer model (Vaswani et al., 2017). Since the PDE components have a crucial impact on the range of the output, we scale and shift X based on the deep equation conditions, and then linearly project X.V to obtain the final output as predictions of Y.
E {scale, shift, select), and X"
3.3. PDE-Conditional Transformer
We propose a conditional Transformer to adaptively fuse deep PDE conditions which are embedded from the com- plete equation components, into hidden representations of inputs within decoupled subspaces.
Subspace decoupling We evenly split the hidden repre- sentations of the inputs along the channel dimension, with one half influenced by dosain-wise deep conditions and the other half by point-wise deep conditions. In multi-head attention ( (Vaswani et al., 2017), our proposed subspace de- coupling 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 representa- tion learning for both categories, and minimized interference between deep PDE conditions from two categories.
4. Experiments
We conduct extensive experiments to evaluate Unisolver on three challenging large-scale benchmarks, covering wide range of PDE components and generalization scenarios.
Benchmarks As summarized in Table 2, three experi- mental large-scale benchmarks cover varied dimensions, resolutions and PDE components. The HeterNS is an ex- tension of the NS dataset from FNO (2021a), incorporating multiple viscosity coefficients and external forces to en- hance diversity. The ID time-dependent PDEs, introduced by PDEformer (2024), is a large-scale dataset containing
Deep condition aggregation We use MLPs to individu- ally project domain- wise and wise conditions into cor-
5

---

Unisolver: PDE-Conditionll Transformers Towards Universal Neural PDE Solvers
three million structured ID PDE samples and evaluate the zero-shot generalization performance on PDEBench (2022). The 2D mixed PDEs, collected by DPOT (2024), include 12 diverse datasets from four well-established benchmarks. More details can be found in Appendix G.
coefficients v and three force terms differentiated by fre- quency w, resulting in 15 combinations of PDE components and 15,000 training samples. As depicted in Figure 4, we evaluate the model performance on in- in-distributton test with only unseen initial conditions and zero-shot generalization involving both unseen initial conditions and variations in viscosity coefficients or force terms.
Table 2. Summary of benchmarks. #GPU hours are calculated by averaging the training time of all models on A100 GPU. Detailed compute resources can be found in Appendix 1.7. Besides, v indicates the PDE component will change among different samples, while x refers to unchanged ones.
T=0
T=20
T=0
T=20
Benchmarks
HeterNS
ID time-dependent PDEs 2D mixed PDEs
#Dim
2D+Time
ID+Time
2D+Time (128, (128,128.10)
#Resolution
(64,64,10)
(256,100)
Unseen IC (In-distribution test)
Unseen IC Viscosity (zeroishot)
#Samples
1Sk
3M
74.1k
T=0
T=20
T=0
T-20
#GPU hours
60h
3300hh
~-800h
Symbols
Coefficient
Force
Geometry
X
Boundary
Unseen IC Force (zero-shot) Unseen IC & Viscosity Force (zero-shot]
Figure 4. Visualization of evaluation scenarios on the HeterNS,
Baselines We compare Unisolver with six advanced base- lines the HeterNS to demonstrate its generalizability under varied PDE components: the well-established FNO (2021a), PINO (2021b) and ViT (2020) and current state- of-the-art methods FactFormer (2023b), ICON (2023) and MPP (2023). We augment these baselines by providing sufficient physics information ensure a fair comparison, either by concatenating inputs with varied PDE compo- nents, providing prompting trajectories (ICON), or applying physics-inforned loss (PINO).
Table 3. Performance comparison on HeterNS. Note that the re- ported In-dist and Zero-shot results are averaged over multiple sub-configurations with different viscosity coefficients and force terms. Full results can be found in Appendix D.
Viscosity
Force
Method
#Params
In-dist Zero-shot
In-dist
Zero-shot
FNO
4.7M
0.0210
0.0574
0.0975
0.4316
PINO
4.7M
0.0364
0.0679
0.1544
0.7155
ViT
4.8M
0.0156
0.0514
0.0593
0.3158
Furthermore, we compare Unisolver with two generalizable Salvers-PEEformee (2024) and DPOT (2024) on zero-shot generalization performance. We refrain from including ad- ditional baselines on these two benchmarks due to the sub- stantial computational cost of using million-scale samples.
FactFormer
5.1M
0.0201
0.1162
0.0654
0.3272
ICON
4.5M
0.0216
0.061
0.0788
0.7621
MPP
4.9M
0.0347
0.0840
0.0997
0.3435
Unisolver Promotion /
4.1M
0.0098
0.0374
0.0428
0.1053
37.2%
27.2%
27.7%
66.7%
Implementations All methods in the HeterNS benchmark are trained for 300 epochs using relative 1.2 loss and the ADAM optimizer (Kingma & Ba, 2015) with an initial learning rate of 0.0005 and a cosine annealing learning rate scheduler (Loshchilov & Hutter, 2016). The batch size is set to 60. For the ID time-dependent PDEs and 2D mixed PDEs, we follow the training strategies from the original papers of PDEformer (2024) and DPOT (2024) to ensure fair comparison. Relative 1.2 is used as the evaluation metric. See Appendix H for full implementation details and hyper-parameter configurations of each model.
Results As shown in Table 3, Unisolver achieves the best performance in both in-distribution test and zero-shot gen- eralization settings. It is worth noting that external force generalization is highly difficult task, the force term fundamentally determines the fluid evolution patterns. Still, Unisolver surpasses other methods in this challenging task, with significantly greater promotions in zero-shot general- ization settings (average 66.7%) than in-distribution test settings (average 27.7 %)., demonstrating the effectiveness of our design capturing generalizable physics relations between external force and model inputs. Even though explicitly concatenate the varied PDE components with the model inputs, most advanced neural operators perform poorly on HeterNS. We also include experiments in Ap- pendix I.1, where both viscosity and force are unseen. Uni- solver still achieves considerable improvement (average 41.3% 3%) ) on this challenging double unseen setting.
4.1. Heterogeneous 2D Navier Stokes Equation: HeterNS
Setups We introduce HeterNS, an extension of the widely used 2D NS dataset (Li et al., 2021a), to assess how mod- els handle diverse PDE components, particularly viscosity coefficients and force terms. It comprises five viscosity
6

---

Unisolver: PDE-Conditional Transformers Towards Universal Neural PDE Solvers
Table Comparison (relative L2) in-distribution test and zero-shot generalization on ID deme-dependent PDEs, advection velocity & are dominated components of target PDEs.
Viscosity
ID Time-dep endent PDEs
Tasks
In-distribution Test
Zero-shot Burgers
7ero shot Advection
Params
0.1
0.01
0.001
0.1
PDEformer 22M 19M
0.0225 0.0108
0.00744
0.0144
0,0393
0,0178
Unisolver
0.00513
0.00995
0,0299
0,0138
Promotion
52.0%
31.0%
30.9%
23.9%
22.5%
Table 5. Performance comparison (relative L2 (x ))) across 12 in-distribution test sets. We use "source-PDE" format to denote different tasks (e.g. FNO-NS). The second row lists the primary PDE components each dataset. See Appendix G.3 details.
2D Mixed PDEs
Tasks
FNO-NS-v
CDEBench-CNS-(MM.
PDEBench
PDEArena
CFDBench-NS
Average Error
Params
le-5
le-4
lo-3
(1.0.1 (I. 0.01.
0011, (0.1, 0.01)
DR
SWE
NS
NS-Torce
Geometry
DPOT Unisolver
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
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
4.2. 1D Time-Dependent PDEs
4.3. 2D Mixed PDEs
Setups This benchmark contains three million high- quality ID time-dependent PDE samples with varying equa- tion formulations, coefficients, force terms and boundary conditions. We perform joint training dataset, where the input for the training task includes all relevant PDE components, and the output records full space- time fields. After training, the model is evaluated across multiple test settings as depicted in Figure 5, including in- distribution test as well as zero-shot generalization on Burg- ers and Advection equations from PDEBench (Takamoto al., 2022), which is another unseen dataset.
Setups This benchmark involves 12 datasets from four prominent benchmarks, covering a wide range of PDEs. After joint training on these diverse datasets, we perform in-distribution tests each dataset shown in Figure 6. Notably, the in-distribution test set also involves challenging variations in the PDE components. Moreover, unlike the balanced data in HeterNS, these datasets exhibit significant imbalances across different PDE components. To mitigate this issue, we adopt the balanced data sampling method from DPOT (Hao et al., 2024); however, it still poses considerable challenges managing such diverse PDE samples.
this extensive
Distribution Test
Distribution Test
Varying Geometry
Varying Force
PDEBench-Burgeen (zero-shot) PDEeench-Advectionn (zero-shot)
Figure 5. Showcases of various evaluation scenarios ID PDEs.
Varying Boundary Condition
Varying Coefficients
Results The results in Table 4 presents that Unisolver achieves significantly better in-distribution test performance compared PDEformer. This indicates that our architec- tural design of incorporating complete PDE components into Transformer model more effective than the computa- tional graph utilized by PDEformer in representing intricate physical relations between inputs and solutions. Addition- ally, Unisolver achieves better performance four zero-shot generalization scenarios, with an average improvement of 27.1 % over PDEformer, even with fewer parameters.
Figure 6. Showcases of in in-distribution test 2D mixed PDEs.
Results As shown in Table 5, Unisolver outperforms DPOT (Hao et al., 2024) in 11 of 12 in-distribution test sets with an remarkable average promotion of 17.5% (5.50 4.54), except for the small Diffusion-Reaction (DR) dataset whose relative L2 tends to be saturated (less than 5%), veri- fying the effectiveness of our design. Unisolver shows con- sistently superior performance several PDE component- dominated tasks, including coefficient generalization in

---

Unisolver: PDE-Conditional Transformers Towards Universal Neural PDE Solvers
Table Ablation results the LLM embeddings and the Condition Modeling. Variants of the former include without I.I.M embeddings (w/o LLM) and replacing by orthogonal random vectors (Random Vector), and variants of the latter include without subspace decoupling (w/o Subspace) and directly concatenating components (Concat). "Unchanged" means no changes to the default design.
Relative L2
Zero-shot Generalization
LLM Embeddings
Condition Modeling
In-distribution Test
ID Time-dep endent PDEs
Burgers
Burgers
Burgers
Advection B = 0.1
0.1
0.01
0.001
w/o LLM
Unchanged
0.0295
0,0189
0.0692
0.1432
0.0637
Unisolver
Random Vector
Unchanged
0.0290
0.0185
0.0675
0.1471
0.0632
Ablations
Unchanged
w/o Subspace
0.0287
0.0187
0.0675
0.1478
0.0625
Unchanged
Concat
0.0317
0.0236
0.0802
0.1586
0.0732
*final
Unchanged
Unchanged
0.0277
0.0176
0.0659
0.1350
0.0603
FNO (incomplete)
FNO (complate)
Unisolver (incomplete)
Unisolver (complate]
5011
Relative L2
6:35
1307
1332
cate
camy
(30
1315
500
1305
0.00
1.
2=
==3
3
60
66
6
== ==3
(a] HeterNSy 1e55
(b) HaterNSV 5e.5
[c] HeterNS 1a-4
[d) HaterNS - Se-4
(o) HaterNS 163
Figure 7. Capability of Unisolver to handle partial-observed data simulated on the HeterNS.
FNO (2021a), force generalization in PDEArena (2023), and geometry generalization in CFDBench (2023), high- lighting its ability to capture generalizable representations from complete PDE components.
drop). which indicates the benefits of our design.
Incomplete component scenario In real-world applica- tions, we may lack complete PDE components. To demon- strate Unisolver's capability to handle such situations, we randomly replace each PDE component with learnable to- kens at 30% probability in the HeterNS benchmark to simulate partially observed real-world data, resulting in 49% of samples with complete components. During inference, we can flexibly choose whether to provide PDE components as inputs. The results in Figure 7 shows Unisolver's perfor- mance under two scenarios: incomplete means no compo- nents available, while complete means full components. The results demonstrate that even with incomplete components, Unisolver surpasses FNO (2021a) in most cases, especially in complex cases with lower viscosity coefficients. More- over, complete PDE information further boosts the model's performance (average 21.6%), supporting our motivation that complete information is essential for PDE solving.
4.4. Model Analysis
Ablation Studies We investigate the effect of LLM em- beddings and condition modeling modules on 50,000 sam- ples from the ID time-dependent PDEs benchmark. The results are summarized in Table 6.
Firstly, in the LLM ablations, without LLM embedding, performance is the worst among all cases, even worse than replacing by orthogonal random vectors. I.LaMA-3 brings 5.76% % averaged promotion compared to models without LLM embedding, indicating its essential role in learning PDEs. Notably, since the I.IM only encodes one of six components, the equation symbols, a promotion of around 5% is a significant margin. Moreover, we compare the Unisolver's performance across different language models in Figure including LLaMA-3, LLaMA-2 and T5. The results are comparable, indicating each model possesses sufficient ability to encode prior mathematical information.
Visualization of learned PDE embeddings Firstly, fo- cusing on the effect of LLMs in Unisolver, we compared various LLMs and found comparable performances among them in Figure 8(a). Going further, we analyze the PDE embeddings in Figure 8(b-c), where principal component analysis (PCA) (Jolliffe & Cadima, 2016) is applied to intu- itively visualize the LLM embeddings of equation symbols and deep PDE conditions learned by Unisolver for 1D time- dependent PDEs. In Figure 8(b), we observe that PDEs with similar complexity are encoded into similar embeddings,
Secondly, in condition modeling ablations, removing sub- space decoupling introduces interference between different groups of PDE conditions, significantly impairing perfor- mance in zero-shot generalization settings, with an average drop of 5.45% %. Moreover, direct concatenation of PDE com- ponents severely hinders relation learning (21.0% average
8

---

Unisolver: PDE-Conditionll Transformers Towards Universal Neural PDE Solvers
Indiatribation leat
Relative
Leamed PDE conditions with varying
Leamed PDE conditim with varying Coz
Ll
conditions vorying Cas
15
Burgers sorp
Advecton
sinas
Aisht wieee
Bargers
Burgars
0.01
0.001
(a) Comparison different LMs
PCA visualization LLM embeddings
[c] PCA visualization of learned PDE conditions
Figure (a) Comparison utilizing different language models to encode equation symbols. (b) PCA visualization LLM embeddings. The considered PDE family contains six coefficients, such as Coz, Cos. well as source and viscosity terms. Different colors represent number non-zero coefficients, intuitively indicating the complexity PDEs. A zero coefficient results the removal of term from the equation, impacting the representations embedded by the LLM. (c) PCA visualization of learned deep PDE conditions, T in Eq. (2). We vary only one coefficient at and keep the others fixed zero, forming shown parabolic-like trajectories.
FNO Ground Truth
[1e-5, 1)
(5e-5, 2)
(1e-4, 2)
[5e44, 3)
(1e-3, 2)
(8e-, 2)
(8e-4, 2)
(1e-5, 1.5) (1e5, 3.5) (2e-5, 0.7) (2e.5, 2.7) (2e-5, 3.7]
Unisolver
Figure Error map comparison on HeterNS, where all cases share same initial condition differ viscosity (v) and force (w) (shown in the first row pairs ()). The left panel shows in in-distribution tests, while the right shows zero-shot generalization.
highlighting that L.M can indeed effectively capture prior mathematical knowledge. Figure 8(c), the trajectories of deep conditions resemble parabolas with varying degrees of curvature, indicating that the learned deep conditions successfully capture the variations of PDE components.
Case Study To provide a clear comparison, we provide showcases on the HeterNS in Figure 9, All presented jectories are generated from the same initial condition but exhibit distinct final fields, underscoring the determining role of PDE components. Further, it can be observed that Unisolver significantly outperforms FNO under more com- plex conditions, such as smaller viscosity v and larger force coefficient w, particularly in zero-shot generalization set- tings. More showcases can be found in Appendix
Scalability Scalability is crucial for building universal neural PDE solver. Figure 10 illustrates 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, achiev- ing better performance with increased data and parameters, posing the potential for a universal neural PDE solver.
5. Conclusion
To break the generalization bottleneck, this paper presents Unisolver PDE-conditional Transformer, which stems from the theoretical analysis of the PDE-solving process. Concretely, Unisolver identifies and systematically encodes complete set of PDE components into domain wiie and point-wise deep conditions separately and specifically. By integrating these conditions with Transformers through a decoupled mechanism, Unisolver can handle universal PDE components and achieve consistent state-of saute-of-the-aee art results across three challenging, large-scale benchmarks. Exten- sive analyses provided to demonstrate the effectiveness, generalizability and scalability of our model.
0.0608
0.608
Ralaaive 12
0.0407
0.0407
0.0273
0.0275
0.085
0.023
50k 10tk 200k
3M
3M
10M
19M
43M
(a) Data Scalability (Samples)
(b) Model Scalability (Porameters)
Figure Data scalability (60x) and model scalability (21x) on ID PDFs. Relative L.2 results are plotted log-log scale.

---

Unisolver: PDE-Conditional TraatiiioonaTTaaaoooonnntt Transformers Towards Universal Neural PDE Solvers
Acknowledgement
operator transformer for operator learning. In ICML, 2023.
This work was supported by the National Natural Science Foundation of China (U2342217 and 62021002), the BNRist Innovation Fund (BNR2024RC01010, and the National Engineering Research Center for Big Data Software.
Hao, Z., Su, C., Liu, S., Berner, J., Ying, C., Su, H., Anand- kumar, A., Song, J., and Zhu, J. Dpot: Auto-regressive denoising operator transformer for large-scale pde pre- training. In ICML, 2024.
Impact Statement
He, K., Chen, X., Xie, S., Li, Y., Dollár, P., and Girshick, R. Masked autoencoders are scalable vision learners. In CVPR, 2022.
This paper presents Unisolver towards universal neural PDE solving. By conditioning Transformers with PDE informa- tion, Unisolver is able to solve a wide scope of PDEs, which verifies its potential serve an efficient surrogate model for traditional numerical PDE solvers. Besides, the scalabil- ity and generalizability of Unisolver could further inspire future research on foundation models for scientific machine learning. Since we only focus on the scientific problem in this paper, there are no ethical issues regarding this work.
Herde, M., Raonić, B., Rohner, T., Käppeli, R., Molinaro, R., de Bézenac, E., and Mishra, S. Poseidon: Efficient foundation models for pdes. In NeurlPS, 2024.
Jolliffe, I. T. and Cadima, J. Principal component analysis: a review and recent developments. Philosophical transac- tions of the royal society A: Mathematical, Physical and Engineering Sciences, 2016.
References
Jumper, J., Evans, R., Pritzel, A., Green, T., Figurnov, M., Ronneberger, O., Tunyasuvunakool, K., Bates, R., Zídek, A., Potapenko, A., et al. Highly accurate protein structure prediction with alphafold. Nature, 2021.
Ames, W. F. Numerical methods for partial differential equations. Academic press, 2014.
Aroo'd, V. Mathematical methods of classical mechanics. Springer Science & Business Media, 2013.
Karniadakis, G. E., Kevrekidis, I. G., Lu, L., Perdikaris, P., Wang, S., and Yang, L. Physics-informed machine learning, Nat. Rev. Phys., 2021.
Brandstetter, J., Worrall, D., and Welling, M. Message passing neural pde solvers. In ICLR, 2022.
Kingma, D. P. and Ba, J. Adam: A method for stochastic optimization. ICI.R, 2015.
Brown, T., Mann, B., Ryder, N., Subbiah, M., Kaplan, J. D., Dhariwal, P., Neelakantan, A., Shyam, P., Sastry, G., Askell, A., et al. Language models are few-shot learners. In NeurlPS, 2020.
Kovachki, N., Li, Z., Liu, B., Azizzadenesheli, K., Bhat- tacharya, K., Stuart, A., and Anandkumar, A. Neural operator: Learning maps between function spaces with applications to pdes. JMLR, 2023.
Devlin, J., Chang, M.-W., Lee, K., and Toutanova, K. Bert: Pre-training of deep bidirectional transformers for lan- guage understanding. In NAACI, 2019.
Li, Z.. Kovachki, N. B., Azizzadenesheli, K., liu, B., Bhat- tacharya, K., Stuart, A., and Anandkumar, A. Fourier neural operator for parametric partial differential equa- tions. In ICLR, 2021a.
Dosovitskiy, A., Beyer, L., Kolesnikov, A., Weissenborn, D., 7hai, X., Unterthiner, T., Dehghani, M., Minderer, M., Heigold, G., Gelly, S., al. An image is worth 16x 16 words: Transformers for image recognition at scale. In ICLR, 2020.
Li, Z., Zheng, H., Kovachki, N., Jin, D., Chen, H., Liu, B., Avizzadenesheli, K., and Anandkumar, A. Physics- informed neural operator for learning partial differential equations. J. Data Sci., 2021b.
Elfwing, S., Uchibe, E., and Doya, K. weighted linear units for neural network function approximation in reinforcement learning. Neural networks, 2018.
Li, Z., Huang, D. Z., Liu, B., and Anandkumar, A. Fourier neural operator with learned deformations for pdes on general geometries. JMLR, 2023a.
Evans, L. C. Partial differential equations. American Math- ematical Society, 2022.
Gupta, J. K. and Brandstetter, Towards multi- spatiotemporal-seat generalized pde modeling. TMLR, 2023.
Li, Z., Shu, D., and Barati Farimani, A. Scalable transformer for pde surrogate modeling. In NeurIPS, 2023b.
Liu, Y., Zhang, Z., and Schaeffer, H. Prose: Predicting operators and symbolic expressions using multimodal transformers. arXiv preprint arXiv:2309. 16816, 2023.
Hao. Z., Ying, C., Wang. Z., Su, H., Dong, Y., Liu, S., Cheng, Z.,, Zhu, J., and Song, Gnot: general neural
10

---

Unisolver: PDE-Conditional Transformers Towards Universal Neural PDE Solvers
Liu, Z., Lin, Y., Cao, Y., Hu, H., Wei, Y., Zhang, Z.,, Lin, S., and Guo, B. Swin transformer: Hierarchical vision transformer using shifted windows. In ICCV, 2021.
Takamoto, M., Alesiani, F., and Niepert, M. Learning neural pde solvers with parameter-guided channel attention. In ICML, 2023.
Lorsung. C., Li, Z., and Barati Farimani, A. Physics in- formed token transformer for solving partial differential equations. Mach. Learn.: Sci. Technol, 2024.
Touvron, H., Lavril, T., Izacard, G., Martinet, X., Lachaux, M.-A., Lacroix, T., Rozière, B., Goyal, N., Hambro, E., Azhar, F., et al. Llama: Open and efficient foundation lan- guage models. arXiv preprint
2023.
Loshchilov, I. and Hutter, F. Sgdr: Stochastic gradient descent with warm restarts. In ICLR. 2016.
Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, u., and Polosukhin, I. Atten- tion is all you need. In NeurIPS, 2017.
Lu, L., Jin, P., Pang, G., Zhang, Z., and Karniadakis, G. E. Learning nonlinear operators via deeponet based on the universal approximation theorem of operators. Nat. Mach. Intell, 2021.
Wang, H., Fu, T., Du, Y., Gao, W., Huang, K., Liu, Z., Chandak, P., Liu, S., Van Katwyk, P., Deac, A., al. Scientific discovery in the age of artificial intelligence. Nature, 2023.
Luo, Y., Chen, Y., and Zhang, Z. Cfdbench: A comprehen- sive benchmark for machine learning methods fluid dynamics. arXiv preprint arXiv:2310.05963. 2023.
Wen, G., Li, Z., Azizzadenesheli, K., Anandkumar, A., and Benson, M. U-fno-an enhanced fourier neural operator- based deep-learning model for multiphase flow. Adv. Water Resour, 2022.
McCabe, M., Régaldo-Saing Blancard, B., Parker, L. H., Ohana, R., Cranmer, M., Bietti, A., Eickenberg, M., Golkar, S., Krawezik, G., Lanusse, F,, al. Multiple physics pretraining for physical surrogate models. In NeurlPS Al for Science Workshop. 2023.
Wu, H., Luo, H., Wang, H., Wang. J., and Long, M. Tran- solver: A fast transformer solver for pdes on general geometries, ICML, 2024.
Park, T., Liu, M.-Y., Wang, T.-C., and Zhu, J.-Y. Semantic image synthesis with spatially-adaptive normalization. In CVPR, 2019.
Yang, L., Liu, S., Meng, T., and Osher, S. J. In-context oper- ator learning with data prompts for differential equation problems. PNAS, 22063)::2331112210
2023.
Peebles, W. and Xie, S. Scalable diffusion models with transformers. In ICCV, 2023.
Ye, 7., Huang, X., Chen, I.., Liu, H., Wang. 7., and Dong, B. Pdeformer: Towards foundation model for one-dimensional partial differential equations. In ICLR Al4Differential Equations In Science Workshop, 2024.
Perez, E., Strub, F,, De Vries, H., Dumoulin, V., and Courville, A. Film: Visual reasoning with general conditioning layer. In AAAI, 2018.
Ying. C., Cai, T., Luo, S., Zheng. S., Ke, G., He, D., Shen, Y., and Liu, T.-Y. Do transformers really perform badly for graph representation? NeurIPS, 2021.
Rahman, M. A., Ross, Z. E., and Azizzadenesheli, K. U-no: U-shaped neural operators. TMLR, 2023.
Raissi, M., Perdikaris, P., and Karniadakis, G. E. Physies- informed neural networks: A deep learning framework for solving forward and inverse problems involving nonlinear partial differential equations. J. Comput. Phys., 2019.
Singh, R., Shukla, A., and Turaga, P. Polynomial implicit neural representations for large diverse datasets. In CVPR, 2023.
Subramanian, S., Harrington, P., Keutzer, K., Bhimji, W., Morozov, D., Mahoney, M. W., and Gholami, A. To- wards foundation models for scientific machine learning: Characterizing scaling and transfer behavior. In NeurIPS, 2023.
Takamoto, M., Praditia, T., Leiteritz, R.. MacKinlay, D., Alesiani, F., Pflüger, D., and Niepert, M. Pdebench: An extensive benchmark for scientific machine learning. NeurlPS, 2022.

---

Unisolver:
PDE-Conditionll TraniiiooonTTTTaaooornnnec Towards Universal Neural PDE Solvers
A. More Ablations
To further verify the role of LLM embeddings in encoding PDE information, we conduct three more additional ablation experiments. In the first experiment, the LLM only encodes the number of non-zero terms in the ID PDE. In the second experiment, the LLM encodes the "wrong" PDE information. Specifically, we replace "** with "/" and adjust polynomial orders their reciprocals. For example, the original latex code u4 + Cu1 u Cu2 u² + s(x) (C11 + C13 =²) 0 is transformed into + Col /u + co2/ /22 + s(x) + (c11/v /uu //a/w//3) = 0. In the third experiment, we manually construct a one-hot vector for each PDE term and combine them to represent full PDE. Then the combined one-hot vector is directly used by Unisolver without being encoded by an LLM. The results of these three ablation studies are shown in Table 7.
Table 7. More ablations about LLM embeddings. We include three more ablations to further demonstrate the rationale for using LLM embeddings. Relative L2 loss reported.
In-distribution Test
Burgers
Burgers V = 0.01
Burgers V = 0.001
Advection = 0.1
Abaltion type
V = 0.1
Number of non zero terms encoded by LLM
0.0285
0.0180
0.0665
0.1391
0.0618
"Wrong" expression encoded by LLM
0.0289
0.0181
0.0672
0.1361
0.0619
Manually constructed representation
0.0282
0,0184
0.0675
0.1386
0.0679
Ours
0.0277
0.0176
0.0659
0.1350
0.0603
The results indicate that the model indeed obtains additional information beyond merely the count of non-zero terms from the LLM embeddings. Moreover, embedding "wrong" mathematical information generally leads to a decline in performance, highlighting the importance of accurately embedding the PDE information. While we cannot definitely claim that the LLM "undeerstadds" mathematical knowledge, we can confirm that the use of LLM enables us encode useful mathematical information into deep representations. Besides, we observe that the LLM embedding case consistently outperforms the manually constructed representation case in both in-distribution tests and four zero-shot generalization settings, showing 4.23% average improvement. Although the manually constructed representation aims to preserve the mathematical structure of the PDE as much as possible, the handcrafted features struggle to perfectly capture the mathematical structure provided by LLMs visualized in Figure 8, leading to a decrease in performance.
B. Fine-tuning Performance
Zero-shot generalization serves valuable metric, but in scenarios where datasets differ substantially from the training set, the model's 's zero-shot performance may be limited. In such instances, fine-tuming performance is critical, since it reflects the model's 's ability to learn fundamentally generalizable representations through large-scale training. We present Unisolver's fine-tuning performance on ID time-dependent and 2D mixed PDFs in Figures 11-12, with 100 epochs for 1D time-dependent PDEs and 200 epochs for 2D mixed PDEs, both amounting to 20% of the total training epochs from scratch, demonstrating fast adaptation.
For 1D time-dependent PDEs, shown in Figure 11, fine- tuning 100 epochs on the Burgers and Advection equations from PDEBench (Takamoto et al., 2022) significantly enhances Uni- solver's performance, reducing error by 61% % compared to zero- shot results and achieving a % improvement over PDEformer (Ye et al., 2024) under the same fine-tuning conditions. These results prove the condition modeling in Unisolver more effective than the computational graph proposed by PDEformer, especially for fast adaptation. For 2D mixed PDEs, as shown in Figure 12, after 200 epochs of fine-tuning for each dataset, Unisolver reduces error by more than 12% compared to zero-shot generalization performance and outperforms DPOT (Hao et al., 2024) under the same fine-tuning conditions by 14%, showcasing its ability to extract generalizable knowledge from diverse training datasets.
0.000
Aactemment (013001
0.055
Unioa ver
Unvio ver (FT.100)
Relative L2
0.025
0.000
0,055
0.00
0.000
Burgersy - 0.1
Burgersy - 0.01 Burgers y- 0.001 Advedion # 0.1
Figure 11. Fine-tuning performance on ID PDEs. "FT-100" means finc-tuning on cach datasct for 100 epochs.
12

---

Unisolver: PDE-Conditioall Transformers Towards Universal Neural PDE Solvers
(901-301 ver (1320)
Relafive L.2
0:20
940 1e0 MM0 1e4 PNO 163
9h (101) ONd (183) 06 81.3 Chi [.E0n]
DK
1Wt
Cribeeck
#05Avmn
FOtArer ||| PDEAreno & Averoge
Awerye
(c) FNO-SSV
[b] PDEBanch & CFDbanch
Figure 12. Performance comparison (relative L2) on 2D mixed PDEs after 200 epochs of fine-tuning.
C. More Experiments about Generalizability
We conduct two additional experiments to evaluate the generalization capability of Unisolver: first, verify the benefits of joint training different types of PDEs rather than training on them independently; second, we evaluate Unisolver's capability to generalize to new types of PDEs.
C.1. The Benefit of Joint Training
We design a new experiment to evaluate the benefit of joint training on the ID time-dependent PDE benchmark. As stated in Appendix G, the general equation formulations used in this benchmark include two polynomials, fo and f1. both with a maximum order of 3. We construct three distinct sub-datasets, each with 10,000 samples, to test the impact of joint training. The polynomials each dataset are fixed orders of 1, 2, and 3, respectively, ensuring that the PDEs contained in these three datasets do not overlap. For instance, in the dataset with polynomials of order 3, only ooa and e13 are non-zero terms, while C01, Co2, C11 and e12 are fixed to zero. We conduct both joint training and independent training for 500 epochs on these 3 subdatasets. The results are shown in the Table 8.
Table 8. The benefit of joint training. We consider three distinct subset, where the polynomials are fixed to orders of 1, 2 and 3, respectively. The performance (relative 1.2) the joint training model compared against the same model trained on each subset independently.
Polynomial order
2
3
Independent Training
0.0792
0.1 161
0.1236
Joint Training Promotion
0.0555
0.0738
0.0695
29.9%
36.5%
43.7%
C.2. Equations Generalization via Finetuning
We design a equation generalization scenario based on the ID PDEs benchmark. As stated in Appendix G, the general equation formulations used in this benchmark include two polynomials, fo and f1, both with maximum order of 3. We pretrain Unisolver on 50,000 samples of PDEs with polynomial orders of up to 2, and fine-tune it for 200 epochs on PDEs with polynomial orders of 3. The fine-tuned model is compared against the same model trained from scratch for 500 epochs, with relative 1.2 error reported in Table 9. Results indicate that Unisolver pretrained on equations of polynomial order up to 2 can be efficiently fine-tuned to handle equations polynomial order Unisolver demonstrates strong generalization capabilities to unseen PDEs, significantly reducing the need for large training datasets when addressing new equations.
Table Generalization to unseen equations. Unisolver is initially trained equations with polynomial order of to and subsequently fine-tuned for 200 epochs on equations with polynomial order of 3. The performance (relative L2) of the fine-tuned model compared against the same model trained from scratch for 500 epochs.
Finetuning Examples
5000
10000
20000
0.3308
0.1913
0.1327
0.1624
0.1036 45.8%
0.0891
Promotion
50.9%
32.9%
D. Full Results for the HeterNS Benchmark
The results in Table 3 are averaged over several sub-configurations containing various viscosity coefficients or external forces. Here we provide the detailed results for each sub-confiuuration in Tables 10 and 11. Unisolver achieves the best
13

---

Unisolver: PDE-Conditional Transformers Towards Universal Neural PDE Solvers
performance in 10 of tasks, covering both in-distribution test and zero-shot generalization settings. is worth noting that external force generalization is highly difficult task, the force term fundamentally determines the fluid evolution patterns. Specifically, all compared neural operators fail to solve the case of w 0.5 in Table 11 with the relative error exceeding 0.5, while Unisolver achieves relative eтror of 0.098 on this task, further highlighting the generalizability of Unisolver.
Table Relative performance on HeterNS with different viscosity and fixed force frequency coefficient w = 2. For clarity, the best result is in bold and the second best underlined. Promtotion refers the relative improvement over the second-best method.
Viscosity
In-distribution Test
Zero-shot Generalization
HeterNS
Params
1e-5
5w-5
le-4
Se-4
1e-3
8e-6
3e-5
8e-5
3e-4
8e.4
20-3
FNO
4.7M
0.0669
0.0225
0.0114
0.0031
0.0011
0.0702
0.0373
0.0141
0.0088
0.0084
0.2057
PINO
4.7M
0.1012
0.0443
0.0263
0.0073
0.0031
0.1014
0.0646
0.0299
0.0142
0.0081
0.1894
VIT
4.8M
0.0432
0.0206
0.0098
0.0031
0.0015
0.0458
0.0353
0.0119
0.0100
0.0174
0.1878
Factformer
5.1M
0.0571
0.0259
0.0148
0.0018
0.0010
0.0489
0.0642
0.0167
0,1808
0,0639
0.3224
ICON
4.5M
0.0585
0.0267
0.0144
0.0054
0.0029
0.0606
0.0387
0.0169
0.0246
0.0110
0.2149
MPP
4.9M
0.0775
0,0496
0.0321
0.0098
0.0043
0.0796
0.0648
0.0376
0.0387
0,0236
0.2595
Unisolver
4.1M
0.0321
0.0094
0,0051
0.0015
0.0008
0.0336
0.0178
0.0064
0.0066
0.0096
0.1504
Promotion
25.7%
54.4%
48.0%
16.7%
20.0%
26.6%
49.6%
46.2%
25.0%
19.9%
Table II. Comparison (relative L2) on HeterNS with varied force and fixed viscosity / = 10
Force
In In-distribution Test
Zero-shot Generalization
HeterNS
Params
w
w 2
w =
w 1.5
= 2.5
w 3.5
FNO
4.7M
0.0640
0.0661
0.1623
1.1
0.1742
0. 1449
0.2974
PINO
4.7M
0,0914
0,1012
0.2707
1,0570
0,5010
0.4660
0.8380
VIT
4.8M
0.0348
0.0432
0.1000
0.7900
0.1412
0.1240
0,2080
Factformer
5.1M
0.0409
0.0570
0.0982
0.8591
0.1207
0.1243
0.2047
ICON
4.5M
0.0435
0,0585
0.1345
1.1950
0.5295
0.5009
0.8231
MPP
4.9M
0.0596
0.0775
0.1620
0.5532
0.2224
0,2180
0,3803
Unisolver
4.1M
0.0244
0.0321
0.0720
0.0980
0.0770
0.0720
0.1740
Promotion
29.9%
25.7%
26.7%
82.3%
36.2%
41.9%
15.0%
E. More Showcases
We provide additional showcases here to supplement the numerical results in main text. First, we visualize the in-distribution test and zero-shot generalization cases on HeterNS in Figure 1.3 and Figure 14 respectively. Next, we present visualizations for ID dime-dependent PDEs in Figure 15, Finally, we illustrate the 12 diverse datasets from 2D mixed PDEs Figure 16.
[te5.1
(te5,
[tes, 3)
[Sw5, 1)
[5e55, 2)
(Be5, 3)
[te4 1)
[te4 2) [1m4, aj
15e4,
[5ad
[5md, 3] [le3, II
[1
[le3,
INO
Facformar
Unisower
Figure Error maps (absolute difference between model predictions and ground truth) in-distribution tests with top three baselines on HeterNS. See Table 10 and 11 for numerical comparison (relative L2). All data has the same initial condition and differs in viscosity (v)
and force (w) (shown the first row by pairs (v, w)). Unisolver achieved the best visual performance among the compared baselines.
14

---

Unisolver: FDE-Conditional Transformers Towards Universal Neural PDE Solvers
[Be-5
Bev4,
[te5, 0.d)
[le5, (ted,
(1ed,
(Ped, 0.7)
(Red,, [Red, [Bed, [Ped, 4.7)
FNO
Foatormm
Unieelver
Figure 14. Error maps for zero-shot generalization setrings with top three baselines the HeterNS dataset with the same initial conditions and differs in viscosity (v) and force (w) (shown in the first row by the pairs (v, w)). See Table 10 and 11 for numerical comparison.
Burgers 0.01
Periodic
Robis
0.1
0.001
0.1
PoSherer Gyoed
Trust
Unioheer
P0S Groad
Uniddwr
Figure 15. Error maps on the in-distrihution test and zero-shot generalization (Burgers and Advection equation from PDEBench (Takamoto al., 2022)) settings in ID time-dependent PDEs. See Table for numerical comparison. We visualize cases: periodic boundary conditions and Robin boundary conditions in in distribution tests. The number in the Burgers columns is the diffusion coefficient while the number the Advection column the advection speed B.
INO
FDEBench-CNt
DR
SWE
PDEAvero
CrO
Unisover
Unisolver
Figure 16. Unisolver predictions and error maps 2D mixed PDEs. See Table for numerical results. Both predictions and error maps are provided. As shown in the CFDBench-NS columns, Unisolver presents an impressive ability to handle different geometry conditions.
15

---

Unisolver: PDE-Conditional Transformers Towards Universal Neural PDE Solvers
F. Analytical Solution for the String Vibration Equation
The solution of Eq. (Ia) with boundary conditions (lb) and initial conditions (Ic) is
((x,t)= 1 ((((z + at) + Φ at))
1 (" dr +
(zta(t--))
(C)) dg
(5,
(3)
2
2a
2a
(a-u(t--)
Initial position
Force
Initial velocity
Geometry
where (x), (z) and F(x, t) are odd, periodic functioos with period 2L defined on the upper half plane, extended from 6(x), (x) and f(z,t) t). The boundary conditions will be explicit by extending the equation the upper half plane and solving it by operator splitting and characteristic lines.
Detailed proof can be found in (Evans, 2022) or other relevant books.
G. Benchmarks
We provide detailed description of the three large-scale benchmarks in our experiments here: challenging, self-generated heterogeneous 2D Navier-Stokes Equations dataset (HeterNS) and two large-scale benchmarks, one proposed by PDEformer (Ye et al., 2024), and the other collected by DPOT (Hao et al., 2024). These benchmarks cover wide range of PDEs and diverse generalization scenarios, which can test the generalizability of PDE solvers well.
G.1. HeterNS
Similar to FNO (Li et al., 2021a), consider the 2D Navier-Stokes equation in vorticity formulation for the viscous, incompressible fluid on a unit torus. We consider both in-distribution test and zero-shot generalization settings on HeterNS. See Figure 13 and 14 for a visual representation.
w,w(r t) + u(æ,t) Vw(x, t) = vAw(2, f(x),
x (0,1)², te (0,7). (0, 1)²; t [0,T]. x (0, 1)².
(4a)
V u(x 0,
(4b)
w(x, 0) wo(x),
(4c)
Train set The problem involves two key PDE components: the viscosity coefficient and the force term. We experiment with viscosity coefficients E [8x10-5 10 2x 10 3] and force terms in the form f(x (0.1ssin(ww((( (22)) (2)))
Specifically, our training set consists of v {1 x 10-5 5 x 10-5 x 10 10-4, 5 x 10 x 10-33 and w E (1,2,3) resulting in 15 unique combinations of PDE components. For each combination, we generate 1000 samples, yielding a total of 15,000 training samples. The dataset can be accessed at the following anonymous link.?
In-distribution test set For testing, we first evaluate the in-distribution test sets, each containing 200 samples. In this setting, only the initial conditions differ from the training dataset, while all other PDE components remain the same.
Zero-shot generalization set Zero-shot generalization settings present much greater challenges, as both the initial conditions and the viscosity coefficient or force terms may be entirely unseen during training. We assess the model's zero-shot performance on 200 samples, offering a more rigorous test of its ability to learn generalizable representations.
G.2. 1D Time-dependent PDEs
This benchmark is proposed by PDEformer (Ye et al., 2024). contains million high-quality 1D time-dependent PDEs with various equation components for training and then evaluates the model performance using in distribution test sets and zero-shot generalization performance on Burgers and Advection equation from PDEBench (Takamoto al., 2022), which another distinct benchmark. See Figure 15 for visual representation.
"https://drrs www.e.ooolle.... e. wtt./ooo//vveee http:///wwi.eeeeeeeeeee/e//////r htt.:p/nnonn.ckkppppeaasaapsappa
16

---

Unisolver: PDE-Conditionll Transformers Towards Universal Neural PDE Solvers
Train set The training dataset is generated by the following PDE family:
du fol) + ((((())) -x()),, = 0, (x,t) E E-1, x 1. u(0, = g(x), x 1-111
(5a)
(5b)
where [u)) = Gru G22 + Cggu ; 0, 1. Each coefficient cu is set to zero with a probability of 0.5, and otherwise uniformly sampled from the interval [-3. The variables x(x)) and s(x) can be zero, constant or physical fields, which are all randomly sampled from pre-defined distributions, as detailed in PDEformer's original paper (Ye et al., 2024). The initial condition (xx) is randomly generated within the family of trigonometric functions, super- puper-position of sinusoidal waves as, Mo (x) = Lke=ki..... A, sin + Φi). where k; = 2 L-g are wave numbers and n,, E N are selected randomly in [1, Mmax] which is same as the zero-shot generalization tasks from PDEBench (Takamoto et al., 2022).
The dataset includes both periodic and non- 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 three pre-defined types: Dirichlet, Neumann, and Robin. The Dirichlet conditions specify the solution value the boundary, while the Neumann conditions set the derivative value at the boundary, and the Robin conditions are a linear combination of the Dirichlet conditions and Neumann conditions. Therefore, Dirichlet and Neumann boundary conditions are regarded as corner cases of the Robin conditions.
We now provide a summary from the perspective of the complete PDE components. The domain-wise components of the training dataset include equation symbolic expression, i.e. Eq. (5), boundary condition types, and coefficients in two polynomials fi while the point-wise components include the physical fields s(x) and <(x), which are considered force terms and boundary value functions. The input observations the initial conditions, discretized spatially at resolution of 256. The output is the final solution u(æ t), discretized spatially at 256 and temporally at 100.
Symbolic variations Additionally, there is one important aspect to consider regarding the symbolic variations of equation symbols. A zero coefficient in the two polynomials fi results the removal of a term from the equation. If the physical fields k(x)) or s(x) are zero, the corresponding term is removed from the prompt. When x(x) is constant, is replaced by more accurately reflect the constant value, and the same applies to s(x). These symbolic variations directly affect the equation formulations further embedded by the LLM, resulting in 26 x 3 x 3 576 types of LLM embeddings, corresponding to 576 distinct equation types.
In-distribution test set We generate 10,000 samples strictly following the configurations of the training dataset to ensure that all PDE components are within the same distribution. However, being the same distribution does not mean that they have been seen before. Given to the multitudinous PDE family, all PDE components, besides the equation symbols, can still exhibit significant variations, making in-distribution tests is also highly challenging task.
Zero-shot generalization set We employ the following two 1D PDE datasets from PDEBench (Takamoto et al., 2022) as zero-shot generalization tasks. All zero-shot generalization tasks follow periodic boundary conditions and the same initial condition family as the training dataset. The resolution of these samples is 1024 x 201. For each dataset, we use 1000 test samples. We downsample the spatial resolution of these datasets 256 and maintain the temporal resolution unchanged. The zhro-shot PDEs consist of the Burgers equation and the Advection equation.
(1) Burgers equation Burgers equation, as the fundamental equation in fluid mechanics, models the non-linear behavior and diffusion process of fluid dynamics as:
a,((t,z) 0, (u(t,x))/2) v/TO,.u E (0, 1) e (0, u(0,x:) uz) x E (0.
(6a)
(6b)
where v is the diffusion coefficient. In our zero-shot generalization settings, the Burgers equation dataset consists of three subsets, distinguished by the diffusion coefficient: v = 0.1, 0.01, 0.001. The diffusion coefficient represents the intensity of fluid variation, with smaller values corresponding to more complex fluid dynamics.
17

---

Unisolver: PDE-Conditionll Transformers Towards Universal Neural PDE Solvers
(2) Advection equation formalized as:
The advection equation models pure advection behavior without non-linearity, which can be
(,u(t, + 30, = 0, E 1), E (0, w(0,z) uz(), x E (0, 1),
(7a)
(7b)
where the constant advection speed /3 and equation symbols are considered domain-wise components in this dataset. In our zero-shot generalization settings, we use an advection speed of 3 0.1. worth noting that the advection equation has an analytic solution, given by u(t, x) = uo(r 31).
Fine-tuning We also provide tin-itniin results on 1D time-dependent PDEs in Appendix B. Compared to zero-shot generalization, fine-tune the model using an additional 9,000 samples while testing on the same 1,000 samples.
Domain alignment Notably, the spatiotemporal domain of the equations in PDEBench is [0, 1] [0.2], whereas the training dataset uses the domain[-1,1] x [0, 1J. To directly infer from the model trained on ID time-dependent PDEs, we need to align the spatiotemporal domains through spatial-temporal coordinate transformations, which will result in corresponding changes to the PDE components. Technically, the zero-shot PDEs after the coordinate transformation are given by:
Burgers equation: u. + (2u²) Bu -e' 1 O, where 1' = - = 2x 1. 2,
Advection equation: u.u 0g (43u) 0, where t' = :/2 2, x 2x 1.
G.3. 2D Mixed PDEs
This benchmark is collected by DPOT (Hao et al., 2024), which consists of the following 12 diverse subsets from benchmarks. We only conduct in-distributtin tests in the 2D mixed PDEs. Notably, the in- distribution test set also involves challenging variations in the PDE components. See Figure 16 for visual representations.
FNO-v (Li et al., 2021a) This well-established benchmark considers the 2D Navier-Stokes equation for viscous, incompressible fluid in vorticity form on the unit torus. The task is to estimate the vorticity field of the future ten timesteps on regular 64 x 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 takes values from the set {1 x 10 10-5, x 10-4, x 10 10-33. We use 1,000 instances for the viscosity value 10-5, 9,800 instances for x 10-4, and 1,000 instances for 10-3 3 to pre-train or fine-tune our model. The remaining 200 instances are used for testing its performance. In in-distribution tests, the initial conditions vary across samples.
PDEBench (Takamoto et al., 2022) The following three subsets are derived from PDEBench (Takamoto et al., 2022), encompassing three distinct equations: the compressible Navier-Stokes equation (CNS), the liffusion-reaction equation (DR), and the shallow-water equation (SWE). All datasets considered in PDEBench adhere to periodic boundary conditions. The spatial resolution of this benchmark is 128 x 128.
(1) The compressible Navier-Stokes equation models compressible fluid dynamics, including phenomena such as shock wave formation and propagation. In this dataset, two dominant nomain-wise components are considered: the Mach number (M) and shear viscosity (C). The dataset includes four combinations of these components, represented coefficient pairs (M, Q): (1, 0.1), (1, 0.01), (0.1, 0.1), (0.1, 0.01). Each combination provides 9,000 instances for training and 200 for testing. The task involves predicting the next 11 timesteps of multiple physical fields fielii-vorticity pressure, and density- seritty--vveen the initial 10 timesteps of observations. In in-distribution tests, the initial conditions vary across samples.
(2) The shallow-wate water equation, derived from the general Navier-Stokes equations, models free-surface flow problems like coastal tides, storm surges, and shallow lake flows. This equation is formalized as,
O,h + V (hu) O, 1 0.(hu) V. (1 ghe²+ 29.11 Vg.
(8a)
(8b)
18

---

Unisolver: PDE-Conditional Transformers Towards Universal Neural PDE Solvers
where h describes the water depth, b describes a spatially varying bathymetry, or describes the gravitational acceleration, and V (hu) can be interpreted as the directional momentum, A key characteristic of this dataset is its long prediction horizon. The task of interest is to predict the future 91 timesteps of water depth based on the first 10 timesteps of observations. In in-distribution tests, the initial conditions vary across samples.
(3) The 2D Diffusion-Reaction Equation involves two non-linearly coupled variables, namely the activator u == =(t,, x; y) and the inhibitor e = v(t, æ, y). It is primarily applicable for modeling biological pattern formation, such as the development of animal coat patterns, skin pigmentation and cellular organization. This equation is formalized as.
d,u Dudgou + Du + Ru.
(9a)
0,0 = DrOrnt + Dr0ygu + R..
(9b)
where D,, = x 10 3 and D, = 5 x 10-3 3 are the diffusion coefficient for the activator and inhibitor, respectively, and Ru, = R. (u, 2) and R2 = (u, v) are the corresponding reaction functions for the activator and inhibitor, which are defined by the Fitzhugh-Naguma equation as,
(u, v) = u1 k U Re (u, = u U,
(10a)
(10b)
where k = 5 x 10-3 The initial condition is generated as standard normal random noise u(0,x,,)) ~ (O,00, for a E (-1,1) and y E (-), The dataset is temporarily discretized into N; = 101. A key characteristic of this dataset is its long prediction horizon. The task of interest to predict the future 91 timesteps of u and given the initial 10 timesteps of observations. In in-distribution tests, the initial conditions vary across samples.
PDEArems (Gupta& & Brandstetter, 2023) This well-established benchmark considers the velocity function formulation of the incompressible Navier-Stokes equations, which is widely used in real-world applications, such as fluid flow in pipes, aerodynamic simulations, and weather prediction models. This equation is formalized as,
a,u= - - Vo+ - -
(11a)
v.
(11b)
where v Ve represents convection, meaning the rate of change of along its own direction, 2 v is the viscosity, i.e. the diffusion or net movement of v, Vp corresponds to the internal pressure, and f represents the external buoyancy force. The inclusion of the incompressibility constraint V u 0 ensures mass conservation within the equations.
The spatial resolution of PDEArena is 128 x 128. This benchmark includes two subsets: one with a fixed external force and another with a varied external force. In the exeeeforce subset, the initial conditions vary across samples and the task is to predict the next timesteps of velocity based on the initial 10 timesteps of observations, with 3,100 samples used for training and 200 samples for testing. In contrast, the more complex varied-force subset, where the initial conditions and force terms vary across samples, requires predicting 46 future timesteps, with 6,500 samples for training and 650 samples for testing.
CFDBench (Luo et al., 2023) We consider three important and representative fluid dynamics problems that provide a comprehensive evaluation of method's 's ability to generalize to unseen PDE components. These problems are: (1) flow in a lid-driven cavity, (2) flow through a circular tube, and (3) flow around a cylinder. The equation is formalized as follows:
(pu) + V. (pu²) Vp + V M( Vu²), V (pu)
(12a)
(pu)= = O.
(12b)
where p is the constant density, μ is the dynamic viscosity, u = (u,v) v)T is the velocity field, and p is the pressure.
In in-distribution test settings, flows are generated for each problem with different PDE components, which are combination of three types: (1) boundary conditions, (2) fluid physical coefficients such density and viscosity, and (3) the geometry of the field. The boundary conditions refer to the inlet velocity or movement velocity, depending on the specific case. Each type of PDE component corresponds to a distinct subset. In each subset, the corresponding PDE components are varied while other parameters remain constant. We mix the three subsets following DPOT's configuration (Hao et al., 2024), resulting in 9,000 training samples and 1,000 testing samples. The initial resolution is x 64, which is then interpolated to 128 128. The task is to predict the next 10 timesteps of velocity given the first 10 timesteps of observations.
19

---

Unisolver: PDE-Conditional Transformers Towards Universal Neural PDE Solvers
Fine-tuning We also provide fine-tuning results 2D mixed PDEs in Appendix B. Given the significant diversity across the 12 subsets, we fine-tune the model using specific training subset allow it focus on the target subset and achieve improved performance.
H. Implementation Details
In this section, we provide a detailed description of the implementation, covering three key aspects: metrics, implementations for each benchmark and LLM embeddding details.
H.1. Loss and Metrics
Relative L2 for physics fields follows:
We can calculate the relative L2 distance between ground truth u and model prediction á as
|/ u||
Relative L2 of (u, ú)
(13)
|| 1.a
where ||u ú|| La is the Li-distance between the predicted solution ú and the ground-truth solution and ||u]|za is the [3-norm of the ground-truth solution. Relative L2 is used as both training loss and evaluation metric.
Relative Promotion Given the error of our model Cours and the error of the second best model Esecund best model. we can calculate the relative promotion as follows:
Cours
Relative Promotion
(14)
Esccond-best model
Relative promotion is widely used in the comparison and analytical experiments across the three large-scale benchmarks to measure the improvement of the Unisolver relative to the base models.
Relative Drop Given the error of our model Cours and the error of the ablation model cablation model, can calculate the relative drop to quantify the extent of performance degradation in the ablation experiments follows:
cablation
Relative
Drop
1.
(15)
Cours
Relative drop is only used in the ablation experiments in Section 4.4 to quantify the performance loss caused by removing or replacing a specific module.
H.2. Implementations for Each Benchmark
HeterNS As outlined in Section 4, all the bascline models are trained under the same training strategy. We train the model using one-step predictions and test the model in an autoregressive manner. Specifically, all the models are trained for 300 epochs using the relative L2 loss and the ADAM optimizer (Kingma & Ba, 2015) with an initial learning rate of 0.0005, along with a cosine annealing learning rate scheduler (Loshchilov & Hutter, 2016). The batch size is set 60. After the training process, use the checkpoint from the last epoch evaluate the model performance.
We also provide the detailed model architecture hyperparameters in Table 12. We configure each model to align their model parameter numbers to ensure fair comparison. Note that for MPP (2023), We utilize the parameter configuration of the tiny version containing approximately five million trainable parameters, which comparable to Unisolver and other baselines.
The varying PDE components in this benchmark include the viscosity coefficient and external force. This physies information provided each baseline in an explicit or implicit way ensure a fair comparison. For FNO (2021a), ViT (2020), FactFormer (2023b), and MPP (2023), explicitly concatenate the viscosity coefficient and the external force to the model input along the channel dimension to ensure fair comparison. As the viscosity coefficient is essentially scalar, we repeat it along the spatial dimensions and then perform the channel-coneatnnatin which takes prompting trajectories as additional inputs to implicitly extract the physics information. Consequently, instead of providing the PDE components, we augment the input to ICON with five additional prompting trajectories with the same viscosity and external force as the target trajectory. Note that ICON also needs additional prompting trajectories when
process. ICON (2023) is a special baseline
20

---

Unisolver: PDE-Conditional Transformers Towards Universal Neural PDE Solvers
Table 12. Model hyperparameters Unisolver and all baselines 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 The number of Fourier Layers in the neural network
depth
PINO
modes
12
The truncation number of Fourier modes The number of channels in the hidden layers The number of Fourier Layers in the neural network
channels
64
depth
VIT
Attention dim
256
The hidden dimension of the transformer attention layer The hidden dimension of the transformer FFN layer
MLP dim
256
pitch_size
The height and width of the VIT patches
n_head
The number of attention heads
dim_head
32
The hidden dimension of each attention heads
depth
12
The number of Transformer Blocks in the neural network
FactFormer
dim
128
hidden dimension of the transformer The number of attention heads hidden dimension of each attention heads
n head
12
dim_head
64
depth
8
The number of Transformer Blocks in the neural network
ICON
Attention dim
256
The hidden dimension of the transformer attention layer The hidden dimension of the transformer FFN layer
MLP dim
256
patch_size
4
The height and width of the VIT patches
n_head
8
The number of attention heads
dim head
32
The hidden dimension of each attention heads
depth
12
The number of Transformer Blocks in the neural network number of prompting trajectories
prompting numbers
MPP
Embed dim
192
Dimension of internal representation The number of attention heads The number of Transformer Blocks the neural network The height and width the VIT patches
n_head
3
depth
patch_size
Unisolver
Attention dim
256
The hidden dimension of the transformer attention layer The hidden dimension of the transformer FFN layer The height and width of the Unisolver patches
MLP dim
256
patch_size
n_head
The number of attention heads
dim head depth
32
The hidden dimension of each attention heads
8
The number of Transformer Blocks in the neural network
conducting evaluation. For PINO (2021b), we follow the experiment setting in the original paper and train the model with physics-informed loss as soft regularization. The proportion of physics-inforee informed loss with regard to data loss is set to 0.
ID Time-dependent PDEs evaluating their in-distribution test and zero-shot generalization performance. We also report the model performance after fine-tuning in Appendix B. The pre-training and fine- tii-tuing configurations for Unisolver and the fine-tuing configurations for PDEformer are listed in Table 13.
We compare Unisolver with PDEformer-L in the ID time-dependent PDEs benchmark,
21

---

Unisolver: PDE-Conditional Transformers Towards Universal Neural PDE Solvers
Following PDEformer's s training strategies, train the model to predict the solution specific spatial-temporal coordinates through INR. After the pre-training process, use the checkpoint from the last epoch evaluate the model performance for the in-distribution test and zero-shot generalization test in Section 4.2. For fine-tuning tasks, we utilize the fine-tuning script provided in the original repository of PDEformer and set the finetuning epochs to 100 for a fair comparison.
The model we use to compare with PDEformer-L. contains 19M trainable parameters, which is comparable to the 22M parameters PDEformer-L. The model scalability experiments in Section 4.4 also show model configurations with different number of trainable parameters. We progressively increase the Unisolver parameter from 3M to 63M, resulting in 4 different model configurations. We present the detailed configurations of these models Table 14. Note that in this benchmark, we utilize adapted version PolyINR (Singh et al., 2023) to decode the encoder output from the Transformer backbone.
Table 13. Pre-trainn training and finetuning configurations on the ID time- dinendependent PDE benchmark.
Parameter
Value
Description
Unisolver Training
batch_size
1024
Total batchsize used in one iteration
learning-rate
6e-4
The initial learning rate for the optimizer The total number of training epochs Use relative L2-Norm for pretraining The optimization algorithm
epochs
500
loss_type
Relative-12 Adam
optimizer
Ir Ir.scheduler
Cosine Annealing
The learning rate scheduler
Unisolver Finetuning
batch size
256
Total batchsize used in one iteration The initial learning rate for the optimizer The total number of training epochs Use relative I.2-Norm for finetuning The optimization algorithm The learning rate scheduler
learning-rate
le-5
epochs
100
loss_type
Relative-12
optimizer
Adam
Ir scheduler
Cosine Annealing
PDEformer Finetuning
batch.sizz
80
Total batchsize used in one iteration The initial learning rate for the optimizer The total number of training epochs Use relative L2-Norm for finetuning
learning_rav rate
5e-6
epochs
100
loss sype optimizer Ir scheduler
Relative-12 Adam
The optimization algorithm
Cosine Amnealing 10
The learning rate scheduler
warmup.epochs
Epochs to linearly increase the learning rate
Table 14. Model configurations of Unisolver with different sizes.
Parameter Count
Attention dim
MLP dim
Layers (Backbone)
Heads
Layers (INR)
3M
256
256
6
4
4
10M
384
384
8
8
8
19M
512
512
8
8
63M
768
768
12
12
12
2D Mixed PDES We compare Unisolver with DPOT-S with comparable model parameters in the 2D mixed PDEs benchmark. The training hyperparameter and model configurations are presented in Table 15. Similar to the HeterNS benchmark, We train the model using one-step predictions and test the model in autoregressive manner.
This benchmark includes multiple diverse PDEs, each including its unique PDE components as illustrated in Appendix G. For example, the viscosity coefficient is the varying PDE components in the FNO- benchmark, while the shallow-water equation does not include this PDE component. Therefore, we must notice Unisolver whether PDE component exists in a certain benchmark. To do so, Specifically, we introduce a binary masking channel to represent the existence of a certain
22

---

Unisolver: PDE-Conditional Transformers Towards Universal Neural PDE Solvers
PDE component. For example, when a PDE component exists in benchmark, concatenate "1" with this component, indicating that this component is valid one. When this PDE component does not exist in benchmark, we concatenate "0" with it, indicating that this is an invalid one. While the LLM embedding can provide some indication of this information, it does not serve as the input to the encoders of other components, This binary mask, however, aids the encoders' learning and further clarifies the information without introducing significant computaiioal overhead.
Table 15. Training configurations on the mixed PDE benchmark.
Parameter
Value
Description
Unisolver Training Configurations
size
320
Total batchsize used in one iteration
learning-rate rate
lc-3
The initial learning rate for the optimizer
epochs
1000
The total number of training epochs
loss_type
Relative-12
Use relative L2-Norm for pretraining
optimizer
Adam W
The optimization algorithm
Ir.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 The hidden dimension of the transformer FFN layer The height and width of the ViT patches
MLP dim
768
patch size
n_head _ead depth
8
The number of attention heads
96
The hidden dimension of each attention heads
The number of Transformer Blocks in the neural network
H.3. Details of the LLM embeddings
Here give detailed description of the prompts of expressing the same PDE using different notations or mathematically equivalent transformations.
encode the equation symbols. We will also diseuss the impact
Note that the pre-training dataset of PDEformer (Ye et al., 2024) contains the PDE family following the formulation:
0.u + fo(u) + s(x) + (z) k(x)).u)) = 0, (x, t) E [-1,1] x [0,1]
(16)
where fi(u) Cau + 222u² + C2yu3, i 0, 1. Each Cty can be zero or non-zero. The source term s(x) and the viscosity term x(z) can be zero, a non-zero constant or a non-uniform function. As stated in Section 3.2, we use the LaTeX code of the equation as a prompt, and the output from the last Transformer block the LLM serves as the symbol embedding the
equation. Table gives some concrete samples of the LaTex code we use. There are 576 different equation symbols in total in the PDEformer benchmark.
Note that a differential equation may have multiple equivalent representations, and different people may express the same equation differently. A potential solution is to design targeted prompts and employ advanced prompting techniques, such as chain of thought, to standardize these variations into unified form, which is clearly within the capabilities of modern LLMs. This standardized form can then be used to enhance the learning of the solver.
I. Additional Analyses
L1. Unseen Viscosity and Unseen External Force on HeterNS
In addition to Tables 10 and 11, we further assess Unisolver's generalization on HeterNS compared to other baselines under more challenging conditions, where both the viscosity coefficient and external force are unseen. Specifically, we generate nine different component pairs (v,w), each with 200 testing samples. Notably, one case features w = 6, which significantly exceeds the maximum value of = 3 used during training, making it particularly difficult. The full results are presented in Table 17. Unisolver consistently outperforms all baselines, especially in the most challenging case with = 6, with relative promotion of 37.1%.
23

---

Unisolver: PDE-Conditional Transformers Towards Universal Neural PDE Solvers
Table 16. Sample LaTeX codes for different equations used the PIEformer benchmark.
LaTeX Code of Differential Equations
Problem Description
u.t + (c-(12) * u22)
Inviscid Burgers Equation
u.t (c.(12) * u=2 kappa uxx)
Viscid Burgers Equation
u.t + (a (11} * u) 0
Advection Equation
u.t (c.(11} * kappa u.x) X
Advection-Diffusiie Equation
u.t + c.(01}
(kappa u.x) x = 0
Reaction-Diffusion Equation
ust c.(01)
c_(02)
u (c.(12)
Fisher-KPP Equation
u kappa
u_x)
u_t c_(01)
c_(02] u + c_103] 3 (x) (c (11)
More General 1D Equations
c {12} u02
{13) 3 + kappa (x) * xUx)
Table 17. Performance comparison (relative L2) on zero-shot generalization settings with unseen viscosity (v) and unseen force (w). The pairs in first row are form (1, w). For clarity, the best result bold the second-best is underlined.
L2RE
(2e-5, 0.7)
(2e-5,
(2c-5, 2.7)
3.7)5.377
(2e-5, 4.7)
(4c-5,
(4c-5, 1.4)
(4c-5, 2.3)
(4c-5,
FNO
0.1862
0.0640
0.1176
0.2404
0.4226
0.0873
0.1516
0.0655
1.3
PINO
0.7002
0.2887
0.4776
0,8991
0,9187
0.3793
0.5596
0.3349
0.9634
VIT
0.1961
0.0690
0.1075
0,2057
0,2226
0.0488
0.1305
0.0772
0,2276
FactFormer
0,2070
0.0720
0.0891
0.1594
0.1868
0.0892
0.1456
0.0618
0,2465
ICON
0,4729
0.3693
0,5202
0,8719
0.7891
0.2212
0,5112
0.3652
0,9058
MPP
0.4532
0.4029
0,5155
0.8421
0.8484
0.2961
0.4084
0.4801
1.0240
Unisolver 0.0781
0.0378
0.0471
0.1421
0.1364
0.0399
0.0433
0.0374
0.1431
Promotion 58.06%
40.94%
47.71%
10.85%
26.98%
18.24%
66.82%
39.48%
37.13%
1.2. More Ablation Studies on PDE Components and Conditional Modeling
In addition to the ablation experiments presented in Table 6, we further conduct ablations on HeterNS to assess whether the proposed PDE information set is essential and whether the condition modeling is effective for the solver's learning. This is demonstrated by removing specific components and replacing Unisolver's condition modeling with direct concatenation of PDE information.
As shown in Tables 18 and 19, removing the information leads to significant drop in performance compared vanilla Unisolver, and concatenating the information directly also results in huge decline. It is worth noting that the absence of external force information or its improper use (c.g. via direct concatenation) significantly degrades performance even in zero-shot viscosity generalization tasks, and vice versa, further highlighting the importance of including complete PDE components.
Table 18, Ablations with different viscosity coefficient v and fixed force w = 2 on the HeterNS on removing some PDE components (W/)), and replacing domain-wise or point-wise conditions from our design directly concat (Concat).
Viscosity
In-distribution Test
Zero-shot Generalization
HeterNS
Params
5e-5
le-4
Se-4
le-3
pr
3e-4
W/o viscnsity
4.1M
0.0388
0.0127
0.0084
0.0031
0.0015
0.0410
0.0367
0.0099
0.0068
0.01
W/o force
4.1M
0.0353
0.0123
0.0074
0.0027
0.0017
0.0378
0.0198
0.0086
0.0096
0.0124
Concat viscosity
4.1M
0.0343
0.0107
0.0058
0.0017
0.0011
0.0359
0.0192
0.0071
0.0278
0.0243
Concat force
4.1M
0.0331
0.0103
0.0061
0.0018
0.0010
0.0357
0.0191
0.0071
0.0104
0.0101
Unisolver
4.1M
0.0321
0.0094
0.0051
0.0015
0.0008
0.0336
0.0178
0.0064
0.0066
0.0096
24

---

Unisolver: PDE-Conditional Transformers Towards Universal Neural PDE Solvers
Table 19. Ablations with different force w and fixed viscosity y = 10 on the HeterNS on removing some PDE components (Wo), and replacing domain-wise wise or point-wiss wint conditions from our design to directly concat (Concar).
Force
In-distribution Test
Zero-shot Generalization
HeterNS
Params
w=
=
3
0.5
1.5
2.5
= 3.5
W/o viscosity
4.1M
0.0310
0.0388
0.0926
0.261
0.250
0.258
0.424
Wo force
4.1M
0.0267
0.0353
0.0804
0.553
0.618
0.657
0.913
Concat viscosity
4.1M
0.0265
0,0343
0.0786
0.1267
0,2057
0,2771
0.2689
Concat force
4.1M
0.0259
0.0331
0.0764
0.5386
0.3392
0.2841
0.2753
Unisolver
4.1M
0.0244
0.0321
0.0720
0.0980
0.0770
0.0720
0.1740
1.3. Long Trajectory Prediction
We extend the temporal evolution steps of HeterNS to 30 steps, corresponding to 30 seconds of complex fluid dynamics, and report the zero-shot performance comparison between Unisolver and the baselines in the Table 20. We present the performance on the subdataset with a viscosity coefficient v = 1 x 10 and a force coefficient of = This is a particularly challenging task, as these models have never seen such long trajectories in the training data (at most 20 seconds). Despite this, Unisolver still achieves the best performance compared with the top three baselines.
Table 20. Zero-shot performance comparison (relative L2) with top three baselines
Unisolver
FNO
VIT
FactFormer
Relative 1.2
0.1956
0.3105
0.2527
0.2962
1.4. Full Scalability
As a supplement to Figure in the main text, we also conduct experiments on different zero-shot generalization tasks from (Takamoto et al., 2022) and record the concrete data Table 21 for clarity.
Table 21. Scalability results on in-distribution and zero-shot generalization tasks, as depicted in Figure 10.
L2RE
Data Scalability (Samples)
Model Scalability (Parameters)
Scale
50k
100k
200k
3M
3M
10M
19M
63M
In-distribution test
0.0232
0.0202
0.0170
0.0106
0.0342
0.0226
0.0202
0.0156
Zero-shot Burgers v
0.0161
0.0116
0.0081
0.0051
0.0143
0.0134
0.0116
0.0091
Zero-shot Burgers v = 0.01
0.0649
0.0412
0.0260
0.0144
0.0552
0.0421
0.0412
0.0351
Zero-shot Burgers = 0..01
0.1399
0.1003
0.0689
0.0299
0.1188
0.0976
0.1003
0.0889
1.5. Efficiency Analysis
We provide the inference time and memory consumption for each model to predict single frame on the HeterNS, along with the calculation time and memory consumption of the numerical solver, which is pseudo-spectral solver adopted by FNO (Li et al., 2021a) used to generate the HeterNS dataset, to calculate the next frame, as summarized in the Table 22. The results are measured on an A A100 GPU with a batch size of 1. Unisolver demonstrates comparable inference speed to FNO, while consuming less memory. Besides, all neural PDE solvers are approximately 1,000 times faster than the numerical solver, highlighting their potential as efficient surrogate models.
1.6. Standard Deviations
We repeat the experiments three times on the HeterNS benchmark and provide standard deviations here. As shown in Table 23-24, Unisolver surpasses the previous state-of-the- models with high confidence. Note that we compare Unisolver
25

---

Unisolver: PDE-Conditional Transformers Towards Universal Neural PDE Solvers
Table 22. Efficiency Analysis. The inference (calculation) time and memory consumption for each model and numerical solver to predict single frame on the HeterNS dataset.
FNO
PINO
ViT
FactFormer
ICON
MPP
Unisolver
Numerical Solver
Average Inference
0.0042
0.0042
0.0045
0.0103
0.0057
0.0120
0.0054
7.26
(Calculation) Time s
Average Memory Usage MB
730
730
558
758
784
1200
554
524
with the second-bets model, which a strong baseline as it is not achieved by single model. The results demonstrate that Unisolver significantly outperforms baseline models, with the second- best result falling more than three standard deviations behind, except in the case of viscosity Se 4.
Table 23. Standard Deviations the HeterNS benchmark with different viscosity coefficients and fixed force frequency coefficient 2.
In-distribution Test
Zero-shot Generalization
Viscosity v
le.5
5e-5
le-4
Sc-4
1e-3
8c-6
8c-4
Second-best model
0.0432
0.0206
0.0098
0.0018
0.0010
0.0458
0.0353
0.0119
0.0088
0.0081
Unisolver
0.0321
0.0094
0.0051
0.0015
0.0008
0.0336
0.0178
0.0064
0.0066
0.0096
Standard Deviation
+0.0005
+0.0003
+000001
+0.0001
+0.00006
+0.0008
#000002
+:0.0004
+000007
=0.00007
Confidence Level
99%
99%
99%
99%
99%
99%
99%
99%
99%
Table 24. Standard Deviations on the HeterNS benchmark with different force (wi and fixed viscosity coefficient 2.
distribution Test
Zero shot Generalization
Force
w 0.5
1.5
2.5
= 3.5
Second-best Model
0.0348
0.0432
0.0982
0.5532
0.1207
0.1240
0.2047
Unisolver
0.0244
0.0321
0.0720
0.0980
0.0770
0.0720
0.1740
Standard Deviation
+ 0.0003
+ 0.0002
0.0000
+ 0.0015
0.0048
- 0.0051
0.0021
Confidence Level
99%
99%
99%
99%
99%
99%
99%
I.7. Detailed Compute Resources
Our models were trained on servers with 32 NVIDIA A 100 GPUs, each with 40GB memory. Here we present the compute resources in terms GPU hours, where one GPU hour represents the time spent training on single A100 GPU for one hour. This metric reflecting the resources required to reproduce the experimental results are shown Table 25.
Table 25, Computational costs GPU hours, measured on NVIDIA A100 GPUs (40 GB memory).
Benchmarks
HeterNS
ID Time demendennn PDEs
2D Mixed PDEs
Models
FNO
FactFormer
VIT
PINO
ICON
MPP
Unisolver
Unisolver
Unisolver
#GPU hours
12
100
24
12
24
30
24
3000
800
J. Full Trajectory Visualizations
To better understand the temporal evolution of the benchmark, we visualize the full trajectory of the ground truth and Unisolver predictions on HeterNS and 2D mixed PDEs Figure and 18.
26

---

Unisolver: PDE-Conditionll Transformers Towards Universal Neural PDE Solvers
Ground Truth
FNO-NS
Unisolver Prediction
PDEBenck-CNS
Ground Truth
Unisolver Prediction
PDEBench-Di
Ground Truth
Unisolver Prediction
PDEBench-SWF
Ground Truth
Unisolver Prediction
Ground Truth
PDEArena-NS
Unisolver Prediction
Ground Truth
CFDBench-NS
Unisolver Prediction
Figure Visualization of full trajectories the 2D mixed PDEs, with the names of subsets displayed on the right. Ground truth and Unisolver predictions are presented, visually highlighting the complexity and diversity of the 2D mixed PDEs.
K. Limitations and Future Work
This paper presents Unisolver solve PDEs under universal PDE components, which achieves impressive performance supported by extensive analyses and visualizations. However, method is currently limited to grid data due to the patchifying process during the embedding of wise components. Actually, this limitation is shared in all the generalizable PDE solvers, such as MPP (McCabe et al., 2023), Poseidon (Herde et al., 2024), PDEformer (Ye et al., 2024) and DPOT (Hao et al., 2024). One fundamental reason is the lack of suitable and large-scale irregular-mesh PDE datasets, which will require extremely high computation costs for generation and massive resources for collection. Since our primary focus in this paper on the study of model architecture design and generalization capabilities, we would like leave the irregular- mesh PDE dataset future work. Also, the capability to handle irregular meshes of Unisolver be achieved by replacing the canonical Transformer with the latest ceometry-geneete PDE solver: Transolver (Wu al., 2024).
27

---

Unisolver: PDE-Conditional Transformers Towards Universal Neural PDE Solvers
Initial Observations
Ground Truth & Unisolver Prediction
T=0
T= 10
T- 11
T= 12
T- 13
T- 14
T- 15
T- 16
T 17
18
T 19
T 20
(1e33, 2.0)
Unisolver Prediction
[1e4, 2.0)
Unisolvar Prediction
[1e5, 1.0)
In-distribution Test
Unisolver Prediction
[1e5,, (.
Unisolver Prediction
(1e5, 3.0)
Unisolver Prediction
[Se5, 2.0)
Unseen Viscosity
Unisolver Prediction
Unseen Force
Unisolver Prediction
12e-5, 2.7)
Unseen Viscosity & Force
Unisolver Prediction
14e5, 1.4
Unisolver Prediction
Figure Visualization trajectories the HeterNS, where all trajectories share same initial condition differ in viscosity (v) and force (w) (shown beside each case by the pairs w)).
28


Reviews:
[{"review_id": "dAKVhfYS4X", "reviewer": "Reviewer_Z6hm", "summary": "This paper introduces Unisolver, a framework that conditions a transformer model on various physical parameters relevant to PDEs. The framework distinguishes between domain-wise components (such as equation symbols and coefficients) and point-wise components (such as external forcing). These are incorporated via adaptive layer normalization, with domain-wise components either extracted from a large language model (LLM) or modeled using an MLP, while point-wise components are patchified for compatibility with the transformer. The method is evaluated on three benchmarks: HeterNS, 1D time-dependent PDEs, and 2D mixed PDEs, where it is compared against baselines.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely thank Reviewer Z6hm for providing valuable feedback and suggestions.\n\n> **Q1:** About the incomplete scenario setup and incomplete ratio.\n\n**(1) Clarify our setting.**\n\nSorry for the confusion. We clarify the incomplete component scenario setup:\n- **During training**, each PDE component (viscosity and force in HeterNS) is independently masked with 30% probability, resulting in 49% samples with full components.\n- **During evaluation**, the **\"incomplete\"** setting in Figure 7 means **no components available**, while **\"complete\"** means full components.\n\nThe results demonstrate that our model can inference with no PDE components, and providing components boosts performance.\n\n**(2) New experiment with a larger incomplete ratio.**\n\nWe perform additional ablation where **each component is masked with 80% probability** during training, resulting in **only 4% data with full components**. As shown in the results below (averaged across multiple forces), our model outperforms FNO and maintains strong performance.\n\n**No components available: (80% masked)**\n|Viscosity|1e-5|5e-5|1e-4|5e-4|1e-3|\n|-|-|-|-|-|-|\n|FNO|0.1039|0.0485|0.0305|0.0097|0.0047|\n|Unisolver|0.0647|0.0237|0.0147|0.0043|0.0022|\n\n**Full component: (80% masked)**\n\n|Viscosity|1e-5|5e-5|1e-4|5e-4|1e-3|\n|-|-|-|-|-|-|\n|FNO|0.1009|0.0473|0.0288|0.0093|0.0044|\n|Unisolver|0.0644|0.0232|0.0136|0.0039|0.0021|\n\n> **Q2:** The method assumes full knowledge of the governing equation, which may not always be realistic.\n\nAccording to $\\underline{\\text{Incomplete component scenario section}}$, our model does not rely on complete knowledge of governing equations and can be trained with partial PDE information, supporting inference with no PDE components. This enables application to real-world settings with incomplete PDE knowledge.\n\nBesides, we want to note that beyond real world, simulation in CAE software is also valuable where complete information is easy to obtain and Unisolver can serve as an efficient surrogate.\n\n>  **Q3:** About contribution of LLM embedding, interpretation of learned PDE embeddings, and whether the model is simply performing pattern recognition.\n\nAs shown in $\\underline{\\text{Table 6}}$, incorporating LLM embeddings yields **an average improvement of 5.76%**, indicating they provide meaningful information. $\\underline{\\text{Table 7}}$ further shows that LLM embeddings outperform manually constructed symbolic embeddings on both in-distribution and downstream tasks, with **over 10%**  improvement on Advection equation, highlighting better generalization performance of LLM embeddings and benefits beyond simple pattern recognition. \n\nRegarding efficiency, as the LLM has been heavily optimized, generating embeddings incurs negligible computational overhead.\n\nTo improve interpretability, we provide additional visualizations: https://anonymous.4open.science/r/rebuttal-4EC7/visualization.png. We filter out equations with viscosity and force for clarity and annotate each cluster with a representative PDE formula. The visualization clearly shows the well-structured latent space of LLM embeddings.\n\n> **Q4:** About statement on existing methods.\n\nWe agree that \"fail to incorporate all PDE information\" may be too strong. We will revise the statement to **\"do not fully utilize all available PDE information\"** to soften our claim.\n\n> **Q5:** MPP and Poseidon should be discussed earlier; CODA , CAPE, and Zebra should be mentioned.\n\nWe will discuss MPP and Poseidon earlier in Related Works for clarity. We have already cited CAPE, and we will cite CODA and Zebra to better contextualize Unisolver’s novelty.\n\n> **Q6:** About distinction between neural PDE solvers and neural surrogates.\n\nWe use the term *neural PDE solver* to refer to PINNs and neural operators following prior works like *Message Passing Neural PDE Solver* and *Transolver*. We also acknowledge that under a narrower definition by [1], only PINNs are considered neural solvers. Considering the mixing concept, we prefer to follow the paper *Message Passing Neural PDE Solver* with more relative topic and maintain the \"neural PDE solver\".\n\n[1] Physics-informed machine learning: A survey on problems, methods and applications.\n\n> **Q7:** About conditioning approach and repeating encoding.\n\nToken repeating is an implementation trick similar to tensor broadcasting. We experiment with cross attention on HeterNS. As shown in the table below, our design outperforms cross attention in PDE information conditioning.\n\n|Viscosity Generalization (Relative L2)|In-Dist|Zero-shot|\n|-|-|-|\n|Unisolver (Cross Attn)|0.01078|0.0416|\n|Unisolver (Ours)|0.0098|0.0374|\n\n> **Q8:** About difference between domain and point-wise embeddings.\n\nThe importance of domain-wise vs. point-wise components varies by dataset. In HeterNS, point-wise components (e.g., force terms) are more critical due to their strong influence on fluid patterns. For 1D and 2D mixed PDEs, both types contribute significantly in guiding the simulation.", "claims_and_evidence": "The authors claim that Unisolver achieves strong performance in both in-distribution and out-of-distribution settings. The experimental results suggest that:  \n- The comparisons with baselines appear fair, with most methods receiving similar input information (except for ICON and PINO, which differ in conditioning).  \n- However, the \"incomplete scenario\" setup is somewhat unclear. Does it refer to partial information during training or only at inference?  \n- The interpretation of the learned PDE embeddings is ambiguous, making it difficult to assess the quality and significance of the extracted representations.", "methods_and_evaluation_criteria": "The proposed approach is conceptually interesting, as it attempts to create a unified conditioning mechanism for neural PDE solvers. However, a few concerns remain:  \n- The idea of conditioning a transformer on all available PDE information is relevant, but the method assumes full knowledge of the governing equation, which may not always be realistic.  \n- The use of an LLM for encoding equation symbols seems questionable. Does the LLM contribute meaningful information, or does it merely introduce additional complexity? The results in Table 7 seem to indicate that the LLM does not significantly improve performance, which raises concerns about the validity of this design choice.  \n- The introduction suggests that prior approaches fail to incorporate all available information, but this framing might be misleading—previous methods likely did not attempt such exhaustive conditioning because it is not always necessary.", "theoretical_claims": "This is an experimental paper.", "experimental_designs_or_analyses": "- The experimental setup appears sound, and the evaluation is conducted on diverse PDE scenarios.  \n- However, the interpretation of learned embeddings is not particularly insightful, making it difficult to assess whether the model genuinely understands PDE structure or is simply performing pattern recognition.  \n- The partially observable setting is not very convincing: the model still relies on 70% fully observed data, which is a relatively mild missing-data scenario.", "relation_to_broader_scientific_literature": "This work aligns with research on generalizable PDE surrogate models and foundation models for PDEs.", "essential_references_not_discussed": "Key references are missing:  \n- MPP and Poseidon should be discussed earlier in the text.  \n- Generalization methods such as CODA [1], CAPE [2], and Zebra [3] should also be mentioned to provide a clearer contextualization of Unisolver’s novelty.  \n\n[1]Generalizing to New Physical Systems via Context-Informed Dynamics Model, Kirchmeyer et al, 2022.\n[2]Learning Neural PDE Solvers with Parameter-Guided Channel Attention, Takamoto et al, 2023.\n[3]Zebra: In-Context and Generative Pretraining for Solving Parametric PDEs, Serrano et al, 2024.", "other_strengths_and_weaknesses": "The distinction between neural PDE solvers and neural surrogates is not well discussed in the introduction. PDE solvers aim to directly solve the equation, whereas surrogates approximate numerical solutions efficiently.", "other_comments_or_suggestions": "1. Clarify the contribution of the LLM: How much does it actually help? If it provides limited gains, should it be removed?  \n2. Improve the \"Neural PDE Solvers\" section in Related Work: The discussion should better differentiate between neural solvers (PINNs, which enforce physics constraints) and neural surrogates (which approximate solutions efficiently).  \n3. Rephrase certain claims: For instance, stating that \"most methods fail to incorporate all PDE information\" is an overstatement—rather, existing methods prioritize different aspects of the problem based on their intended applications.  \n4. Consider an alternative conditioning approach: Instead of repeating equation information at each location, why not use a single token and cross-attention to encode the PDE domain information?", "questions_for_authors": "1. What PDE knowledge does the LLM actually encode? Would removing the LLM affect performance?  \n2. Why use repeated encoding for PDE parameters? Would a single token with cross-attention be a more efficient alternative?  \n3. Did you find a difference between domain and pointwise embeddings ? Which ones are the most important ?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "kWJYKdVVx9", "reviewer": "Reviewer_8zJe", "summary": "The paper introduces Unisolver, a universal neural PDE solver that can handle a wide range of PDEs, unlike traditional neural solvers that are limited to specific equations and coefficients. Instead of merely scaling up data and parameters, Unisolver leverages theoretical insights into PDE structures, embedding key components (e.g., equation symbols, coefficients, and boundary conditions) into a Transformer-based model. This approach integrates physics and deep learning, achieving state-of-the-art performance and superior generalization across diverse PDEs.", "questions": "", "limitations": "", "rating": 2, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely thank Reviewer 8zJe for providing valuable feedback and insightful questions.\n\n> **Q1:** \"'Universal' is too big to say.\"\n\nThanks for this rigorous review, which is very instructive for us.\n\n**(1) We adopt \"Universal\" in the context of deep learning to highlight model's flexibility and broad applicability.**\n\nWe appreciate the reviewer’s concern. Our use of the term \"universal\" is not meant to claim generality across all kinds of PDEs, but rather to highlight the **flexibility and broad applicability** of our model in **incorporating diverse PDE information and handling large-scale PDE datasets**. Similar usage also exists in the previous work as \"Universal Physics Transformers\" [1].\n\nSpecifically, our method allows the deep model to incorporate all available PDE information including PDE types, coefficients, boundary conditions, domain geometries and force terms. This allows the model to flexibly handle large-scale and diverse PDE datasets, as demonstrated by our experiments on three challenging benchmarks. \n\nIn particular, we have trained a single unified model for all considered 1D PDE datasets, and likewise a single model for 2D PDE datasets containing all kinds of PDE components, each covering a wide spectrum of PDE variations. The experiment setup demonstrates **a certain degree of universality**, as the models generalize well across diverse PDE families within their respective domains.\n\n[1] Universal Physics Transformers: A Framework For Efficiently Scaling Neural Operators, NeurIPS 2024\n\n**(2) We will change the title to \"Towards Universal Neural PDE Solvers\" for scientific rigor.**\n\nThanks for the reviewer's kind reminder, we acknowledge the potential overstatement of \"universal\" and we promise to revise our title to **\"Unisolver: PDE-Conditional Transformers Towards Universal Neural PDE Solvers\"** to more accurately reflect the scope and position of our work. Similar usage also exists in previous work [2]. This revision will help clarify that the goal of our work is to make progress towards more generalizable and practical neural PDE solvers.\n\n[2] Towards Foundation Models for Scientific Machine Learning: Characterizing Scaling and Transfer Behavior, NeurIPS 2023\n\n> **Q2:** \"For such simple examples... How do you justify?\" \n\n**(1) Not all the PDE-solving tasks require large computation resources.**\n\nWe apologize for the confusion. While we conduct experiments on a 32-GPU server, we would like to clarify that **not all our models were trained on 32 A100 GPUs**. The HeterNS dataset is a relatively small dataset and models were trained with a single GPU. For the 2D mixed PDEs dataset, we use 8 A100 GPUs to train our model, aligning with the training configuration reported in DPOT. For the 1D time-dependent PDEs dataset, our model was trained on 32 GPUs, which is justified by its large scale and high diversity, comprising over **3 million training samples** in total. The computational cost of Unisolver is shown below, which can also be found in $\\underline{\\text{Appendix I.7}}$.\n\n| Benchmarks | HeterNS | 1D Time-dependent PDEs | 2D Mixed PDEs |\n| - | - | - | - |\n| GPU Hours  | 24      | 3000                   | 800           |\n\nBesides, we also want to highlight that although Unisolver requires some time for training, once trained, it can generalize to new PDEs without retraining and efficiently generate solutions, as can be seen in the efficiency analysis in $\\underline{\\text{Appendix I.5}}$. This allows Unisolver to serve as an efficient surrogate of numerical solvers, significantly reducing the computational overhead, which further justifies the training cost.\n\n**(2) Our PDE-solving benchmarks are among the hardest ones in current research community.** \n\nRegarding the concern about simplicity, we respectfully argue that the 2D datasets are highly non-trivial, covering a wide range of complex PDE types and diverse PDE components. While the 1D dataset is relatively simpler from a numerical solving perspective, it is specifically constructed to cover a large variety of time-dependent PDEs, posing significant challenges for training and generalization from a deep learning perspective. Specifically, the 1D PDE family contains six polynomial coefficients, various viscosity terms and force terms, three types of boundary conditions, all of them can vary simultaneously, imposing intricate challenges for the model to capture the complex relationship between PDE component inputs and the corresponding solutions.\n\nFinally, we would like to note that the datasets used in our paper **are not intended as end goals, but as a foundation for systematically studying the performance of deep surrogate models** across increasingly challenging PDE settings, paving the way towards more practical and powerful PDE foundation models.", "claims_and_evidence": "'Universal' is too big to say.", "methods_and_evaluation_criteria": "Yes", "theoretical_claims": "No Such.", "experimental_designs_or_analyses": "Yes, all of them.", "relation_to_broader_scientific_literature": "Foundational model is open question to the SciML community. This paper provides a tangent solution.", "essential_references_not_discussed": "No such.", "other_strengths_and_weaknesses": "'Our models were trained on servers with 32 NVIDIA A100 GPUs, each with 40GB memory.' for such simple examples... How do you justify?", "other_comments_or_suggestions": "No such.", "questions_for_authors": "No such.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "BdontfjuTk", "reviewer": "Reviewer_GsAC", "summary": "The paper proposes a method for solving various types of PDEs by leveraging a pretrained LLM alongside known parameterizations in the form of equations and values. The symbolic equations are embedded using the pretrained LLM, while numerical values and boundary/initial conditions are incorporated separately through conditioning in a Transformer. The approach is evaluated in both in-domain settings (where equations and general conditions remain similar) and out-of-domain settings (where equations and/or their coefficients and parameterizations differ).", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely thank Reviewer GsAC for providing the insightful review and valuable suggestions.\n\n> **Q1:** \"The claim of universality may be overstated. The equations considered remain a relatively small subset of possible PDEs.\" \"The domain remains discretized on a grid, even in CFDBench, which further weakens the claim of universality.\" \"Soften their claim on universality.\"\n\nWe thank the reviewer for providing valuable feedback. \n\n**(1) Regarding handling irregular geometry.**\n\nWe acknowledge the limitation of our current method in handling irregular geometries and have discussed this in $\\underline{\\text{Appendix K}}$. We highlight that this limitation is shared by existing PDE foundation models such as DPOT, MPP, and Poseidon, all of which focus on data defined on regular grids. One fundamental reason behind this is the lack of suitable large-scale PDE datasets on irregular meshes. To extend Unisolver to irregular geometries, one possible approach is to replace the current canonical Transformer with geometry-general PDE models like Transolver.\n\n**(2) Regarding the potential overstatement of \"universality\".**\n\nWe agree that the claim of \"universal\" may overstate the current scope of our model, as our method does not handle all possible kinds of PDEs, but rather focuses on several diverse sets of PDEs. Our definition of a \"universal\" neural PDE solver should be able to incorporate all possible PDE information and to generalize across the wide family of PDEs, and our method is one step beyond existing approaches by systematically encoding available PDE information. We promise to revise our title to **\"Unisolver: PDE-Conditional Transformers Towards Universal Neural PDE Solvers\"** to more accurately reflect the scope and goal of our work. \n\n>  **Q2:**  Although the paper claims to provide theoretical analysis, I did not find any substantial theoretical justification in the text.\n\nWe acknowledge the reviewer's concern. Our model design is inspired by theoretical insights on how PDE components influence the solutions, as shown in the motivating example in $\\underline{\\text{Section 3.1}}$. However, we would like to clarify that we do not attempt to provide formal theoretical analysis or guarantees in this work.\n\n> **Q3:** Suggestions on improving visualization of embedded PDE conditions.\n\nWe sincerely thank the reviewer for providing the valuable suggestions. We provide an additional visualization of the learned PDE condition embeddings, which can be found through this anonymous link: https://anonymous.4open.science/r/rebuttal-4EC7/visualization.png. The new 3D plot provides more insightful visualization of the learned embeddings, clearly reflecting how each coefficient impact the embedded PDE conditions.\n\n> **Q4:** Unisolver exhibits more artifacts.\n\nWe thank the reviewer for the detailed observation. $\\underline{\\text{Figure 14}}$ displays the error maps of each model, which is the absolute difference between the model predictions and the ground truth. Although there are visual artifacts in the error map, it is important to note that the model predictions does not display such significant artifacts. This can be demonstrated in the Full trajectory visualization in $\\underline{\\text{Figure 17 and 18}}$​. Actually, the absolute error is much smaller than the ground truth values, thus making the artifacts in error maps more noticeable.", "claims_and_evidence": "- Generalization ability: The proposed method appears to achieve strong performance in both in-domain and out-of-domain settings. This claim is supported by the quantitative results presented in Section 4.  \n\n- Universality claim: While the model demonstrates some OOD generalization capabilities through extended experiments on challenging benchmarks, the claim of universality may be overstated. The equations considered remain a relatively small subset of possible PDEs. To substantiate this claim further, additional validation with a broader range of equations, such as those used in mesh-based simulations (e.g., Pfaff et al., 2021), would be necessary.  \n\n- Theoretical analysis: Although the paper claims to provide theoretical analysis, I did not find any substantial theoretical justification in the text.  \n\nReference:\n- Pfaff et al. (2021), Learning Mesh-Based Simulation with Graph Networks", "methods_and_evaluation_criteria": "- The proposed method of incorporating prior information by encoding it within a sufficiently flexible framework is conceptually sound. I find it particularly interesting that separating the equation skeleton from the specific numerical values leads to improved performance, highlighting certain limitations of pretrained LLMs.\n- The evaluation benchmarks are adequate for demonstrating the model’s generalization capabilities within the family of equations considered, particularly in a squared domain setting.", "theoretical_claims": "Not applicable. Although the abstract mentions a \"theoretical analysis of the PDE-solving process,\" I did not find any substantive theoretical claims in the paper.", "experimental_designs_or_analyses": "I reviewed the experimental settings and did not find any specific issues.", "relation_to_broader_scientific_literature": "The key contributions of this paper align with existing research on neural solvers, particularly in the context of enhancing OOD generalization by leveraging privileged information about the equations.", "essential_references_not_discussed": "To the best of my knowledge, there are no essential related works missing from the discussion.", "other_strengths_and_weaknesses": "Strengths:\n- The results clearly demonstrate OOD generalization when leveraging privileged information about the equations.  \n- The approach of using LLMs to embed symbolic equations is an interesting direction.\n- The inclusion of a CFD benchmark is valuable for showcasing the method in an applied setting, in which the model should have access to privileged information on the equations.\n\nWeaknesses:\n- The domain remains discretized on a grid, even in CFDBench, which further weakens the claim of universality.", "other_comments_or_suggestions": "- I found the PCA analysis of embedded PDE conditions interesting but in need of improvement. I suggest visualizing it in 3D, as there are three varying conditions, which would better illustrate their distinctions across additional axes. Additionally, adjusting the color shading for each condition could help assess whether the condition embeddings follow a consistent trajectory that aligns with the order of prior values.\n- I recommend that the authors soften their claim on universality, as it appears too strong given the current scope of the experiments.", "questions_for_authors": "Despite strong quantitative results, I noticed that Unisolver exhibits more artifacts in its solutions compared to other baselines, even compared with ViT (e.g., visible in Appendix E, Fig. 14). Do the authors have any insights into why this occurs and what might be causing more often artifacts?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Q3ocZJ3LI1", "reviewer": "Reviewer_WyYo", "summary": "This paper proposes Unisolver, a universal neural PDE solver designed as a “foundation model” for solving a broad range of PDEs. Unisolver leverages Transformer architectures pre-trained on a diverse set of PDEs. The model make use of many PDE components, incorporating information such as coefficients, boundary conditions, and notably LLM-based embeddings of PDE expressions. The model achieve better performance than existing models on a diverse dataset.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely thank Reviewer WyYo for providing a detailed review and insightful questions.\n\n> **Q1:** \"The proposed techniques may be limited in its applicability. Though it's a good combination of existing techniques.\"\n\nThank you for the feedback. Regarding the applicability, we would like to clarify that in the $\\underline{\\text{Incomplete Component Scenario}}$ section, the \"incomplete\" setting in $\\underline{\\text{Figure 7}}$ actually means inference with no PDE components provided. Therefore, our method can be trained with partial PDE information, and supports inference with no PDE components, which strengthens its applicability in real-world scenarios with incomplete PDE knowledge.\n\nWe acknowledge that our method builds upon existing some existing components; however, the novelty of our method lies in the systematic integration of PDE information into neural surrogate modeling through conditional embedding, which to our knowledge has not been explored in prior works. \n\n> **Q2:** \"Could you clarify which numerical solver was used in Table 22?\"\n\nThank you for raising this question. In $\\underline{\\text{Table 22}}$, the numerical solver used for comparison is the pseudo-spectral solver adopted by Li et al. in the original FNO paper, which is not extensively optimized for the specific task. As noted in the paper you cited, FNO is reported to be up to 1,000 times faster than **a pseudo-spectral solver**, which is consistent with the result in $\\underline{\\text{Table 22}}$. We promise to cite the mentioned paper and discuss the efficiency improvement more rigorously.\n\n> **Q3:** \"How does the neural solver handle extrapolation to future time steps?\"\n\nThank you for your valuable suggestions. We have provided the extrapolation behavior comparison between our model and baseline models in $\\underline{\\text{Appendix I.3}}$. By extending the prediction horizon to time steps beyond the training range, we observe that all model's performance drops, while our method still outperforms all other baselines.\n\n> **Q4:**  How does the model generalize, with and without fine-tuning, to new type of PDEs, especially PDEs with distinct behaviors, such as inviscid burger's equation, wave equation, or Cahn–Hilliard equation, etc.\n\nWe appreciate the reviewer's valuable suggestions on providing more PDE type generalization experiments.\n\nNote that we have provided a new PDE generalization analysis in $\\underline{\\text{Appendix C.2}}$, where Unisolver trained on equations with polynomial order up to 2 demonstrates strong generalization capabilities to equations of polynomial order 3 via fine-tuning.\n\nAs per your request, we test our model's performance on **2D wave equation** using 200 samples for training and 20 samples for evaluation. The 2D wave equation exibits significantly differnt behaviors from the PDEs of the training dataset. The zero-shot and fine-tuning performance of our model and a FNO model trained from scratch is shown in the table below. Relative L2 is reported.\n\n| Unisolver (Zero-shot) | Unisolver (Fine-tuned) | Unisolver (From scratch) | FNO (From scratch) |\n| -------------------- | --------------------- | ----------------------- | ----------------- |\n| 0.774                | 0.0078                | 0.0406                  | 0.0667            |\n\nThe zero-shot performance of our model is not very impressive, showing that zero-shot generalization to PDEs with significantly different behavior is rather hard. However, when fine-tuned with only 200 samples, our model is able to achieve a relative error smaller than 1% on evaluation samples, which is far better than FNO trained from scratch and also better than Unisolver trained from scratch, demonstrating its strong generalization capability to different types of PDEs and its effectiveness in adapting to new PDEs with limited data.", "claims_and_evidence": "The claims are supported by the evidence.", "methods_and_evaluation_criteria": "The method and the evaluation criteria make sense.", "theoretical_claims": "Experimental work. No theoretical claim.", "experimental_designs_or_analyses": "The experiments and the analyses are sound.", "relation_to_broader_scientific_literature": "This work make progress toward \"foundation model\" for solving PDE.", "essential_references_not_discussed": "Not to the knowledge of the reviewer.", "other_strengths_and_weaknesses": "Strength:\n- Extensive experimental evaluation across diverse PDE benchmarks.\n- Detailed and thorough appendices.\n- Well-written and easy to follow.\n\nWeakness\n- The proposed techniques may be limited in its applicability. Though it's a good combination of existing techniques. \n- See Questions.", "other_comments_or_suggestions": "N/A", "questions_for_authors": "1. Could you clarify which numerical solver was used in Table 22? According to McGreivy & Hakim (2024) [1], after controlling for accuracy and resolution, FNO is only approximately 7× faster than traditional numerical solvers. Understanding the specific solver used would help contextualize the reported speedup.\n\n2. For time dependent PDEs, how does the neural solver handle extrapolation to future time steps? Given that neural solvers are typically trained on fixed time intervals, it would be insightful to discuss the model’s ability to generalize beyond the training range.\n\n3. For 2D mixed PDE, the dataset are mostly variants of Navier-Stokes equations with diffusion. How does the model generalize, with and without fine-tuning, to new type of PDEs, especially PDEs with distinct behaviors, such as inviscid burger's equation, wave equation, or Cahn–Hilliard equation, etc. \n\nQ2 and Q3 can be demonstrated without extensive training. While it is expected that performance will degrade for out-of-distribution samples, providing such examples would be valuable for understanding the model’s limitations and guiding future research directions.\n\n[1] McGreivy, N., Hakim, A., 2024. Weak baselines and reporting biases lead to overoptimism in machine learning for fluid-related partial differential equations. Nat Mach Intell 6, 1256–1269.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: The paper proposes a framework for training PDE surrogates conditioned both on field state values and on prior information, including PDE coefficients, parameters, the algebraic expression of the PDE (via an LLM embedding of this expression), and boundary conditions. The prior information conditions a transformer architecture through adaptive layer-wise scale and shift operations. The method is evaluated on three benchmarks involving 1D and 2D PDEs.

The reviewers appreciate the novelty of the approach and the flexibility of the framework. They requested clarifications and additional experiments. The authors' responses were found to be relevant, and they included some complementary experiments. A key question remains regarding the practical usefulness of the framework, as pointed out by the reviewers, given that the authors make several assumptions about the availability of prior information—assumptions that are rarely met in practice. Nevertheless, the approach may inspire new ideas in the community, and I propose acceptance. The authors are encouraged to take into account the reviewers’ comments, particularly their concerns about the overstatements regarding the method’s supposed “universality,” by moderating their claims and better highlighting its limitations.
Decision: Accept (poster)