Wavelet Diffusion Neural Operator

Thank you for recognizing our work as interesting, technically rigorous, well-articulated, and convincing. Below are our responses to your comments.

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

• Thanks for the suggestion. To make the structure of the paper clearer and easier to follow, at the beginning of the Method section in the revised manuscript, we add a specific introduction to each subsection. Additionally, we further organize the content in Section 3.1 and 3.2 of Method by adding more detailed subheadings.

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

• In the ablation study (Figure 5a) in the original submission, we have already provided the errors of U-Net and WDNO at different timesteps, pointing out that the MSE of U-Net increases much faster than that of WDNO. To further support this, we add the errors of other methods at different timesteps. The results in the table below show that WDNO exhibits the slowest error growth, confirming its ability to capture long-term dependencies. We have updated Figure 5a and analysis in Section 4.7 of the updated manuscript accordingly. | Time step | 1 | 11 | 21 | 31 | | --- | --- | --- | --- | --- | | FNO | 0.0019 | 0.0094 | 0.0076 | 0.0034 | | OFormer | 0.0113 | 0.0226 | 0.0311 | 0.1138 | | MWT | 0.0110 | 0.0242 | 0.0209 | 0.0706 | | WNO | 0.1060 | 0.1252 | 0.1158 | 0.0826 | | U-Net | 0.0038 | 0.0053 | 0.0099 | 0.0306 | | DDPM | 0.0018 | 0.0034 | 0.0133 | 0.0173 | | WDNO | 0.0011 | 0.0012 | 0.0031 | 0.0049 |
• To better demonstrate the capabilities of WDNO, we add experiments on the challenging real-world weather prediction dataset ERA5 [1], which has a higher dimensionality compared to PTB-XL. Compared to the standard task of predicting the next 12 hours based on 12 hours of initial conditions [2], we increase the difficulty by extending the prediction horizon to 20 hours. The experimental results in the table below show that WDNO still achieves the best performance. Notably, its relative $L_2$ error for the 20-hour prediction is as low as 0.0161, highlighting its superior capabilities. Here we experiment with different parameters for WNO, but all configurations fail to converge. In the updated manuscript, we incorporate these results into Section 4.5 and Table 1. | | MSE | | --- | --- | | WNO | - | | MWT | 21.85750 | | OFormer | 18.26230 | | FNO | 14.38638 | | U-Net | 15.51342 | | DDPM | 15.21103 | | WDNO | 12.83291 |
• First, in the revised manuscript, we have cited these two references to the first paragraph of Related Work. Second, we compare with MSVI [3] to compare WDNO with neural ODE based approaches. The results of WDNO on 1D compressible Navier-Stokes equation are presented below. It is obvious that WDNO has superior performance. The results have been added to Table 5 in the updated manuscript. | | MSE | MAE | $L_{\infty}$ error | | --- | --- | --- | --- | | MSVI | 1.7063 | 0.6047 | 17.0386 | | WDNO | 0.2195 | 0.1049 | 13.0626 |

[1] Kalnay E, Kanamitsu M, Kistler R, et al. The NCEP/NCAR 40-year reanalysis project[M]. Renewable energy. Routledge, 2018: Vol1_146-Vol1_194.

[2] Tan C, Li S, Gao Z, et al. Openstl: A comprehensive benchmark of spatio-temporal predictive learning[J]. Advances in Neural Information Processing Systems, 2023, 36: 69819-69831.

[3] Iakovlev V, Yildiz C, Heinonen M, et al. Latent Neural ODEs with Sparse Bayesian Multiple Shooting[C]. The Eleventh International Conference on Learning Representations.

Thanks for the detailed comments. Here are responses to your concerns.

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

• Achieving "good model performance" is inherently challenging, which is why special design is necessary. Learning complex dynamics, such as systems with abrupt changes, is particularly difficult. This is evident in two aspects:
  • First, existing literature indicates that deep neural networks tend to prioritize low-frequency signals, leading to poor learning of high-frequency signals like abrupt changes [1].
  • Second, as shown in our experiments through tables and many visualizations in the original submission, other models, including DDPM, struggle to achieve "good model performance." Figures 2a, 6, 7, and 9 clearly illustrate how WDNO outperforms DDPM in capturing signals with abrupt changes.
