PDEformer: Towards a Foundation Model for One-Dimensional Partial Differential Equations

Authors: zhanhong  ye
xiang  huang
leheng  chen
hongsheng  liu
zidong  wang
bin  dong
Venue: ICLR_cc_2024_Workshop_AI4DiffEqtnsInSci
Published: 2024-02-07 07:17:48

Abstract: This paper introduces PDEformer, a neural solver for partial differential equations (PDEs) capable of simultaneously addressing various types of PDEs. We propose to represent the PDE in the form of a computational graph, facilitating the seamless integration of both symbolic and numerical information inherent in a PDE. A graph Transformer and an implicit neural representation (INR) are employed to generate mesh-free predicted solutions. Following pretraining on data exhibiting a certain level of diversity, our model achieves zero-shot accuracies on benchmark datasets that is comparable to those of specifically trained expert models. Additionally, PDEformer demonstrates promising results in the inverse problem of PDE coefficient recovery.

Main paper: 
 introduction and related work: 
 The efficient solution of PDEs plays a crucial role in various scientific and engineering domains, from simulating physical phenomena to optimizing complex systems. In recent years, many learning-based PDE solvers have emerged. Some methods (Raissi et al., 2019; Sirignano & Spiliopoulos, 2018; Ee & Yu, 2017; Zang et al., 2020) represent the approximate PDE solution with a neural network, and are tailored to individual PDEs. Other approaches, such as neural operators like Fourier Neural Operator (FNO) (Li et al., 2021) and DeepONet (Lu et al., 2021) , tackle parametric PDEs by taking the PDE parameters (coefficient fields, initial conditions, etc.) as network inputs. While these methods exhibit a higher level of generality, their capability is still limited to solving a specific type of PDE.Drawing inspirations from successful experiences in natural language processing and computer vision, we aim to develop a foundation PDE model with the highest generality, capable of handling any PDE in the ideal case. Given a new PDE to be solved, we only need to make a direct (zero-shot) inference using this model, or fine-tune it only for a few steps using a relatively small number of solution snapshots. By leveraging the power of generality, foundation models have demonstrated great potential in capturing the similarity inherent in a wide range of tasks, and producing highquality feature representations that are beneficial to various applications (Bommasani et al., 2022; Zhou et al., 2023) . Specific to the realm of scientific computing, we anticipate such a foundation PDE model can achieve high solution accuracy, which is comparable with or even surpass expert models that are trained to solve a specific type of PDE. Besides, it should be easily adapted to tackle with down-stream tasks, including inverse problems, inverse design, optimal control, etc.A PDE to be solved would involve two parts of information: one is the symbolic part specifying the mathematical form of the PDE, and the other is the numeric part that includes the PDE coefficients, initial and boundary values, etc. Typical neural operators like FNO and DeepONet deal with a specific form of PDE, and only need to take the numeric information as the network input. However, in order to construct a foundation model generalizable to different PDEs, the symbolic information has to be integrated seamlessly. Some existing approaches towards this direction (Lorsung et al., 2023; Yang et al., 2024; Liu et al., 2023 ) employ a language model, where the mathematical expression of the PDE serves as the input. These methods may struggle to fully capture the complex interaction between the symbolic and the numeric information. Other strategies avoid explicit input of the PDE forms, opting to encode it implicitly in the numeric input to the model. For example, Yang et al. (2023) and Yang & Osher (2024) use several parameter-solution pairs of the target PDE. Specific to time-dependent PDEs, McCabe et al. (2023) trains a model to predict the next time-step based on a few history solution snapshots, which contain the information of what the underlying dynamics is. Subramanian et al. (2023) specifies the PDE to be solved by the location of the nonzero input channels. Such implicit input methods could be insufficient for encoding classes of PDEs with greater variaty and complexity, and may have to be accompanied with another solver to prepare the additional solution snapshots.In this paper, we introduce PDEformer. Different from previous approaches, we propose to express the symbolic form of the PDE as a computational graph, ensuring that the resulting graph structure, along with its node types and feature vectors, encapsulate all the symbolic and numeric information necessary for solving the PDE. A graph Transformer and an INR are utilized to generate meshfree predicted solutions. After pretraining on PDEs with a certain level of diversity, evaluation on benchmark datasets shows that PDEformer exhibits higher zero-shot prediction accuracy compared with baseline expert models, or can achieve this after fine-tuning with limited data. The potential of application to various down-stream tasks is primarily validated by the PDE coefficient recovery inverse problem. Although our experiments are currently limited to one-dimensional PDEs, we believe it would serve as a noteworthy milestone towards building a foundation PDE model. 
