On conditional diffusion models for PDE simulations

We thank the reviewer for the thoughtful consideration of our work. We have taken the time to address all the points raised in the review.

Novelty - While we acknowledge that the techniques per-se are not new, we were the first to:

1. quantitatively evaluate the decomposition on physics data;

2. propose AR sampling instead of AAO sampling, which works much better in scarce observed data regimes and is key to performing forecasting;

3. use both reconstruction guidance and amortisation on different sets of variables simultaneously, which improved model performance and flexibility.

Flexibility vs performance - We agree with the reviewer that flexibility does not necessarily have to be orthogonal to performance. One of the goals of the paper is to quantitatively evaluate the current status of diffusion models for PDE simulation compared to other deep learning approaches by highlighting pros and cons of each such approach on the considered tasks (forecasting, DA, and a mix thereof). Unsurprisingly, we found that in the case of pure forecasting, our methods lag behind SOTA methods that have specifically been designed for forecasting. However, we argue that even so, the forecasting performance of our methods is close to other commonly used ML-based benchmarks, while benefiting from a much more flexible conditioning framework.

Universal amortised training for PDE-Refiner - Indeed, the “universal” amortized training and larger window size can be easily incorporated into PDE-Refiner. However, reconstruction guidance and data assimilation cannot be applied out-of-the-box, as PDE-Refiner first produces an estimate of the next state, and then refines this estimation using noise. This is in contrast to plain diffusion models that gradually build this estimation from pure noise, rather than from the neural operator prediction.

NFE - We provide a more detailed analysis of the NFE below in the KS and Kolmogorov cases.

For KS we simulate $640$ states. For the joint AR and amortised models, we use $128$ diffusion steps in the paper with a window $9$ model, and the conditioning scenario $P \mid C$ depends on the task. Overall, the computational cost is $(\frac{640-9}{P} + 1) \times 128$. In contrast, PDE-Refiner first produces the NO estimate, and then performs $3$ refining steps, leading to $4$ FE for each state. They always predict one state at a time, leading to NFE=$640 \times 4$, indeed leading to a lower computational cost than our approaches. However, subsequent analysis on the amortised model indicates that we can achieve similar performance with the universal amortised model by performing as few as 16 diffusion steps (please also see Fig. 7 in the attached PDF) (This is not true for the joint AR approach). This would lead to NFE = $(\frac{640-9}{P} + 1) \times 16$, which becomes comparable (and depending on $P$, even lower) than PDE-Refiner cost.

For Kolmogorov, the same analysis holds, but we only need to simulate $64$ states, and found that for the universal amortised model we can also achieve good performance with $16$ diffusion steps.

Thus, the universal amortised approach is comparable to PDE-Refiner in terms of computational time, while we acknowledge that the joint AR approach is more computationally intensive.

MSE-based U-Net for offline DA - In the online data assimilation case, we assume that in the first step we do not observe the initial condition (lines 346-348), but only some sparse observations within the $s$ first time steps. Hence, plain autoregressive U-Net trained with MSE target cannot be directly applied to this task, as it expects to have fully observed initial state.

However, we have included a comparison in the offline DA scenario, where we report the RMSD for the MSE-trained U-Net based on the initial conditions alone (please see Fig. 8 in attached PDF).

We hope we have successfully addressed your concerns.

We thank the reviewer for the positive consideration of our work and the useful feedback on it.

Turbulent dynamics + metrics - Thanks for your suggestion, we will add a mention about the turbulent dynamics of the data. For a discussion about metrics and spectrum, please refer to common answer.

Classical solvers vs ML-based + runtime - This is a great point and we will include a discussion in the updated manuscript. We believe ML-based approaches for PDE modelling are still fairly new and thorough comparisons between them and classical solvers will emerge. In the case of our paper we considered implementing such methods is out of the scope, since the focus is on ML-based solutions, with emphasis on diffusion models. But we agree that we should mention the possibility of implementing the classical approach.

