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
Primary area: applications to physical sciences (physics, chemistry, biology, etc.)
Keywords: PDE, physics, simulation, control, diffusion model, wavelet, abrupt changes, multi-resolution

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.

Decision: accept (poster)
Reviewer scores: 6, 8, 6, 6, 8

Reviews and discussion:
<thread>
creation date:
1730034386573

thread_type:
official_review

<thread_object>
creator: reviewer_rcjb

<content>
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.

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. 

The 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:
1. 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.
2. 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.
3. 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. 

Figures 1 and 2 could be improved for clarity and relevance. 
In 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.
Figure 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.
A 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. 

In 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.

Additionally, the work presented in Appendix C is noteworthy and should at least be mentioned briefly in the main text. 

Please 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. 

Please add an arrow indicating the direction of time on the time axis in Figure 4.

questions:
Line 194, “ we combine the use of both guidance methods” Where is the classifier based guidance method used?

What is the impact of measurement noise? An ablation study for datasets with increasing measurement noise would be great.

How many trajectories are required? Can you show the impact of the available number of trajectories on the performance?

What happens if the dynamical changes due to interventions/perturbation upon the system? Is the learned model still useful?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730427759965

thread_type:
official_review

<thread_object>
creator: reviewer_nmdb

<content>
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.

strengths:
- This paper aims to solve an interesting problem regarding the abrupt changes in spatiotemporal dynamics. 

- This paper has shown comprehensive details of the experimental setup, model details, and results. 

- 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. 

- 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?

- 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].

- 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. 

- 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.

**References:**

[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. 

[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.

[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.

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. 

- Some minor typos: 

    - On Page 2, the paragraph name “wavelet domain” should be “Wavelet domain”. 

    - On Page 2,  it seems not common to see “contribute the following”.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730459169674

thread_type:
official_review

<thread_object>
creator: reviewer_ivpn

<content>
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.

strengths:
1) Multi-resolution training

2) 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

2) Training time, inference speed, and memory usage are not compared for WDNO and baselines.

3) Lacks novelty as compared to previous works such as:

* Benchmarking Autoregressive Conditional Diffusion Models for Turbulent Flow Simulation (arXiv:2309.01745)

* DiffusionPDE: Generative PDE-Solving Under Partial Observation (arXiv:2406.17763)

questions:
1)  How does WDNO compare to baselines regarding training time, memory usage, and inference speed?

2) How was the basis wavelet chosen for the benchmarks?

3) Have the authors investigated using diffusion in the Fourier domain for comparison with WDNO, considering the computational efficiency of Fourier transforms?

4) How were hyperparameters for baselines chosen? Were they optimized fairly compared to WDNO?

5) The paper needs to elaborate on why WDNO outperforms DDPM, considering Parseval's identity, which states that information content remains constant during transformations.

6)  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)?

7) The paper should cite relevant related work and add as baselines, including:

* Benchmarking Autoregressive Conditional Diffusion Models for Turbulent Flow Simulation (arXiv:2309.01745)

* DiffusionPDE: Generative PDE-Solving Under Partial Observation (arXiv:2406.17763)
</content>
</thread_object>

</thread>
<thread>
creation date:
1730514197553

thread_type:
official_review

<thread_object>
creator: reviewer_34r3

<content>
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.

strengths:
* The main idea of combining WNO with diffusion is simple, novel and well justified.
* Every design choice is justified with multiple tests and ablation studies.
* 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.
* The generalization to more complicated tasks or boundary conditions might not be trivial.

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.
* 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.
* 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?

Final 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.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730644513533

thread_type:
official_review

<thread_object>
creator: reviewer_fypm

<content>
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.

strengths:
1. A wavelet-based neural operator that surpasses the performance of exising methods.
2. Scaling of the input system reduces the solution manifold that the model needs to learn.
3. 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.

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)?
</content>
</thread_object>

</thread>
<thread>
creation date:
1732158430794

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General Response

comment:
We thank the reviewers for the thorough reviews and constructive suggestions. We acknowledge the positive comments such as novel (Reviewer rcJb, 34r3), solid (Reviewer rcJb, 34r3, 34r3), technically rigorous (Reviewer rcJb), well-written (Reviewer rcJb, FYpM), and interesting (Reviewer rcJb). We also believe that our proposed WDNO method would significantly contribute to the community. 