methodology: 
 We consider 1D time-dependent PDEs on (t, x) ∈ [0, 1]×[-1, 1] with periodic boundary conditions, of the general form F(u, c 1 , c 2 , . . . ) = 0, u(0, x) = g(x), where c 1 , c 2 , • • • ∈ R are real-valued coefficients, and g(x) is the initial condition. Here, we assume the operator F has a symbolic expression, which may involve differential and algebraic operations. The goal is to construct a surrogate of the solution mapping (F, g, c 1 , c 2 , . . . ) → u that essentially takes the form of the operator F as its input. We illustrate the overall network architecture in Figure 1 . A primary intepretation of the elements involved will be presented in the following text, with further details left for the appendix. Graph Construction We first represent F, i.e. the symbolic information specifying the PDE form, as a computational graph. In such a computational graph, a node may stand for an unknown field variable (denoted as UF), a scalar coefficient (SC), the initial condition (IC), as well as a differential or algebraic operation, and a directed edge can be used to specify the operands involved in an operation. This would constitute of a directed acyclic graph (DAG) with heterogeneous nodes and homogeneous edges.Then, in order to include the numeric information, we endow each graph node with a feature vector in R d f . For a scalar coefficient c, the value is repeated d f times to form the feature vector of the corresponding SC node. In terms of the initial condition g(x), we assume it is given at an equi-spaced grid with n x points. Inspired by ViT (Dosovitskiy et al., 2021) , we divide these grid values into N = n x /d f patches, yielding N vectors g g g 1 , . . . , g g g N ∈ R d f . These will be used as the feature vectors of the N newly-introduced "patch" nodes, whose types are denoted as p 1 , p 2 , . . . , p N , respectively. We shall connect these patch nodes with the corresponding IC node. The feature vectors of all the remaining nodes are set as zero.Moreover, we introduce L additional nodes with type m 1 , m 2 , . . . , m L , and connect them to the corresponding UF node. These nodes will be used to decode the predicted solution as explained below.Encoding Graph Data The symbolic and numeric information encapsulated in the graph data is integrated into a latent code µ µ µ = [µ1 , . . . , µ L ] T ∈ R L×de . This is accomplished by the graph Transformer, a class of modern Transformer-based graph neural networks with impressive representational capabilities. An adapted version of Graphormer (Ying et al., 2021) is utilized as the specific architecture in the experiments, while more potential alternatives can be found in Min et al. (2022) .For ℓ = 1, . . . , L, we let µ ℓ ∈ R de be the embedding vector assigned to the node with type m ℓ in the output layer of this graph Transformer.Decoding the PDE Solution We employ an INR that takes the coordinate (t, x) as input, and produces the mesh-free prediction û(t, x) according to µ µ µ. Various INR architectures with such an external condition have been adopted in neural operators (Yin et al., 2023) , data compression (Dupont et al., 2022) and generative models (Singh et al., 2023) . In the experiments, we utilize an adapted version of Poly-INR (Singh et al., 2023) with L hidden layers due to its efficiency, and the modulations of the ℓ-th hidden layer is generated based on µ ℓ . 