A motivating application of our work is ML-based approaches for weather modelling. For a more detailed description of the weather forecasting setup please see reply to Reviewer D4vs. Works in this field, such as Aardwark Weather [1] show that there is potential for ML-based methods to reduce the cost of medium-range weather forecasting (which is currently based on classical solvers) by orders of magnitude, while still achieving competitive performance. There have also been works proposing to tackle the forecasting component using diffusion-based models [2], where the probabilistic treatment is appealing - it allows one to predict a range of probable weather scenarios, which can be crucial for decision making, disaster prevention, etc.

Regarding runtime, as pointed out in the common response, the universal amortised model can achieve a significant speed-up (8x faster) during sampling compared to what we report in Figure 3 in the main paper.

We will incorporate a subsection on classical solvers in the related work section, mentioning their applicability in the tasks considered. We will also add a discussion on how ML-based methods compare to classical solvers.

Theoretical justification - We acknowledge that classical methods benefit from a much more mathematically-rigorous framework. Unfortunately, guarantees for ML-based methods are rarely available (not only in the context of PDE modelling). Diffusion models have mathematical foundations (e.g. [3] or [4]), but it is still non-trivial to develop guarantees when taking into account the complex network dynamics within diffusion models. As such, this is probably more of a limitation of the field rather than of this paper in particular, and it is unclear whether such guarantees will ever be available for ML-based methods.

W1minor - Thanks for pointing this out. The dependence on $u$ should be included as well.

W2minor - L126: We omit putting (0) as time dependency and just use $x_{1:L}$. The initial Markov assumption is introduced on the sequence coming from the original data distribution.

$\mu_t$ - Yes, $\mu_t$ is a scalar value and then we divide the numerator component-wise.

Denoising score matching - Indeed, computing $\nabla_x \log p(x)$ is not easy as we do not know $p(x)$. The workaround is to use the denoising score matching loss (Eq. 3), where we use $\nabla_x(t) \log p(x(t)|x(0))$ instead of $\nabla_x \log p(x)$. For a justification of why this is the case, see Appendix of [5]. Then, the quantity of interest becomes $p(x(t)|x(0))$, which is referred to as the noising kernel in diffusion models. This can be derived from the governing noising SDE. In particular, if the drift and the diffusion coefficient are affine transformation of $x(0)$, then $p(x(t)|x(0))$ is Gaussian. In our case, we are using an Ornstein-Uhlenbeck process (see Eq. 1) as noising process and therefore this satisfies the above conditions making $p(x(t)|x(0))$ Gaussian.

Offline DA setup - The initial states are fully observed in the offline DA setting - we found this to be a convenient setting to analyse the performance of the model in a variety of settings, going from almost pure forecasting (the sparsest setting) to a very dense observation regime. We understand that in the context of classical solvers, where the treatment is deterministic, this just collapses to plain forecasting. However, sparse observations can potentially be helpful in reducing accumulating approximation errors (resulting from e.g. space/time discretisation). Nevertheless, this was not the focus of the experiment, but rather to see whether our approaches can “correct” their forecasting mistakes (i.e. make use of the observations to bring the predictions closer to ground truth), and alleviate error accumulation which is an inherent issue of autoregressive approaches.

In the online DA experiment, the initial states were not fully observed, and thus, could not be tackled by a classical numerical solver. This is the task we believe best reflects the advantages of our approaches, and acted as the main motivation when developing them. The previous investigations on forecasting and offline DA were useful investigations to understand the workings of the models in a variety of conditioning scenarios, but we consider the online DA task to be the one with most relevance to upstream applications, such as in weather modeling.

References

[1] Vaughan, A. et al. (2024). Aardvark weather: end-to-end data-driven weather forecasting.

[2] Price I. et al. Gencast: Diffusion-based ensemble forecasting for medium-range weather, 2024.

[3] De Bortoli, V. Convergence of denoising diffusion models under the manifold hypothesis. TMLR (2022).

[4] Benton, J. et al. (2024). Nearly d-linear convergence bounds for diffusion models via stochastic localization. ICLR

[5] Vincent P. A connection between score matching and denoising autoencoders, 2010.

Thank you for your review and your comments, we aim to address them below.