• Besides, to address the fundamental requirement for neural PDE solvers to generalize across different resolutions, we also introduce the multi-resolution training seamlessly integrated with wavelets, demonstrating superior performance compared to previous super-resolution baselines in experiments. Additionally, we unify simulation and control tasks within one single framework, highlighting the utility of diffusion models in the PDE domain.

[1] Xu, Zhi-Qin John, et al. "Frequency principle: Fourier analysis sheds light on deep neural networks." arXiv preprint arXiv:1901.06523.

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

• Firstly, both in simulation and control, studies have highlighted that diffusion models can help capture long-range dependencies. In simulation, the noise-learning mechanism enhances temporal stability, a conclusion widely validated in fluid prediction [1], long-term human motion generation [2] and weather forecasting [3]. Some works even propose, inspired by DDPM, introducing an adapted Gaussian denoising step for stable rollouts in PDE simulation [4]. In control, many RL-related studies [5, 6, 7] suggest using diffusion models to model entire trajectories from a global perspective, enabling trajectory-level optimization. Thus, this approach facilitates non-greedy planning and achieves near-globally optimal plans more effectively.
• Secondly, diffusion models are widely recognized for their strong modeling capabilities in complex and high-dimensional systems, including 3D fluids [8], weather [3], videos [9], 3D shape generation [10] and real-world robots control [5], among others. This advantage stems from its denoising and noising mechanism, which not only transforms modeling the original data into predicting simpler noise, but also decomposes the task of modeling a complex system into multiple simpler subtasks through multi-step denoising.
• Thanks for the suggestion. In the updated manuscript, we have included the references in the second paragraph of Introduction and the last paragraph of Related Work.

[1] Kohl G, et al. Benchmarking autoregressive conditional diffusion models for turbulent flow simulation[C]. ICML AI for Science Workshop, 2024.

[2] Yang Z, et al. Synthesizing long-term human motions with diffusion models via coherent sampling[C]. Proceedings of the 31st ACM International Conference on Multimedia, 2023.

[3] Rühling Cachay S, et al. Dyffusion: A dynamics-informed diffusion model for spatiotemporal forecasting[J]. Advances in Neural Information Processing Systems, 2024, 36.

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

[5] Janner M, Tenenbaum J, et al. Planning with Diffusion for Flexible Behavior Synthesis[C]. International Conference on Machine Learning. PMLR, 2022.

[6] Chi C, et al. Diffusion policy: Visuomotor policy learning via action diffusion[J]. The International Journal of Robotics Research, 2023.

[7] Wei, Long, et al. A Generative Approach to Control Complex Physical Systems.[C] Advances in Neural Information Processing Systems, 2024.

[8] Li T, Biferale L, et al. Synthetic Lagrangian turbulence by generative diffusion models[J]. Nature Machine Intelligence, 2024: 1-11.

[9] Ho J, et al. Video diffusion models[J]. Advances in Neural Information Processing Systems, 2022.

[10] Vahdat A, et al. Lion: Latent point diffusion models for 3d shape generation[J]. Advances in Neural Information Processing Systems, 2022.

We thank the reviewer for the comments and questions. Below we address the reviewer's concerns.

Q1: The evaluation omits comparisons with state-of-the-art operators like Transolver, GNOT, LSM, DPOT, and CNO etc.