pretraining stage: 
 We generate a dataset containing 500k samples, distinguished by equation types, coefficients and initial conditions. Specifically, the addressed PDEs follow the form 1 u t + f 0 (u) + f 1 (u) x -νu xx = 0, (t, x) ∈ [0, 1] × [-1, 1] with periodic boundaries and initial condition u(0, x) = g(x), x ∈ [-1, 1], where f i (u) = 3 k=0 c ik u k for i = 0, 1. The corresponding PDEs are solved using randomly generated c ik , ν, and g(x) with the Dedalus package (Burns et al., 2020) to create the pretraining dataset. Pretraining involved 1,000 epochs on 90% of the data, reserving the remaining 10% for testing. PDEformer achieves a relative L 2 error of 0.0104 on the training dataset, and 0.0128 on the test dataset. Figure 2 illustrates the pretrained PDEformer's predictions on the test dataset, emphasizing its high accuracy and proficiency in learning representations across diverse PDEs. The pretrained PDEformer is highly versatile in handling various equations. Its performance on forward problems is evaluated using parametric PDEs from PDEBench (Takamoto et al., 2022) , including Burgers', Advection, and 1D Reaction-Diffusion PDEs. Comparative analysis is conducted with neural operator models tailored for individual PDEs. In this context, neural operators receive initial conditions and predict the entire solution field. Notably, all previous methods as well as PDEformer-FS are trained from-scratch and tested separately on different datasets, while the PDEformer model shows zero-shot inference across all test datasets post pretraining. Furthermore, the PDEformer-FT model involves an additional fine-tuning process on the corresponding PDEBench dataset.In Table 1 , the pretrained PDEformer model showcases zero-shot proficiency by attaining reasonable accuracy in all in-distribution tests2 . Remarkably, for Burgers' equation with ν = 0.1 and 0.01, the zero-shot PDEformer outperforms all the baseline models trained specifically on these datasets. Such superior performance can be partially attributed to the network architecture we have utilized, as PDEformer already exhibits competitive performance when trained from-scratch. We believe that the additional improvement of the pretrained PDEformer stems from its exposure to diverse PDEs during pretraining, from which the model may learn a generalizable law to outperform models trained specifically for individual PDEs. Results of the out-of-distribution tests can be found in Table 3 in the Appendix. The fine-tuned PDEformer consistently excels in all in-distribution and out-of-distribution tests, further highlighting the robustness and versatility of our approach in solving a wide range of PDEs.Table 1 : Test relative L 2 error on PDEBench, in which the PDE coefficients lie within the range of the pretraining data. We format the first and second best outcomes in bold and underline, respectively.Burgers Advection ν = 0.1 ν = 0.01 β = 0.1 U-Net (Ronneberger et al., 2015) 0.1627 0.2253 0.0873 Autoregressive U-Net 0.2164 0.2688 0.0631 DeepONet (Lu et al., 2021) 0.0699 0.1791 0.0186 FNO (Li et al., 2021) 0.0155 0.0445 0.0089 PDEformer-FS (Ours) 0.0135 0.0399 0.0124 PDEformer (Ours) 0.0103 0.0309 0.0119 PDEformer-FT (Ours) 0.0046 0.0146 0.0043Thanks to the high-quality initialization obtained after pretraining, general foundation models are known to exhibit efficient adaptation to new tasks (Bommasani et al., 2022; Zhou et al., 2023) . To compare the efficiency of fine-tuning PDEformer with training traditional expert models, we conduct a comparative analysis on the Advection equation (β = 1, OoD) dataset with a limited number of 100 training samples. As depicted in Figure 3 , PDEformer rapidly reaches convergence in about just 100 iterations. Conversely, the FNO model, trained from scratch, results in a higher test error even after several thousands of iterations. Indeed, it is possible for traditional neural operators to start from a better initialization. However, designed for a specific type of PDE, they cannot be pretrained on 500k data samples containing diverse PDEs as PDEformer does. The valid option left for us is to pretrain them on one different PDE, and then transfer to the target setting, which could be much less efficient. Post pretraining on 9k samples of the Advection equation with β = 0.1 for 1k iterations, the FNO-FT model only exhibits a limited improvement over the corresponding fromscratch version, as can be seen in the figure. This contrast highlights the pretrained PDEformer's swift and accurate adaptability, marking a significant advancement over existing expert models. In addition to forward problems, we can leverage the pretrained PDEformer to the PDE coefficient recovery inverse problem based on one noisy observed solution instance. For each PDE, we feed the current estimation of PDE coefficients into the pretrained PDEformer to get the predicted solutions, and minimize the relative L 2 error against the observations to obtain the recovered coefficients. As this optimization problem exhibits a lot of local minima, the particle swarm optimization algorithm (Wang et al., 2018) is utilized. Figure 4 illustrates the outcomes involving 40 PDEs from the test set explained in Section 3.1. The number of coefficients to be recovered varies for each equation (ranging from 1 to 7). In the absence of noise, the recovered coefficients closely align with the ground-truth values, with scattered points primarily distributed along the y = x line. The existence of a few outliers could be attributed to the intrinsic ill-posed nature of this inverse problem. Even under high noise levels, the majority of the PDE coefficients can be effectively recovered. In the Appendix, we offer comprehensive supplementary materials to enhance the understanding of our study and support the reproducibility of our results. Appendix A delves into the details of our computational graph representation, elucidating its design and the rationale behind its structure. In Appendix B.1, we present an overview of the datasets employed during the pretraining stage of our models, including data sources and preprocessing steps. A DETAILED INTEPRETATION OF THE COMPUTATIONAL GRAPH REPRESENTATIONg(x) u u x u t c cu x u t + cu x 0 ∂t IC ∂x = IC UF dx dt SC × + = 0Figure 5: Illustration of how the form of a PDE can be represented as a computational graph, taking the advection equation u t + cu x = 0, u(0, x) = g(x) as the example. The left panel shows the logical meaning of the nodes and edges, and the right panel illustrates the formalized data structure that is taken as the input of PDEformer. We also note that, different from textual representations, this formalization of DAG is independent of the choice of symbols and the order of addition or multiplication. For example, the equation βv x + v t = 0, v| t=0 = v 0 also corresponds to the DAG shown on the right panel.Figure 5 gives an illustration of the semantic meanings of the computational graph. We also have the following remarks on the computational graph representation of the PDE form:• Only a small number of node types are involved in this computational graph: UF (unknown field variable), SC (scalar coefficient), IC (initial condition), dt, dx (differentiation with respect to t and x, respectively), + (sum), × (product), -(negation), (•) 2 (square), = 0 (being equal to zero). • Note that "-" stands for negation rather than subtraction in the computational graph, making it a unary rather than a binary operation. This is because nodes in this computational graph are not ordered, and the edges are homogeneous. If the binary subtraction operation is involved in the computational graph, we cannot differentiate the subtrahend and the minuend. • For the same reasons, we do not include a power operation node. However, we note that powers with a positive interger exponent can be expressed. For example, u 3 = u × u 2 , andu 11 = ((u 2 ) 2 ) 2 × u 2 × u since 11 = 2 3 + 2 1 + 2 0 .• Although not involved in our experiments, node types representing special functions such as sin, cos, exp, log can be introduced as well, depending on the form of the PDEs involved. This also enables expression of the general power operation, since we have a b = exp(b × log(a)).• Disregarding the auxiliary nodes, nodes with type UF and SC would have a zero in-degree, + and × have an in-degree that is equal to or greater than two, and the in-degrees of all the remaining nodes would be exactly one. • In terms of the auxiliary nodes, we let each patch node p i receive an edge from the corresponding IC node, and each latent modulation node m ℓ emanate an edge towards UF. We adopt such convention of edge direction in order to improve the connectivity of the final DAG, since we shall mask out the attention between disconnected node pairs in the graph Transformer module (see Appendix C.1).In the experiments, the initial condition g(x) is discretized on a spatial grid with n x = 256 points, and we divide the values into N = 16 patches of length d f = 16. This exhibits better solution accuracy compared with the case N = 4 or N = 1. For the external condition of the INR decoder, we observe that using different latent vectors µ 1 , . . . , µ L ∈ R de for each hidden layer leads to improved performance compared with the case of using a shared one µ ∈ R de . The introduction of the auxiliary nodes is therefore deemed meaningful.The dataset constitutes of solutions to PDEs of the formu t + f 0 (u) + f 1 (u) x -νu xx = 0, (t, x) ∈ [0, 1] × [-1, 1], u(0, x) = g(x), x ∈ [-1, 1],where f i (u) = 3 k=0 c ik u k for i = 0, 1. Each coefficient c ik is set to zero with probability 0.5, and drawn randomly from U ([-3, 3]) otherwise. 3 The viscosity ν satisfies log ν ∼ U ([log 10 -3 , log 1]).For the case of a linear flux, i.e. when c 12 = c 13 = 0, we set ν = 0 with probability 0.5. Note that terms with a zero coefficient will be excluded in the computational graph of the PDE. The random initial condition g(x) is generated in the same way as the PDEBench dataset, as will be explained in Appendix B.2.The numerical solutions are obtained using the open-source Python package named Dedalus v3 (Burns et al., 2020) , which is a flexible solver based on spectral methods. To generate the data samples, we use a uniform spatial grid with 256 grid points. The solver proceeds at a time-step of δt solver = 4 × 10 -4 , and the solution snapshots are recorded with time-step δt data = 0.01, yielding a total of 101 temporal values for each data sample. When Dedalus fails to solve the PDE, or when the L ∞ -norm of the solution exceeds 10, the corresponding data sample will be discarded, and not included in the final dataset.As the PDEs have a periodic boundary condition, and are discretized on uniform grid points, we introduce data augmentation by a random translation along the x-axis during the pretraining stage. For each data instance, a total of 8192 spatial-temporal coordinate points are randomly sampled from the 101 × 256 grid. These sampled points are taken as the input of the solution decoder INR, and we compare the model predictions with the ground-truth numerical values to compute the loss.In this subsection, we present an overview of three 1D PDE datasets derived from PDEBench, which we employed in our experimental analysis. Each dataset, tailored to a specific PDE type and coefficient configuration, encompasses 10k instances. For our training purposes, we utilized 9k samples from each dataset, reserving the remaining 1k samples for testing. It is crucial to note that all these PDEBench datasets adhere to periodic boundary conditions.• Burgers' equation 4 :∂ t u + ∂ x (u 2 ) = ν π ∂ xx u for (t, x) ∈ [0, 2] × [-1, 1],where ν ∈ {0.1, 0.01, 0.001}. This equation is a fundamental partial differential equation from fluid mechanics.• Advection equation:∂ t u + β∂ x u = 0 for (t, x) ∈ [0, 2] × [0, 1],where β ∈ {0.1, 1}. The equation models the transport of a quantity u without alteration in its form.• Reaction-Diffusion equation:∂ t u = ν∂ xx u + ρu(1 -u) for (t, x) ∈ [0, 1] × [0, 1],where we only consider ν = 1, ρ = 1. This equation represents a process combining chemical reaction and diffusion dynamics.The initial conditions for each dataset are given by u 0 (x) = ki=k1,...,k N A i sin (k i x + ϕ i ), with frequency numbers k i = 2πni Lx , where n i are integers randomly selected within a pre-determined range and L x is the length of the spatial domain, amplitudes A i are random numbers within [0, 1], and phases ϕ i are chosen randomly from the interval (0, 2π). The absolute value function with a random signature, as well as restriction to a random sub-interval by multiplying a window function, are applied afterwards with 10% probability each. For the Reaction-Diffusion equation, the range of the initial condition is rescaled to the unit interval [0, 1].In order to utilize the pretrained PDEformer model to make predictions, we rescale the spatialtemporal coordinates to the range (t ′ , x ′ ) ∈ [0, 1] × [-1, 1], and the resulting PDEs taken as the input of PDEformer have the following form:• Burgers' equation:∂ t ′ u + ∂ x ′ (2u 2 ) -2ν π ∂ x ′ x ′ u = 0, where t ′ = t/2, x ′ = x. • Advection equation: ∂ t ′ u + ∂ x ′ (4βu) = 0, where t ′ = t/2, x ′ = 2x -1. • Reaction-Diffusion equation: ∂ t ′ u -4ν∂ x ′ x ′ u + (-ρ)u + ρu 2 = 0, where t ′ = t, x ′ = 2x -1.To ensure equitable comparisons among the baseline models, we standardize the resolution of all PDEBench samples to 256 × 256. More specifically, the original PDEBench datasets have a spatial resolution of 1024, which is downsampled to 256. The original number of recorded time-steps is 201 for the Burgers and Advection datasets and 101 for the one-dimensional Reaction-Diffusion dataset, and a linear interpolation is utilized to obtain a temporal resolution of 256. It is important to note that PDEformer makes mesh-free predictions, enabling us to set the temporal resolution to 101 for the pretraining dataset, and 256 for the PDEBench dataset. For FNO and U-Net (nonautoregressive case), the initial value is repeated 256 times to form the two-dimensional data with resolution 256 × 256, and then taken as the network input.The specific graph Transformer architecture employed in our experiments is based on Graphormer (Ying et al., 2021) , with some adaptations to fit our setting. The details are presented as below.Initial Embedding Vector In the graph Transformer, the initial embedding vector of node i is given ash (0) i = x type(i) + Feat-Enc(f i ) + z - deg -(i) + z + deg + (i), where x, z -, z + ∈ R de are learnable embedding vectors specified by the node type type(i), indegree deg -(i) and outdegree deg + (i), respectively. In order to encode the node feature vector f i ∈ R 16 that is not involved in the original Graphormer (Ying et al., 2021) , we utilize a featureencoder Feat-Enc, which is a three-layer multi-layer perceptron (MLP) with ReLU activations and 256 neurons in each hidden layer. Attention Bias Denote ϕ(i, j) to be the shortest path length from node i to node j. If such a path does not exist, or has a length greater than 14, we shall set ϕ(i, j) = 14. For each attention head involved in the graph Transformer, the attention bias corresponding to the node pair (i, j) is given asEQUATIONHere, b + ϕ(i,j) and b - ϕ(j,i) are learnable scalars indexed by ϕ(i, j) and ϕ(j, i) respectively, and shared across all layers. The additional term d ij , which does not appear in the original Graphormer, is introduced to mask out attention between disconnected node pairs. More specifically, when node i and node j are connected in the graph, i.e. there exists a path either from i to j or from j to i, we take d ij = 0, and set d ij = -∞ otherwise. We observe in our experiments that the overall prediction accuracy can be improved with such an additional masking operation. Moreover, since our graph has homogeneous edges, we do not introduce the edge encoding term that appears in the original Graphormer.The structure of the graph Transformer layer is the same as the original Graphormer, and we include it here for convenience to the readers. Each layer takes the formh(l) = Attn(LN(h (l-1) )) + h (l-1) h (l) = FFN(LN( h(l) )) + h(l) ,where FFN represents a position-wise feed-forward network with a single hidden layer and GeLU activation function, and LN stands for layer normalization. In terms of the self-attention block Attn, we shall follow the convention in the original Graphormer paper, and only present the single-head case for simplicity.Let H = [h ′ 1 , • • • , h ′ n ]T ∈ R n×de denote the input of the self-attention module involving n graph nodes, the self-attention is computed as (Lu et al., 2021 ), HyperDeepONet (Lee et al., 2023) for neural operators, as well as SIREN (Sitzmann et al., 2020) , WIRE (Saragadam et al., 2023) , MFN (Fathony et al., 2021) , Poly-INR (Singh et al., 2023) and others (Ramasinghe & Lucey, 2022; Chen & Wang, 2022; Jun & Nichol, 2023) in computer vision. In the experiments, we utilize an adapted version of Poly-INR (Singh et al., 2023) , which exhibits better prediction accuracy and training stability compared with other candidates in our setting. Inspired by COIN++ (Dupont et al., 2022) , we also employ L hypernets, in which the ℓ-th hypernet takes µ ℓ ∈ R de as its input, and generates the scale-and shift-modulations for the ℓ-th hidden layer of our Poly-INR.Q = HW Q , K = HW K , V = HW V , A = QK T √ d e + B, Attn(H) = softmax(A)V, where W Q , W K , W V ∈ RThe intricate architecture of our INR decoder is illustrated in Figure 6 , with the mathematical framework detailed below. We take h 0 = 1 to be the vector with all entries equal to one. For ℓ = 1, 2, . . . , L, we computeg ℓ = W in ℓ t x + b in ℓ , s scale ℓ = MLP scale ℓ (µ ℓ ), s shift ℓ = MLP shift ℓ (µ ℓ ), q ℓ = s scale ℓ ⊙ W h ℓ (h ℓ-1 ⊙ g ℓ ) + b h ℓ + s shift ℓ , h ℓ = σ (q ℓ ) ,and the network output is given as û(t, x) = W Last h L + b Last . Here, the activation function σ(•) is a leaky-ReLU operation with a slope of 0.2 at the negative input range, followed by a clipping operation into the interval [-256, 256] to improve training stability. The hypernets correspond to MLP scale ℓ and MLP shift ℓ . Note that in the original Poly-INR, the hypernets are utilized to generate W in ℓ and b in ℓ . Compared with our practice of generating s scale ℓ and s shift ℓ , this method exhibits better accuracy, but deteriorates the training efficiency, and is therefore not adopted in our experiments. The experimental settings, including model hyperparameters and configurations, are outlined in Table 2 . For a comprehensive understanding of the baseline models employed in our experiments, we provide an overview of all models:• DeepONet: DeepONet employs a unique architecture with two sub-networks: a branch net and a trunk net. The branch net processes a fixed number of sensor observations (256 points from the initial condition in our case), while the trunk net handles coordinate inputs for inference, akin to PDEformer's input mechanism. The outputs from both networks are combined to produce the solution value. Each sub-network consists of a six-layer MLP with 256 hidden neurons and utilizes the ReLU activation function. Notably, DeepONet's mesh-free nature allows for training with scattered data points, enabling us to sample 8192 points per iteration from 256 × 256 grids for each data sample during both DeepONet's training and PDEformer's fine-tuning processes.• FNO: The Fourier Neural Operator (FNO) operates on a mesh-dependent yet resolutionindependent principle. It initially transforms regular grid data into multi-channel hidden features through a pointwise fully connected layer, followed by processing through several Fourier Layers, and finally map to the solution grid. In Fourier Layer, the FNO keeps the lowest 12 Fourier modes. In our experiments, the FNO2D model is utilized, with the initial condition (256 spatial points) extended to form a 256 × 256 input grid, allowing for simultaneous full field output.• U-Net: U-Net adopts a CNN-based encoder-decoder framework, distinguished by its 4 layers of downsampling and upsampling convolutions, bridged by intermediate residual connections. Analogous to FNO2D, both the input and output dimensions are set to 256 × 256. Unlike the mesh-free DeepONet or PDEformer, FNO and U-Net require training data organized in regular grids.• PDEformer: The Transformer-based Graphormer is configured with 9 layers, a 512dimensional embedding space, and 32 attention heads. The Poly-INR part employs L = 8 hidden layers with 256 neurons, and each hidden layer is dynamically modulated using separate scale and shift hypernets, each comprising of a 3-layer MLP with independent parameters.In the pretraining stage of PDEformer, we employ the normalized root-mean-squared-error (nRMSE) loss function due to its effectiveness in improving training efficiency. A learning rate schedule is implemented, progressively reducing the learning rate at predetermined epochs to improve the stability of the training process. Moreover, a warm-up period is utilized at the start of training to mitigate the risk of early training failures by gradually increasing the learning rate from zero to the initial pre-scheduled value. Throughout this study, we quantify performance using the relative L 2 error as our primary metric for testing. The relative L 2 error is mathematically represented by the loss function:EQUATION)where ∥u -û∥ L 2 is the L 2 -distance between the predicted solution û and the ground-truth solution u, and ∥u∥ L 2 is the L 2 -norm of the true solution. This metric offers a normalized measure of the error, thereby enabling consistent comparisons across datasets with varying scales and magnitudes.All the experiments are conducted using MindSpore5 2.0, and the pretraining involving 1, 000 epochs takes about 79 hours on 8 NPUs (84 hours if the internal testing evaluations is taken into account). Figure 7 illustrates the pretraining process of PDEformer.In our investigation of the forward problem, Table 3 compares the prediction results on PDEs with coefficients lying outside the range of the pretraining data. Note that all baseline methods as well as PDEformer-FS do not involve a pretraining process. We also embarked on a detailed investigation to assess the model's learning efficiency with limited data. Specifically, we reduced the training dataset size from 9k to 100 and 1k samples. As depicted in Figure 8 , the fine-tuned PDEformer model notably excels, outperforming all other methods in the test. Moreover, the zeroshot PDEformer establishes a commendably high benchmark, demonstrating robust performance without any fine-tuning. It is particularly noteworthy that under OoD conditions, such as in the Advection (β = 1) and Reaction-Diffusion scenarios, the fine-tuned PDEformer rapidly attains superior results. This highlights the model's few-shot learning6 ability in adapting to unfamiliar scenarios.In terms of the inverse problem, the additive noise value at each grid point is randomly sampled fromU ([-r∥u∥ L ∞ , r∥u∥ L ∞ ]), where u is the true solution without noise, and r is the noise level.Table 4 shows the recovered coefficients for three PDEs out of the 40 random samples, with the corresponding noisy observations and PDEformer predictions illustrated in Figure 9 . Note that the input of PDEformer is the recovered PDE coefficients rather than the ground-truth values. The results implies the promising accuracy of PDEformer in both forward and inverse problems, even in the case of noisy observations. Table 5 showcases a comparison of the number of parameters, per-sample inference time and prediction accuracy for a range of models, including DeepONet, FNO, U-Net, and PDEformer. We also include the results of two traditional numerical solvers. The former is based on the first-order upwind finite-difference (FD) scheme, utilizing the solve_ivp function provided by the SciPy Python package, and the latter being Dedalus, the spectral-method-based solver employed in generating our ground-truth solution data. The evaluation was conducted using the 1D Advection equation (β = 1.0) on a 256 × 256 spatial-temporal grid as a test case, with neural network mod- 4 .els tested on a single NPU and traditional solvers executed on a CPU. The neural network models are adequately trained on the corresponding dataset, and the batch size is standardized to 10 during the test. We average the total time consumption of each model across all samples to show the per-sample inference time. As the FD solver exhibits lower accuracy, the spatial grid resolution is refined to 16 × 256 = 4096 in its solution process. While the comparison reveals a significantly longer inference time for Dedalus, it's essential to acknowledge the inherent differences in the computational platforms and the nature of the models themselves. This juxtaposition, though not strictly fair, aims to illustrate the potential efficiency of machine learning methods in solving PDEs.The U-Net model exhibits unsatisfactory performance in our experiments, and some may speculate that the practice of predicting the entire spatial-temporal solution is not suitable for U-Nets.To address these concerns, we also implement an autoregressive variant of the U-Net model. Following PDEBench Takamoto et al. (2022) , this model takes ℓ consecutive timesteps as the input, and predicts the next unknown timestep. In other words, the model approximates the mapping [u(t -ℓ, •), . . . , u(t -1, •)] → û(t, •). The model architecture is analogous to the non-autoregressive U-Net, except that it now operates on one-dimensional data.During training, we randomly select ℓ consecutive timesteps from a data sample, feed it into the U-Net model, and rollout to predict the next K timesteps:[u(t -ℓ, •), . . . , u(t -2, •), u(t -1, •)] → û(t, •), [u(t -ℓ + 1, •), . . . , u(t -1, •), û(t, •)] → û(t + 1, •),• • • [û(t -ℓ + K -1, •), . . . , û(t + K -2, •)] → û(t + K -1, •).The loss function is a weighted average of the prediction error, in the formL(θ) = K-1 k=0 λ k • nRMSE(u(t + k, •), û(t + k, •)).In the implementation, we select ℓ = 4, K = 16, λ 0 = 1, λ 1 = • • • = λ 15 = 0.1.For the inference phase, we feed the first ℓ timesteps into the model, and rollout until we obtain the entire spatial-temporal solution. Note that all the other models involved in our experiments only takes the initial value (i.e. the first timestep) as the network input. Figure 10 illustrates the predictions of the autoregressive U-Net model. We notice that the model successfully captures the overall dynamics inside the spatial interval, and exhibits a high per-step prediction accuracy. However, small error would appear near the boundary points, and is then amplified during the rollout prediction process, leading to unsatisfactory spatial-temporal prediction results. Indeed, such boundary errors might be mitigated if we modify the network architecture to enforce periodicity, but the resulting network design would then be equation-specific, and is not applicable to more general PDEs with non-periodic boundary conditions. Terms with zero coefficients are removed, and not taken as the input of PDEformer. The PDEs involved have different equation types in this sense.Here, "in-distribution" and "out-of-distribution"(OoD) refer to the range of the PDE coefficients. Indistribution PDEs have coefficients lying within the range of the pretraining data, whereas OoD samples fall outside this range. Note that PDEBench datasets labeled as "in-distribution" are not utilized during pretraining. Therefore, we term the corresponding PDEformer inference as "zero-shot".The value of c10 has no effect on the PDE solutions, and PDEformer can learn such redundancy during training. We exclude this term in the inverse problems.The convection term is ∂x(u 2 ) rather than ∂x(u 2 /2) due to an implementation issue in the PDEBench data generation code. See https://github.com/pdebench/PDEBench/issues/51 for more details.https://www.mindspore.cnIn this work, we use the term few-shot learning to describe the model's proficiency in adapting to and learning from new data that falls outside the distribution of the training set, using only a small number of examples for fine-tuning. 
 
 Appendix:  