Evaluation metrics - Thank you for your suggestions, when we chose the evaluation metrics we followed [1], but we acknowledge that some other metrics can be used to assess the model. We provide some examples of energy spectra for the trajectories in the attached PDF (Fig. 3 and Fig. 9), showing that they are in agreement with the ground-truth spectra for low frequencies, and the differences mostly come from high frequency components, as also noted in [1].

Algorithm/Pseudocode - Thank you for the suggestions, we agree that this would make it easier to compare the different methods and will include them in the updated manuscript should the paper be accepted.

Amortised Kolmogorov - Thank you for your suggestion. These models are still training at the moment because the cluster is currently very busy and they have been spending some time in the queue. We hope to be able to provide the results by the end of the discussion period. But this is definitely something we will include in the final version of the manuscript.

Ablation window size - We launched models for window 3 and 7 for Kolmogorov for both the joint and amortised models. However, please note that because of the time constraints, their hyperparameters are not tuned and from the training curves, we can clearly see that the window 7 model would benefit from more training. We provide a few example results in terms of high correlation time below, as the other scenarios are still currently being queued on our cluster:

• Amortised $1 \mid 2$ (window $3$) - $9.09 \pm 0.50$s
• Amortised $1 \mid 4$ (window $5$) - $8.85 \pm 0.53$s - from paper
• Amortised $1 \mid 6$ (window $7$) - $8.63 \pm 0.50$s
• Amortised $3 \mid 4$ (window $7$) - $8.01 \pm 0.44$s
• Joint $1 \mid 2$ (window $3$) - $7.63 \pm 0.41$s
• Joint $1 \mid 4$ (window $3$) - $8.16 \pm 0.62$s - from paper
• Joint $1 \mid 6$ (window $7$) - $7.73 \pm 0.50$s

For the joint model, we can see that the performance is comparable between all window sizes, although the best one is achieved with the window $5$ model. We believe the window $7$ one does not outperform it because of convergence issues.

For the universal amortised, we see a decreasing trend in performance with increasing window size. However, the performance of all models with $P=1$ is within error (and hence, comparable). We believe that this result might be influenced by hyperparameter tuning.

Architecture - No, the MSE-trained U-Net and the modern U-Net architecture are not the same. We have experimented with several architectures and found that the architecture inspired by [2] (SDA) was the most stable one in the context of joint prediction. The MSE-trained U-Net is inspired by [1] and we found it better suited for amortised-style of models. As such, we optimised the architecture for each model and chose the optimal one, meaning that we used the SDA architecture for the joint and for the amortised models (for a one-to-one comparison), and the PDE-Refiner architecture for the baselines used in forecasting.

We will include a more detailed description of the MSE-trained U-Net in the Appendix in the updated version of the manuscript, thanks again for pointing this out.

Non-universal vs. universal - You are correct and we also thought this finding is interesting. We believe the reason why the universal amortised models outperforms the plain amortised models (where we need to train a model for each $P \mid C$ combination) is because the training task becomes more complex, which encourages generalisation. In effect, the universal amortised model also amortises over the predictive horizon $P$. We note, however, that this comes with the downside that in the KS case we had to train the universal amortised model for longer than the plain amortised models (which would start overfitting if they had been trained for the same amount of epochs as the universal amortised model).

Universal amortised for non-diffusion models - This is indeed possible, as the framework we propose is not specific to diffusion-based models. We find your suggestion very interesting, but potentially not directly relevant to the topic of the paper, which aims to study diffusion-based models for PDE modelling.

References

[1] Lippe, P. et al. (2023). PDERefiner: Achieving Accurate Long Rollouts with Neural PDE Solvers. In Thirty-Seventh Conference on Neural Information Processing Systems.

[2] Rozet, F. and Louppe, G. (2023). Score-based Data Assimilation. In Thirty-Seventh Conference on Neural Information Processing Systems.

We thank the reviewer for the time taken to review our paper and the positive consideration of it. We address the comments below.

Alternative metrics

Thanks for suggesting alternative metrics, such as the energy spectrum, as well as providing more qualitative examples. We agree that these are indeed helpful in understanding the quality of the generated PDE trajectories, Please find in the attached PDF:

