Wavelet Diffusion Neural Operator
Authors: Peiyan Hu
Rui Wang
Xiang Zheng
Tao Zhang
Haodong Feng
Ruiqi Feng
Long Wei
Yue Wang
Zhi-Ming Ma
Tailin Wu
Venue: ICLR 2025 Poster (2025)
Keywords: PDE
physics
simulation
control
diffusion model
wavelet
abrupt changes
multi-resolution
Primary area: applications to physical sciences (physics, chemistry, biology, etc.)
TL;DR: We propose Wavelet Diffusion Neural Operator (WDNO), a novel method for generative PDE simulation and control, to address diffusion models' challenges of modeling system states with abrupt changes and generalizing to higher resolutions.

Abstract: Simulating and controlling physical systems described by partial differential equations (PDEs) are crucial tasks across science and engineering. Recently, diffusion generative models have emerged as a competitive class of methods for these tasks due to their ability to capture long-term dependencies and model high-dimensional states. However, diffusion models typically struggle with handling system states with abrupt changes and generalizing to higher resolutions. In this work, we propose Wavelet Diffusion Neural Operator (WDNO), a novel PDE simulation and control framework that enhances the handling of these complexities. WDNO comprises two key innovations. Firstly, WDNO performs diffusion-based generative modeling in the wavelet domain for the entire trajectory to handle abrupt changes and long-term dependencies effectively. Secondly, to address the issue of poor generalization across different resolutions, which is one of the fundamental tasks in modeling physical systems, we introduce multi-resolution training. We validate WDNO on five physical systems, including 1D advection equation, three challenging physical systems with abrupt changes (1D Burgers' equation, 1D compressible Navier-Stokes equation and 2D incompressible fluid), and a real-world dataset ERA5, which demonstrates superior performance on both simulation and control tasks over state-of-the-art methods, with significant improvements in long-term and detail prediction accuracy. Remarkably, in the challenging context of the 2D high-dimensional and indirect control task aimed at reducing smoke leakage, WDNO reduces the leakage by 78% compared to the second-best baseline. The code can be found at https://github.com/AI4Science-WestlakeU/wdno.git.

Published as a conference paper at ICLR 2025
WAVELET DIFFUSION NEURAL OPERATOR
Peiyan Hu 28- Rui Wang 38+
Haodong Feng' Ruiqi Feng' Long Wei Yue Wang" Zhi-Ming Ma Tailin Wu
Xiang Zheng Tao Zhang'
Department of Artificial Intelligence, School of Engineering, Westlake University, 2 AAcademy of Mathematics and Systems Science Chinese Academy of Sciences, Fudan University, 'South China University T'Thhnologyy "Mierosoft Al4Science (hupei yan, wutaill lin) @west .aw.aaee........ edu. cn, ruiwang @0 fudan. edu.ch
ABSTRACT
Simulating and controlling physical systems described by partial differential equa- tions (PDEs) are crucial tasks across science and engineering. Recently, diffusion generative models have emerged as a competitive class of methods for these tasks due their ability capture long-term dependencies and model high-dimensional states. However, diffusion models typically struggle with handling system states with abrupt changes and generalizing higher resolutions. In this work, we pro- pose Wavelet Diffusion Neural Operator (WDNO), novel PDE simulation and control framework that enhances the handling of these complexities. WDNO com- prises two key innovations. Firstly, WDNO performs diffusion-basee generative modeling in the wavelet domain for the entire trajectory to handle abrupt changes and long-term dependencies effectively. Secondly, to address the issue of poor generalization across different resolutions, which is one of the fundamental tasks modeling physical systems, we introduce multi-resolution training. We validate WDNO on five physical systems, including 1D advection equation, three challeng- ing physical systems with abrupt changes (ID Burgers' equation, 1D compressible Navier-Stokes equation and 2D incompressible fluid), and a real- world dataset ERA5, which demonstrates superior performance on both simulation and control tasks over state-of-the art methods, with significant improvements in long-term and detail prediction accuracy. Remarkably, in the challenging context of the 2D high- dimensional and indirect control task aimed at reducing smoke leakage, WDNO reduces the leakage by 78% compared to the second-best baseline. The code can be found at atthtps:: // thub. wwt.eccomaaaa AIIccc/A ence West l waeee//wdr.ggg
INTRODUCTION
Many systems across science and engineering are described by partial differential equations (PDEs). Simulating and controlling these PDE systems are fundamental tasks with numerous applications, in- cluding weather forecasting (Lynch 2008), controlled nuclear fusion (Carpanese 2021 astronomical
simulation (Courant et al. 1967). 7, and aviation (Paranjape et al. 2013).
With developments of neural networks, deep saariing-based methods have emerged to address this problem (Li et al. 202T Lu et al. [20TT] Tripura & Chakraborty [2022] Hu et al., [2022). Among them,
diffusion generative models (Ho al. (002bb) achieve impressive results in both simulation (Cachay 2023 Price et al. 2023 Rühling Cachay et al.) 2023) and control (Ajay et al. 2022 Chi et al.
ct al.
2023
Wei al. 2024h. On the one hand, simulation and control tasks are typically long-term, where small variations in the early stage can have a long-term term-termmm impact on the full trajectory, making their accurate prediction and control difficult. Diffusion models alleviate the long- challenge by the noise-learning mechanism and recovering the full trajectory from a Gaussian distribution as a whole. Therefore, they can better capture long-term dynamics and generate coherent plans for certain goals (Janner et al. [2022] Chi et al. 20233; Wei et al., (224)). On the other hand, PDE dynamics are typically
digh-simeniio and nonlinear, and the diffusion model demonstrates strong capabilities in modeling
complex high-ding sigh-dimensional data (Ho et al. 2022 Harvey et al. 2022 Vahdat et al. 2022 et al.
2024). See Appendix for more related works.
"Equal contribution. Work done an intern Westlake University: 'Corresponding author.

---

Published as a conference paper at ICLR 2025
Original Domain
Wavelet Domain
Original Domain
Transform
Denolse
Condition
Downsample
Training
Inference
Figure 1: Overview of WDNO. The figure shows the training and inference of the Base-Resolutie Model (BRM) and Super- Resolution Model (SRM).
However, for PDE simulation and control with diffusion models, two key challenges arise. Firstly, the evolution of physical systems is often accompanied by abrupt changes, which reflect key mechanisms of the system (Ben-Dor Ben-Dor 2007: Rassweiler al. [011]. Due to their rapid and intense local variations, and even discontinuitics, these changes are difficult to capture, Secondly, existing resolution, and cannot generalize to
diffusion models typically operate on a fixed ppatial-tempon
finer resolutions (Croitoru et al. 2023 Yue et al. 2024 (Shan et al. 2024), which is a fundamental
requirement of neural PDE solvers (Li et al. 2021 Boussif al.] 2022] Yin et al.] 2022).
In this work, we introduce Wavelet Diffusion Neural Operator (WDNO) to address the above two challenges. Our WDNO method consists of two key innovations: (1) Generation in the wavelet domain. Since the wavelet transform is both space and frequency localized and excels at approximating functions with abrupt changes (Tripura & Chakraborty [2022), generation in the wavelet space endowed by the wavelet transform is ideal for modeling abrupt changes. Besides, due to the linearity and locality of the wavelet transform, it can integrate seamlessly with the multi-resolution training. (2) Multi-resolution training. To enable generalization to finer resolutions, we prepare training datasets across multiple spatial and temporal resolutions utilizing the approximate scale invariance. Since changes of the equation forms approximately the same across resolutions, the model is trained to generalize to finer resolutions conditioned on coarser resolutions, which opens up the capability to generalize even finer resolutions not seen during training.
Concretely, our contributions include the following: (1) We introduce the WDNO method that comprises diffusion in the wavelet space, addressing the challenges of modeling states with abrupt changes in simulation and control. (2) We propose multi-resolution training to address the issue of poor generalization higher-resolution simulations, which is fundamental task in PDE modeling. (3) We evaluate our method on ID advection equation, complex PDEs with abrupt changes including ID Burgers' equation, ID compressible fluid, and 2D incompressible fluid, and a real-world dataset ERA5. Compared with strong baselines in physical simulation and control, our method shows competitive performance. Particularly, the 2D experiments are extremely challenging as they involve indirect control with 3,584 spatial control variables at each time step. for total of 32 time steps. It is noteworthy that WDNO reduces 79% of the leaked smoke compared to the prior state-of thet--.
2 PRELIMINARY
2.1 PROBLEM SETUP
We consider PDE on
[0,
D C R Rdd with the following form
Ou F
Ou 8u
Øt u, Ox dz
f(1,2), (1,2) E [0,7] x D,
(1)
u( x) vo(x),
D,
B[ul(t,x)=0, 0, (t,x) [o, x 0D.
where u : [0. x D + Rm is the solution, with the initial condition ux(x)) at time t = 0 and boundary condition (t, 0 on the boundary. F is function and f(t, x:) is the force term
For such PDE systems, there are two fundamental tasks: simulation and control. The former involves learning a mapping from certain parameter functions a, such initial conditions and boundary
2

---

Published as a conference paper at ICLR 2025
conditions, to the solutions that represent mapping from another definite-dimensiin function space. The latter task involves identifying the external control f for specific objective Ju, f) which is function of and f, aiming finding f that minimizes J.
function space to
2.2 DIFFUSION MODEL
A representative instance of diffusion models is the Denoising Diffusion Probabilistic Model (DDPM)
(Ho et al. 2020b), which contains forward and reverse process to generate samples. In the forward process, noise is progressively added to clean data Xo until it is corrupted into Gaussian noise XK ~ I). This process follows the Gaussian transition kernel g(xx-1 Xxx = NN(*+11 VarXk, (1 0g ax)D), where fax [i-1 denotes the variance schedule. In the reverse process, data sampled from Gaussian noise (O, and denoising model са gradually removes the noise from the data until it returns the original clean data distribution. The model predicts the mean (4g (Xx)) of Xg-1 and the reverse process is defined with the transition pe (xx-1 (xx) N (xxxx11 (0 k), T...).
K
To train the denoising model 6g. the training loss is defined as follows, which optimizes a simplified variant of the variational lower-bound for the data's log-likelihood (Ho et al. 2020b).
- + v1 - cge, 1)]3 where dg. :=
II 04. (2)
Guided Diffusion Generation. Modeling the conditional distribution g(xly) enables controllable
sample generation. Methods for conditioning in diffusion models include classifier-bases guidance
(Du et al.] (023) and classifier-free guidance (Ho & Salimans) 2022 Ajay et al. (022). The former employs an additional classifier model trained on clean data to directly modify the denoising direction of the data during generation. The classifier-free conditioning simplifies the architecture and enables guided generation without an explicit classifier. It trains the model to learn both conditional and unconditional probabilities (x @) Vx log (xx) and cx( Vx log qxlly), where 2 is an identifier that tells the model eg to output p(x) instead of p(aly) (Ho & Salimans 2022). During
sample generation, it combines noise terms following (x, Ø3) (x.(x.)) <g(x, Ø 2)), where
w E [0, is the weight. In this paper, we combine the use of both guidance methods.
3 METHOD
In this section, we detail our proposed WDNO from two perspectives: Section 3.1 describes how we perform the generative process within the wavelet domain, including basic concepts and practical implementation of wavelet transforms, and algorithms for applying WDNO to simulation and control problems. Section 3.2 presents the approximate scale invariance of PDE systems and our proposed multi-resolution training based on this property. The overall algorithm is presented in Figure
3.1 GENERATION IN THE WAVELET DOMAIN
The WDNO performs generative control and simulation in the wavelet domain. Compared to the Fourier transform, the wavelet transform features locality while simultaneously retaining information in both space-time and frequency domains, allowing more accurate modeling for abrupt changes.
Wavelet basis. Intuitively, we use wavelet analysis to represent signals with basis functions localized in both space-time and frequency domains, taking values only within finite intervals. Specifically, this set of basis functions can be divided into two categories: one type is the scaling function & used to represent the general outline (low-frequeng
information) of the original signal, and the other type is the mother wavelet vh, which is used to depict the detailed information (high-frequency information) of the original signal (Alpert et al. 2002 Selesnick et al. 2005).
By scaling the function e and mother wavelets vh, we obtain 6i.m and v/t,m
Φr,m (r) 2 (22xx m), (1) = 22'22 m).
where m adjusts the position of the wavelet along the z-axis and / represents the level of the basis. When increases, the wavelet narrows, and its frequency increases. Then, the entire space can be spanned by Φr.m at a particular level lg and a/r... at levels greater than or equal to lg. which can be presented as follows:
3

---

Published as a conference paper at ICLR 2025
u(x) Σ Ea (m) (x)
de de)
1=fa
Thus we get the coarse wavelet coefficients (m) and the detail wavelet coefficients di(m).
Practical implementation. However, due to the discrete nature of real-world data, the levels of 4/2. will have an upper bound L, meaning there exists a minimum length interval for 1.. As mentioned in the introduction (Sec. D. preserve the locality of the data for integration with the multi-resolution training, we choose lo T. So, the decomposition can be presented as:
u(x))
cz (L(m))L, (x) + Σ (m))).mm...
To verify the reliability of the wavelet decomposition's implementation, in Appendix A we conduct tests of the reconstruction error on training data and discover that the relative l2 errors of such reconstructions are on the order of 10-7, indicating that there nearly no information loss. Further details about the wavelet transform can be found in Appendix A
WDNO for simulation. For simulation, as introduced Section 2.1 the objective to learn a mapping from the equation parameter function a to the solution function M[o.T]- We can view the learning of this mapping as learning a conditional probability p(ujo. However, we consider the conditional probability in the wavelet space, p(W "to. (Wa). where W.. and Wa are the wavelet- transformed values of "[o,T] and Here, adopt classifier-free conditioning to guide the sampling process in diffusion models. Specifically, to ensure that the generated wavelet-transfonee Wao.rr align with the corresponding Wa. we include Wa as a conditioning factor. Specifically, we initialize an optimization variable wr(k) Wjo.r with Gaussian noise (O, I), and iteratively update it via: ((k-1) = ((()) "jo,r] - (W "porr Wa, k) + & ~~ ~~ o
values
(3)
where i denotes the denoising step. is the scaling factor and o% represents the noise schedules. Repeatedly applying this denoising procedure from k = M down k yields the final solution
w Besides, during inference, we follow the Denoising Diffusion Implicit Model (DDIM) (Song
et al. 2020). which can largely speed up the sampling process.
WDNO for control. For the control problem, in a task aimed to minimize J, our goal is to find the optimal f1o,7 based
environment determined by parameter function a, such the initial condition. Consequently, this problem can be naturally modeled as learning p(foo. (ffoorla)) la). Here, we also transform it into the wavelet domain, thus learning pWW fp.m WWa) Similar to the simulation, we employ conditional diffusion model. However, a challenge arises in that we can only model zr (Wa) as represented in the training set, where f is typically not optimal. To address this issue, we view the control problem from an energy optimization perspective, and thus during inference, we enhance the denoising process with guidance J to steer the generation of towards a smaller J. Note that without this term, the model can only generate control sequences that follow the same distribution as the dataset, without optimizing for the control objectives. Specifically, initializing (+( fo.21 from Gaussian noise NO, I), we iteratively update
and train p(Wiio.z)
(K(k-) = ((())
(w/()) Marr Wa, kt) + XV (WW/() )) +5. o), (4) 10,z1
fo.z: Mp.T]
where 0% and 17 are the noise schedule and the scaling factor respectively, and À is the weight of guidance. Here (K) is the approximate noise-free (() estimated from (+) by:
foor
fiom
fpo.7
(k-(k)
foo.r: (w-(k)
-I V1- (W Wa, Wa,k))/Vor.
(5)
We calculate I in Eq.
fja.ry because otherwise noise in (Ak could bring errors to J. Repeatedly applying this denoising procedure yields the final fjo.r]
4 based (7(())
instead of directly using w-(k)
foo.
solution (()))
fJa.Ty" Similar the simulation, we also employ the DDIM accelerate the denoising.
3.2 MULTI-RESOLLUTION FRAMEWORK
Next, to enable the diffusion model to generalize across different resolutions, we will introduce our multi-resolution framework based on the approximate scale invariance, which we will introduce in
4

---

downsample it Published as a conference paper at ICLR 2025
the following. In contrast to the model mentioned in the previous section, which we refer to the Base-Resolution Model (BRM), we will introduce a super-Resolution Model (SRM) in this section. The framework integrates seamlessly with the wavelet transform technique and enables zero-shot super-resolution, which is one of the fundamental requirements of a neural operator.
Approximate scale invariance. We first introduce the approximate scale invariance. For simplicity, let us first assume that the spatial domain of the PDE in Eq. D is = [o, 1]. Given the high-resolution data d4 of size N M, and low-resolution data d_ of size (N/2) x (M/2), although both are originally defined over the same spatiotemporal domain [0, 7 x D, we can rescale the low-resolution data into a new spatiotemporal domain [0, T T/2 x D. where the spatial domain D is scaled to [0. 1/2]. In this case, d_. and d_ can be aligned to the same precision. However, note that the coordinates of d_ are now scaled, meaning that the system no longer follows the original equation.
For any arbitrary spatial domain, we can always achieve such alignment through linear transforma- tion. We denote the linear transformations of time and space as a,t + b and 02 + b2 respectively. Then the stretched function actually satisfies the transformed version of the original equation:
du az "adx' "zzxx" az02²'' ........... + (t,z) E [0, T/2 x D.
du pu
(6)
Note that if consistently consider the same factor of resolution change, this linear transformation remains constant, meaning that the coefficients aj and d2 are fixed. Therefore, the pattern of change between different resolutions is consistent. Additionally, since the wavelet transform is linear and localized, this pattern remains consistent in the wavelet domain.
Correspondingly, in practical operations, we consider that each refinement of the discrete observations of the physical system follows the same pattern, which inspires us to develop the idea of multi- given resolution,
resolution training. Specifically, based the training dataset to create a multi-resolution training dataset and then use this dataset for training to learn this pattern. Thus, during inference, we can naturally follow this pattern to achieve zero-shot super-resolution.
Multi-resolltion training data. In practical implementation, we introduce the Super-Resolutee Model, which is a conditional diffusion model. Assuming the resolution of the original training dataset is N x M, that is, N time steps and M spatial points, we obtain data at the resolution of (N/2) (M/2) through downsampling, which means we do not need finerrresolution data. We thus get the data pairs of sizes N x M and (N/2) x (M/2). This downsampling process can be repeated to obtain data pairs of (N/2) (M/2) and (N/4) x (M/4), (N/4) x (M/4) and (N/8) (M/8), and so forth, to compose the multi-resolution training dataset for training the Super-Resolution Model.
Training. We take the conditional diffusion model (Ho & Salimans 2022) to model the conditional probability pWW, Wi, W.,), where h and I respectively present high- and low-resolutin data of data pairs in the multi-resolution training dataset, aj is the resh-resottttin equation parameter, and Wn. W7 and Wa, the corresponding savelet-transformen values. In detail, to align low- www-resouution with high-resolution data, we duplicate the low-resolution data to match the size of high high-resolution data. During training, each batch randomly selects data pairs from given resolution.
Inference.
During the inference process, when super-resolution is required, we first downsample the high-resolution equation parameters to the same resolution N x M as the training data and perform a wavelet transform. Then, using the Base-Resolution Model, we first generate the wavelet coefficients of the base low resolution. Subsequently, we utilize the Super-Resolluton Model to generate the data based on both the wavelet coefficients of lower-resolution results with size N x M and the wavelet coefficients of an the post-super-resolution resolution 2N 2.M. This process is iterated, allowing us to ultimately generate results with the same resolution as the original a.
4 EXPERIMENTS
In this section, we aim to test (1) the advantages of WDNO in handling complex long-term dynamics with abrupt changes simulation and control problems, (2) the effectiveness of multi-resolution training in performing zhoo super-resolution, and (3) the benefits of integrating wavelet transform.
We report the Mean Squared Error (MSE) measured entire state sequences excluding initial conditions for the simulation tasks, and the control objective T for control problems. Besides, we consider stateoof-the-an baselines from different fields. For control tasks, the following methods are
5

---

Published as a conference paper at ICLR 2025
compared: (1) the classical control algorithm Proportional- Integral- Dntegral-Derriative (PID) (Li al. 2006
(2) Supervised Learning method (SL) (Hwang et al. 2022), reinforcement learning and imitation
learning methods including (3) Soft Actor-Critee (SAC) (Haarnoja et al. 2018). (4) Behavior Cloning
(BC) (Pomerleau 1988). . (5) Behavior Proximal Policy Optimization (BPPO), and (6) DDPM (Zhuang
et al. 2023). For simulation,
consider (1) DDPM (Ho et., 2020a), (2) Wavelet Neural Operator
(WNO) (Triprra & Chakraborty 22222). (3) Multiwavelet Neural Operator (MWT) (Gupta et al.
2021), (4) Fourier Neural Operator (FNO) (Li et al.) 2021), (5) CNN (Hwang al. 2022). (6)
2023). and (7) U-Net (Ronneberger et al. 2015). Details Operator Transformer (OFormer) (Li et al.
can be referenced in Appendix [ and
For reproducibility, the code is available here
4.1 1D BURGERS' EQUATION
Experiment setting. We first consider the ID Burgers' equation, a fundamental equation describing shock waves and turbulence in fluid dynamics, with the Dirichlet boundary condition and external
force f, which follows previous works (Hwang et al. 2022 Mowlavi & Nabi] 2023) and is more difficult due to the long time horizon of 81 steps. The visualizations are presented in Figure 6 More details about the setting are in Appendix [F] The simulation task is to learn the mapping from the
initial condition uo and force term f to the entire trajectory Ho,T], while the control objective J corresponding to the target state u² (z) and the fixed weight a is
J J= [u(T.)) (x)) dr + a
/o.rI [5(,, didd..
(7)
Data preparation. We perform a 2D wavelet transform on the original data using the bior2.4 wavelet basis and the 'periodization' mode, implemented using the pytorch_ cee_wwleeeeeesseeeeeeeeeeeeeeeee package (Cotter 2019). Since the initial condition and the target state are ID, we take the 1D wavelet transform, repeat the coefficients, and then concatenate them with the 2D coefficients.
Results We report results of simulation and control tasks in Table T and Table 2 From Table is evident that WDNO and DDPM achieve results that far surpass other baselines in simulation, demon- strating the capability of diffusion models for long-term predictions. In this particular simulation experiment, the performance of WDNO and DDPM is quite similar, while advantages of WDNO over DDPM are detailed in Section 4.6
(.6 and Section 4.7 For the control problem, WDNO achieves the best results, which clearly illustrates the superior performance of WDNO.
Table 1: Results of simulation. Bold font denotes the best model and the runner-up is underlined.
1D
2D
Methods
Burgers'
Advection
Navier-Stokes
Fluid
ERA5
WNO
0.00572
4.216c-02
6.5428
0,07975
MWT
0.00052
3,468c-04
1.3830
0,01556
21.85750
OFormer
0.00023
1.858e-04
0.6227
0.04303
18.26230
FNO
0.00015
9.712e-04
0.2575
0.00684
14.38638
CNN (ID) U-Net (2D)
0.00198
5.033e-04
12.4966
0.00737
15.51342
DDPM
0.00013
4.209e-05
5.5228
0.01578
15.21103
WDNO (ours)
0.00014
2.898c-05
0.2195
0.00231
12.83291
4.2 ID ADVECTION EQUATION
Experiment setting. Next, we consider the advection equation, which models pure advection behavior without nonlinearity. This dataset, sourced from PDEBench (Takamoto et al., 2022). is set up to predict 80 timesteps of evolution based on the one-time-ste step initial condition. The system exhibits relatively smooth and simple dynamics. We aim to observe the performance various methods system without abrupt changes using this dataset.
Data preparation. Since the data shape is similar to that of the 1D Burgers' equation, the data preparation process is consistent with that of the first experiment.
Results. From results in Table we can observe that most models achieve low prediction errors. However, WDNO still delivers the best results.

---

Published as a conference paper at ICLR 2025
OPFN
t=
21
28
531e
(a) Results of WDNO and DDPM the ID Navier- Stokes equation.
(b) Results of WDNO on the 2D indirect control. The objective is to navigate the ycllow smoke to get around grey obstacles and reach the target bucket located the top center.
Figure 2: Visualizations of ID Navier-Stokes equation and 2D incompressible fluid.
4.3 1D COMPRESSIBLE MAVIER-STORE EQUATION
Experiment setting. We also consider the important 1D compressible Navier-Stokes equation which can describe complex phenomena, such as shock wave formation and propagation in aerodynamics around airplane wings and interstellar gas dynamics. We consider particularly challenging scenario from the ID CFD dataset in PDEBench (Takamoto et al. (002). We select extremely small viscosity coefficients, 7 10-3 and Ç 10 10-8 The initial conditions are shock-tube fields consisting of piecewise constant values generating shocks and rarefactions. Boundary conditions allow waves to exit the domain. Since this pre-existing dataset does not include time-varying control terms, we only perform the simulation task on it. We provide more details in Appendix
Data preparation. The data preparation process is also similar to the above ones.
Results. From Table we can observe that WDNO still gains the best performance among strong baselines. It is particularly noteworthy that the MSE of DDPM exceeds that WDNO by over times. In Figure 2a and Figure 7 we further present the detailed prediction results of DDPM and WDNO. It can be seen that for physical dynamics with abrupt changes, DDPM struggles to model shocks and loses many fine details. This highlights the necessity of introducing the wavelet transform. More results, including MSEs, MAEs, and Loo, and other baselines can be found in Appendix C.I
Table 2: Results of control tasks. Bold font denotes the best model and the runner-up is underlined. (a) 1D Burgers' equation.
(b) 2D incompressible fluid.
Methods
T
Methods
J
PID (surrogate-solver)
0.6645
BC
0.3085
SAC (pseudo-online)
0.1376
BPPO
0.3066
SAC (offline)
0.3210
SAC (pseudo-online)
0.3212
BC (surrogate-solver)
0.2998
SAC (offline)
0.6503
BC (solver)
0.1879
DDPM
0.3124
BPPO (sureegate-solver)
0,3075
WDNO (ours)
0.0679
BPPO (solver)
0.1867
SI.
0.0235
DDPM
0.0272
WDNO (ours)
0.0205
4.4 2D INCOMPRESSIBLEF FLUID
Experiment setting. Next, we experiment on 2D fluid problems following the incompressible Navier-Stokes equation. The experiment setting, a complex scenario close to real-world, follows previous works (Wei et al. 2024), where the control can only be exercised out of the frame as shown in Figure 26 The boundary condition at obstacles is the no-slip condition, meaning that the velocities set to at the boundary. This experiment thus includes fluid-solid coupling, where functions have discontinuities and are hard to model. The different data trajectories share the same initial velocity field; the variations in initial smoke positions, specifically the smoke's initial density, and control
7

---

Published as a conference paper at ICLR 2025
sequences. The simulation task is to predict the smoke's density, velocity field, and the percentage of smoke passing through the target bucket based on the initial smoke density and control sequences.
For the control problem, our goal is to move the smoke from its initial position, located beneath the central obstacle, into the middle bucket at the top. To more specific, is defined as the percentage of smoke not passing through the target bucket. Firstly, this objective presents considerable challenges due to the restriction that forces can only be applied the peripheral regions. This problem requires the model plan ahead in the middle of the entire trajectory avoid entry into the wrong opening. Furthermore, we need to generate 3,584 control parameters over a time span of 32 steps in these peripheral zones indirectly control the velocity field in the central region.
Data preparation. We perform 3D wavelet transform on original data using biorl.3 wavelet basis and 'zero' mode, implemented through Pytorch Wavelet Too Toolbox (ptwt) (Wolter et al.) 2024). Since the initial condition and percentage of smoke are 2D and 1D respectively, we take the 2D and ID wavelet transform and repeat the coefficients to concatenate them.
Results. Table I are the simulation results, showing our method is far superior to DDPM and exceeds all the baselines. is worth noting that the prediction error of WDNO is an order of magnitude lower than that of DDPM. As for the results of the control problem shown in Table 26 our method can make more than 90% of the smoke pass through the target bucket, and its J is 22% of the next best method's J, showing model's superiority under complex dynamics with abrupt changes.
4.5 ERA5
Experiment setting. The ERA5 dataset (Kalnay et al. 2018), provided by ECMWF, is a challenging real-world dataset for weather forecasting. offers hourly atmospheric estimates with a 0.25 latitude-longitude resolution from the Earth's surface to 100 km altitude, spanning from 1979 to the present. We conduct simulation experiments on this dataset to demonstrate the superior performance of WDNO. The selected variable temperature, and the specific task involves predicting the system's cvolution over the next 20 hours based on its state over the past 12 hours.
Data preparation. Due to similar data size, the process of data preparation is similar to Section 4.4
Results. We present the results Table I Here we experiment with different parameters for WNO, but all configurations fail to converge. is clear that WDNO still achieves the best performance, with a relative Le eиror as low as 0.0161, demonstrating its outstanding capability on challenging datasets.
4.6 ZERO-SHOT SUPER RESOLUTION
No upsampling
Single upsampling Double upsampling
In this subsection, we will present the super-resolution simulation results for the ID Burgers' equation and 2D incompress- ible fluid. For the ID experiments, the resolution of the training dataset is the time-space resolution 80 120. We demon- strate the results of single, double, and triple super-resolution steps on both time and space, with the corresponding unseen resolutions of 160 240, 320 480, and 640 x 960 respectively. For the 2D experi- ments, the training dataset has resolution of 64 x 64, and we transfer to the reso- lution 32 128 128. The visualization of ID zero-shot super resolution is presented in Figure
Figure 3: 1D zero-shot super-resolution. The first shows WDNO's simulation results with super resolution, one-level super resolution, and two-level super resolution. The second row the ground truth, and the third row is the difference between the first and second rows. As resolution increases, WDNO's output gets closer to the ground truth, demonstrating its zero-shot super resolution capability.
To evaluate the performance across dif- ferent resolutions. we interpolate the out- comes of each super-resolution step to the highest resolution level. This allows us to assess whether the model can accurately
8

---

Published as a conference paper at ICLR 2025
WUNDSimaar -NOlincar) WNCD)
WoB Insn
Logitstt
LogjMSSE
Lag/MSE)
Number Upsampling
Number of Upsampling
Number Upsampling
(a)
(b)
(c)
Figure 4: Results in Section 4.6 and Section 4.7 after super-resolution steps. All MSEs
calculated at the finest resolution through linear or nearest interpolation. (a), (b) are ID Burgers' and 2D results in Section 4.6 respectively, and (c) the results in Section
0.12
WDNC (ours)
150
3.5
MSR
WNG
3.0
50
0.2195
MWT
0.08
+FFT
Denoiser
0.06
FNO
U--Net
2.0
(c) Comparison with Fourier transform.
Models
DDPM
(ours)
1.0
1x10
0.5
0.00
100
15
20
25
30
20
40
60
80
100
15
14
Time Step
Percentage of Trajectories Used (%)
3.61
0,0001
(a) Long-term dependencies.
(b) Number of training samples.
(d) Measurement noise.
Figure 5: Results of ablation studies.
generate data on finer grid points beyond the resolutions encountered during training. We consider linear interpolation and nearest interpolation, taking the mesh-invariant model FNO and WNO as the baselines. Due to WNO's implementation, it can only perform spatiotemporal super-resolution simul- tancously, making unsuitable for 2D super-resolutin experiments. As shown in Figure 4aj Figure 4b Table 16 and Table 17 in both ID and 2D scenarios, our method surpasses results of interpolation by achieving significantly improved outcomes with each super-resolution step. It can effectively reconstruct the values the newly added grid points the highest resolution, outperforming the mesh-invariant FNO and WNO.
4.7 ABLATION STUDY
Abrupt changes. We first verify whether the wavelet transform enhance DDPM's ability model abrupt changes. To this end, we present the system's states and prediction errors of WDNO and DDPM time in Fig- and Fi Note that, although WDNO and DDPM have similar overall MSEs in Table we can observe that at moments when the state exhibits abrupt changes in space, WDNO achieves lower prediction error compared to DDPM. This demonstrates that the wavelet transform helps to model dynamics with abrupt changes that are otherwise difficult to learn.
MOSO
0.0
0.00
Figure 6: MAE and state trajectories.
Combination of wavelet and multi-resolution training. To assess the efficacy of integrating wavelet transform with multi-resolution training due to the wavelet transform's locality, we provide outcomes from DDPM combined with multi-resolution training by applying the framework directly in the space-time domain, as depicted in Figure 4c Notably, in the ID experiment, as the number of super-resolution steps increases, evaluations at the highest resolution reveal that the disparity between WDNO and the application of the training in the original space-time domain becomes more pronounced, verifying the efficiency of utilizing wavelet transforms for super resolution.

---

Published as a conference paper at ICLR 2025
Comparison with Fourier transform. We also evaluate the diffusion model in the Fourier domain (Diffusion + FFT). The implementation strictly follows WDNO, except for replacing the wavelet transform with Fourier transform. The MSEs the ID compressible Navier-Stokes equation are shown in figure 5cp While the Fourier transform also provides some improvement over DDPM, its performance is significantly inferior that of the wavelet transform, which verifies that wavelet transforms inherently decompose information into low-frequency components and high-freegg f details across different directions, making them more effective for learning complex system dynamics, such as those with abrupt changes. In addition, we take the FNO as the noise prediction model (FNO Denoiser), but the results indicate inferior performance. This may be because FNO tends filter out high-frequency information, which is crucial for a noise prediction mode.
Long-term dependencies. Long-time predictions tend to perform poorly due to error accumulation and prediction instability. Therefore, capturing long-term dependencies allows WDNO to grasp the dynamics over extended periods better, naturally improving WDNO performance. To further verify it, in Figure Sal provide errors of baselines and WDNO at different time steps in the 2D simulation experiment. is obvious that WDNO exhibits the slowest error growth, confirming its ability to capture long-term tergtermm dependencies.
Measurement noise. To evaluate on datasets with increasing mcasurement noise, we add noise to both the training and testing datasets of ID Burgers' equation, sampled as Gaussian noise scaled by the original data's standard deviation multiplied by a scale factor. We test scale factors of 0.01, 0.001, and 0.0001. As shown in the Figure 5d WDNO 's results exhibit minimal variation with changes in scale, demonstrating its robustness to noise.
Number of training samples. We reduce the training dataset size 0.2, 0.4, 0.6, and 0.8 times the current size (9000 samples) and measure WDNO's 's MSE on the ID compressible Navier-Stokes equation. The results in Figure 5b show that even when the dataset size reduced 0.4 times, WDNO 's стrоr remains within a relatively small range. When the dataset size is reduced to 0,2 times, the eπrоr shows a noticeable increase.
Additional results. Due to space constraints, we provide additional details in Appendix which include sensitivity analysis of key hyperparameters, verifying approximate scale invariance, evaluating the sensitivity of baselines and WDNO to noise in control sequences, comparing computational resource usage between baselines and WDNO, and analyzing the impact of the guidance parameter.
LIMITATION AND FUTURE WORK
Firstly, although we do not conduct real-world experiments, WDNO is not limited to the specific environments, which means that it can be applied to real scenarios, such as turbulence, structural materials and plasma, which we will leave as future work. Secondly, due to the wavelet transform and denoising model U-Net, WDNO only applicable static, uniform grid data. We are considering applying WDNO to irregular data by using geometric wavelets (Xu et al.] 2018) combined with diffusion models designed for graph structures (Vignac et al.] (023), or projecting data from irregular grids onto regular uniform grids (Li al. 2020b;Lin et al. (002).. among others. Finally, current approach does not yet incorporate information from equations, such as adding chysics-informe loss based on the PDEs, which can enhance the model's 's accuracy, robustness, and generalizability.
6 CONCLUSION
In this paper, we have introduced Wavelet Diffusion Neural Operator (WDNO), method for simulation and control of PDE systems. By introducing two innovations of generation in the wavelet domain and multi-resolution training, WDNO addresses the challenges of modeling states with abrupt changes and generalizing across resolutions typical in PDE systems. Experiments settings including the ID Burgers' equation, 1D Adevection Equation, 1D CFD, 2D incompressible fluid and ERA5 demonstrate WDNO's superior performance and its ability to generalize to much finer spatial and temporal resolutions than in training. We believe that WDNO will be useful for complex physical simulation and control in wide range of scientific and engineering domains.
challenging
10

---

Published as a conference paper at ICLR 2025
7 ACKNOWLEDGMENT
We thank Tengfei Xu and Tao Zhang for discussions and for providing feedback on our manuscript. We also gratefully acknowledge the support Westlake University Research Center for Industries of the Future; Westlake University Center for High-pon pet-ormarfoeeee Computing. The content is solely the responsibility of the authors and does not necessarily represent the official views of the funding entities.
REFERENCES
Anurag Ajay, Yilun Du, Abhi Gupta, Joshua B Tenenbaum, Tommi S Jaakkola, and Pulkit Agrawal. Is conditional generative modeling all you need for decision making? In The Eleventh International Conference Learning Representations, 2022.
Beylkin Alpert, Gregory Beylkin, David Gines, and Lev Vozovoi. Adaptive solution of partial differential equations in multiwavelet bases. Journal of Computational Physics, 182(1):149-19))) 2002.
Gerben Beintema, Alessandro Corbetta, Luca Biferale, and Federico Toschi. Controlling rayleigh- bénard convection via reinforcement learning. Journal of Turbulence, 21(9-10):585 605, 2020.
Gabi Ben-Dor and Gabi Ben-Dor. Shock wave reflection phenomena, volume 2. Springer, 2007.
Oussama Boussif, Yoshua Bengio, Loubna Benabbou, and Dan Assouline. Magnet: Mesh agnostic neural pde solver. Advances in Neural Information Processing Systems, 35:3 1992-399-3,,,, 2022.
Johannes Brandstetter, Rianne van den Berg, Max Welling, and Jayesh K Gupta. Clifford neural layers for pde modeling. arXiv preprint arXiv:2209. 04934,, 2022a.
Johannes Brandstetter, Daniel E Worrall, and Max Welling. Message passing neural pde solvers. In International Conference on Learning Representations, 2022b.
Salva Rühling Cachay, Bo Zhao, Hailey James, and Rose Yu. Dyffusion: A dynamics-informed diffusion model for spatiotemporal forecasting. arXiv preprint arXiv:2306.0/990 2023.
Qianying Cao, Somdatta Goswami, and George Em Karniadakis. Laplace neural operator for solving differential equations. Nature Machine Intelligence, 6(6):631-640, 2024.
Francesco Carpanese. Development of free-boundary equilibrium and transport solvers for simulation and real-time interpretation of tokamak experiments. Technical report, EPFL., 2021.
Boyuan Chen, Kuang Huang, Sunand Raghupathi, Ishaan Chandratreya, Qiang Du, and Hod Lipson. Automated discovery of fundamental variables hidden in experimental data. Nature Computational Science, 2(7):433 442, 2022a.
Peter Yichen Chen, Jinxu Xiang, Dong Heon Cho, Yue Chang, GA Pershing, Henrique Teles Maia, Maurizio M Chiaramonte, Kevin Carlberg, and Eitan Grinspun. Crom: Continuous reduced-order modeling of pdes using implicit neural representations. arXiv preprint arXiv:2206.02667, 2022b.
Ricky T.Q. Chen, Yulia Rubanova, Jesse Bettencourt, and David Duvenaud. Neural ordinary differential equations. Advances in Neural Information Processing Systems, D518-Deccm:6571- 6583, 2018. ISSN 10495258.
Sihao Cheng, Rudy Morel, Erwan Allys, Brice Ménard, and Stéphane Mallat. Scattering spectra models for physics. PNAS nexus, 3(4):pgae1 2024.
Cheng Chi, Zhenjia Xu, Siyuan Feng, Eric Cousineau, Yilun Du, Benjamin Burchfiel, Russ Tedrake, and Shuran Song. Diffusion policy: Visuomotor policy learning via action diffusion. The International Journal of Robotics Research, pp. 02783649241273668 2023.
Fergal Brian Cotter. Uses of Complex Wavelets in Deep Convolutional Neural Networks. PhD thesis, Trinity College, 2019.
11

---

Published as a conference paper at ICLR 2025
Richard Courant, Kurt Friedrichs, and Hans Lewy. On the partial difference equations of mathematical physics. IBM journal of Research and Development, 11(2):215-22 -244, 1967.
FA Croitoru, V Hondru, RT lonescu, and M Shah. Diffusion models in vision: A survey. IEEE Transactions on Pattern Analysis and Machine Intelligence, 2023.
Jonas Degrave, Federico Felici, Jonas Buchli, Michael Neunert, Brendan Tracey, Francesco Carpanese, Timo Ewalds, Roland Hafner, Abbas Abdolmaleki, Diego de Las Casas, et al. Magnetic control of tokamak plasmas through deep reinforcement learning. Nature, 602(7897)4414411 2022.
Prafulla Dhariwal and Alexander Nichol. Diffusion models beat gans image synthesis. Advances in neural information processing systems, 34:8780-8794, 2021.
Jie Ding, Min Wu, and Min Xiao. Nonlinear decoupling control with pi À neural network for mimo systems. IEEE Transactions on Neural Networks and Learning Systems, pp. 1 1-8, 2022. doi: 10.1009/NNN.S
Yilun Du, Conor Durkan, Robin Strudel, Joshua B Tenenbaum, Sander Dieleman, Rob Fergus, Jascha Sohl-Dickstein, Arnaud Doucet, and Will Sussman Grathwohl. Reduce, reuse, recycle: Composi- tional generation with energy-based diffusion models and memc. In International conference on machine learning, pp. 8489-8510. PMLR, 2023.
MA Elhawary. Deep reinforcement learning for active flow control around a circular cylinder using unsteady-mode plasma actuators. arXiv preprint urXiv:2012. 2020.
Soheil Esmaeilzadeh, Kamyar Arizzadenesheli, Karthik Kashinath, Mustafa Mustafa, Hamdi A Tchelepi, Philip Marcus, Mr Prabhat, Anima Anandkumar, al. Meshfreeflownet: A physics- constrained deep continuous space-time super-resolution framework. In SC20: International Conference for High Performance Computing, Networking, Storage and Analysis, pp. 1-15. IEEE, 2020.
Haodong Feng. Yue Wang, Hui Xiang, Zhiyang Jin, and Dixia Fan. How control hydrodynamic force fluidic pinball via deep reinforcement learning. Physics of Fluids, 35(4), 2023.
Han Gao, Luning Sun, and Jian-Xun Wang. Phygeonet: Physics-informed geometry- anometty-adaptive convolutional neural networks for solving parameterized steady-state pdes irregular domain. Journal of Computational Physics, 428: 10079, 2021.
Tiantong Guo, Hojjat Seyed Mousavi, Tiep Huu Vu, and Vishal Monga. Deep wavelet prediction for image super-resolution. In Proceedings of the IEEE conference on computer vision and pullern recognition workshops, pp. 13,, 2017.
Gaurav Gupta, Xiongye Xiao, and Paul Bogdan. Multiwavelet-basee operator learning for differential equations. Advances in neural information processing systems, 0124048-240663 2021.
Florentin Guth, Simon Coste, Valentin De Bortoli, and Stéphane Mallat. Wavelet score-based generative modeling. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho (eds.), Advances in Neural Information Processing Systems, 2022.
Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with a stochastic actor. In International conference machine learning, pp. 1861-1870. PMLR, 2018
Elie Hachem, Hassan Ghraieb, Jonathan Viquerat, Aurélien Larcher, and P Meliga. Deep reinforce- ment learning for the control of conjugate heat transfer. Journal of Computational Physics, 436: 110317, 2021.
Paul Hagemann, Sophie Mildenberger, Lars Ruthotto, Gabriele Steidl, and Nicole Tianjiao Yang. Multilevel Diffusion: Infinite Dimensional Score- Diffusion Models for Image Generation, November 2023. URL /tRl://tt////av..rr///////////// math, stat].
arXiv:2303.0477. [cs,
12

---

Published as a conference paper at ICLR 2025
Zhongkai Hao, Zhengyi Wang, Hang Su, Chengyang Ying. Yinpeng Dong, Songming Liu, Ze Cheng, Jian Song, and Jun Zhu. Gnot: A general neural operator transformer for operator learning. In International Conference Machine Learning, pp. 12556- 12569. PMLR, 2023.
William Harvey, Saeid Naderiparizi, Vaden Masrani, Christian Weilbach, and Frank Wood. Flexible diffusion modeling of long videos. In S. Koyejo, S. Mohamed, (eds.), Advances in Neural Infor-
A. Agarwal, ottpr////ppppoeeeeenoooooooooeee
Belgrave, Cho, and A. mation Processing Systems, volume 35, pp. 27953-27965. Curran Associates, Inc.,
2022. URL https::/prooeeedings neur wttppp.ccpppppppee cwtpcccppppppeeeeeeeeee httpreerriilssspprr//aaae Ei Eile/ Fieeeeeeee.eeccccc00022222233833 Coner .df pdf
Paper Conference
Jonathan Ho and Tim Salimans. Classifier-free diffusion guidance, 2022.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. In H. Larochelle. M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin (eds.), Advances in Neural Information Processing Systems, volume 33, pp. 6840-6851. Curran Asso-
ciates, Inc., 2020a. URL http::///ppeeeooo.oeeeeeee.o///s
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in neural information processing systems, 1855440-6851, 2020b.
Jonathan Ho, Tim Salimans, Alexey Gritsenko, William Chan, Mohammad Norouzi, and David Fleet. Video diffusion models. arXiv:2204.03455 2022.
Philipp Holl, Nils Thuerey, and Vladlen Koltun. Learning to control pdes with differentiable physics. In International Conference on Learning Representations, 2020.
Peiyan Hu, Qi Meng, Bingguang Chen, Shiqi Gong, Yue Wang, Wei Chen, Rongchan Zhu, Zhi-Ming Ma, and Tie-Ya Liu. Neural operator with regularity structure for modeling dynamics driven by spdes. arXiv e-prints, pp. arXiv-2204, 2022.
Huaibo Huang, Ran He, Zhenan Sun, and Tieniu Tan. Wavelet-smet: A wavelet-based cnn for multi- scale face super resolution. In Proceedings of the IEEE international conference on computer vision, pp. 1689-1697, 2017.
Jiahe Huang, Guandao Yang, Zichen Wang, and Jeong Joon Park. Diffusionpde: Generative pde- solving under partial observation. In ICML 2024 Al for Science Workshop, 2024.
Ka-Hei Hui, Ruihui Li, Jingyu Hu, and Chi-Wing Fu. Neural Wavelet-domain Diffusion for 3D Shape Generation, September 2022. URL http://araav.cc//bb/2200.. arXiv:2209.08725 [es].
Rakhoon Hwang, Jae Yong Lee, Jin Young Shin, and Hyung Ju Hwang. Solving Donstonntr Control Problems Using Operator Learning. AAAI, 36(4):4504 June 2022. ISSN 2374-3468, 2159-5399, doi: 134.6009//aaii..6440223.3. 20373. URL httt:ttts///sssaaaee
org/ wwt//ndex....ooooooooo...o.mm... php/
Valerii lakovlev, Cagatay Yildiz, Markus Heinonen, and Harri Lähdesmäki. Latent neural odes with sparse bayesian multiple shooting. In The Eleventh International Conference on Learning Representations, 2022.
Valerii lakovlev, Markus Heinonen, and Harri Lähdesmäki. Learning space-time time continuous neural pdes from partially observed states. arXiv preprint arXhv:2307.04110, 2023a.
Valerii lakovlev, Cagatay Yildiz, Markus Heinonen, and Harri Lähdesmäki. Latent neural ODEs with sparse bayesian multiple shooting. In The Eleventh International Conference on Learning Representations, 2023b.
1b
Devendra K Jangid, Neal Brodnik, Michael G Goebel, Amil Khan, SaiSidharth Majeti, McLean Echlin, Samantha H Daly, Tresa M Pollock, and BS Manjunath. Adaptable physics-based super- resolution for electron backscatter diffraction maps. npj Computational Materials, 8(1):255, 2022.
13

---

Published as a conference paper at ICLR 2025
Michael Janner, Yilun Du, Joshua Tenenbaum, and Sergey Levine. Planning with diffusion for flexible behavior synthesis. In Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvari, Gang Niu, and Sivan Sabato (eds.), Proceedings of the 39th International Conference on Machine Learning, volume 162, pp. 9902-9915. PMLR, 17-23 Jul 2022.
Eugenia Kalnay, Masao Kanamitsu, Robert Kistler, William Collins, Dennis Deaven, Lev Gandin, Mark Tredell, Suranjana Saha, Glenn White, John Woollen, et al. The neep/ncar 40-year reanalysis project. In Renewable energy, pp. Vol1_146 101_146-Vool_me 194. Routledge, 2018.
Diederik P Kingma and Jimmy Ba. Adam: method for stochastic optimization. arXiv preprint arXiv: wet..:112666.0 2014.
Georg Kohl, Liwei Chen, and Nils Thuerey. Benchmarking autoregressive conditional diffusion models for turbulent flow simulation. In ICML 2024 Al Science Workshop, 2024.
Nikola Kovachki, Zongyi Li. Burigede Liu, Kamyar Azizzadenesheli, Kaushik Bhattacharya, Andrew Stuart, and Anima Anandkumar. Neural operator: Learning maps between function spaces with applications to pdes. Journal of Machine Learning Research, 24(89): 1-77, 2023.
Kai Lagemann, Christian Lagemann, and Sach Mukherjee. Invariance-basee learning of latent dynamics. In The Tweifth International Conference on Learning Representations, 2023.
Qiufu Li, Linlin Shen, Sheng Guo, and Zhihui Lai. Wavelet integrated cnns for noise-robust image classification. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 7245-7254, 2020a.
Tianyi Li, Luca Biferale, Fabio Bonaccorso, Martino Andrea Scarpolini, and Michele Buzzicotti. Synthetic lagrangian turbulence by generative diffusion models. Nature Machine Intelligence, pp. 1-11, 2024.
Yun Li, Kiam Heong Ang, and G.C.Y. Chong. Pid control system analysis and design. IEEE Control Systems Magazine, 26(1):32-41, 2006. doi: 10. 2009
Zijie Li, Kazem Meidani, and Amir Barati Farimani. Transformer for partial differential equations' operator learning. Transactions on Machine Learning Research, 2023, ISSN 2835-8856. URL
https://cpentevview wwt het/forumiiddEPPe 3uERT
Zongyi Li, Nikola Kovachki, Kamyar Azizzadenesheli, Burigede Liu, Kaushik Bhattacharya, Andrew Stwart, and Anima Anandkumar. Fourier neural operator for parametric partial differential equations. arXiv preprint arXiv:2010.08895 2020b.
Zongyi Li. Nikola Borislavov Kovachki, Kamyar Azizzadenesheli, Burigede liu, Kaushik Bhat- tacharya, Andrew Stuart, and Anima Anandkumar. Fourier neural operator for parametric partial differential equations. In International Conference Learning Representations, 2021. URL atttps://openreviee net wtw.leewwweee/oouomeeeeme
Timothy P Lillicrap, Jonathan J Hunt, Alexander Pritzel, Nicolas Heess, Tom Erez, Yuval Tassa, David Silver, and Daan Wierstra. Continuous control with deep reinforcement learning. arXiv preprint 150.0099, 2015.
Jae Hyun Lim, Nikola B. Kovachki, Ricardo Baptista, Christopher Beckham, Kamyar Azizzade- nesheli, Jean Kossaifi, Vikram Voleti, Jiaming Song. Karsten Kreis, Jan Kautz, Christopher Pal, Arash Vahdat, and Anima Anandkumar. basee-based Diffusion Models in Function Space, Novem- ber 2023. URL
07400 00Xiv:2302.07400 [cs, math, stat].
Haitao Lin, Lirong Wu, Yongjie Xu, Yufei Huang, Siyuan Li, Guojiang Zhao, and Stan Li. Non- equispaced fourier neural solvers for pdes. In ICLR 2023 Workshop Physics for Machine Learning, 2023.
Phillip Lippe, Bas Veeling, Paris Perdikaris, Richard Tumnrr, and Johannes Brandstetter. Pde- refiner: Achieving accurate long rollouts with neural pde solvers. Advances in Neural Information Processing Systems, 36, 2024.
14

---

Published as a conference paper at ICLR 2025
Dong C. Liu and Jorge Nocedal. On the limited memory BFGS method for large scale optimization. Mathematical Programming, 45(1):503 -28, August 1989. ISSN 1436-4646. doi: 10.1007/ BF01589116. URLhttps:: //doi http:///.....r......00////111///
Pengju Liu, Hongzhi Zhang, Wei Lian, and Wangmeng Zuo. Multi-level wavelet convolutional neural networks. IEEE Access, 7:74973-74985, 2019.
Lu Lu, Pengzhan Jin, and George Em Karniadakis. Deeponet: Learning nonlinear operators for identifying differential equations based on the universal approximation theorem of operators. arXiv preprint (EI.:39900003999 2019.
Lu Lu, Pengzhan Jin, Guofei Pang, Zhongqiang Zhang, and George Em Karniadakis. Learning nonlinear operators via deeponet based on the universal approximation theorem of operators. Nature machine intelligence, 3(3):118 (23):21--229.210 2021.
Bethany Lusch, Nathan Kutz, and Steven L Brunton. Deep learning for universal linear embeddings of nonlinear dynamics. Nature communications, (1))4995 2018.
Peter Lynch. The origins of computer weather prediction and climate modeling. Journal of computa- tional physics, 227(7):343-33444 13444, 2008.
Saviz Mowlavi and Saleh Nabi. Optimal control of pdes using physics-informe informedd neural networks. Journal of Computational Physics, 473:111731, 2023.
Guido Novati, Siddhartha Verma, Dmitry Alexeev, Diego Rossinelli, Wim M Van Rees, and Petros Koumoutsakos. Synchronisation through learning for two self- seff-propelled swimmers. Bioinspiration & biomimetics, 12(3):036001, 2017.
Yi Pan, Amir-massoud Farahmand, Martha White, Samira Nabi, Pulkit Grover, and Daniel Nikovski. Reinforcement learning with function- manluen-vvaludd action spaces for partial differential equation control. In Jennifer Dy and Andreas Krause (eds.), Proceedings of the 35th International Conference on Machine Learning, volume 80, pp. 3986339869999 Stockholm, Sweden, 10-15 July 2018. PMLR.
Aditya Paranjape, Jinyu Guan, Soon-Jo Chung, and Miroslav Krstic. Pde boundary control for flexible articulated wings on robotic aircraft. IEEE Transactions on Robotics, 29(3):625-640, 2013.
Jakiw Pidstrigach, Youssef Marzouk, Sebastian Reich, and Sven Wang. Infinite-Dimensionnal Diffusion Models, October 2023. URL httte///arrvv..orgg/////22222..1 10130 [cs, math, stat].
Dean A Pomerleau. Alvinn:
An autonomous land vehicle in neural network. Advances in neural information processing systems, 1, 1988.
Ilan Price, Alvaro Sanchez-Gonzalee Ferran Alet, Timo Ewalds, Andrew El-Kadi, Jacklynn Stott,
Shakir Mohamed, Peter Battaglia, Remi Lam, and Matthew Willson. Gencast: Diffusion-bais based ensemble forecasting for medium-range weather. arXiv preprint
/5796, 2023.
Jean Rabault, Miroslav Kuchta, Atle Jensen, Ulysse Réglade, and Nicolas Cerardi. Artificial neural networks trained through deep reinforcement learning discover control strategies for active flow control. Journal of fluid mechanics, 865:281-302, 2019.
Maziar Raissi, Paris Perdikaris, and George Karniadakis. Physics-informed neural networks: A deep learning framework for solving forward and inverse problems involving nonlinear partial differential equations. Journal of Computational physics, 378:686-707, 2019.
Bogdan Raonic, Roberto Molinaro, Tim De Ryck, Tobias Rohner, Francesca Bartolucci, Rima
Alaifari, Siddhartha Mishra, and Emmanuel de Bézenac. Convolutional neural operators for robust and accurate learning of pdes. Advances in Neural Information Processing Systems, 36, 2024.
Jens J Rassweiler, Thomas Knoll, Kai-Uwe Köhrmann, James A McAteer, James E Lingeman, Robin O Cleveland, Michael R Bailey, and Christian Chaussy. Shock wave technology and application: update. European urology, 5955-444477 2011.
15

---

Published as a conference paper at ICLR 2025
Pu Ren, Chengping Rao, Yang Liu, Zihan Ma, Qi Wang. Jia Wang, and Hao Sun. Physr: Physics-inforned deep super-resolution for spatiotemporal data. Journal of Computational Physics, 192:1,, 2023.
Olaf Ronneberger, Philipp Fischer, and Thomas Brox. U-net: Convolutional networks for biomedical image segmentation. In Medical image computing and compuler-ussisted InterveniionMIICAA. 2015: 18th international conference, Munich, Germany, October 5-9, 2015, proceedings, purt III 18, pp. 234-241. Springer, 2015.
Salva Rühling Cachay, Bo Zhao, Hailey Joren, and Rose Yu. Dyffusion: dynamics-informed diffusion model for spatiotemporal forecasting, Advances in Neural Information Processing Systems, 36, 2023.
Alvaro Sanchez- Gonazllzz Jonathan Godwin, Tobias Pfaff, Rex Ying, Jure Leskovec, and Peter Battaglia. Learning simulate complex physics with graph networks. In International conference on machine learning, pp. 8459-8468. PMLR, 2020.
Ivan W Selesnick, Richard G Baraniuk, and Nick C Kingsbury. The dual-tree complex wavelet transform. IEEE signal processing magazine, 22(6): 2005.
Shuyao Shang, Zhengyang Shan, Guangxing Liu, LunQian Wang, XingHua Wang. Zekai Zhang, and Jinglin Zhang. Resdiff: Combining enn and diffusion model for image super-resolution. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 38, pp. 8975-8983, 2024.
Dule Shu, Zijie Li, and Amir Barati Farimani. A physics-informed diffusion model for high-fidelity flow field reconstruction. Journal of Computational Physics, 478: 11972, 2023.
Sabrine Slama, Ayachi Errachdi, and Mohamed Benrejeb. Neural adaptive pid and neural indi- rect adaptive control switch controller for nonlinear mimo systems. Mathematical Problems in Engineering. 2019, 2019.
Jiaming Song, Chenlin Meng. and Stefano Ermon. Denoising diffusion implicit models. arXiv:2010.02502, October 2020. URL/httss:/aaxxi.. wwwp:/wwwww.rraa////bbbeaa 02502
Makoto Takamoto, Timothy Praditia, Raphael Leiteritz, Daniel MacKinlay, Francesco Alesiani, Dirk Pflüger, and Mathias Niepert. Pdebench: An extensive benchmark for scientific machine learning. Advances in Neural Information Processing Systems, 35: 15:15161161,, 2022.
Tapas Tripura and Souvik Chakraborty. Wavelet neural operator: neural operator for parametric partial differential equations. arXiv preprint arXiv:2205.02191 2022.
Arash Vahdat, Francis Williams, Zan Gojeic, Or Litany, Sanja Fidler, Karsten Kreis, et al. Lion: Latent point diffuusin models for 3d shape generation. Advances in Neural Information Processing Systems, 35:10021-10039 2022.
Siddhartha Verma, Guido Novati, and Petros Koumoutsakos. Efficient collective swimming by harnessing vortices through deep reinforcement learning. Proceedings of the National Academy of Sciences, 115((3):5849 5854, 2018
Clément Vignac, Igor Krawczuk, Antoine Siraudin, Bohan Wang, Volkan Cevher, and Pascal Frossard. Digress: Discrete denoising diffusion for graph generation. In Proceedings of the 1thh International Conference on Learning Representations, 2023.
Rui Wang, Yihe Dong, Sercan Ö Arik, and Rose Yu. Koopman neural forecaster for time series with temporal distribution shifts. arXiv preprint arXiv.2210003675 2022.
Tao Wang, Changhua Lu, Yining Sun, Mei Yang, Chun Liu, and Chunsheng Ou. Automatic ecg classification using continuous wavelet transform and convolutional neural network. Entropy, 23 (1):119, 2021.
ZP Wang, RJ Lin, ZY Zhao, X Chen, PM Guo, N Yang, ZC Wang, and DX Fan. Leam to flap: foil non-pan con-paramttric path planning via deep reinforcement learning. Journal of Fluid Mechanics, 984: A9, 2024.
16

---

Published as a conference paper at ICLR 2025
Long Wei, Peiyan Hu, Ruiqi Feng, Haodong Feng, Yixuan Du, Tao Zhang, Rui Wang, Yue Wang, Zhi-Ming Ma, and Tailin Wu. A generative approach to control complex physical systems. arXiv preprint ar arXiv:2407.066
2024.
Moritz Wolter, Felix Blanke, Jochen Garcke, and Charles Tapley Hoyt. ptwt the pytorch wavelet toolbox. Jourmal of Machine Learning Research, 25(80):1-7, 2024, URL
org/
papers/v25/2--0533 html
Haixu Wu, Huakun Luo, Haowen Wang, Jianmin Wang, and Mingsheng Long. Transolver: fast transformer solver for pdes general geometries. In Forty-first International Conference Machine Learning, 2024a.
Tailin Wu, Takashi Maruyama, Long Wei, Tao Zhang, Yilun Du, Gianluca Iaccarino, and Jure Leskovec. Compositional generative inverse design. In The Twelfth International Confer- ence on Learning Representations, 2024b. URL /ttttp:://opnnreveew.
wmXX0CFFdd]
Bingbing Xu, Huawei Shen, Qi Cao, Yunqi Qiu, and Xueqi Cheng. Graph wavelet neural network. In International Conference on Learning Representations, 2018.
Yuan Yin, Matthieu Kirchmeyer, Jean- Yves Franceschi, Alain Rakotomamonjy. and Patrick Galli- nari. Continuous pde dynamics forecasting with implicit neural representations. arXiv preprint arXiv:2209. 14855, 2022.
Zongsheng Yue, Jianyi Wang, and Chen Change Loy. Resshift: Efficient diffusion model for image super-resolution by residual shifting. Advances in Neural Information Processing Systems, 36, 2024.
Mykhaylo Zayats, Malgorzata Zimoó, Kyongmin Yeo, and Sergiy Zhuk. Super resolution for turbulent flows in 2d: stabilized physics informed neural networks. In 2022 IEEE 61st Conference on Decision and Control (CDC), pp. 3377-3382. IEEE, 2022.
Yi Zhu, Fang-Bat -Bo Tian, John Young, James Liao, and Joseph CS Lai. A numerical study fish adaption behaviors in complex environments with deep reinforcement learning and immersed boundary-lattice boltzmann method. Scientific Reports, 11(1):669, 2021.
Zifeng Zhuang, Kun Lei, Jinxin Liu, Donglin Wang, and Yilang Guo. Behavior proximal policy optimization. arXiv preprint 133. 2023.
17

---

Published as a conference paper at ICLR 2025
A DETAILS OF WAVELET DECOMPOSITION
In this section, we provide detailed introduction wavelet transforms. Let V be the space spanned by scaling functions @t.m m E E, and Wi be the space spanned by wavelets Vi.m. m E Z. The scaling functions possess two fundamental properties:
The scale function orthogonal for its integer translation.
2. The wavelet spaces satisfy nested and increasing sequence of spaces:
V-oc C ... C V-1 C Vo c Vi C... C. Voc.
As for the space Vi and Wi, they have the following relationship:
V++1 = V22 @ Wi.
Intuitively, the spaces Wi spanned by the wavelet functions complement the missing information between the scaling function spaces of different levels.
Then, the process of wavelet transform can be viewed convolving the scaling and wavelet functions of a certain level with the original signal, effectively splitting the signal into low-frequency and high-frequency components. Subsequently, the low-frequency part further decomposed. This results in obtaining coefficients ct, and dig dio+1 di, +2,
There are numerous types of wavelet bases that have different waveforms. Here provide further insights into the criteria used for wavelet selection. For the wavelets we consider (bior, db, sym), bior and sym wavelets offer symmetry, which reduces phase distortion during processing and allows for more accurate reconstruction compared to db, also reflected in the reconstruction loss table. Regarding the choice of wavelet scale, despite the higher smoothness higher-order wavelets, they generally have larger value ranges. Therefore, for data with small spatiotemporal size, high-order wavelets may not be suitable. For example, in the ID data with size of N 81 120, we choose bior2.4, while for the 2D data with a size of N x 32 x 64 x 64, we select bior lower-order wavelet. Using wavelets with excessively large support lengths may distort coefficients near the boundaries and fail to effectively decompose signal details, hindering the effectiveness of multi-resolltion decomposition.
Specifically, select bior2.4 and biorl.3 from the Biorthogonal wavelet family for experiments on the 1D Burgers' equation and 2D incompressible fluid, respectively. And we we use the 'peri- odization' mode in ID and the 'zero' mode in 2D. Due the presence of the temporal dimension, we perform a two-dimensional wavelet transform on data from the ID Burgers' equation and three-dimensional wavelet transform on data from the 2D incompressible fluid.
In Table we report the reconstruction eπrоrs of wavelet transforms using different wavelet bases on the ID Burgers' equation and 2D incompressible fluid, the results show that the reconstruction is significantly low.
Table 3: Reconstruction relative L2 errors on 1D Burgers' equation and 2D incompressible fluid.
Types of wavelet
1D
2D
bior1.3
1.09e-07
3.32e-07
bior2.4
8.32e-08
2.65e-07
db4
1.39e-07
4.39e-07
sym4
1.17e-07
3.74e-07
Besides, in Table we provide the total time consumption for Fourier and wavelet transforms on the training set of the ID compressible Navier-Stoker equation. The Fourier transform is implemented using PyTorch's ' 2D Fast Fourier Transform function. Both times are recorded on an A100 GPU with batch size of 2000. From the results, observe wavelet transform's efficiency.
18

---

Published as a conference paper at ICLR 2025
Table Total time consumption for Fourier and wavelet transforms on the training set of the ID compressible Navier-Stokes equation.
Wavelet transform
Fourier transform
Time (s)
1.0171
1.3810
B VISUALIZATION OF EXPERIMENT RESULTS
B.1 VISUALIZATIONSS OF ID COMPRESSIBLE NAVIER-STOKE EQUATION
In Figure we present visualizations of predictions from WDNO and DDPM. It is clear that WDNO far better at modeling states with abrupt changes. While DDPM can not capture details, WDNO successfully predict these precise changes.
DCPN
Sheend truit
Ground thahh
NCMC ICPM
CDM
WOND
WCNC Greund
21210
Mtrte
DOPN CourE DE27
Ground
stare
Figure 7: Visualizations of WDNO's and DDPM's performance on simulation of the ID com- pressible Navier-Stokes equation.
19

---

Published as a conference paper at ICLR 2025
B.2 VISUALIZATIONS OF 2D INCOMPRESSIBLE FLUID
We provide visual results of WDNO on challenging 2D control tasks in Figure It can easily be observed that, for many trajectories, our method successfully guides the smoke to pass essentially through the target bucket, which is not achieved by other baselines.
t
t 14
t t= 21
t 28
T
T
T
Figure 8: Visualizations of WDNO's performance on the 2D incompressible fluid control task.
20

---

Published as a conference paper at ICLR 2025
C ADDITIONAL RESULTS OF EXPERIMENTS
C.I MORE COMPARISONS ON ID COMPRESSIBLE NAVIER-STOKES EQUATION
Here, we provide more results on simulation of ID Compressible Navier-Stokes Equation. We further compare WDNO with Transolver (Wu et al. 2024a), CNO (Raonic et al. 2024). MSVI (lakovlev et al. (022), ACDM (Kohl al., 2024). ) and DiffusionPDE (Huang 2024). We also add comparisons with diffusion models in Fourier domain and FNO denoiser. The results in Table 5 demonstrate that WDNO still achieves the best performance on MSE. can be observed that the trends of MAE align closely with MSE. However, the L error values across different methods are relatively similar because this metric only considers the maximum value across the entire spatiotemporal domain, thus capturing less information.
Table 5: Comparison of Various Models Based Error Metrics
Model
MSE
MAE
Togo Error
Transolver
4.9984
0.4025
4.87284
CNO
0.3987
0.2765
9,9169
MSVI
1.7063
0.6047
17.0386
ACDM
4.6574
0.8946
60.9370
DiffusionPDE
5,5936
0.9792
16.0514
WNO
6.5428
1.1921
21.3860
MWT
1.3830
0.5196
11.3677
OFormer
0.6227
0.4006
30,9019
FNO
0,2575
0.1985
11.1495
CNN
12.4966
1.2111
17.6116
DDPM
5.5228
0.9795
16.0532
Diffusion FFT FNO Denoiser
3.0258
0.8498
14.6670
145.0469
6.6406
31.7515
WDNO (ours)
0.2195
0.1049
13.0626
C.2 ABRUPT CHIANGES
Here we provide more visualizations of the comparison between WDNO's and DDPM's MAE of different time steps. Figure verifies that WDNO can better model abrupt changes due to the wavelet transform.
21

---

Published as a conference paper at ICLR 2025
WODOD
0.010
WDkD DOPN
100
2
0.004
0.002
0.000
20
90
10
20
30
40
50
70
40
-1-
0.0
0.2
0.4
0.6
0.8
1.0
X
0.020
0.015
DOPH
0.000
DOPN
0.006
0.004
0.005
60
1.0
1.0
0.5
0.0
00.
-1.0
0.0
0.2
0.4
0.6
0.0
1.0
0.0
0.6
0.6
0.8
1.0
0.015
WOND
0.000
WORD
100
0.004
3.025
0.002
0.000
0.000
20
20
30
50
40
3.5
0.5
3.0
0.5
-10
0.0
0.4
0.5
1.0
0.0
0.2
0.6
1.0
0.15
0.815
D0PM
0.030
0.009
0.000
40
0.0
0.2
0.4
0.5
0.5
1.0
Figure Visualizations WDNO's and DDPM's MAE on the 1D Burgers' equation.
22

---

Published as a conference paper at ICLR 2025
C.3 APPROXIMATE SCALE INVARIANCE
conduct experiments to verify approximate scale invariance by training FNOs on original, once- downsampled, dodie-dowssappled, and mixed datasets, then testing these three resolutions. From 6 it is evident that the model trained on the mixed dataset performs better than those trained at specific resolutions.
Table 6: Approximate scale invariance on the ID Burgers' equation.
Original
Once-downsampled
Twice-downssmpled
Mix
4.72e-04
4.69e-04
4.95e-04
Individual
4.94e-04
5.42e-04
4.94e-04
C.4 SENSITIVITY ANALYSIS
To conduct a sensitivity analysis of key guidance weight DDIM sampling steps, and coefficient 7) ID simulation and control. As shown in Table 7 and Table g WDNO is not sensitive to hyper-parameters.
we analyze the impact of wavelet type,
Table 7: Results of simulation on ID Burgers' equation.
(a) DDIM step.
(b) DDIM 77.
(c) Wavelet type.
MSE
MSE
MSE
20
0.00022
0.2
0.00020
0.2
0.00020
40
0.00017
0.5
0.00020
0.5
0.00020
50
0.00014
0.8
0.00017
0.8
0.00017
100
0.00015
0.00014
0.00014
200
0.00017
Table 8: Results of control on 2D incompressible fluid.
DDIM step.
(b) DDIM 1).
(c) Guidance weight (le4).
Results
Results
Results
20
0.0223
0.2
0.2285
9
0.0213
40
0.0215
0.5
0.0694
10
0.0207
50
0,0205
0.8
0.0244
11.5
0,0205
100
0.0200
0.0205
12.5
0.0205
200
0.0217
13
0.0215
C.5 ROBUSTNESS
In addition, to test the robustness of WDNO, we first conduct 1D experiments with a 0.1 probability of noise in the control sequence. Table shows that WDNO outperforms other learning-base based methods, showing its robustness. We also give the results (mean + std) of ID control averaged over 50 testing samples in Table 10 showing that WDNO's std relatively low.
23

---

Published as a conference paper at ICLR 2025
Table 9: ID control experiments with 0.1 probability of noise in the control sequences.
Methods
J
PID
(surrogate-solver)
0.6644
SAC (pseudo-onlinn)
0.2166
SAC (offline)
0.3979
BC
0.2457
BPPO
0.2392
SL
0.0348
DDPM
0,0701
WDNO (ours)
0.0305
Table 10: Results of control tasks (mean-istd). Bold font denotes the best model and the runner-up
underlined.
Method
Results
PID (surrogate)
0.6645 + 0.5940
SAC (pseudo-online)
0.1376 + 0.1729
SAC (offline)
0.3210 + 0.2733
BC (surrogate)
0.2998 + 0.1137
BPPO (surrogate)
0.3075 = 0,1178
SL
0.0235 = 0.0171
DDPM
0.0272 + 0.0198
WDNO (ours)
0.0205 + 0.0198
C.6 COMPUTATIONAL RESOURCES
We provide inference times for a batch size of 1 on A100 in Table [1] It is evidence that WDNO's runtime is moderate, and its relatively large parameter count is due the U-Net base model, which can be replaced with smaller models. In addition, we test WDNO's total training and inference times. As shown in Table T2 WDNO has reasonable spatial and temporal costs. Notably, for the 1D Burgers' equation, WDNO achieves the lowest training time, as shown Table 13
Table Number of parameters and inference time (s) results of 1D Burgers' equation.
(a) Control task.
(b) Simulation task.
Methods
Parameters
Time
Methods
Parameters
Time
PID (surrogate)
4034952
0.081
WNO
153408
0.0105
SAC (pseudo-online)
89011116
0.503
MWT
2733978
0.0093
SAC (offline)
88854770
0.466
OFormer
2556321
0.0338
BC
1543408
0.075
FNO
4769601
0.003
BPPO
6579171
0.079
CNN
156346
0.277
SL
156346
181.35
DDPM
140703746
1.884
DDPM
140703746
1.903
WDNO (ours)
140748553
0.966
WDNO (ours)
140748553
1.131
24

---

which reveal that our method is Published as a conference paper at ICLR 2025
Table 12: Total training and inference time of WDNO.
Control
Simulation
Time
Space (MB)
Time
Space (MB)
ID train (A100)
2.4h
7165
2.5h
7165
ID inference (A100)
455s
3691
5.2s
3761
2D train (2A100)
7.8h
8043+8039
7.9h
8043+8039
2D inference (A100)
2676s
30255
70.9s
16621
Table 13: Training time (h) results 1D Burgers' equation.
Methods
Time (h)
WNO
4.5
MWT
6.5
OFormer
19.7
FNO
10.5
CNN
63.8
DDPM
7.8
WDNO (ours)
2.5
Moreover, in Table 14, provide the time and space required for WDNO to generate a batch of size 5 during inference on the ID Burgers' equation experiment, without super-resolution, and with one, two, and three levels of super-resolution. As shown in Table, with each increase in the level of resolution, the required time and space increase, and the growth rate increasing. We can infer that the level of super-resolution increases, the spatiotemporal costs also rise, indicating potential areas for further algorithm optimization.
Table 14: Inference time and space of 1D super resolution.
Level of super resolution
0
2
3
Time (8)
1.7
1.8
6.9
28.1 1
Space (MB)
1711
1831
3503
10631
C.7 IMPORTANCE OF GUIDANCE
Since À in Eq. is a hyperparameter, it can be set to zero, which means not including this term during denoising. In practice, we have selected the hest-performing A. To further show the effectiveness, we We see that without
have also provided the 1D control results with and without this term Table 15 this term, the performance drops lot.
Table 15: Results of control on the 1D Burgers' equation.
Results
0
0.3360
120000
0.0205
C.8 ZERO-SHOT SUPER-RESOLLUTION
In Table T6 and Table 17 we provide results in Section 4.6 outstanding zero-shot super-resolution.
25

---

Published as a conference paper at ICLR 2025
Table 16: Mean squared error of zero-shot super-resolution on the 1D Burgers' equation.
Methods
times
I times
times
3 times
WNO (linear)
0.01 10
0.6284
1.4474
2.0588
FNO (linear)
0.00312
0.02463
0.04458
0.05610
WDNO (linear)
0.00259
0,00074
0.00036
0.00035
WNO (nearest)
0.0079
0.6491
1.5007
2.0588
FNO (nearest)
0.00439
0.02473
0.04450
0.05610
WDNO (nearest)
0,00382
0.00110
0.00046
0.00035
Table 17: Mean squared error of zero-shot super-resolution the 2D incompressible fluid.
Methods
0 times
times
FNO (linear)
0.09497
0.01692
WDNO (linear)
0.09309
0.00765
FNO (nearest)
0.11963
0.01692
WDNO (nearest)
0.
0.00765
D RELATED WORK
PDE simulation. Solving a family of PDEs can be regarded as approximating nonlinear operators in the functional space, where neural operators have recently proved effective (Kovachki et al. 2023),
such as DeepONet (Lu et al. 2019). FNO (Li et al., 2021). GNOT (Hao et al. 2023). L LNDDE (lakovlev et al. 2023a), CROM (Chen et al. 2022b). DINo (Yin et al. 2022). MagNet (Boussif ct al. 2022). Transolver (Wu et al.) 2024a) and CNO (Raonic et al. 2024). Additionally, scientific legge such Clifford algebras (Brandstetter al.) (022a) and Koopman theory (Wang et al.)
2022) has been incorporated into neural networks to improve neural operators' performance. There also neural ODE based approaches able simulate PDE systems (lakovlev et al. 2022 Lagemann ct al. 2023). Among them, some models explicitly learn inside functional spaces, such as the Fourier domain (Li et al., (22T)) and the wavelet domain (Gupta al., [021] Tripura & Chakraborty [0022]
Cheng et al. 2024). However, most existing works mainly focus on simulating physical systems, lacking physical system control problems. Our work proposes a wavelet diffusion neural operator that excels in simulating physical systems and can naturally manage both control and super-resolution
simulation tasks.
Super-resolution tasks. Super-resolution tasks aim to reconstruct high-resolution data from low- resolution data. In recent years, many studies on physical system simulation have focused on super resolution tasks. Some methods transfer the learning of dynamies into function space, naturally enabling zero-shot super resolution capabilities (Li et al, 2021 Tripura & Chakraborty 2022 Cao al. 2024). Additionally, some studies attempt incorporate physical information, such equation forms, into the model's learning process to achieve super-resolution (Gao et al., 2021 Jangid et al. 2022 Zayats et al. 2022; Jangid al., 2022). Other studies primarily achieve super-resolution by injecting high-resolution information into super-resolution model (Esmaeilzadeh et al.) 2020 Ren et al. 2023 Shu et al., (023). Given the importance of the super-resolution task, we propose leveraging approximate scale invariance to enable diffusion models to achieve super-resolution capabilities.
Wavelet transform. The wavelet transform, a powerful tool for signal processing and analysis, is widely utilized in designing deep learning algorithms. Due its ability to decompose signals into high-frequency and low- eowffreueency components. it is used to enhance robustness (Li et al. 2020a). improve accuracy (Li et al. 2020a Liu et al. 2019), enable super-resolution (Guo al. 2017 Huang al. 2017), and extract features (Wang al., (222), among other applications. Two closely related works (Hui et al. 2022 Guth et al. 2022) incorporate the wavelet transform into the diffusion model, correlations. Also, our paper focuses different tasks and emphasizes the operator characteristics in PDE systems, such as mapping between infiiie-diienssinn function spaces.
but these works do not involve space- space-time
26

---

Published as a conference paper at ICLR 2025
Long-term predictions. Error accumulation is common challenge for transient PDE predictions, and several methods have been proposed to address it. Some works suggest training prediction models over multiple steps rather than single step to enhance robustness in multi-step predictions (Lusch et al., 2018 Brannstetter et al. 2022b). Techniques such as noise injection into training data and adversarial training are employed to improve the model's resilience to small disturbances (Sanchez-Gonzalez et al.] 20200] Lippe et al. (2024), while other works use geometric manifold learning identify the intrinsic dimensions of observed systems, enabling robust predictions of
underlying dynamics (Chen et al. 2022a).
PDE control. For the task of controlling physical systems governed by PDEs. various deep learning-
based techniques have been proposed (Feng et al.) 20233 Zhhu et al. [021] Degrave et al., 2022) ).
prominent class of methods is supervised learning (SL.) (Holl et al.] 2020, Hwang et al. 2022)
which optimizes control input via backpropagation through a neural surrogate model. Unlike these
methods, our approach does not rely on auto-regressive surrogate models but instead learns both entire state trajectories and control sequences. Besides, deep reinforcement learning (DRL) has been applied to various physical problems such as drag reduction (Rabault et al.] 20199] Elhwwrry
2020; Feng et al., 2023] Wang al. 2024). heat transfer (Beintema et al., 2020) Hachem al. 2021). and swimming Novati et al., 2017 Verma et al.) 2018). These methods often implicitly
incorporate physical information and make decisions sequentially. In contrast, our approach generates entire trajectories, facilitating trajectory-level optimization while emhedding physical insights learned by models. Additionally, physies-informed neural networks (PINNs) (Raissi et al., 2019) have recently been used for control (Mowlavi & Nabi 2202)), but they require explicit formulations of PDE dynamics. In contrast, our method is data-driven and can address a broader spectrum of complex physical system control problems without knowledge of the explicit PDE dynamics.
Diffusion models. The diffusion model (Ho et al., 2020b) proficient in learning high- sigh-dimensional distributions and has succeeded in image and text generation (Dhariwal & Nichol (021). It has also demonstrated remarkable ability, including strong modeling capabilities complex and high- dimensional systems and temporal stability. in scientific or engineering problems such as robot
control (Janner et al. 2022 Ajay et al. 2022). fluid prediction (Li et al. 20244) Kohl et al. 2024).
weather forecasting (Price et al. 2023). 3D human motion generation (Vahdat et al.) 2022) ) PDE simulation based on sparse observation (Huang et al., 2024) and PDE control (Wei et al. 2024
2024) is an autoregressive model for fluid Wu et al.] 2 Among them, ACDM (Kohl al.
simulation, which means both training and inference are performed sequentially. The model predicts the next k steps U[k,2k-1] based on the previous steps Wc0 Then, using U[k,2k-1], it predicts 321,33-11-- and this process continues iteratively until the entire trajectory "[0,T-1] of T steps is generated. In contrast, our proposed method predicts the entire trajectory u[o directly in single inference step based on the given k initial steps, significantly reducing computational overhead. Bisides, generalization across different resolutions and modeling states with abrupt changes remain challenging, and our WDNO proposes a promising direction to tackle the challenges. Many works have focused on improving the well-posedness functional space diffusion generation (Pidstrigach) et al., 2023 Hagemann et al., 2023 Lim et al., (003).. Previous works commonly choose the Fourier space as the functional space (Lim et al.] [023 Hagemann et al., 2023). Our method differs by using the wavelet transform since is better approximating the important abrupt changes.
E PSEUDOCODE
To help understand the entire algorithm, provide the pseudocode of WDNO's training and inference in Algorithm [ and the visualization of WDNO's entire training and inference on ID Burgers' equation in Figure 10)
27

---

Published as conference paper at ICLR 2025
ReeeRRestttion Model (BRM) Training
super-Resolution Model (SRM) Training Downsample
diffuse
denoise
20 30
10 15 A
40
BRM
Super-resolution
Wavelet Transform
120
160 240
320 680
Inverse Wavelet Transform
BRM
domain
Super-resclution
Space
PDE conditions:
Condition
Intal conditon, force sem,,
R7
Objective guidance
120
160 340
(only control);
Inference (simulattionconnrrol]
Output
-N(O.))
Figure 10: Overview of WDNO. The figure illustrates the training of Base-Resoluttin Model (BRM, top left), training of Super-Resolution Model (SRM, top right), and inference (bottom) of WDNO ID Burgers' equation. Through multi-resolution training and generation in wavelet space, WDNO is capable of generating superior simulation and control trajectories and conducting zero-shot super-resolution.
Algorithm Training and Sampling for WDNO
Require Diffusion models e (y Wa. k). objective J()) for control task, covariance matrix 2(()), I, condition Wa, schedule ãn, hyperparameters X, K
Training:
repeat
2: "(o,7] (
3: Apply the discrete wavelet transform to u to get W:
(0)
4:
Uniform(1....,. K)
5:
6~ (~,10,
6: Take gradient descent step on Volle - celVan Wa. + VI âge, k)
7: until converged
Sampling:
1: W I.,
2: for k = K, K,...,dd.
3: { {,,N I) if
clse 8 =0
4: w(k-1)
wilk)
17 ( "j W.. k) + & for simulation task
(k)
(((kl)) = w(()
v. T) (W uj.r)::
Wa. k) + (WW(WWl (4)
U1o,21 + & for control task
5: end for
6: Apply the inverse discrete wavelet transform to (0) "4..r to get u 10,77
(0)
7: return u² u "
28

---

Published as a conference paper at ICLR 2025
F ADDITIONAL DETAILS FOR 1D BURGERS' EQUATION CONTROL
F.1 EXPERIMENT SETTING
The equation takes the form
@u(t.z) -u(t 81 u(t, x) =0 u(0, X) uo(a)
Oult,z) dx
Pu(t,)) Ou + x(t,x) in [0 D,
+
on [0,7] x aD, at (t= 0},
(8)
where Up is the initial condition, the diffusion coefficient 0.01, T- 8 and [0, 1].
During inference, alongside the control sequence f(t, z), our diffusion model generates states ptt, x). Besides, some models produce surrogate states f(t, x:) when fed with the control ftt, æ). However, the state deviation JD [u(T, z) - u" (x)) dz in our reported evaluation metric J is always based on the output u(T', æ) = Ug.. (2", x) of the ground-truth solver given the control force f(,
The solver solves the Burgers' equation (Eq, 8) described in Appendix F.2 where the internal grid size of the ground-truth numerical solver is consistently at high-resolution (80 x 16, 120 x 16). When model outputs are a lower resolution, we linearly interpolate them before feeding them into the solver.
F.2 DATA GENERATION
We use the finite difference method (referred to as the solver or ground-truth solver henceforth) to solve the Burgers' equation of Eq. [ and generate the training data for the 1D Burgers' equation. Specifically, the initial state un(2)) and the control force f'(1, x) are both randomly generated, and then the state's evolution u(t, r) is numerically simulated using the solver. In the numerical simulation using the ground-truth solver, a domain of æ = [0, 1], t = [0, 8) is simulated. The space is discretized into 120 16 grids and time discretized into 4800 x 16 steps. However, only 80 time stamps are stored in the dataset, and the control sequence f is kept constant between two time stamps.
After simulation, we downsample by 16 times both spatially and temporally before saving the dataset. Therefore, the data size of each trajectory is [81, 120] for the state and [80, 120] for the force f. As for the super-resolution dataset for super-resolution simulation, we downsample the original data with shapes [80 x N + 1, 120 x N] for u and [80 x N, 120 x N] where N 2,3, corresponding to 1, 2, times super-resolution in Table 16
(r-ba)2
The initial value u(0, x) ) is superposition of two Gaussian functions z0, = Σ are
20
Li=|
where ae by, of are all randomly sampled from uniform distributions: a1 (0, 2), a2 U(-2,0) 0), b ~ U(0 0.4), by ~ U(066, 0.8), ~ U0.05 0.15), ~ U(0.05 0.15). Sim- ilarly, the control sequence f(x,t) t) is also a superposition of 8 Gaussian functions f(t, x) a bg.22
is 4,02 201.
E E aRe are 1), b2d U(O,1), 1), 01s (0..1, σ2 10.1, 0.4), while d1 U(-1.5, and for i 2 2, a, ~ U(-15, 1.5) or 0 with equal probabilities. u(t, x), (t + 0) is then numerically simulated (using the ground-truth touud-truthhao solver) given u(0,x)) x)) and f(t, x) based on Eq.
E
205.
where each parameter independently generated as follows: bi, ~
8 The dataset
generation setting is based on previous works Hwang et al. b; WWi et al. (004)).
generate 40000 trajectories for the training set. For the Burgers' equation control task, generate another 50 trajectories for testing. For the super-resolution task, we generate another 2000 trajectories for testing in the 0x super-resolution task (which is the original resolution). In 1, 3x super-resolution tasks, another 100 samples are generated and shared across the three different super-resolution settings.
F.3 DATA PREPARATION FOR WDNO
We perform 2D wavelet transform on the original data using the bior2.4 wavelet basis and the 'periodization' mode, implemented using the pytorch htteeh_aaaeeeeeeeeeeeeeeeeeeeeee package (Cotter) 2019). This transform the data, originally sized 81 x 120, into four sets of wavelet coefficients, each sized 41 x 60.
29