Decision: This paper is relevant to the paper since it investigates developing a Transformer-based foundation model for PDEs. In particular, it tests how to generalize across various PDEs, which is very important. This was a borderline paper. I am voting for acceptance but in the camera ready paper the authors should address the valid points from Reviewer 6gPK by correcting the abstract and double checking their U-Net baseline.

Reviews: review: title: title: Torn between recommending acceptance and rejection; interesting architecture but multiple concerns
main_review: main_review: **Explanation of paper** 

This paper introduces a method for solving 1D time-dependent PDEs which they call PDEformer. The unique (or at least unusual) feature of the PDEformer architecture is that it represents the PDE being solved via a computational graph (see graph construction, figure 1). The symbolic form of the PDE is thus an input to PDEformer, and by doing by PDEformer is able to solve multiple different classes of PDEs. The goal of the paper is to develop a “noteworthy milestone” towards a so-called foundation model for PDEs, which can *generalize* over many different types of PDEs. This is certainly a worthy goal and of interest to the workshop.

I have three primary concerns about this paper.

**Result in abstract is inconsistent with results in section 3**

The abstract states the following:
> Following pretraining on data exhibiting a certain level of diversity, our model achieves zero-shot accuracies on benchmark datasets that surpass those of adequately trained expert models.

Table 1 compares five models across six tests. The three expert models are U-Net, DeepONet, and FNO. These are compared to the zero-shot PDEformer and the fine-tuned PDEformer. Table 1 shows that, on four of the six tests, the zero-shot PDEformer model does worse than the expert models. It is incorrect to say that the model does better than expert models when it does worse on the majority of benchmark problems. The abstract of the paper should instead say