• Thanks for your suggestions. We previously selected baselines that are diverse, high-performing, and highly relevant. These include methods based on wavelet transforms (WNO, MWT), Fourier transforms (FNO), Transformer architectures (OFormer), convolutional networks (CNN, U-Net), and diffusion models (DDPM).
• Based on your valuable comments, first, we have cited these methods in Related Work in the revised manuscript. Second, we have added the mentioned baselines of Transolver and CNO to further evaluate the effectiveness of our method. These are chosen because Transolver has been shown to outperform GNOT and LSM across various experiments in its paper [1], while DPOT, a pre-trained foundation model, requires 10 timesteps to predict the next frame, which is not applicable to our setting where 1 timestep is used to predict the entire trajectory. In experiments on the 1D compressible Navier-Stokes equation, WDNO continues to show the best performance. We have incorporated the results in Table 5 in the updated manuscript. | | MSE | MAE | $L_\infty$ error | | --- | --- | --- | --- | | Transolver | 4.99843 | 0.40253 | 4.87284 | | CNO | 0.3987 | 0.2765 | 9.9169 | | ACDM | 4.6574 | 0.8946 | 60.9370 | | DiffusionPDE | 5.5936 | 0.9792 | 16.0515 | | WNO | 6.5428 | 1.1921 | 21.3860 | | MWT | 1.3830 | 0.5196 | 11.3677 | | OFormer | 0.6227 | 0.4006 | 30.90186 | | FNO | 0.2575 | 0.1985 | 11.1495 | | CNN | 12.4966 | 1.2111 | 17.6116 | | DDPM | 5.5228 | 0.9795 | 16.0532 | | WDNO | 0.2195 | 0.1049 | 13.0626 |

[1] Wu H, Luo H, Wang H, et al. Transolver: A Fast Transformer Solver for PDEs on General Geometries[C]. Forty-first International Conference on Machine Learning.

Q2:

1. 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)
2. Lacks novelty as compared to these works.

• Firstly, our work differs significantly from these two works. The first work, ACDM, uses a diffusion model for autoregressive system state prediction, which fundamentally differs from our approach of modeling the entire trajectory distribution simultaneously. As we have elaborated in our original submission, this global modeling allows the model to better capture long-range dependencies, leading to higher accuracy in prediction tasks and enabling globally optimal planning in control tasks. Moreover, our work introduces wavelet transforms to capture abrupt changes, a multi-resolution framework for super-resolution, and demonstrates how the entire framework seamlessly applies to control tasks with superior performance.
• As for the second work, DiffusionPDE focuses on the scenarios where there lacks full knowledge of the scene. Its algorithmic innovations based on diffusion models are specifically designed for the partially observed scenario, which means this work focuses on a different aspect, and its innovations do not overlap with ours.
• Finally, to verify our statement, we conduct the comparisons with these two methods. As shown in the above table and the Table 5 in the updated manuscript, the results on 1D compressible Navier-Stokes equation simulation demonstrates that WDNO outperforms them. We also have included them in the discussion of Related Work.

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

• First, we emphasize that Fourier transforms are not more computationally efficient than wavelet transforms. We note that the wavelet transform we adopt is parallelizable and can run on GPUs. To demonstrate this, we provide a speed comparison in the table below, showing the total time for Fourier and wavelet transforms on the training set (9000 samples) of the 1D compressible Navier-Stokes equation. The Fourier transform is implemented using PyTorch's 2D Fast Fourier Transform function, while the wavelet transform is implemented via PyTorch Wavelets [1] as already mentioned in the Appendix. Both times are recorded on an A100 GPU with a batch size of 2000. It can be observed that the wavelet transform requires even less time than the Fourier transform. We have incorporated the results in Table 4 of Appendix A in the updated version. | | Wavelet transform | Fourier transform | |---|---|---| | Time (s) | 1.0171 | 1.3810 |
• Following your suggestion, we conduct experiments in the Fourier domain. The implementation strictly follows WDNO, except for replacing the wavelet transform with Fourier transform. The MSE and MAE results on the 1D compressible Navier-Stokes equation are shown in the table below. We see that while the Fourier transform also provides some improvement over DDPM, its performance is significantly inferior to that of the wavelet transform. As already mentioned in the Section 3.1 in the original submission, this is because wavelet transforms inherently decompose information into low-frequency components and high-frequency details across different directions, making them more effective for learning complex system dynamics, such as those with abrupt changes. These results have been incorporated in Figure 5c and Ablation Study of the updated manuscript. | | MSE | MAE | | --- | --- | --- | | DDPM | 5.5228 | 0.9795 | | In the Fourier domain | 3.0258 | 0.8498 | | WDNO | 0.2195 | 0.1049 |