---

Published as a conference paper at ICLR 2025
Among these four sets of coefficients, there is set of coarse coefficients and three sets of detail coefficients. For the multi-resolution dataset used to train the Super- Resolution Model, obtained through downsampling,
conduct wavelet transforms data sizes 41 60, 21 31, and 11 x 15, which correspond to four sets of wavelet coefficients sized 21 x 30, 11 x 15, and 6 x respectively.
Notably, since the initial condition and the target state 1D, take the 1D wavelet transform, repeat the coefficients, and then concatenate them to other data.
When aligning the sizes of low-resolution wavelet coefficients with high- nigh-resolution ones by duplica- tion, special handling is required at the boundaries due to the presence of odd numbers. Specifically. duplicate the last temporal dimension of the high-resolution data once to ensure that the sizes match perfectly.
F.4 MODEL
The model architecture in this experiment follows the Denoising Diffusion Probabilistic Model (DDPM) (Ho et al., 2020b). In simulations, the Base-Resolution Model conditions on ug and f to predict No,r- For control tasks, we condition on My, ur and apply guidance related to I to generate J. The Super-Resolutie
Model conditions the same variables as the Base- Rass-Reeolution Model, with the addition of conditioning the initial condition and target state better, we involve the loss between the inverse wavelet transform of the initial condition and the target state's wavelet coefficient channel and the ground truth into the guidance. The hyperparameters on WDNO are recorded in Table
the low-resolution data. Besides, to make the generation meet
Table 18: Hyperparameters the UNet architecture and training for the results of 1D Burgers' equation in Table 2a and Table
Hyperparameter name
Base-Resolution Model
Super- Super-Resolution Model
UNet cs())
Initial dimension
128
128
Poausampling/UUsampliinnn layers
4
4
Convolution kernel size
3
3
Dimension multiplier
[1,2, 8j
[1,
4, 8]
Resnet block groups
8
8
Allention hidden dimension
32
32
Attention heads
4
4
UNet eg (u, f)
Initial dimension Cranssapliin/UUssappliii layers
128
128
4
4
Convolution kernel size
3
3
Dimension multiplier
[1,
4, 8]
[1, 2, 4,
8]
Resnet block groups
8
8
Atlention hidden dimension
32
32
Attention heads
4
4
Training
Training batch size
16
16
Optimizer
Adam le-4
Adam
Learning rate Training steps Learning rate scheduler
le-4
190000
290000
cosine annealing Inference
cosine annealing
DDIM sampling iterations of DDIM Sampling
50
50
1
1
Intensity of guidance in control Scheduler of guidance
120000
0
cosine
cosine
30

---

Published as a conference paper at ICLR 2025
G ADDITIONAL DETAILS FOR 1D COMPRESSIBLE NAVIER-STOKES EQUATION
G.1 EXPERIMENT SETTING
This fluid dynamic equation takes the form
O.p+ V. (pv) = O,
p(0,++v v Vv) = -Vp++ A (S+ 3) V(V.v) 01 c+ 5 +V. [c+p+ - v 0.
(9)
where p is the density, v is the velocity, p is the pressure, c = p/(T - 1) is the internal energy with I= = 5/3, o' is the viscous stress tensor, and S are the shear and bulk viscosity, respectively. The sound velocity is defined as:
C
r P
(10)
and the Mach number is:
M v
(11)
C2
The velocity field for turbulence is initialized as:
v -= 0) Σ A sin + Φi),
(12)
where A, Tejz. d 1,2 for 2D and 3D cases, and ē = c,M.
The shock-tube field is initialized as:
Qn),
(13)
where Q (p, v. p), with random constants Qz and Qu-
In detail, we use a ID compressible Navier-Stokes equation dataset provided by PDEBench. The initial conditions include random fields, turbulent fields, and shock-tube fields. The ran- dom and turbulence fields are prepared by adding perturbations to a uniform background, while the shock-tube setup consists of piecewise constant values generating shocks and rarefactions. Boundary conditions allow waves to exit the domain, and numerical solutions are computed using second-order HLLC and central difference schemes. And choose the most challenging dataset "ID_CFD Shock Etal. e d__etall e trans Train hdf5'
Since the original data of quite high resolution, downsample and the final resolution of the used data is 81 120, the same as the 1D Burgers' equation.
G.2 DATA PREPARATION FOR WDNO
Since the data size is the same as Appendix F the data preparation is almost the same. We also perform 2D wavelet transform the original data using the bior2.4 wavelet basis and the 'periodization' mode, implemented using the pytorch wavelets package (Cotter) 2019). As for the initial condition, take the ID wavelet transform repeat the coefficients, and then concatenate them to other data.
G.3 MODEL
The model architecture in this experiment follows Appendix F The hyperparameters on WDNO are recorded in Table 19
31

