CAPE: Channel-Attention-Based PDE Parameter Embeddings for SciML
Authors: Makoto Takamoto
Francesco Alesiani
Mathias Niepert
Venue: Submitted to ICLR 2023 (2023)
Keywords: machine learning
partial differential equation
attention
generalization

Abstract: Scientific Machine Learning (SciML) designs machine learning methods that predict physical systems governed by partial differential equations (PDE). These ML-based surrogate models substitute inefficient and often non-differentiable numerical simulation algorithms and find multiple applications such as weather forecasting, molecular dynamics, and medical applications.
While a number of ML-based methods for approximating the solutions of PDEs have been proposed in recent years, they typically do not consider the parameters of the PDEs, making it difficult for the ML surrogate models to generalize to PDE parameters not seen during training. 

We propose a new channel-attention-based parameter embedding (CAPE) component for scientific machine learning models and a simple and effective curriculum learning strategy. The CAPE module can be combined with any kind of ML surrogate model, which can adapt to changing PDE parameters without harming the original model's ability to find approximate solutions to PDEs. The curriculum learning strategy provides a seamless transition between teacher-forcing and fully auto-regressive training. 
We compare CAPE in conjunction with the curriculum learning strategy using a PDE benchmark and obtain consistent and significant improvements over the base models. The experiments also show several advantages of CAPE, such as its increased ability to generalize to unseen PDE parameters without substantially increasing inference time and parameter count.
An implementation of the method and experiments are available at \url{https://anonymous.4open.science/r/CAPE-ML4Sci-145B}.

Under review as a conference paper at ICLR 2023
CAPE: CHANNEL-ATTENTION-BBAE PDE PARAME- TER EMBEDDINGS FOR SCIML
Anonymous authors
Paper under double-blind review
ABSTRACT
Scientific Machine Learning (SciML) is concerned with the development of ma- chine learning methods for emulating physical systems governed by partial differ- ential equations (PDE). ML-based surrogate models substitute inefficient and of- ten non-differentiable numerical simulation algorithms and find multiple applica- tions such as weather forecasting, molecular dynamics, and medical applications. While number of ML-based methods for approximating the solutions of PDEs have been proposed recent years, they typically do not consider the paramcters of the PDEs, making it difficull for the ML surrogate models generalize to PDE parameters not seen during training. We propose a new channel-attention- based parameter embedding (CAPE) component for scientific machine learning models and simple and effective curriculum learning strategy. The CAPE module can be combined with any neural PDE solvers allowing them adapt to unseen PDE parameters without harming the original models' ability to find approximate so- lutions. The curriculum learning strategy provides a seamless transition between teacher-forcing and fully auto-regressive training. We compare CAPE in conjunc- tion with the curriculum learning strategy using a PDE benchmark and obtain consistent and significant improvements over the base models. The experiments also show several advantages of CAPE, such as its increased ability to generalize to unseen PDE parameters without substantially increasing inference time and pa- rameter count. An implementation the method and experiments are available at 4open. wwwenceee/eeeeeccoocc //ceee////PPPPP M
INTRODUCTION
Many real-world phenomena, ranging from weather forecasts to molecular dynamics and quantum systems, can be modeled with partial differential equations (PDEs). While for some problems the mathematical description of these equations is available, finding its solutions is complex and usually needs some numerical treatments. Numerical simulation methods have been developed for many years and have achieved high level of accuracy in solving these equations. However, numeri- cal methods are resource intensive and cone-consuming even when run on larger supercomputers to obtain sufficiently accurate results. Especially high-resolution and high-dimensional hydrodynamic- type field equations are computationally demanding. The situation becomes even worse it is nec- essary perform simulations with various PDE parameters since numerical simulation is required for each of the initial conditions and for each PDE parameter's ' configurations.
Recently, there has been rapidly growing interest in machine learning methods for the problem of solving PDEs due to their various applications in science and engineering Guo et al. (2016); Lusch
al. (2018); Sirignano & Splliopoulos (2018): Raissi 2018); Kim al. (2019]; Hsich et al. (011))
Bar-Sinai et al. (2019); Bhatnagar et al. 2019 Plaff et al. (2020): Wang et al. (2020); Khoo et al. (2021). For example, several prior studies reported that ML. models can estimate solutions more efficiently than classical numerical simulators (Li et al. 2021a: Stachenteld et al. (02T). Moreover, using neural networks as surrogate models allows us to compute derivatives with respect to the input variables. Differentiable surrogate models allow one to use backpropagation and automatic differentiation to solve the so-called inverse problems which have numerous real- real-world applications but are difficult to solve using traditional numerical methods (Coros al.] 2013 Allen al. 2022 A considerable number of papers have shown the advantage ML-based surrogate models (Li et al.
2020 2021a Stachenfeld et al. 2021 The majority of these methods, however, are purely data-

---

Under review as a conference paper at ICLR 2023
Autoregressive
CAPE
aands
hase
Time
Step
Step 2
Figure 1: The standard autoregressive approach (left) and the proposed CAPE approach (right) which consists of two interdependent steps.
driven, which does not allow us to change PDE parameters. Although a few models are taking into account PDE parameters, they are tailored specific neural networks and cannot be used with other state-of-the-an methods. This makes it difficult for the SciML community to develop models with high generalization capability only for the initial conditions but both for different types of PDEs and PDE parameters.
To overcome the shortcomings of existing data-driven SciML models, a straightforward approach would include the PDE parameters as additional input. However, this naive method requires modi- fication of the BASE network which is potentially harmful to its accuracy. An alternative approach attaches an external parameter embedding module to the network. However, there are too many possible module structures and methods to provide the embedded parameter information to the base network, and it is in general non-trivial to select the best one. In this paper, propose a new and effective parameter embedding module by utilizing the channel-attention method inspired by nu- merical solver with the implicit method and style transfer in ML (see See. 2.3 and Sec. 2.4) The crucial idea is that a neural network generates intermediate (approximated) field data for future time steps which then interpolated by BASE model such the FNO (Li et al.) 2021a) predict the field data for the next time step. CAPE can be combined with any existing autoregressive neural PDE solvers. Fig. Tillustrates the proposed CAPE framework.
We make the following contributions. First, we propose a CAPE module which can be combined with any existing neural PDE solvers and effectively transfer the PDE parameter information to the base network (BASE). Second, we propose simple but effective curriculum learning strategy that seamlessly bridges the teacher-foreing and auto-regressive methods. Third, we perform exten- sive experiments using various PDEs with large number of different parameters evaluating the effectiveness and efficiency of the proposed method in comparison with ofate-of-the-art methods.
CAPE: A FRAMEWORK FOR NEURAL PDE SOLVERS
2.1 BACKGROUND: PARTIAL DIFFERENTIAL EQUATIONS
Following the notation by Brandstetter al. 2022), we consider Partial Differential Equations (PDEs) over time dimension t E [0,T] and over spatial dimensions Ei,..., FD] X C R which can be written as
0,u = Ftt, x, u. Uw Uu.uu...... (t. E [0, x X
(1)
(0, x)) = (x), x New X, (lu] = O,
(2)
where :00, 2]] x X + R is the solution of the PDE, where c is the field dimension, used to describe various field quantities such velocity, pressure, and density, while u²(x) is the initial condition at time t = O. and B (t, x) = 0 are the boundary conditions at a in 0x, which is the boundary of the domain X. Here, Ozu, u, a.n u, are the partial derivatives of the solution u with respect to the domain, while 0,u is with respect to the time. The functional F describes the possibly non-linear interactions between the PDE's terms.
2.2 PROBLEM DEFINITION
We consider PDEs Sec. 2.I whose solution described as a temporal sequence of field data #t) k=U..... w u N where / is the field data at time step tg. that is, the state of the
2