[1] Cotter, Fergal. "Uses of Complex Wavelets in Deep Convolutional Neural Networks". Apollo - University of Cambridge Repository, 2019, doi:10.17863/CAM.53748.

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

• Firstly, for different models, the total input information remains constant, yet their performance varies significantly. This discrepancy arises from differences in model design, which impact the effectiveness of learning.
• Secondly, existing studies have shown that deep neural networks tend to prioritize learning low-frequency information [1]. As a result, they naturally struggle to capture dynamics with abrupt changes, which are common in PDEs, necessitating special designs to address this limitation. To tackle this, as stated in the penultimate paragraph of Section 1 and the first two paragraphs of Section 3.1 in the original submission, wavelet bases are localized in both space-time and frequency domains, decomposing the system state into low-frequency and high-frequency information. This naturally helps models capture the high-frequency components that are typically harder to learn.
• Moreover, to analyze why WDNO outperforms DDPM, we have provided extensive visualizations of the prediction results for DDPM and WDNO on two equations in Figures 2a, 6, 7, and 9 in the original submission. Our analysis clearly demonstrates that WDNO achieves significant improvements in capturing dynamics with abrupt changes, which are otherwise difficult to learn.

[1] Xu, Zhi-Qin John, et al. "Frequency principle: Fourier analysis sheds light on deep neural networks." arXiv preprint arXiv:1901.06523 (2019).

Thank you for your response and constructive discussion. We are glad that we have addressed most of your questions. Below, we provide further clarification on the points you raised.

• ACDM is an autoregressive model, which means both training and inference are performed sequentially.
  • The model predicts the next $k$ steps $u_{[k,2k-1]}$ based on the previous $k$ steps $u_{[0,k-1]}$. Then, using $u_{[k,2k-1]}$, it predicts $u_{[2k,3k-1]}$, and this process continues iteratively until the entire trajectory $u_{[0,T-1]}$ of $T$ steps is generated. In contrast, our proposed method predicts the entire trajectory $u_{[0,T-1]}$ directly in a single inference step based on the given $k$ initial steps.
  • For example, in our 1D experiment, where $k=1$ and $T=80$, the model needs to perform inference 80 times to predict the entire trajectory, while our proposed method predicts the entire trajectory in a single inference step, significantly reducing computational overhead.
• We note that we have already made the corresponding modifications to the revised manuscript in response to your questions, as mentioned in our previous replies. Regarding the new clarification about ACDM, we have incorporated this content into the 'Diffusion models' part of Related Work. We have also uploaded the revised manuscript.

If you have any further points to discuss, please feel free to reach out at any time.

Thank you for complimenting our paper as novel, solid, and high-performing. We have summarized and categorized your questions. Here are our responses to your comments.

Q1:

1. The method is constrained to static uniform grid data.
2. 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.

• A regular grid refers to a uniform grid that is structured in a square or rectangular arrangement. Using a regular grid means that our architecture cannot learn data with irregular structures like graphs.
• Currently, the wavelet transform we use does make it challenging to handle data on irregular grids. However, as noted in Section 5 (Limitation and Future Work), our method could potentially be extended to irregular data. There are several possible approaches, including using geometric wavelets [1] combined with diffusion models designed for graph structures [2], or projecting data from irregular grids onto regular uniform grids [3, 4], among others.
• In Section 5 in the updated manuscript, we have added explanation and offered approaches to generalize our method to more complex geometries.

[1] Xu B, Shen H, Cao Q, et al. Graph Wavelet Neural Network[C]. International Conference on Learning Representations. 2018.

[2] Vignac C, Krawczuk I, Siraudin A, et al. DiGress: Discrete Denoising diffusion for graph generation[C]. The Eleventh International Conference on Learning Representations.

[3] Li Z, Huang D Z, Liu B, et al. Fourier neural operator with learned deformations for pdes on general geometries[J]. Journal of Machine Learning Research, 2023, 24(388): 1-26.