---

Published as a conference paper at ICLR 2025
Table 19: Hyperparameters of the UNet architecture and training for the results of ID com-
pressible Navier-Stokes equation in Table 1
Hyperparameter name
Base-Resolution Model
UNe c (f
Initial dimension
128
Downsampling/ Upsampiinggi layers
Convolution kernel size
3
Dimension multiplier
1, 4, 8]
Resnet block groups
8
Attention hidden dimension
32
Attention heads
UNet en (u,
Initial dimension wwansaapling/ppsaapppii layers
128
4
Convolution kernel size
3
Dimension multiplier
[,
Resnet block groups
8
Attention hidden dimension
32
Attention heads
4
Training
Training batch size
16
Optimizer
Adam
Learning rate
le-4
Training steps Learning rate scheduler
190000
cosine annealing
Inference
DDIM sampling iterations
850
n of DDIM Sampling Scheduler of guidance
cosine
H ADDITIONAL DETAILS FOR 2D INCOMPRESSIBLE FLUID
H.1 EXPERIMENT SETTING
The equation takes the form
8v(t.)) + v(l,2). Vv(l,x) -(4,))))VVv(((a)) Vp(t,x) f(1,2), Bt
v-v(t,) = 0,
(14)
v(0 - vo(x),
where f is the external force, p denotes pressure, v is the velocity and v is the viscosity coefficient.
H.2 DATA PREPARATION FOR WDNO
We performed 3D wavelet transform on the original data using the biorl.3 wavelet basis and 'zero' mode, implemented through the Pytorch Wavelet Tool box (ptwt) (Wolter et al. 2024). The data of size x 64 64 are transformed into eight sets of wavelet coefficients, each sized 18 34 34. Among these, there is one set of coarse coefficients and seven sets of detail coefficients. For the multi-resolltion dataset used to train the Super- Resolution Model, obtained through downsampling,
do wavelet transforms corresponding to wavelet coefficient with sizes 18 x 18 x 18, 18 x 10 x 10 respectively.
data sizes 32 x 32 32 and 32 16 x 16
Specifically, since the initial condition and the percentage of smoke through the target bucket 2D and ID respectively, we take the 2D and ID wavelet transform and repeat the coefficients to concatenate them.
32

---

Published as a conference paper at ICLR 2025
Similarly, to align the duplicated low-resolution data with the high-resolution data, duplicate the boundary values on both sides of the dimensions requiring super-resolution in the high-resolution data.
H.3 MODEL
In this paper, the architecture of the three-dimensiondl U-net we employ is inspired by the previous work (Ho et al. 2022). In this experiment, we use there are three main modules in our U-net: downsampling encoder, middle module. and an upsampling decoder.
3D convolutions. Specifically,
In the simulation, the diffusion model conditions the initial density and control to generate the entire trajectories of density, velocity, and percentage smoke passing the target bucket. In the control problem, the diffusion model conditions smoke through the target bucket at the last time step as the guidance. Same as the ID case, to satisfy the initial condition better, the guidance also involves loss between the inverse wavelet transform of the initial condition's 's wavelet coefficient channel and the ground truth initial condition. The hyperparameters of the 3D-Unet architecture in the Table 20
the initial density and takes the negative percentage of
Table 20: Hyperparameters of 3D-Unet architecture 2D experiments.
Hyperparameter Name
Value
Number of attention heads
4
Kernel size of conv3d
(3. 3)
Padding of conv3d
(1,1,1)
Stride of conv3d
(1,1,1)
Kernel size of downsampling
(1,4,4)
Padding of downsampling
(1,2,2)
Stride of downsampling
1,
Kernel size of upsampling
(1,4,4)
Padding of upsampling
(1,2.2)
Stride of upsampling
1,
DDIM sampling iterations
100
y of DDIM Sampling
Intensity of guidance in control
100
I ID CONTROL BASELINES
1.1 PID
Propercentageal Integral Derivative (PID) control (Li et al. 2006) is versatile and effective method widely employed in numerous control scenarios. It operates by using the error, i.e., the difference between the desired target and the current state system. Due to its simplicity and effectiveness, PID control is often the default choice for many control problems. However, despite its widespread use, PID control faces challenges such as parameter adaptation and limitations in Single Input Single Output (SISO) systems.
In our study, the ID Burgers' Equation Control problem presents a Multiple Input Multiple Output (MIMO) scenario, rendering direct application of PID control infeasible. Inspired by early works Slama et al. 2019 Ding et al. 2022) that employed neural networks as PID parameter adapters, we integrated deep learning with PID control to address the MIMO control problem. As illustrated in Figure [T] the ANN (artificial neural network) PID uses a neural network to adapt PID parameters, enabling multiple sets of SISO PID control.
33

---

Published as a conference paper at ICLR 2025
Err
[Kpe, Kig, Kdzl
ft
NN
controller
Figure 11: The architecture of ANN PID Controller. To use the MIMO PID controller to control state de target state ud, we train neural-network-haar PID parameter planner to output MIMO PID parameters based on Err, then use the PID controller to output the control sequence ft-
The neural network generating the PID parameters consists two ID convolutional layers, two fully connected layers, and four activation layers. We utilize the L1 loss between the current and target states as the training loss, and the Adam optimizer (Kingma & Ba 2014) to train the model. Detailed architecture information is provided in Table 21
Table 21: Hyperparameters of network architecture and training for ANN PID.
Hyperparameter name
Full observation
Kernl size of conv ld Padding of convld
3
Stride of conv ld Activation function
Softsign
Batch size
16
Optimizer
Adam
Learning rate
0.0001
Loss function
MAE
Given that PID is inherently SISO control method, the ANN PID employs neural network to derive multiple PID parameter sets, facilitating multiple SISO PID controls for MIMO control in the context of the Burgers' equation. However, ANN PID requires that the input and output dimensions match, thus it can only address problems with full observation and full control, or partial observation and partial control.
Additionally, the ANN PID controller has two training setups: one involving direct interaction with the solver, and the other involving interaction with the 1D surrogate model.
1.2 SAC
The Soft Actor-Critic (SAC) algorithm, developed by Haarnoja et al. (2018), represents a significant advancement in reinforcement learning techniques. Designed an enhancement of the conventional Actor-Critic frameworks, SAC sets itself apart by incorporating an entropy regularization term in its loss function. This addition promotes more effective exploration by simultaneously maximizing the expected cumulative reward and the entropy of the policy itself, leading to improved decision-man making processes in complex environments.
Compared to the Deep Deterministic Policy Gradient (DDPG) algorithm (Liliicrap et al. 2015 Pan
et al. 2018), the Soft Actor-Critic (SAC) algorithm introduces entropy regularization that promotes more effective exploration and avoids premature convergence to suboptimal policies, common drawback DDPG's deterministic nature. Furthermore. SAC's twin Q-networks counteract the overestimation bias that can affect DDPG's value updates, resulting in more stable learning processes. The automatic adjustment of the temperature parameter in SAC also eases the balancing act between exploration and exploitation, minimizing the necessity for careful hyperparameter tuning, As result, these features make SAC generally more sample-efficient and robust, especially in complex and continuous action spaces.
34

---

Published as a conference paper at ICLR 2025
During training, experiences stored in replay buffer and sampled randomly to update the networks. Initially, the entire training set is loaded into the replay buffer. For offline SAC, this replay buffer remains unchanged. In contrast, online SAC alternates between gathering experiences through environment interaction and updating the networks using the replay buffer. Offlin SAC, however, utilizes surrogate model trained the training real environment. The policy network is optimized to maximize the expected return by considering both the Q-value and the entropy term. The critic networks are trained minimize the error between their Q-value predictions and the target Q-values. To further stabilize training, SAC employs target critic network, which is slowly updated with the weights from the main critic network. For inference, SAC uses the policy network to select the action with the highest probability.
gather experiences instead of the
To effectively guide the system towards the target state with accuracy and speed, it essential to incorporate the distance between the state at each time step and the target state into the reward function. Therefore, the reward function for given time step i, state uy target state uy, and action wris defined follows:
rt,, Ues VT, ws) = --/// |u,, ual dx a S |we|? dx,
(15)
where the space domain and a is the weight of energy. We take the Adam optimizer (Kingma & Ba 2014) to train the networks and update the temperature parameter. The detailed values of hyperparameters are provided in Table 22
Table 22: Hyperparameters of 1D SAC.
Hyperparameter name
Value
Hyperparameters for ID Burgers' equation control:
Discount factor for reward Target smoothing coefficient Learning rate of critic loss Learning rate of entropy loss Learning rate of policy loss Training batch size Number of episodes
0.5
0.05
0.0003
0.003
0.003
8192
1500
Number of model updates per simulator step
50
Value target updates per step
15
Size of replay buffer
1000000
weight of energy cost
0.00002
Number of trajectories interacted with the environment per step
Number of layers of critic networks
Number of hidden dimensions of critic networks
4096
Number layers the policy network
Number of hidden dimensions of the policy network
4096
Activation function
ReLU [e 20,e2
Clipping's range policy network's standard deviation output
1.3 SUPERVISED LEARNING
The paper (Hwang et al. 2022) proposes supervised-learnnn--he based control algorithm that takes neural operator as surrogate model solve control problems. contains two stages. In the first stage, we take neural operator to learn the PDE constraint as Hwang et al. (2022). Two VAEs based on CNN learn to project state u, control signal f into the latent space, and a CNN learns the transition from u, to M,+1 in the latent space. In the second stage, these three neural networks used as surrogate models to calculate the gradient of the objective function with respect to the control input and optimize the control signal
35

---

Published as a conference paper at ICLR 2025
During optimization, the reconstruction loss for the control force is also included to guide it out of the adversarial mode of the surrogate model. We consider the control f as a learnable parameter and update it with the LBFGS optimizer (Liu & Nocedal, 1989). The hyperparameters of this supervised learning method are recorded in Table 23
Table 23: Hyperparameters of inference the 1D supervised learning method.
Hyperparameter name
Value
Hyperparameters for ID Burgers' equation control
Learning rate of updating
0.1
Number of epochs
100
Weight of average objective function loss
Weight of average reconstruction loss
0.01
Termination tolerance on first-order optimality of LBFGS optimizer
1e-5
Termination tolerance on parameter changes LBFGS optimizer
le-5
1.4 BPPO
The Behavior Proximal Policy Optimization (BPPO) algorithm, introduced by (Zhuang et al. 2023 an advanced reinforcement learning method that combines the strengths of Proximal Policy Optimization (PPO) with elements of behavior cloning. BPPO is an offline algorithm designed to monotonically improve the behavior policy a manner akin to PPO. Due to the inherent conservatism of PPO, BPPO restricts the ratio of the learned policy the behavior policy within specific range, similar to other offline RL methods, which ensures the Iearned policy closely aligns with the behavior policy. By leveraging the conservatism of online on-policy algorithms,, BPPO effectively addresses the overestimation issue often encountered offline RL settings.
The algorithm begins by estimating a behavior policy using behavior cloning and then iteratively improves a target policy using the PPO objective with a behavior constraint. This process of policy improvement, advantage estimation, and policy update enables BPPO refine the target policy while ensuring it remains close to the behavior policy. By integrating the strengths of online on-policy methods with tailored offline RL techniques, BPPO has demonstrated promising results on the D4RL benchmark, surpassing oftte-of-theart offline RL algorithms.
During training, BPPO first initializes the behavior policy Ty and the target policy Te. The behavior policy Tg is then estimated using behavior cloning to replicate the behavior demonstrated in the offline dataset. Subsequently, the target policy Ty is optimized using the PPO objective with behavior constraint, ensuring the target policy remains close to the behavior policy. The advantage function A then estimated using the behavior policy ny to evaluate the quality actions taken by the target policy. Finally, the target policy is updated by maximizing the PPO objective with the estimated advantage function, and adjusting the policy parameters to enhance performance. In the implementation, a state value network and Q value network are pre-trained using the state, action, and reward data from the offline dataset.
In practice, to enable the system to approximate the target state accurately and swiftly, it is essential to incorporate the distance between the state at each time step and the target state into the reward function. Thus, the reward function time step I, given the state u,, target state ugs and action ft, is defined as follows:
rtt, Ue, uas ft) --)=-1 |uee ux
Sn Ifr
where 52 is the space domain and is the weight of energy. We the Adam optimizer (Kingmal Ba BaBa[201) 2014) to train the networks and update the temperature parameter. The specific values of the hyperparameters used are detailed in Table 24
36

---

Published as a conference paper at ICLR 2025
Table 24: Hyperparameters of ID BPPO.
Hyperparameter name
Value
Hyperparameters for ID Burgers' equation control:
State value network:
Learning rate of value network
1 x 10-4
Steps of value network
2 x 10
Number of layers of value network
3
Batch size of value network
512
Number of hidden dimensions of value network
512
Q value network:
Learning rate of Q network
1 x 10-4
Steps of Q network
2 x 106
Number of layers of Q network
2
Batch size of Q network
512
Number of hidden dimensions of Q network
1024
Target Q network updates per step Soft update factor
2 0,005
Discount factor for reward
0.99
Behavior cloning:
Learning rate of BC
1 x 10 512 5 x 106
Training batch size of BC
Number of episodes of BC
BPPO:
Number of episodes of BPPO
1 x 102 2 1024
Number of layers of policy networks
Number of hidden dimensions of policy networks
Learning rate of BPPO
1 1 10-3
Training batch size of BPPO Clip ratio of BPPO Weight decay factor Weight of advantage function Size of replay buffer
512
0.25
0.96
0.9
2 x 106
Activation function
ReLU
Table 25: Hyperparameters of ID BC.
Hyperparameter name
Value
Hyperparameters for ID Burgers' equation control:
Learning rate Training batch size Number of episodes Size of replay buffer
I x 10-4
512
5 105
2 x 106
Number of layers of policy networks
2
Number of hidden dimensions of policy networks
1024
Activation function
ReLU
I.5 BC
The Behavior Cloning (BC) algorithm, introduced by (Pomerleau) 1988 is foundational technique imitation learning. BC is designed to derive policies directly from expert demonstrations, utilizing supervised learning to associate states with corresponding actions. This method eliminates the necessity for exploratory steps commonly required in reinforcement learning by replicating the actions observed in expert demonstrations. One of the significant advantages of BC is that it does
37

---

Published as a conference paper at ICLR 2025
involve interacting with the environment during the training phase, which streamlines the learning process and diminishes the demand for computational resources.
In this approach, policy network is trained using standard supervised learning strategies aimed at reducing the discrepancy between the actions predicted by the model and those performed by the expert in the dataset. The commonly used loss function for this purpose is the mean squared error between the predicted actions and expert actions. The dataset for training comprises state-action pairs harvested from these expert demonstrations. During inference, the model is evaluated using the same objective function as used in SAC. The specific hyperparameters utilized are detailed in Table 25
J 1D SIMULATION BASELINES
J.1 FNO
FNO represents a deep learning framework capable of mapping between infinitediimesiioall spaces. By parameterizing the integral kernel in Fourier space, FNO processes input through sequence of Fourier layers, performing linear transformations in the Fourier domain for efficient convolutions. This architecture supports zero-shot super-resolution, allowing models trained on lower resolutions to predict higher resolutions without retraining.
In ID experiments, we train the FNO model using the initial state and all controls as the input and using the rest states as the output. The parameters are outlined Table 26
Table 26: Hyperparameters of 1D FNO.
Hyperparameter name
Value
Hyperparameters for 2D Burgers' equation:
Number of modes to keep in Fourier Layer
16
Width of the FNO (i.e. number of channels)
64
Number of input channel
3
Number of output channel
1
Number of hidden channels of the lifting block
256
Number of hidden channels of the projection block Number of Fourier Layers Expansion parameter of MLP layer
256
4
0.5
Non-Linearity module
Gelu
Rank of the tensor factorization of the Fourier weights
1.0
Mode of domain padding
one-sided
Learning rate
1 x 10-4
Optimizer
Adam
Training epochs
1000
Learning rate scheduler
Cosine
Training batch size
50
J.2 WNO
The Wavelet Neural Operator (WNO) (Tripura & Chakraborty 2022) a novel operator learning algorithm that blends integral kernel with wavelet transformation. We record the hyperparameters of it in ID simulation in Table 27 On ID Burgers' equation and ID compressible Navier-Stokes equation, we choose 'sym4', 'bior2.4" wavelet respectively.
38

---

Published as a conference paper at ICLR 2025
Table 27: Hyperparameters of ID WNO.
Hyperparameter name
Value
Hyperparameters of the model architecture
Type of wavelet
sym4
Level of wavelet decomposition
5
Uplifting dimension
Number of wavelet layers
4
Training
Training batch size
100
Optimizer
Adam
Learning rate
le-3
Training epochs
1000
Learning rate scheduler
StepI.R
J.3 CNN
Convolutional Neural Network is the key block in deep learning. For ID simulation, our CNN model based on Hwang al. (022). Details of model architecture and training can found in Table 28
Table 28: Hyperparameters of ID CNN.
Hyperparameter name
Value
Autoencoder of state
Convolution kernel size
Convolution padding
Activation function
ELU
Latent vector size
256
Autoencoder of force
Convolution kernel size
5
Convolution padding
Activation function
ELU
Latent vector size
256
Training
Training batch size
5100
Optimizer
Adam
Learning rate
le-3
Training epochs
500
Learning rate scheduler
cosine annealing
J.4 MWT
To compare with other wavelet-based methods, we mainly implement the ID baseline adapted from Gupta et al. (2021). We select 'legendre" wavelet here, following the original work. More configurations can be found in Table 29
Table 29: Configuration of ID MWT.
Hyperparameter name
Value legendre
Wavelet basis
Number of Fourier modes
10
Kernel size Training batch size Training epochs
4
256
300
Optimizer
Adam
Learning rate scheduler
MultiStepLR
39

---

Published as a conference paper at ICLR 2025
J.5 OFORMER
The Operator Transformer (OFormer [i et al. (2023) is novel framework built upon self-attention. cross-attention, and a set point-wise multilayer perceptrons. Details of model architecture and training be found in Table 30
Table 30: Hyperparameters of ID OFormer.
Hyperparameter name
Value
Encoder
Type
SpatialEncoder21>
Input Channels Embedding Dim of Token
96
Embedding Dim Encoded Sequence
256
Heads
Depth
Resolution Dropout of Embedding
120
0.05
Decoder
Type
PontWiseDecooee2DSSimpp
Latent Channels
256
Out Channels
Scale
0.5
Res
120
Training
Training batch size
Iteration Learning rate
50000
le-4
J.6 LADID
In the dynamics trajectory prediction community, predicting trajectories of dynamical systems is of interest (Chen et al. 2018). Among them, MS-L-NODE (Iakovlev et al. 2023b) a represen- tative method that is dedicated to learning the system invariant dynamics. Since MS-L-NODE effectively operates on sigh-dimensional spatial- data, we include as baseline for more comprehensive empirical comparison.
Since the original work only considers spatially 2D input, we modify the encoder and decoder from stacked 2D CNNs 1D CNNs, tune the latent dimension in (4, 8, 16, 64), and the CNN base output dimension (16, 64, 128). Important hyperparameters are reported in Table [3] and the rest are kept the same as in Iakovlev al. (2023b).
K 2D SIMULATION BASELINES
K.1 FNO
In 2D experiments, we train the FNO model using the density, velocity, control, and percentage of smoke through the target bucket of the previous step the input and using the rest step's 's density, velocity, and percentage of smoke through the target bucket as the output. The parameters are outlined in Table 32
40

---

Published as a conference paper at ICLR 2025
Table 31: Hyperparameters of MS-L-NODE.
Hyperparameter name
Value
Encoder
Encoder CNN channels
128
Latent dimension
8
Decoder
Encoder CNN channels
128
Latent dimension
8
Aggregation Network
Heads
16
Static layers
4
Dynamical layers
Training
Training batch size
64
Training iterations
37500
Learning rate
le-3
Table 32: Hyperparameters of 2D FNO.
Hyperparameter name
Value
Hyperparameters for 2D incompressible fluid:
Number of modes to keep in Fourier Layer
16
Width of the FNO (i.e. number of channels)
64
Number of input channel
6
Number of output channel
4
Number of hidden channels of the lifting block
256
Number of hidden channels of the projection block Number of Fourier Layers
256
4
Expansion parameter of MLP layer
0.5
Von-Linearity module
Gelu
Rank of the tensor factorization of the Fourier weights
1.0
Mode of domain padding
one-sided
Learning rate
1 x 10
Optimizer
Adam
Training epochs
1000
Learning rate scheduler
Cosine
Training batch size
50
Table 33: Hyperparameters of 2D WNO.
Hyperparameter name
Value
Hyperparameters of the model architecture
Type of wavelet
db4
Level of wavelet decomposition
Uplifting dimension
Number of wavelet layers
3
Training
Training batch size
50
Optimizer Learning rate
Adam
0.05
Training epochs
500
Learning rate scheduler
StepLR
41

---

Published as a conference paper at ICLR 2025
K.2 WNO
The hyperparameters of WNO for the 2D simulation task are in Table And the wavelet we choose is 'biorl.3'.
K.3 MWT
For comprehensive comparison with other wavelet-based approaches, focus primarily implemening the 2D baseline inspired by the work in Gupta et al. (202I). Following the previous work, select the 'Legendre' wavelet. More details the configurations can be found in Table 34
Table 34: Configuration of 2D MWT.
Hyperparameter name
Value
Wavelet basis
legendre
Number of Fourier modes
12
Kernel size Training batch size Training epochs
200
300
Optimizer
Adam
Learning rate scheduler
MultiStepLR
Table 35: Hyperparameters of 2D OFormer.
Hyperparameter name
Value
Encoder
Type
SpariaTTemporalEEnoder222
Input Channels Embedding Dim of Token Embedding Dim of Encoded Sequence Heads
34
96
192
Depth
5
Decoder
Type
PointWiseDecoder2D
Out Channels Embedding Dim of Token
96
Propagate forward
Length of output sequence
32
Propagator depth
Curriculum ratio
0.16
Curriculum steps
10
Training
Training batch size
Iteration
100000
Learning rate
1e-4
K.4 OFORMER
The Operator Transformer (OFormer) Li et al. (023) is an attention- based framework for learning solution operators of partial differential equations using self-attention, cross-attention, and point- wise MLPs, designed to handle various input sampling patterns and query locations. More details on the configurations can be found in Table 35
42

---

Published as a conference paper at ICLR 2025
L BROADER IMPACTS
Our research proposes a method to simulate and control complex physical systems. We believe our research will bring in significant progress for various scientific and engineering domains, including climate forecasting, fluid control, robotic control, et al. However, there is also a potential that the method might be abused incur negative social consequences, upon which we should remain vigilant.
43


Reviews:
[{"review_id": "hTt9FekrLZ", "reviewer": "Reviewer_FYpM", "summary": "Aiming at improving the performance of diffusion generative models, in particular for cases with abrupt changes, the authors propose to construct such diffusion models in the wavelet space. To reduce the solution manifold that the model needs to learn and therefore to ease super-resolution, the authors propose to scale the systems such that they align spatio-temporally, and additionally propose to train the model at multiple scales. The resulting model over-performs the state-of-the-art baselines in various simulation and control tasks.", "questions": "1. It seems that the Fourier neural operator is second to the proposed method in the simulation tasks. I am curious about how would FNO perform if it is also trained through a similar process as the proposed model (since the original FNO is not trained as a denoising model)?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "1. A wavelet-based neural operator that surpasses the performance of exising methods.\n2. Scaling of the input system reduces the solution manifold that the model needs to learn.\n3. The paper is generally well-written, with details carefully reported.", "weaknesses": "The most significant weakness point of the work is the small number of temporal steps in all the test cases, which might actually be one of the limitation of these models, since the usual approach adopted to achieve long roll-out inference without instability is to introduce artificial training noise - it might be difficult to do so when one is already training a denoising model. It should be noted that the authors iteratively try to strengthen the point that they are doing \"long-term dynamics\" in the temproal domain, but their reported cases, e.g., the 2D incompressive fluid case, only consist 32 time steps each, which in the commonsense definition is quite a short forecasting window. I strongly suggest removing such terms from the paper to avoid confusion for future readers of the work.", "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": "", "extra_scores": {}}, {"review_id": "Rjvb8ql4tl", "reviewer": "Reviewer_34r3", "summary": "This paper presents a new method to simulate and control physical systems governed by specific PDEs. The method is based on learning a diffusion generative model in the wavelet domain with multi-resolution. The use of wavelet basis is specially convenient for discontinuous behaviours such as shock waves in fluid dynamics, and the diffusion procedure ensures that the full-trajectory error remains bounded. The method is tested with several fluid dynamic problems such as 1D Bruger's equaion, 1D and 2D Navier-Stokes in both simulation and control tasks.", "questions": "* Can the authors discuss on the implication of using a regular grid? Is this constraint related to the use of a Fast Wavelet Transform? This restriction might be problematic when applying this algorithm to more complex geometries.\n* The examples shown on the paper have changing initial conditions, but the geometry and boundaries remain fixed. Can the authors describe how the WDNO architecture is generalizable to other parameters or boundary conditions? It might be interesting for the reader to have a comment over problems with parametric solutions.\n* Eq. 4: There is no reference to the guidance rate $\\lambda$ until Appendix C3. Can the authors define it next to the equation, to improve readability?\n\nFinal comment: The paper has a solid idea, with well justified design choices and the results outperform existing techniques. Based on the comments above, my rating for this paper is marginally above the acceptance threshold.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "* The main idea of combining WNO with diffusion is simple, novel and well justified.\n* Every design choice is justified with multiple tests and ablation studies.\n* The method outperforms other similar state-of-the-art techniques such as traditional U-Net/CNN to the newest Neural Operator architectures.", "weaknesses": "* The method is constrained to static uniform grid data and low-dimensional toy examples.\n* The generalization to more complicated tasks or boundary conditions might not be trivial.", "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": "", "extra_scores": {}}, {"review_id": "fPpeMbhDGt", "reviewer": "Reviewer_iVpN", "summary": "This paper proposes a wavelet diffusion neural operator (WDNO), which performs diffusion-based generative modeling in the wavelet domain for the entire trajectory of time-dependent PDEs and multiresolution training to generalize across different resolutions.", "questions": "1)  How does WDNO compare to baselines regarding training time, memory usage, and inference speed?\n\n2) How was the basis wavelet chosen for the benchmarks?\n\n3) Have the authors investigated using diffusion in the Fourier domain for comparison with WDNO, considering the computational efficiency of Fourier transforms?\n\n4) How were hyperparameters for baselines chosen? Were they optimized fairly compared to WDNO?\n\n5) The paper needs to elaborate on why WDNO outperforms DDPM, considering Parseval's identity, which states that information content remains constant during transformations.\n\n6)  Why wasn't the Wavelet Neural Operator (WNO) included as a baseline for super-resolution tasks? How does WDNO compare to WNO, DDPM, UNET, etc., for long-range dependencies (Figure 6b)?\n\n7) The paper should cite relevant related work and add as baselines, including:\n\n* Benchmarking Autoregressive Conditional Diffusion Models for Turbulent Flow Simulation (arXiv:2309.01745)\n\n* DiffusionPDE: Generative PDE-Solving Under Partial Observation (arXiv:2406.17763)", "limitations": "", "rating": 6, "confidence": 5, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "1) Multi-resolution training\n\n2) Diffusion in the wavelet domain to capture long-term dependencies and abrupt changes effectively.", "weaknesses": "1) The evaluation omits comparisons with state-of-the-art operators like Transolver, GNOT, LSM, DPOT, and CNO etc\n\n2) Training time, inference speed, and memory usage are not compared for WDNO and baselines.\n\n3) Lacks novelty as compared to previous works such as:\n\n* Benchmarking Autoregressive Conditional Diffusion Models for Turbulent Flow Simulation (arXiv:2309.01745)\n\n* DiffusionPDE: Generative PDE-Solving Under Partial Observation (arXiv:2406.17763)", "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": "", "extra_scores": {}}, {"review_id": "R3aUuLxntX", "reviewer": "Reviewer_NMDB", "summary": "The standard diffusion models cannot learn dynamics with abrupt changes. This paper proposes a new method, Wavelet Diffusion Neural Operator (WDNO), to solve this issue. It combines wavelet transform and diffusion models in the context of learning PDEs from the perspective of neural operators. This method also considers multi-resolution training. Multiple datasets have been tested to evaluate the model performance compared to the baseline models.", "questions": "- For the evaluation metrics, why does this paper only consider MSE? If this paper focuses on the dynamics with abrupt changes, then the Mean Absolute Error (MAE) or infinity norm should be considered. \n\n- Some minor typos: \n\n    - On Page 2, the paragraph name “wavelet domain” should be “Wavelet domain”. \n\n    - On Page 2,  it seems not common to see “contribute the following”.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- This paper aims to solve an interesting problem regarding the abrupt changes in spatiotemporal dynamics. \n\n- This paper has shown comprehensive details of the experimental setup, model details, and results. \n\n- This paper is easy to follow.", "weaknesses": "- This paper shows incremental novelty. It uses diffusion models in the wavelet space. But that is not a big issue as long as the model performance is good. \n\n- The motivation for using diffusion models for PDE learning is not well established. This paper mentions that diffusion models can better capture long-term predictions and model high-dimensional systems. However, many deep learning-based models can do both, such as Fourier neural operators. How do diffusion models differentiate from other deep learning-based models? Also, can you provide a reference paper that shows diffusion models can do a better job of capturing long-term predictions?\n\n- I have some concerns about the datasets used in this paper. 1D Burgers with a viscosity of 0.01 is not particularly challenging, though it presents shock wave phenomena. The standard physics-informed neural networks [1] are also capable of handling such tasks. The authors may consider some other challenging datasets, such as ERA5 [2], since the proposed method can deal with long-term predictions. Second, it would also be interesting to see the model performance on the dynamics without abrupt changes. I think that will also give the audience a broader sense of how this model works. The authors may consider some datasets in PDEBench [3].\n\n- It would be good to have another ablation study on the DDPM + Fourier transform, which changes the wavelet transform in the proposed method. The additional experiment will be trained in Fourier space instead of wavelet space. It will further validate the effectiveness of using wavelet transform to capture local details. \n\n- The presentation for Section 2 Related Work can be improved. It would be better to have several paragraphs to introduce the related works from different perspectives.\n\n**References:**\n\n[1] Raissi, M., Perdikaris, P., & Karniadakis, G. E. (2019). Physics-informed neural networks: A deep learning framework for solving forward and inverse problems involving nonlinear partial differential equations. Journal of Computational Physics, 378, 686-707. \n\n[2] Pathak, J., Subramanian, S., Harrington, P., Raja, S., Chattopadhyay, A., Mardani, M., ... & Anandkumar, A. (2022). Fourcastnet: A global data-driven high-resolution weather model using adaptive fourier neural operators. arXiv preprint arXiv:2202.11214.\n\n[3] Takamoto, M., Praditia, T., Leiteritz, R., MacKinlay, D., Alesiani, F., Pflüger, D., & Niepert, M. (2022). Pdebench: An extensive benchmark for scientific machine learning. Advances in Neural Information Processing Systems, 35, 1596-1611.", "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": "", "extra_scores": {}}, {"review_id": "vFwJ1hravX", "reviewer": "Reviewer_rcJb", "summary": "The authors propose the Wavelet Diffusion Neural Operator , a data-driven framework designed for PDE simulation and control. WDNO uses diffusion-based generative modeling within the wavelet domain, which is claimed to capture entire trajectories of the system dynamics. To address the  challenge of generalizing across different resolutions multi-resolution training approach is introduced, enabling WDNO to work across varying data scales. This framework represents a tool for precise, adaptable modeling in complex, multi-scale systems.", "questions": "Line 194, “ we combine the use of both guidance methods” Where is the classifier based guidance method used?\n\nWhat is the impact of measurement noise? An ablation study for datasets with increasing measurement noise would be great.\n\nHow many trajectories are required? Can you show the impact of the available number of trajectories on the performance?\n\nWhat happens if the dynamical changes due to interventions/perturbation upon the system? Is the learned model still useful?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "The paper presents an interesting and rich contribution to the field and offers a new approach through the Wavelet Diffusion Neural Operator. This work is both technically rigorous and well-articulated, making the methodology and results accessible to the reader. The introduction of wavelets is notable, as it appears to deliver the claimed advantages in handling multi-scale features, abrupt changes, and long-term dependencies in PDE-based simulations. These claims are convincingly supported by a number of experiments, which effectively demonstrate the impact of the wavelet-based approach.", "weaknesses": "The structure of the work can be challenging to follow, partially due to the number of innovations introduced in quick succession, which may benefit from clearer organization or further segmentation for clarity. \n\nThe experimental evaluation, while well-executed and featuring comparisons to other methods, is somewhat limited in the number of test cases. Expanding the range of test cases would enhance the generalizability of the results:\n1. An analysis of the error distribution over time would be valuable, as it could reveal any systematic errors or trends in performance decay over extended simulations.\n2. Lack of a real-world test case. Please include a real-world test case to provide further validation of the approach in practical settings:  a dataset like PTB-XL should be addressable with the proposed method and could offer relevant, real-world complexity.\n3. Additional 2D cases across a broader range of complex spatial-temporal domains would strengthen the  work. [Iakolev et al, Latent neural ODEs with sparse bayesian multiple shooting, ICLR 2023, Lagemann et al, Invariance-based Learning of Latent Dynamics, ICLR 2024] are presenting results for 2D PDE-based test cases. It makes also sense to compare WDNO against these neural ODE based approaches for the simulation part. \n\nFigures 1 and 2 could be improved for clarity and relevance. \nIn my opinion Figure 1 lacks meaningful content and could be more effectively repurposed. The caption offers minimal context, leaving the figure’s purpose unclear, especially as the main text already conveys WDNO’s super-resolution capabilities. Reallocating this space to highlight other aspects of the method might provide more value.\nFigure 2 is informative and well-designed but would benefit from an expanded caption to guide readers through its details, despite space constraints. I would suggest a more general, high-level version of this figure in the main text and the current version of Figure 2 with a detailed caption in the Appendix. This would offer both an accessible overview and a richer, in-depth explanation in the Appendix.\nA detailed explanation tracing the steps from input data to output data of the pipeline and the transformations involved in each stage, with explicit dimensions based on one of the datasets would be great and improve the clarity. \n\nIn my opinion, eq. 9 and 10 appear not important in the main discussion and could be more appropriately relocated to the corresponding Appendix (E, F, G). Instead, a general paragraph on data preparation in the main text would provide readers with a clearer understanding of the processes involved, which is currently lacking.\n\nAdditionally, the work presented in Appendix C is noteworthy and should at least be mentioned briefly in the main text. \n\nPlease reorganize the Appendix to ensure that tables are placed immediately after they are referenced. As there is no page limit for the Appendix, there is no need to conserve space, which would increase the readability. \n\nPlease add an arrow indicating the direction of time on the time axis in Figure 4.", "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": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)