• Example energy spectra (Fig. 3 - KS and Fig. 9 - Kolmogorov).
• More examples of trajectories from each of the evaluated models on forecasting (e.g. Fig. 1 - KS and Fig. 2 - Kolmogorov, besides the examples provided in Fig. 24 for KS and Fig. 25 for Kolmogorov in the Appendix)
• Example of trajectories for DA offline, alongside the masks used and interpolation results, which we included as a baseline for offline DA (Fig. 5 - KS and Fig. 6 - Kolmgorov)

We will make sure to include them in the final version of the manuscript.

For the quantitative evaluation, we decided to use high correlation time for an easier comparison with the results provided in [1]. RMSD/MSE are metrics that we frequently encounter in the literature (e.g. [1], [2]). While we acknowledge they are not perfect in evaluating PDE dynamics, we believe the problem of finding the right metric is a research question on its own and was not the focus of this investigation.

Training times - We found that for the KS dataset, the joint model requires less training (4k epochs) compared to the universal amortized model (15k) steps (please also see Table 3 p.21 for more training details). However, both models, universal amortised and joint, were trained for the same amount of epochs for Kolmogorov. All considered models take from a few days to a maximum of one week to train depending on the dedicated GPU hardware. Since we relied on a GPU cluster with different machines on it, it is challenging for us to report a more specific comparison about the training times.

Standard diffusion validation loss (measured by how well the model predicts the added noise) is not comparable across joint and amortised models, as the amortised model predicts the conditional score, so validation loss is often smaller than the one for the joint model.

Long-term trajectory stability - The KS dataset is already analysing this behavior, since training examples are shorter than validation and test (140 states / 28s for train vs. 640 states / 128s for validation and test). The evolution of MSE error can be found in Fig. 23, where the error starts to linearly increase after around 50 seconds.

To analyse even longer-term behaviour (e.g. does the model still produce plausible states), we generated a trajectory of 2000 steps (400s), clearly showing that the model produces plausible predictions (see Fig. 4 on the attached PDF). Moreover, we computed the energy spectra of a states close to the initial condition, a state in the middle of the trajectory, and of the last state. We observe that the energy spectra of these three states do not show different behaviour, indicating that our model does not start generating physically unplausible states if we roll it out for very long times.

Fourier modes - Please see Fig. 3 and Fig. 9 on the attached PDF for example energy spectra.

Physics match - We believe the energy spectrum + long-term behaviour already address this question, but we are open to investigate other ways of assessing how well the generated trajectories match the physics.

References

[1] Lippe, P. et al. (2023). PDERefiner: Achieving Accurate Long Rollouts with Neural PDE Solvers. In Thirty-Seventh Conference on Neural Information Processing Systems.

[2] Georg Kohl et al. Turbulent flow simulation using autoregressive conditional diffusion models. 2023.

Thank you for addressing my concerns regarding the metrics, and for providing them in the PDF. I think that their addition to the appendix of the paper strengthens it overall.

The weakness regarding inference cost persists (I read the rebuttal to D3JV). I see the analysis and discussion regarding PDE-Refiner, but that just means that both methods are expensive. However, I acknowledge that this is an initial exploration and therefore cannot be expected to solve all problems.

I see from the spectral analysis that PDE-Refiner has a more accurate spectra, and that even the U-Net and FNO perform similarly or better than the proposed method. I agree with D3JV that the denoising process could and should be used to achieve effects similar to PDE-Refiner. While I also agree with you in your rebuttal to D3JV when you point out that the non-forecasting tasks are not trivial to incorporate into PDE-Refiner's framework, I think it is a necessary improvement to the proposed method, as it can at least amortize some of the inference cost.

All in all, you have addressed some of my concerns, and I acknowledge that those that remain unresolved can (and should) be the subject of future work. I am happy to raise my score from 6 to 7 and recommend that this paper be accepted.

We would like to thank the reviewer for their positive consideration. We address the points raised in the Weaknesses and Questions sections below.

Forecasting results