> Following pretraining on data exhibiting a certain level of diversity, our model performs worse than adequately trained expert models on most benchmark datasets.

The mistake in this case was that the result reported in the abstract "[omitted] outcomes which are deemed to be unfavourable or statistically insignificant." This is called [outcome reporting bias](https://catalogofbias.org/biases/outcome-reporting-bias/), a form of [reporting bias](https://catalogofbias.org/biases/reporting-biases/). Reporting biases are a form of scientific misconduct.


**Weak U-Net baseline in table 1**

Table 1 compares five models across six tests. A few days ago, I spent a couple hours going through the literature, looking at other papers that reported the relative L2 error (sometimes called the “percent error”). I don’t have time to go back through the papers that I found and cite them all, but I concluded that the FNO results in table 1 all seemed reasonable and consistent with the relative L2 error reported in previous papers. I think the FNO baseline is implemented correctly and is a strong baseline.

My concern about a possible weak baseline is with the U-Net. I couldn’t find U-Nets that had solved these benchmark PDEs, but I know that on other 1D and 2D benchmark PDEs (such as the Kuramoto-Sivashinsky equation, the incompressible Euler equation, etc) U-Nets were found to work better than FNOs at solving PDEs [2,3,4]. Thus, the terrible performance of U-Nets on every benchmark problem is surprising and suggests that the U-Nets were not implemented correctly and are a weak baseline.

I believe the reason that U-Nets achieved such unexpectedly poor performance was that they were not implemented as a 1D autoregressive model (see, for example, [4]) but as a full-solution model that outputs the entire 2D (space & time) solution all at once. For this reason, the U-Net baseline is almost certainly a weak baseline.

The input should have been a 1D 256-length array, and the output the solution or change in solution at the next timestep. Instead, the input was a 256x256 array, and the output was the solution over every timestep.

Given the results of [2,3,4] where autoregressive U-Nets outperform FNOs on 1D and 2D PDEs, I would expect the same would likely occur in table 1. It is quite possible that autoregressive U-Nets would outperform PDEformer on this problem.

**Failure to identify the sources of empirical gains**

Please read *Troubling trends in machine learning scholarship* [5].  [5] discusses four troubling trends, the first two of which are

> 1. Failure to distinguish between explanation and speculation.
> 2. Failure to identify the sources of empirical gains, e.g. emphasizing unnecessary modifications to neural architectures when gains actually stem from hyper-parameter tuning.

This paper makes two statements which purport to explain, but actually are speculating. I believe that both statements are incorrectly identifying the source of empirical gains, or at minimum failing to provide sufficient evidence for the source of empirical gains.


Statement #1:
> This superior performance stems from the pretrained PDEformer’s exposure to diverse PDEs during pretraining, from which the model learns a generalizable law to outperform models trained specifically for individual PDEs.

The paper fails to provide enough evidence to make this conclusion. PDEformer trains on 500k samples of diverse PDEs, then (in 2 out of 6 instances) outperforms expert models trained on 9k samples. Yet it is very much possible that the training on diverse PDEs is not the reason that PDEformer does better in those 2 instances. Perhaps the paper is emphasizing unnecessary modifications to training procedures when gains actually stem from an improved architecture. 

To test whether it is the training procedure or the architecture that causes these empirical gains, the PDEformer architecture trained only on the 9k samples should be compared to the other expert models trained only on the 9k samples. Performing such an ablation study would allow us to identify the source of empirical gains.

Statement #2:
> This contrast highlights PDEformer’s swift and accurate adaptability, marking a significant advancement over neural operators.

Again, the paper does not provide enough evidence to conclude that PDEformers are more adaptable than neural operators. If neural operators were trained on 500k samples with diverse PDEs, perhaps they too would be more capable of being trained on a small dataset with fewer iterations. It is possible that the neural network initializations, not something inherently better about PDEformer than neural operators, is the reason for the improved adaptability. Perhaps the paper is emphasizing unnecessary modifications to architectures when gains actually stem from improved initialization.

To test whether it is the architecture or the initialization that causes these empirical gains, instead of training an FNO from scratch, compare fine-tuning of a PDEformer to retraining an FNO originally trained for a different task.


**Additional comments:**
* Line 2: is “advocate” really the right word to use here? I don’t think so.
* "By leveraging the power of generality …” What is the power of generality? Can you explain what it is and provide references?
* What does “adequately trained” mean? What does “sufficiently trained” mean? In figure 3, the train and test loss for the FNO are both still decreasing on a logarithmic scale. Shouldn’t they be trained longer?
* How do I know that the “patch nodes” are helpful? Are they? Can you perform an ablation study where you test with and without these patch nodes?


**Conclusion:** 

I am torn between recommending acceptance and recommending rejection. On the one hand, this paper introduces an interesting architecture that is likely of interest to the workshop. On the other hand, I have concerns about reporting bias in the abstract (a form of scientific misconduct), a weak baseline (potentially leading to inaccurate conclusions), and a set of conclusions which aren’t supported by the experiments performed in the paper.

My personal recommendation would be to reject this paper. While reporting bias is sadly quite common in ML research, just because scientific misconduct is common does not mean it should be tolerated. When papers attempt to blatantly mislead readers — doing worse than baselines, but claiming to do better — I think that is clear grounds for rejection.

I will ultimately leave it up to the editor to make the final decision.


[1] Zhuang, Jiawei, et al. "Learned discretizations for passive scalar advection in a two-dimensional turbulent flow." Physical Review Fluids 6.6 (2021): 064605.

[2] Stachenfeld, Kimberly, et al. "Learned coarse models for efficient turbulence simulation." arXiv preprint arXiv:2112.15275 (2021).

[3] Gupta, Jayesh K., and Johannes Brandstetter. "Towards multi-spatiotemporal-scale generalized pde modeling." arXiv preprint arXiv:2209.15616 (2022).

[4] Lippe, Phillip, et al. "Pde-refiner: Achieving accurate long rollouts with neural pde solvers." Advances in Neural Information Processing Systems 36 (2024).

[5] Lipton, Zachary C., and Jacob Steinhardt. "Troubling trends in machine learning scholarship." arXiv preprint arXiv:1807.03341 (2018).
title: title: Good paper
main_review: main_review: The paper is in the direction of creating foundation models for PDEs. I have no comment as the paper is nicely written and has great results.