---

Under review as a conference paper at ICLR 2023
physical system governed by the PDE under consideration at time t4 discretized using At T/N. Each u E X' C R=x=" represents the field tensor data with c, the number of physical variables such as density and velocity, and x; the spatial dimensions of the i-th coordinate. For example, for a 1-d problem we have X' C Rex=; for a 2-d problem X' C Rexa, x#=, and for a 3-d problem X' C TR x Xas, We will often refer to c the channel dimension. We aim to emulate numerical simulators of PDEs which iteratively map M X + X' from u² to u =*1 The emulator (or surro- gate model) a learnable function modeled parameters of neural network as weights to avoid conflict in terminology with the parameters of PDEs. In the following. we denote the emulator's 's prediction at time index it as a² Auto-regressive neural networks predict the next time step's field data based on sequence of field data tensors of length E
neural network NN with weights O. We refer the
ü ++1 NN NN(
ū²;0)
Given the length of the input sequence N E N. and an initial input sequence (u",... ,..., u 1 (ū" -) of length é < N, the ML model auto-regressively generates the remaining sequence (ū i ). The training loss is typically the normalised root-mean-square error (RMSE) between the predicted and the true field data tensors
L( Σ ARMSE (u*, it
1|
where ||u||2 is the L2 Lo-nomm of (vector-valued) variable u. Since we training auto-regressive neural network, the gradients of the above loss can be backpropagated in time in various ways. We discuss this in the following sections. Figure Ieftt) illustrates this auto-regressive approach to solving PDEs. In the vast majority of experimental setups, the assumption is made that > 1, and, therefore, an initial input sequence of length é is available to the model; in practice, this would require a numerical simulation to be run for L - 1-time steps from the initial condition and for each PDE parameter X. The main idea of CAPE to learn generate these sequences based on the current field data and parameter values À and use those as input to an aff-the-shelf neural surrogate model such as an FNO (Li et al. 2021a) or an U-Net (Ronneberger et al. 2013) to perform a complex interpolation.
2.3 COMBINING NEURAL PDE SOLVERS WITH THE CAPE MODULE
The proposed approach is motivated by the need for neural PDE solvers to generalize to PDE param- eters unseen during training. We propose CAPE, novel neural network architecture that takes the prior state of the system n and PDE parameters À as input and predicts the l-intermediate future :==.... CAPE(u t A: ACAPE) The output of CAPE is then used by the BASE network. The overall structure is provided in Fig. rright). The intuition behind this approach is that the intermediate future states capture information about the PDE parameters' impact by attending to the results of the convolutional operations. While we do not change the architecture of the base neural PDE solvers, we propose to use them predict, given the past temporal states and the inter- mediate future states, the state for the next time step. This is contrary to the typical use of neural PDE solvers. The base network is trained jointly with the CAPE module. As shown in Sec. this choice improves the prediction capability of the BASE network.
states { cape b(tt))))))
During training, the output of CAPE is regularized by the additional loss term
N min(t.N kk)
Leapo ( (OCAPE) Σ
nRMSE
(t
u uA+)),
(4)
cape
k=
which forces the CAPE module to predict a temporal sequence of future field data {u fi=1.....
Ji=1....f concatenated with u², the field data at time f.. and given to the base network to make the final prediction. In summary, the CAPE module transforms the input variables {u* }} into ampoora--sqqentiie intermediate field data
Finally, the intermediate sequence { ( (ok=)
( sk= a cape which is then interpolated by the base neural network. Before we
In principle is possible for CAPE to also predict field data of past time steps: ( cape b-t+() J sinthonnte"
3

---

Under review as a conference paper at ICLR 2023
attention weights
Channel attention
Ixl Conv
c channels
kemels
K- U CAPE
u
Conv
CAPE
channels
Figure 2: The CAPE module for one type of convolution (residual connections are omitted).
introduce the inductive bias of the CAPE module, we motivate the general approach from a classical numerical simulation perspective.
CAPE as style interpolation. Let us consider the following simple PDE:
0,u = F(u: A).
(5)
The base neural network can be expressed with the equation
u *+1 Shase (u", {
( cappe ('eape**) Ji=1..,,e OBASK)
(6)
where Shase is the function expressed by the base network and kok-kti CAPE(u", A: OCAPE) are the intermediate future states predicted by the CAPE module generated for
cape Ji=l.....
time index k. The CAPE module can be interpreted tial data u², ( (u=+),
pre-processor network providing sequen-
while Eq. 6 can be interpreted as an interpolation network applying a
particular style of the given data u to {ū k+i cape Ti (see the last figure of Fig. 4).
CAPE as implicit discretization method. On the other hand, those equations be understood from a numerical method perspective, where using implicit discretization (Anderson et al. 2016 Eq. 5 reduces to:
k a AtF(u (t ; Fuu u t1 A))
Now, Eq. with 1 can be interpreted as an approximation of the implicit method of Eq. When CAPE can therefore be generalized MI. MI--based variant of the implicit method.
2.4
CHANNEL-ATTENTUEE -ASEDD PARAMETER EMBEDDING (CAPE MODULE)
CAPE computes 3 different d-dimensional channel attention masks On E Rdd 1,2,3 from the parameters of the PDE using a 2-layer MLP
a. WWi,0 AI,
where is the channel dimension in the feature space and a is the GeLU activation function (Hendrycks & Gimpel 2016 W. [W.... W2..) ) are the weights associated with three operators:
a 11x I-convolution ( a depth-wise convolution (g2), and a spectral convolution (Li al. 202
(gs), that are used to compute the tensor representations 2 B.d...... 2 Da (u"" The
tensors are then multiplied by the attention
using the Hadamard operator (01) over the first dimension (the channel dimension) which is equiv- alent to the broadcast operation of ML programming languages. The three convolutions can be interpreted finite difference method since convolution operations accumulate local information of mesh which, in principle, can simulate local interactions such as advection and diffusion. Intu- itively, channel attention is equivalent to choosing an appropriate physical process for each PDE pa- rameter. A similar mechanism has been proposed for visual tasks, called the aaneeze-and-excciaaiin networks (Hu et al. 2018) which enhances useful channels of the feature vector convolutional networks through attention mechanism. The feature va E dx 1,, 2, 3 are combined to form an intermediate feature A RexxExn."" as
yd -1ix1dd-oxx (a (MMIILOoaal) +
(10)
4

---

Under review as a conference paper at ICLR 2023
PDE
training parameters
testing (unscen) parameters
ID Advection
3 (0.2, 0.4, 0.7, 4.0)
(0.1, 1.0, 7.0)
Burgers
(0.002, 0.007, 0.02, 0.04, 0.2, 0.4, 2.0)
(0.001, 0.01, 0.1, 4.0)
2D NS
= Ç (10 0.000,, 0.004, 0.01, 0.04,
= Ç (0.007,0.07)
Table 1: PDE parameters used in the experiments.
where 11x1, 1 convolutions that adjust the number of dimensions, in particular 11x1,e-dd e x my.... d x ...... while hix1,d--ex d x fh c & x ...... Finally, the sequence of predictions is computed
ko C'cape Ii=1...,f
(u LLyerNorm((")))er (im)
(11)
where ÿ is the i-th element of the data tensor y², selected from the second dimension For the sake of presentation, we omitted the batch dimension. Figure illustrates the architecture of the CAPE module.
2.5 CURRICULUM LEARNING
For each initial condition u., the BASE and CAPE models (referred as NN(u 0)) jointly pre- dict the full temporal sequence {u" )== We propose the following curriculum learning strat- egy. For each training epoch, we split the temporal sequence into two parts. For the first part (ze" ..... u ( ) we use auto--rgressive training using य. the prediction of the model at time index k. as input predict the solution for time index
1, that is, u++1 NN(@"; 9). For the sec- ond part (u f1 ..... u we train the model using teacher-forcing. The teacher-forcing strategy (Williams Zipser 1989 Bengio et al. 2015) computes the prediction for time index h: + using noisy version of the true value at time index k, that is, h NN
0 where is random
noise increasing the stability at inference time (Sanchez-Gonzale: Gonzalez et al. 2018 2020 Plaff et al. 2020 Stachenfeld et al. 2021). The time index ktrans determines the time step where we switch from auto-regressiv training to teacher-forcing and is computed using the following monotonically increasing function of the epoch number
N 2
0.5
ktrans
1 tanh
(12)
A
where N is the total epoch number, and A is a hyper-parameter controlling the steepness of the transition function. A plot of the function is provided in Fig. 6) and a detailed algorithm of the training strategy is provided in Appendix (Algorithm M.
The strategy is based on the following two assumptions: (1) the prediction error decreases as the number of training epochs increases, (2) the accumulated error increases as the number of auto- regressive rollout steps increases. Teacher-forcing training is usually more stable since it avoids the accumulation of prediction errors and should be used exclusively the first phase of training. The auto-regressive strategy simulates the behavior test time and exposes the model to inputs that evolved further from the true data, making it more robust to error accumulation. For the same reasons, however, it tends to be less stable, especially in the early phase of training. The proposed curriculum-eeaming strategy is used to combine the advantages of both approaches.
3 EXPERIMENTS
We used datasets provided by PDEBench (Takamoto et al. 2022 benchmark for SciML from which we selected the following PDEs:
ID Advection Equation. This equation describes the pure-advection of waves
a + (0,u(t, = 0,
(13)
We found that in the case of 2D NS obtain improved results when modifying the right-hand side of Eq. 11 ux" + IN(y.)) and we applied to obtain the 2D NS results in this paper.
5

---

Under review as a conference paper at ICLR 2023
1.00
Advection
1.00
Burgers
1.0
0.75
0.75
MMSE
0.50
0.50
0.8
0.25
0.25
0.6
0.00
FNO
Unet
0.00
FNO
Unet
FNO
Unet
Model
Model
Mode
CAPE
Eave
S-tte
CArE
Figure 3: Plots of the normalized MSE (smaller better) with an error bar for Advection eq. (Left),
Burgers eq. (Middle), and 2D Compressible NS equations (Right).
where / is the PDE parameter describing advection velocity. The exact solution of this equation is: utt, x) = uo((, x - -t) where un the initial condition. Hence, this PDE can be used to check if the ML models understand the property of advection, updating the solution by just advecting the initial profile without changing it.
ID Burgers Equation. Burgers' equation is mathematical model equation simulating the non- linearity and diffusivity in the hydrodynamic equation by a scalar variable
0,)(t,z) + u(t, u),, x) = u/T
(14)
where y is the diffusion coefficient and the parameter of this equation. This PDE can be used to check if ML models can understand the non-linear behavior from the second term in the left-hand side of Eq. and the diffusion process whose strength is controlled by the parameter v.
2D Compressible Navier-Stokes Equations (2D NS). The compressible Navier-Stokes equations (NS eqs., in the following) one of the basic physics equations describing classical fluid dynamics
dep + V (pv) = 0, p(d.r + v. Vv) = Vp+ náv + (5 + V/3)V(V - v),
(15)
V [ pv
de
4 2
v = 0,
(16)
2
where is the mass density, v is the velocity, p is the gas pressure, c= P/ is the internal energy, I'= 5/3, o' the viscous stress tensor, and 1), C are the shear and bulk viscosity, respectively. For 1-dimensional PDEs, we used N = 9000 training instances and 1000 test instances for each PDE parameter with resolution 128. For 2-dimensional NS equations, we used N 900 training instances and 100 test instances for each PDE parameter with spatial resolution 64 64.
Experiment Setup. We evaluated the neural models U-Net (Ronneberger et al. 2015) and FNO al. 2021a) with some datasets provided by PDEBench (Takamoto al.) (022) for various pa- rameters for the ID Advection equation, 1D Burgers equation, and 2D compressible Navier-Stokes equations. We also evaluated the message passing neural PDE Solvers (MPNN) (Brandstetter al. 2022] as a baseline allowing conditional treatment of PDE parameters. We trained each of the neu- ral models (1) Base: without any changes (vanilla model), (2) PINO: with a PINO loss (Li al., dimensini,m,ssson, (4) 2-step: with the field data for the current and previous time-steps as input (u², -² and (5) CAPE: with the CAPE module. Other than case (4), we only provided field data for one time step to the models and, therefore, the models cannot obtain PDE parameters' information from the given data. The PINO loss function regularizes the ML models to satisfy the residuals of the PDEs and might lead to an improved generalization behavior for unseen PDE parameters. The CAPE module predicts intermediate field data for one future time step and this is used as the input the BASE model together with the field data of the current time step The amount of field data provided to the BASE network in cases (4) and (5) is the same: in case (4) the model always obtains field data for time steps & and 1: - 1 as input to predict the field data for time step l: +++1 while in case (5) the BASE model obtains field data for time step l: and intermediate field data for time step k generated by the CAPE model to predict the field data for time step k + 1. Hence, the model in case (4) obtains two true field data for 2 initial steps and, therefore. has the opportunity adapt to different PDE parameter values. Hence, while case (4) is more expensive, we consider it a strong baseline for the problem.
2021b), (3) Conditional: the parameters are added to the input data as new
First introduced in Li et al. 2021a a with twenty steps as input.

---

Under review as a conference paper at ICLR 2023
True
Base
CAPE
CAPE module
Burgers,
0.5
0.0
-1.0
0.0
Figure 4: Visualization of the results: Advection eq. at the final time-step 2.0) (Left), Burgers
eq. at t2 = 20( = 1.0) (2nd-left) at the final time-step, and Vz of 2D NS equations at i. = 5(t = 0.25) (Right). Here "Base" is the vanilla FNO, "CAPE" " is the FNO with CAPE, "CAPE module" is the direct output from CAPE module only; the CAPE module provides a higher frequency proposal the BASE model which then more accurately predicts the field data.
PDE
model
BASE
BASE (PINO)
Conditional
prev. 2-steps
CAPE
Advection
FNO
0.69
0.70 41.6110
0.05+1.2x00
0.13 2
0.04 #3.2x10
Unet
0.72
0.76*1.8x10
0.33
0.33
0.11 08.1x0
MPNN
0.33#2.5x10-2
0.34
Burgers
FNO
0.54#0.40
0.49-1.3x10--
0.23
0.17#2.1110
0.13
Unet
0.53-4.0x10-2
0.85-3.2x10-1
05.100
0.50#6.3x10
0.45
MPNN
0.26
0.28
2D NS
FNO
100
0.86
0.91#0.10
0,80
Unet
0.77#3.6x10-2
0.74
0.74
0.70 13
Table 2: List the normalized RMSE (the smaller, the better) for Advection eq., Burgers eq,, and 2D Compressible NS equations.
Since the solutions of each PDE are not normalized and based on prior results on evaluating PDE solvers Takamoto et al. (2022), we measure the normalized RMSE (nRMSE). We used the normalized RMSE loss function LARMSE with the auxiliary loss function of the CAPE module performed
Leape := LOMSE aLLcape where is the weight coefficient. The optimization with Adam (Kingma & Baj for 100 epochs. The learning rate was divided by 2.0 every 20 epochs. For fair comparison, table with parameter sizes is provided in Tab. 6 the appendix. A more detailed description of the hyper-parameters is provided in Appendix B
made the model size of the different methods similar possible. A
Varying the parameter values. Fig. 3 shows bar plots comparing the BASE models with and without CAPE module, the models with PINO loss, and the models with the 2-steps as input. The CAPR module results the lowest error in all cases. In particular, the CAPE module leads to an impressible error reduction ranging from 20 % (2D NS equation) to 95 % (ID Advection). We partly attribute this to the BASE network's ' ability to capture physical dynamics from the PDE parameter- dependent data provided by the CAPE module. The vanilla FNO is a state-of-the-art model and is superior to the U-net our setting. We hypothesize that the PINO loss heavily affected by and dependent on the time- step size. A more detailed explanation of this observation is given in Appendix C Interestingly, the CAPE module provides either comparable or a little better results than the case with 2-step information. This indicates that the CAPE module succeeded in providing equivalent and even more useful information to the BASE network.
BASE model | Interestingly, the PINO loss provides almost no benefit in
Generalization Ability. using FNO as the BASE network. The parameter the ID Advection PDE controls the advection velocity and the parameters of the remaining equations control the strength of the diffusion process.
5 plots the normalized MSE for each parameter value of the PDEs
'For the 2D NS PDE, the U-net achieves a smaller error than the FNO. We hypothesize that this is partly due to the difference in model size (the number of weights of the U-net is nearly 10 times larger) and partly hecause the U-net typically excels image-to-image mapping problems.
7

---

Under review as a conference paper at ICLR 2023
Advection
Burgers
20 NS (Nm0.1)
2D NS (M=1.0)
OMSE
10"
10P
2006
MMSE
10
Mode
Model
Model
Mode
Figure 5: Plots of the normalized MSE (smaller is better) in terms of each PDE parameter for Advection eq. (Left), Burgers eq. (Middle-Left), 2D NS with M 0.1 (Middle-right) and 2D NS with M = (Right). The results were obtained using the FNO.
PDE
Model
Ablation
nMSE
curriculum strategy
8.0 x 10
pure Autoregressive pure Teacher-Forcing
1.3 x 10
(+0.5)
2D NS
FNO
3.2 x 10+0
(42.4)
curriculum strategy
7.0 x 10-1
pure Autoregressive
1.0 x 10+0 (40.3)
2D NS
Unet
pure Teacher-Forcing
1.0 x 10
(40.3)
Table 3: Ablation study for the 2D CFD equations with FNO and Unet as BASE model.
First, the ID-Adectioon result shows that the CAPE module overfits with the trained parameter (B 0.2,0.4, 0.7, 2.0, 4.0), though it showed a very nice generalization performance to the trained PDE parameters. This could also be indicated by the fact that only the CAPE result for the ID Advection equation showed a much lower error than the approach receiving the 2-step field data which in theory provides similar amount of information the unseen PDE parameters, we expect the difference of these errors to be correlated with the CAPE module overfitting to the trained parameters.
the CAPE module. Hence, evaluate generalization to
On the other hand, the other cases (the parameter describes the diffusion process) showed a good generalization unseen PDE parameters. Note that the plots also indicate that vanilla models show a preference for the parameter regime; in all the cases, the vanilla models exhibit better results on smaller diffusion coefficients but lose accuracy as the diffusion coefficients increase.
Ablation experiments. In this section, we performed an ablation study to separate the impact of the curriculum learning strategy. The ablation study on CAPE structure is provided in Tab. 9 in Appendix. Tab. 3 lists the result for the 2D NS equations using FNO and Unet as the BASE network. The proposed curriculum learning strategy drastically impacts the accuracy of the model in all cases, indicating the effectiveness of seamlessly bridging teacher-forcing and auto-regressive training. The full ablation study results are provided in Tab. 8 in Appendix.
Qualitative analysis of the CAPE module. In Fig. 4 we plot some representative outputs of the vanilla FNO, the CAPE module, and the overall CAPE model, and compare them with the true solutions. Interestingly, imation of the CAPE module into the typical shape (style) of the final solution.
that the BASE network often interpolates higher noise approx-
Inference time. We provide the comparison of the inference time between the hybrid approach (including initial steps as input as done in the FNO paper) and CAPE (only using the initial step obtained from
the input). Here we only consider a scenario where the initial time steps a numerical simulation. [Tb.. 4 lists the inference time in the case of the 2D NS equation. The inference time of the FNO with the CAPE module is much shorter than the hybrid method where the inference time is dominated by the simulation time 5
'The experiments were run using Nvidia GeForce RTX 3090 with CUDA-11.6. The ML. models are implemented using PyTorch 1.1 and the numerical simulations with IAX-0.3.17.
8

---

Under review as a conference paper at ICLR 2023
PDE
Resolution
Total Inference time [see]
simulation time [sec]
Simulation + FNO
512²
582.8
582.6
CAPE + FNO
512²
1.3
Table 4: Inference time comparison of simulation (initial steps=10) FNO and CAPE FNO in the case of 2D CFD (1) = = 0.1). The time-step size is At = 0.05 and the computations were performed until & 1.0 as in this paper's other experiments.
RELATED WORK
Scientific Machine Learning Models Scientific Machine Learning aims at data-driven modeling of physical systems. A notable example is Physics Informed Neural Networks (Cai et al.] 2022 (PINNs) that, having access the PDE of system, learns a neural network over the domain X _d by enforcing small residual, i.e. the error when the solution is evaluated by the PDE or the boundary conditions, over set of sampled points. While PINNS have the capacity to model various physical systems, they need to be trained for each new condition or parameter. Neural Operators (Li et al. 202Ib). as FNO (Li al. 2021a) or Graph NO (Li et al. 2020). models the continuous operators over an infinite space and have shown the ability to generalization at multiple scales, Also, more traditional amage-to-image neural networks such as the U-Net (Ronneberger al., 2015) can be adopted to model NOs. Physics- Informed Neural Operators (PINO) (Li et al., 2021b) improve the representational power PINNs by pre-training a NO but having similar limitations to NOs. Message passing neural PDE solver (Brandstetter et al., 2022) extends the message passing principle to solve PDEs. In signal processing, for the artificial bandwidth extension (ABE) task, the Time- Frequency Network(Dong. of channel attention (CA).
al. 2020) (TFNet) has been proposed, which shares a similar concept
Training Autoregressive Models As was discussed Sec. 2.5 there are two representative train- ing strategies for Scientific MI. that is, teacher-forcing and auto-regressive training. The teacher- forcing strategy was originally developed in natural language processing (Williams & Zipser 1989] Bengio et al. 2015) which predicts n+ 1-th step data (word) using the true n-th step information (word). This method is known to prevent from the error-accumulation in the predicted sequential data during the model training. In the case of Scientific ML, it was found that it is profitable to add a random noise for improving the robustness against the accumulated error at the inference time
al. 2018 2020 Pfaff al. 2020 Stachenfeld et al. 2021 lp. On the other hand, the auto-regressive strategy uses the previous prediction of the model as n-th timestep infor- mation. Because of the error accumulation problem, not so many works were adopted, e.g. (Li al. 2021a). Recently, (Brandstetter et al. 2022) proposed reconciling method for this problem which the so-called "pushforward trick". To increase stability, this method uses an adversarial-style loss which predicts the next timestep data using the previous prediction which is calculated using true data. Note that our method adopted a curriculum training strategy to prevent error accumulation instead of using an additional loss function.
(Sanchez-Gonzalez
5 CONCLUSION
This paper proposed a CAPE module which allows any data-driven SciML models to incorporate PDE parameters. We also proposed a simple but effective curriculum training strategy that allows us to bridge teacher-forcing and auto-regressive learning. We performed an extensive set of experi- ments and showed the effectiveness and efficiency of our method from various aspects: generaliza- tion of seen/unseen PDE parameters during training, parameter efficiency, and inference time.
One of our key findings the behavior of ML models without parameter embedding. We found that the ML models without parameter embeddings either (1) show a poor performance uniformly for all the PDE parameters (1D Advection eq.), (2) overfitting to a preferred parameter regime (ID Burgers eq. and 2D NS eqs.). Hence, the ML model in this case has no generalization power in terms of PDE parameters. On the other hand, the models with CAPE module generalized well for 1D Burgers' eq. and 2D NS eqs. whose parameters govern the physical systems' diffusion behavior. CAPE cannot be generalized for ID Advection eq., and we consider it necessary to formulate CAPE to be more physics- innfore- which consider for future work.

---

Under review as a conference paper at ICLR 2023
REFERENCES
Kelsey R. Allen, Tatiana Lopez-Guevara, Kimberly Stachenfeld, Alvaro Battaglia, Jessica Hamrick, and Tobias Pfaff. Physical Design using Differentiable Learned Sim- ulators. arXiv e-prints, art. arXiv:2202.00728, February 2022.
Peter
Dale Anderson, John C Tannehill, and Richard H Pletcher. Computational fluid mechanics and heat transfer. Taylor & Francis, 2016.
Yohai Bar-Sinai, Stephan Hoyer. Jason Hickey, and Michael P Brenner. Learning data-driven dis- cretizations for partial differential equations. Proceedings of the National Academy of Sciences, 116(31)11334411549) 2019.
Samy Bengio, Oriol Vinyals, Navdeep Jaitly, and Noam Shazeer. Scheduled sampling for sequence prediction with recurrent neural networks. In C. Cortes, N. Lawrence, D. Lee, M. Sugiyama, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 28. Curran As- sociates, Inc., 2015. URL https: httpt::s:s/pproeeeednnn wwwp:/eweee.......ee.eeeeeeeeeee neurieeee httpps/paappee/2/1//11111
pdd
Saakaar Bhatnagar, Yaser Afshar, Shaowu Pan, Karthik Duraisamy, and Shailendra Kaushik. Predic- tion of aerodynamic flow fields using convolutional neural networks. Computational Mechanics, 64(2):525-543 2019.
Johannes Brandstetter, Daniel Worrall, and Max Welling. Message passing neural pde solvers. arXiv preprint arXiv:2202.03376, 2022.
Shengze Cai, Zhiping Mao, Zhicheng Wang. Minglang Yin, and George Em Karniadakis. Physics- informed neural networks (pinns) for fluid mechanics: A review. Acta Mechanica Sinica, pp. 1-12, 2022.
Stelian Coros, Bernhard Thomaszewski, Gioacchino Noris, Shinjiro Sueda, Moira Forberg, Robert W Sumner, Wojciech Matusik, and Bernd Bickel. Computational design mechanical characters, ACM Transactions on Graphics (TOG), 32(4):1-12, 2013.
Yuanjie Dong, Yaxing Li, Xiaoqi Li, Shan Xu, Dan Wang, Zhihui Zhang, and Shengwu Xiong. A time-frequency network with channel attention and non-local modules for artificial bandwidth extension. In ICASSP 2020-22 222002220002 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP), pp. 6954 6958. IEEE, 2020.
Xiaoxiao Guo, Wei Li, and Francesco Iorio. Convolutional neural networks for steady flow ap- proximation. In Proceedings of the 22nd ACM SIGKDD international conference on knowledge discovery and data mining, pp. 1114481999 2016.
Dan Hendrycks and Kevin Gimpel. Gaussian error linear units (gelus). a6Xiv: 08445, 2016.
arXiv preprint
Lun-Ting Hsieh, Shengjia Zhao, Stephan Eismann, Lucia Mirabella, and Stefano Ermon. Learning neural pde solvers with convergence guarantees. arXiv preprint arXiv: 1906.01 1906.01200 2100, 2019.
Jie Hu, Shen, and Gang Sun. Squeeze-and-excitation networks. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 7132-7141, 2018.
Cheng-Zhi Anna Huang, Ashish Vaswani, Jakob Uszkoreit, Noam Shazeer, Ian Simon, Curtis Hawthomne, Andrew M Dai, Matthew D Hoffman, Monica Dinculescu, and Douglas Eck. Music transformer. arXiv preprint arXiv: 1809: 04281, 2018.
Yuehaw Khoo, Jianfeng Lu, and Lexing Ying. Solving parametric pde problems with artificial neural networks. European Journal of Applied Mathematics, 32(3):421-453 2021.
Byungsoo Kim, Vinicius C Azevedo, Nils Thuerey, Theodore Kim, Markus Gross, and Barbara Solenthaler. Deep fluids: A generative network for parameterized fluid simulations. In Computer graphics forum, volume 38, pp. 59-70. Wiley Online Library, 2019,
Diederik P. Kingma and Jimmy Ba. In ICLR (Poster).
10

---

Under review as a conference paper at ICLR 2023
Zongyi Li, Nikola Kovachki, Kamyar Azizzadenesheli, Burigede Liu, Kaushik Bhattacharya, An- drew Stuart, and Anima Anandkumar. Neural operator: Graph kernel network for partial differ- ential equations. arXiv preprint arXiv:2003.0344 2020.
Zongyi Li, Nikola Kovachki, Kamyar Azizzadenesheli, Burigede Liu, Kaushik Bhattacharya, An- drew Stuart, and Anima Anandkumar. Fourier neural operator for parametric partial differential equations. International Conference on Learning Representations (ICLR), 2021a.
Zongyi Li, Hongkai Zheng. Nikola Kovachki, David Jin, Haoxuan Chen, Burigede Liu, Kamyar Azizzadenesheli, and Anima Anandkumar. Physics-informed neural operator for learning partial differential equations. arXiv preprint arXiv:2111.0333 2021b.
Bethany Lusch, Nathan Kutz, and Steven L. Brunton. Deep learning for universal linear embeddings of nonlinear dynamics. Nature communications, 91):1-10, 2018.
Mchdi Mirza and Simon Osindero. Conditional generative adversarial nets. arXiv preprint 1r11v:1411.178% 2014.
Tobias Pfaff, Meire Fortunato, Alvaro Sanchez-Goozalezz, and Peter W Battaglia. Learning mesh- based simulation with graph networks. arXiv preprint arXiv:2010.03409, 2020.
Maziar Raissi. Deep hidden physics models: Deep learning of nonlinear partial differential equa- tions. The Journal of Machine Learning Research, 19(1):932-955, 2018.
Olaf Ronneberger, Philipp Fischer, and Thomas Brox. U-Net: Convolutional Networks for Biomed- ical Image Segmentation, May 2015.
Alvaro Sanchez-Gonzalee Nicolas Heess, Jost Tobias Springenberg. Josh Merel, Martin Riedmiller, Raia Hadsell, and Peter Battaglia. Graph networks as learnable physics engines for inference and control. In International Conference on Machine Learning, pp. 4470-4479. PMLR, 2018.
Alvaro Sanchez-Gonzalez,, Jonathan Godwin, Tobias Pfaff, Rex Ying, Jure Leskovee, and Peter Battaglia. Learning to simulate complex physics with graph networks. In International Confer- ence on Machine Learning, pp.
PMLR, 2020.
Peter Shaw, Jakob Uszkoreit, and Ashish Vaswani. Self-attention with relative position representa- tions. arXiv preprint WrXiv:1803002/55 2018.
Justin Sirignano and Konstantinos Spiliopoulos. Dgm: A deep learning algorithm for solving partial differential equations. Journal of computational physics, 375:1339-1364, 2018
Kimberly Stachenfeld, Drummond B Fielding. Dmitrii Kochkov, Miles Cranmer, Tobias Pfaff, Jonathan Godwin, Can Cui, Shirley Ho, Peter Battaglia, and Alvaro Lanchez-Gonzaler Learned coarse models for efficient turbulence simulation. arXiv preprint 12X 2021.
Makoto Takamoto, Timothy Pradita, Raphael Leiteritz, Dan MacKinlay, Francesco Alesiani, Dirk Pflüger, and Mathias Niepert. PDEBench: A diverse and comprehensive benchmark for scientific machine learning, 2022. URL / uni wnt.:uusssuntttttttttstue pwwppr/eaeeurrll..
xhtml wtwkent 7526 348a 40ed
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, E ukasz Kaiser, and Illia Polosukhin. Attention is all you need. In L. Guyon, U. Von Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Gar- nett (eds.), Advances in Neural Information Processing Systems, volume 30. Curran Asso- httpppp.c//paae/22///1111
ciates, Inc., 2017. URL (aseee24544ddeeeeoooo333eeeee.e
Paper. pdf
Rui Wang, Karthik Kashinath, Mustafa Mustafa, Adrian Albert, and Rose Yu. Towards physics- informed deep learning for turbulent flow prediction. In Proceedings of the 26th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining, pp. 1457-1466, 2020.
Ronald J Williams and David Zipser. A learning algorithm for continually running fully recurrent neural networks. Neural computation, 1(2):270-25 28000022800000 1989.


Reviews:
[{"review_id": "t7PbWlNmZN", "reviewer": "Reviewer_adWD", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "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": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "", "main_review": "", "summary_of_the_paper": "This paper proposed CAPE, which is a channel attention module that embeds PDE parameters. The module can be combined with off-the-shelf PDE solver. CAPE facilitates models to adapt to unseen PDE parameters, and is beneficial for inference efficiency. The authors also propose a curriculum learning strategy to improve model performance. Experiments are provided to demonstrate the effectiveness of the proposed method.", "summary_of_the_review": "See above", "strength_and_weaknesses": "Strength:\n\n* The proposed method is easy to understand, and the writing is clear. \n\n* The proposed method can reduce inference time since the method alleviates the need of running simulations. \n\n\nWeaknesses:\n\n* My main concerns are about experiments.\n\n  * The authors mentioned in the introduction that several approaches can be used to incorporate the PDE parameters into the model, such as adding the PDE parameters as additional inputs, and include the parameters in the embedding module. However, it is not clear why these methods are not good. And the authors should compare with these methods in the experiments.\n\n  * The proposed method contains more parameters than existing methods because of the introduction of the channel attention module. Therefore, the current comparison is not fair. The authors should train all models with the same number of FLOPs (or equivalently same training time).\n\n* I also have concerns about the proposed curriculum learning.\n\n  * The design of the proposed curriculum learning seems rather arbitrary. What is the intuition here? Especially for Eq. 12, I cannot see the benefit of this particular deign. The authors should elaborate more on this part.\n  * From the results in Table 3, it doesn’t seem like the curriculum learning method is effective. For the advection equation, w/o curriculum learning has the best performance. And for the Burgers equation, the curriculum learning method also does not increase model performance. \n\nUpdate:\n\nI still believe the paper is below the acceptance bar. I've decided to keep my rating.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "oH6TDiVcfu", "reviewer": "Reviewer_65qq", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "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": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "", "main_review": "", "summary_of_the_paper": " This paper proposes a CAPE module which allows any data-driven SciML models to incorporate PDE parameters. The key idea is to  transform the input variables $u_k,\\lambda$  into an  intermediate field data, and make the final prediction based on the original input and the intermediate output.", "summary_of_the_review": "Please see Strength And Weaknesses.", "strength_and_weaknesses": "Strength: The problem considered is interesting and the proposed method seems novel.\n\nWeaknesses:\nThe biggest limitation from the reviewer's point is that the proposed CAPE module is not well motivated. More discussion is needed on why this module is introduced. The author mentions that the straightforward  method of including the PDE parameters as additional input will have negative impact on the accuracy of the network. However, no reference  is provided and  no experiments include this method as a baseline. The reviewer believes that the author should provide more references on the similar topic and add  more baseline methods for comparison, for example, $u^{k+1} = Base([u^k,\\lambda])$ and $u^{k+1} = Base([u^{k-1},u^k,\\lambda]).$\n\nAlso, more explanations are needed on the ablation experiments. It is clear that on the Advection and Burgers model, the method without curriculum learning is at least as good as the full method, and what is the advantage of using curriculum learning?\n\nUpdate on Dec. 5: I do not have further questions.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "2FLu9yo6Xz", "reviewer": "Reviewer_WSSM", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "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": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "", "main_review": "", "summary_of_the_paper": "CAPE proposes a channel-attention-based method of generalizing existing neural PDE architectures to support training on equations with varying parameters. This is then combined with a curriculum learning strategy during training time. Evaluations are done on advection, Burgers, and 2d NS equations with Fourier Neural Operator and UNET being used as a baseline.", "summary_of_the_review": "Generalization of PDE parameters is motivated as an important problem and the CAPE method using channel attention + curriculum learning is proposed as a solution which can be combined with existing solvers. However additional experimentation demonstrating the insufficiency of existing methods for generalization + some ablations is absent.", "strength_and_weaknesses": "Strength: A good number of PDEs are chosen to demonstrate the strengths (and weaknesses) of the CAPE method. A need for more parameter general neural PDE solvers is motivated and made clear.\n\nWeakness: While other methods for incorporating PDE parameters are said to be insufficient, it would be useful to be comparisons against such naive methods. Additionally, ablations on the CAPE method would be useful to see which components of the architecture/curriculum learning strategy are the most critical. Also I would be curious to see a CAPE like method being used to augment pre-trained solvers, instead of the joint training required here.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "0ZpJMFZuUjw", "reviewer": "Reviewer_PGvV", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "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": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "", "main_review": "", "summary_of_the_paper": "This paper proposes a channel-attention-based parameter embedding (CAPE) component for PDEs. The module takes as input the state time step t, and the PDE parameter labmda, predicts the state for several intermediate steps, which the main network takes as input (together with state at time t) to predict state at t+1. It details the architecture of the module, and also show in experiments that it results in improved accuracy.", "summary_of_the_review": "In summary, the module the paper proposes is interesting. However, in its current form, there is no enough empirical evidence to justify the benefits of CAPE. Therefore, in its current form, I recommend for reject. On the other hand, if the concerns are addressed, I'm willing to increase my score.", "strength_and_weaknesses": "Strengths:\n\nThe paper addresses an important problem. The method seems to be novel.\n\nWeaknesses:\n\nIt is not well justified or motivated why introduce this module. And why the simplest method of putting the parameter (or a transformation of the parameter) does not work. The improved accuracy of the module can be due to:\n(a) The CAPE module contains the information of the PDE parameter\n(b) It has more parameters than the model without this module, thus have higher expressivity\n(c) The CAPE module needs to predict more future steps, thus having a regularizing effect.\n(d) The inductive bias of the CAPE module is useful and can generalize better\n\nI believe that the simple method of providing the PDE parameter will also benefit from the reason (a). To really demonstrate the usefulness of the CAPE module, the paper needs to rule out (a)(b)(c) (and other potential reasons) and then (d) can then be the likely reason.\n\nTo rule out (a), I think at least 2 more baseline are necessary: \n(I) a base network which takes as input the concatenation of the field data with the PDE parameter (expanded into channel features) concatenated on the channel dimension, i.e.\nu^{k+1} = Base([u^k,\\lambda])\n\n(II) a base network which takes as input the concatenation of the field data with an embedding of PDE parameter:\nu^{k+1} = Base([u^k,MLP(\\lambda)])\nwhere the MLP can be a 2- or 3-layer simple neural network. The layer number and the number of neurons needs to be hyperparameter searched.\n\nIn my previous experience, the above two simple methods works quite well, and can be a strong baseline.\n\nTo rule out (b), add another baseline, \n(III) a larger base network whose number of parameter is similar to the (4) CAPE + base network.\n\nTo rule out (c), the method needs to compare with a baseline where the base network needs to predict the same number of future steps as the CAPE.\n\nThe above baselines need to be hyperparameter searched with similar budget as in CAPE, to ensure a fair comparison.\n\nThere can also be other potential factors, but I think the above are necessary to show the benefit of CAPE.\n\nAlso, the CAPE architecture needs to be better justified\n\n--\nUpdate: \nThrough the rebuttal, the reviewer has addressed my concerns and resolved the misunderstanding. Thus I have increased my score. \n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Decision: Reject