It is true that in the plain forecasting task, our proposed models do not achieve SOTA performance, as we also highlight in the limitations section of our work, with the main reason being that we choose to trade off forecasting performance with conditioning flexibility. Nevertheless, the paper brings two main contributions in the context of forecasting:

1. With the proposed autoregressive (AR) sampling approach, we show that diffusion-based models can achieve results of comparable quality with other benchmarks (e.g., MSE-trained U-Net, MSE-trained FNO), while preserving flexibility. This is in contrast to the all-at-once (AAO) sampling approach proposed by [1] which predicts trajectories that diverge from the ground state early on.

2. We propose a new training strategy for amortised models (which can also be applied to non-diffusion based models), which achieves stable performance for a wide range of conditioning scenarios. This is in contrast to other amortised approaches from the literature which report a decrease in performance with increasing history length [2].

Scenarios where flexibility is important

A setting where flexibility is important is the mix of forecasting and data assimilation (DA) considered in the paper, especially in the online DA case. More specifically, not only do we want to produce forecasts based on some estimated initial states, but we also gather some observations in an online manner that can help the model correct its previous forecasts to better match the observed data.

One important application of the above setting is weather prediction. The goal is to produce a forecast for a certain horizon (e.g. days to weeks for medium-range prediction). At the same time, we continuously acquire new observations from a variety of sources, such as weather stations, radiosondes, wind profilers, and satellites [3]. Thus, we would like to update our forecasts whenever these new observations arrive to better reflect reality. This is currently done in two stages, where the forecasting system first generates the forecast, and then the data assimilation system incorporates the observations. However, they are both very costly systems [4]. Most of the current ML-based approaches for weather modeling only tackle the forecasting component (e.g. PanguWeather [5], GraphCast [6]). A more preferable approach would be to perform both tasks at once, being able to flexibly condition on both initial states, as well as any incoming observations. And this is exactly what we are proposing in this paper, with the proposed online DA task being the one that most closely resembles this scenario. However, we do acknowledge that it is still in a toy-ish setting and a natural future direction involves investigating the potential of the technique to scale.

Finally, another setting where flexibility is important is where the nature of the conditioning information might change over time, implying that during training one does not have access to examples of the conditioning information they might encounter at test time. In our case, we address this by conditioning the models using reconstruction guidance. In terms of practical applications, we can imagine an example in the weather modelling setting. For example, a new sensor might be added to the sources of information. Due to the flexibility of our models, we would not have to re-train the model to account for this new instrument.

Experimental settings vs. AAO

Indeed, in contrast to our work, [1] does not provide any quantitative results (on the Kolmogorov flow).

Their paper performed the following experiments:

• Observing the velocity field every four steps, coarsened to a resolution 8 $\times$ 8 and perturbed by Gaussian noise ($\Sigma_y = 0.1 \mathbf{I}$).
• Observing a regularly sampled sparse velocity field at every time step, with factor $n$ ($n$ implies 1 observation for each $n \times n$ region). They use $n$=2,4,8,16. $n$=16 implies 0.3% observed data.

They did not use the 1D KS dataset, only 2D Kolmogorov flow data. Moreover, in our work, we sample observations at random locations, and at each time step. The fraction of observed variables varies between 0.1% to 30%.

In our experiments, we show that the proposed method is able to perform at least on par with AAO across different settings, and it could, thus, be useful if we expect the fraction of observed variables to vary.

References

[1] Rozet, F. and Louppe, G. (2023). Score-based Data Assimilation. In Thirty-Seventh Conference on Neural Information Processing Systems.

[2] Lippe, P. et al. (2023). PDERefiner: Achieving Accurate Long Rollouts with Neural PDE Solvers. In Thirty-Seventh Conference on Neural Information Processing Systems.

[3] Keeley, S. (2022). Observations. https://www.ecmwf.int/en/research/dataassimilation/observations.

[4] Bauer, P. et al. (2020). The ECMWF Scalability Programme: Progress and Plans.

[5] Bi, K. et al. (2023). Accurate medium-range global weather forecasting with 3D neural networks. Nature, 619(7970):533–538.

[6] Lam, R. et al. (2023). GraphCast: Learning skillful medium-range global weather forecasting.