[4] Gao H, Sun L, Wang J X. PhyGeoNet: Physics-informed geometry-adaptive convolutional neural networks for solving parameterized steady-state PDEs on irregular domain[J]. Journal of Computational Physics, 2021, 428: 110079.

Q2:

1. The method is constrained to low-dimensional toy examples.
2. The generalization to more complicated tasks might not be trivial.

• To address your concern, we add a challenging real-world weather prediction dataset, ERA5 [1], to our experiments. Additionally, we select a more challenging task: predicting 20 hours ahead based on 12 hours of observations, instead of the standard 12-hour prediction task [2]. The results in the table show that WDNO still achieves the best performance, with a relative $L_2$ loss as low as 0.0161. And we have experimented with different parameters for WNO, but all configurations fail to converge. These results are included in Section 4.5 and Table 1 in the updated version. | | MSE | | --- | --- | | WNO | - | | MWT | 21.85750 | | OFormer | 18.26230 | | FNO | 14.38638 | | U-Net | 15.51342 | | DDPM | 15.21103 | | WDNO | 12.83291 |

[1] Kalnay E, Kanamitsu M, Kistler R, et al. The NCEP/NCAR 40-year reanalysis project[M]. Renewable energy. Routledge, 2018: Vol1_146-Vol1_194.

[2] Tan C, Li S, Gao Z, et al. Openstl: A comprehensive benchmark of spatio-temporal predictive learning[J]. Advances in Neural Information Processing Systems, 2023, 36: 69819-69831.

Thank you for your recognition of our paper. Below are our responses to your concerns.

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

• We agree that adding artificial noise can help achieve stable long-term predictions for the model. And we believe this is one of the key reasons why diffusion models are well-suited for long-term predictions.
  • This is because the noise-learning mechanism in the training process naturally equips the model with a stronger ability to handle noise. For instance, as highlighted in [1], inspired by diffusion models, introducing an adapted Gaussian denoising step can enhance the temporal stability of PDE simulations, enabling accurate long roll-outs. Moreover, we point out that the noise addition in diffusion models is global and more comprehensive [2], which may lead to even better temporal stability compared to such specially designed approaches.
  • In addition, many studies have applied diffusion models in specific domains and concluded that they are effective for long roll-out inference. These applications span various fields, including fluid prediction [3], weather forecasting [4], decision-making [5], human motion generation[6] and so on.
  • We have added the mentioned references into Introduction and Related Work in the updated manuscript.
• Regarding the timesteps chosen in our experiments,
  • We use 80 timesteps for both 1D experiments. Notably, in the 1D Burgers experiment, we extend the timesteps from 10 to 80 compared to previous studies [7, 8]. For the super-resolution task, the total timesteps reach $\mathbf{80 \times 2^3=640}$.
  • In the 2D experiment, the dataset contains 32 timesteps, but the actual physical simulation spans 128 timesteps. Another reason for selecting this timestep is that it corresponds to the point where the smoke almost disappears from the observation region.
  • Our proposed method has the ability to perform longer tasks. However, for 2D problems, the presence of two spatial dimensions limits long-term predictions due to memory constraints.
• In Ablation Study in the original submission, we have already compared the prediction errors of WDNO and U-Net at different timesteps. To further illustrate this, we add results of other models. As shown in the table below, WDNO exhibits the slowest error growth over time, demonstrating its suitability for long-term predictions. We have updated Figure 5a and the analysis in Ablation Study in the updated manuscript. | Time step | 1 | 11 | 21 | 31 | | --- | --- | --- | --- | --- | | FNO | 0.0019 | 0.0094 | 0.0076 | 0.0034 | | OFormer | 0.0113 | 0.0226 | 0.0311 | 0.1138 | | MWT | 0.0110 | 0.0242 | 0.0209 | 0.0706 | | WNO | 0.1060 | 0.1252 | 0.1158 | 0.0826 | | U-Net | 0.0038 | 0.0053 | 0.0099 | 0.0306 | | DDPM | 0.0018 | 0.0034 | 0.0133 | 0.0173 | | WDNO | 0.0011 | 0.0012 | 0.0031 | 0.0049 |