Based on the reviewers' valuable feedback, we have conducted additional experiments and revised the manuscript, which hopefully resolve the reviewers' concerns. The major additional experiments and improvements are as follows:
- We add experiments on **1D advection equation** and **a challenging real-world dataset ERA5**. From the results, WDNO surpasses all the baselines, which verifies WDNO's strong modeling capability on both simple and complex dynamics. For more details, see responses to Reviewer NMDB, rcJb, and 34r3.
- To provide a more comprehensive comparison and analysis, we evaluate **other strong baselines**, including Transolver, CNO, MSVI, ACDM, and DiffusionPDE, on the 1D compressible Navier-Stokes equation. Additionally, we included **MAE and $L_\infty$ error** results for all methods. WDNO consistently achieves the best performance in both MSE and MAE. For $L_\infty$ error, while the results across methods are closer, WDNO still shows relatively better results. See responses to Reviewer iVpN and rcJb for more details.
- We provide **WNO**'s performance on **zero-shot super resolution**, from which we can observe that WNO does not generalize well to finer resolution. This further confirms the effectiveness of our proposed multi-resolution training. For details, please refer to responses to Reviewer iVpN.
- We **compare WDNO with Fourier transform**, including DDPM in Fourier domain and take FNO as the noise prediction model. The results show that DDPM in the Fourier domain achieves some improvement over the original DDPM but performs far worse than the wavelet transform, highlighting the superiority of wavelet transforms for complex systems with abrupt changes. Additionally, FNO as a noise prediction model performs poorly. For more details, see responses to Reviewer iVpN, NMDB, and FYpM.
- We add ablation studies including evaluating the **influence of measurement noise and number of training samples**. Also, based on previous results in the original submission, we additionally report **errors at different timesteps** and **training times** of all baselines. Results show that WDNO is robust to noise, good at long-term prediction and trains quickly. Please refer to responses to Reviewer rcJb, iVpN and FYpM for details.
- We expand **Related Work** by adding three new paragraphs, making it more comprehensive. More details can be found in responses to Reviewer NMDB.
- To **improve the clarity and structure of the paper**, we add explanatory statements and subheadings, adjust the layout of tables in Appendix, and reorganize Experiments. Also, we redrew **Figure 1** to include more information, provide a high-level overview, and better summarize our method. Additionally, we move the original **Figure 2** to the appendix, add more details, and align it with the training and inference process of the 1D Burgers' equation.

We now address each reviewer's concerns individually. Please see responses below each review.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732508159135

thread_type:
official_comment

<thread_object>
creator: area_chair_tkdk

<content>
title:
Reviewer responses

comment:
Dear Reviewers,

as the author-reviewer discussion end is approaching, I would strongly encourage you to read the authors' responses and acknowldge so, while also checking if your questions/concerns have been appropriately addressed.

This is a crucial step, as ensures that both reviewers and authors are on the same page, and it also helps us to put in perspective your recommendation.

Thank you again for your time and expertise

Best, 

AC
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper introduces the Wavelet Diffusion Neural Operator (WDNO), a method for simulating and controlling PDE-based physical systems. The authors claim that WDNO addresses limitations of standard diffusion models in handling abrupt changes in trajectory by incorporating wavelet transforms and multi-resolution training. While the proposed approach appears promising, particularly for fluid dynamics problems, the reviewers agree that further benchmarking, with more realistic examples, is necessary to fully assess its capabilities and generalizability.

Also, at least one reviewer raised concerns that the claim regarding long-time statistics was not appropriate, particularly given that recent publications do handle long-term statistics of dynamical systems. Along the same lines, there are some possible issues with the memory footprint for 3D problems with long relaxation times (which require long trajectories to capture the phenomenon of interest). Other reviewers were concerned about how the methodology would generalize to different boundary conditions and geometries, given the difficulty of traditional wavelet methods to handle them. In addition, some of the experimental data seem to be lacking details, and they are not compared against state-of-the-art methods using community metrics (ERA5).

Besides all these issues, the reviewers seem to appreciate the content of the paper, and the authors were responsive and thorough on their responses, so I recommend acceptance.