On conditional diffusion models for PDE simulations

Summary

The paper studies forecastiong temporal dynamics associated with forward PDE problems using diffusion models. Different strategies for approximating the score function of trajectories and how to parameterize them are proposed, ranging from sampling the whole trajectory all at once to sample it step by step autoregressively. In addition, the authors also propose to incorporate information from observation to refine and modulate the prediction process. The proposed methods are tested on two types of PDEs under different settings involving both forecasting and data assimilation.

Strengths

1. The writing of the paper is pretty clear and self-contained. 2. Comprehensive empirical studies of different components and hyperparameters are carried out to support the claims. 3. The proposed universal amortized model, while being straightforward, is quite effective and flexible. It resolves several issues existed in common amortized models, including fixed context windows and performance degradation when given a longer context. 4. Applying reconstruction guidance to the diffusion process can be potentially useful for building hybrid forecasting system, e.g. combing data-driven forecasting model with some physics prior.

Weaknesses

1. Many techniques in the proposed framework are not new. Using the local segment's score to approximate and decompose the score of global dynamics is introduced in Score-based Data Assimilation. Reconstruction guidance has been used in several previous works that applied diffusion model to inverse problems (e.g. Manifold Constrained Gradient). 2. The empirical performance of the proposed method does not outperform many data-driven forecasting models. Nonetheless, as authors have stated, the objective of the proposed framework is not only to do pure data-driven forecasting but do flexible and versatile sampling tasks such as online data assimilation. However, I think flexibility does not necessarily have to be orthogonal to the performance (see question section)

Questions

1. In its current form, the best model variant of the proposed method still falls behind PDE-Refiner. Given PDE-Refiner is also closely related to denoising diffusion model (with velocity prediction, starting from last time step instead of pure noise), it seems that it is also possible to apply many of the proposed techniques to PDE-Refiner. More specifically, according to my understanding, PDE-Refiner uses the amortized architecture (P=1 and C=1) and it is autoregressive, while the paper has showcased the proposed universal amortized method seems superior (can predict more time steps at once, flexible time window and can use longer context). Can the authors provide more comments related to this point? 2. The major limitation of applying vanilla diffusion model to simulate PDE autoregressively is that it increases the NFE by several times (e.g. if the diffusion model needs k steps to sample one step then it will need kT NFE to sample T steps). The proposed method can alleviate this by doing joint sampling and predicting multiple future steps in one sampling loop, yet not too many study and comparison on NFE is presented in the main text. It will be clearer if computational cost statistics regarding the proposed method and existing methods like PDE-Refiner are also reported. 3. Compared to pure data-driven forecasting models like using UNet trained with MSE target, how much performance gain does online data assimilation bring? (the authors have conducted experiments on varying portions of observed data but it can be helpful to see a comparison against method that cannot do data assimilation)

Rating

7

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

The authors have discussed the limitation in the conclusion section.

Summary

The manuscript addresses the challenging problem of conditional diffusion modeling, towards accurate and efficient data assimilation in problems governed by PDE. The authors compare different diffusion modeling approaches in this setting, focusing on the auto-regressive setting, and also propose new approaches. The comparisons are done on relevant PDE benchmarks (chaotic, 1d and 2d spaces).

Strengths

Using conditional diffusion models for data assimilation is an interesting approach, and the results on the given PDE systems are impressive. The paper is relatively easy to read, and contains interesting suggestions for improvements. I very much appreciated the brief introduction to score-based modeling, as I am not working day-to-day with diffusion models (more with PDEs). Scalability of the approaches are also addressed, briefly in the main part and more detailed in the appendix. In general, the appendix is also extensive and contains a lot of details and additional studies related to the main paper.

Weaknesses

1) The PDEs considered as examples have turbulent / chaotic dynamics. There is not discussion about this, but it is important to mention: comparing trajectory data is not a good metric, because the PDE itself is causing sensitivity to the initial conditions, not the models. Rather, the PDE itself could be used as a metric (i.e., evaluate the PDE on the generated trajectory and check if it is zero); alternatively, Lyapunov exponents or density measures on the attractors can be used for comparison. The 1D Burgers equation does not have these problems. The "correlation time" defined by the authors is probably a way to mitigate the issues with chaotic/sensitive dynamics, but is not ideal, because the divergence from the true solution depends on the initial condition. Regardless of this issue with the metric, it seems the trained models are capable of generating trajectories close to the attractor (i.e. they "look reasonable"), which is already impressive. 2) There is no comparison to a classical scheme that solves the inverse problem. Especially looking at the long solution times (fig.3, between 25-100 minutes); I think it is possible to design a classical inverse problem solver that (a) starts with a random initial state of the PDE (b) solves the PDE with this guess, and then (c) updates the guess based on the error to the observations, and does all this iteratively in less than 25 minutes. It is ok to study how diffusion models perform in this setting, but it is not ok to not even mention this possibility of the classical approach in the related work. Many inverse problem solvers exist in this direction (also data-driven ones), a good overview is here: Arridge, S., Maass, P., Öktem, O., Schönlieb, C.-B., 2019. Solving inverse problems using data-driven models. Acta Numerica 28, 1–174. https://doi.org/10.1017/S0962492919000059 It would be good to include a brief discussion of the benefits of the neural network / diffusion model approach compared to the more classical setting. 3) There is no theoretical justification (meaning: no rigorous proofs) of the suggested new architectures and training methods. The new approaches are motivated plausibly in the main paper, with additional computations done in the appendix, but without theory. Minor: 1) l103: the statement $\mathcal{P}(u;\alpha)=\mathcal{P}(\partial_z u,\dots)$ is not correct. Does $\mathcal{P}$ only depend on $u$ or only (?) on its derivatives? Are PDEs like $<v,\partial_z> u + u=0$ not included (because they include $u$ as well as $\partial_z u$)? Probably this is just an abuse of notation here. 2) l126: do the authors mean "...time dependency [...] $=x_{1:L}(t)$" instead of "$x_{1:L}(0)$"?

Questions

1) l96: how can we divide by $\mu_t$ if it is a vector? Is it done component-wise? 2) It is probably my lack of experience with diffusion models, but how do we know the "true" probability density p to compare to (resp, $\nabla \log p$)? For example, in equation (3) and l.138, "the local score network ... is trained by minimizing the following DSM loss", which contains the expression $\nabla \log p$. This is probably not contained in the training data as a value for each x, right? How do we compute it? Equation 7 does not help, because it is also not clear to me how to compute the individual probability densities (for each i). In Equation 8 I can understand that the conditional is assumed to be Gaussian, which then leads to a simple expression, but it cannot be the case for the general density. 3) l305: "We assume to always observe some full initial states" this is very confusing to me. What is "some full"? Does it mean the full initial state is observed? This would invalidate the entire approach, because then there is no need for "data assimilation", the PDE can just be solved with a classical solver, from this given initial condition. Looking at appendix G, it seems that only a portion of the initial state is "always" given, which would be ok - then the classical approach is not really possible - but then the sentence should be rewritten to clarify this.

Rating

6

Confidence

3

Soundness

3

Presentation

3

Contribution

3

Limitations

Some limitations are discussed, but not in the direction of a comparison to classical solvers (as I mentioned in the "weaknesses" section). It is not clear if the diffusion based modeling can be better than classical solution methods for the particular problems, and if not, which other problems it can be used on that are hard to solve with classical methods.

Summary

This paper studies the application of diffusion models (DMs) to PDE forecasting and data assimilation, focusing on different approaches to condition the DMs on initial conditions (forecasting) or sparse observations (data assimilation). The conditioning can either occur explicitly during training by conditioning the score network (amortised model), or at inference time via reconstruction guidance (joint model). The authors present an improved autoregressive sampling technique for the latter, which is shown to be important for forecasting tasks. The methods are evaluated on multiple PDE datasets, including on a hybrid forecasting+data assimilation problem setup.

Strengths

- The studied models can tackle both forecasting and data assimilation jointly. Even if performance lags behind forecasting-only models, there might exist use-cases for this, and performance could be improved in future work. - There are a lot of careful ablations and insightful comparisons of the studied conditional diffusion models variants. - The universal amortised model is an interesting approach (taken from the ML for natural videos literature) that seems to perform promisingly.

Weaknesses

- The forecasting evaluation could be enhanced with probabilistic metrics such as the CRPS and spread-skill ratios, by evaluating ensemble forecast of the diffusion models. It would also be beneficial to see evaluations of the spectra of the predictions. - I think that it would be useful to include algorithms or pseudocode for how the different methods (joint AR, joint AAO, amortised etc.) are trained and how they function at inference time. This will make it easier to compare the different approaches in a more detailed and self-contained way. - In fig. 1 for Kolmogorov: It would be nice to include some results for the amortised model, especially since you observe differences between the datasets. - I might have missed an ablation for the window size W=P+C of the universal amortised model. This would be especially interesting to see on the Kolmogorov dataset (i.e. basically repeat Fig. 1 right but with W \in {3, 7, 9} etc.)

Questions

- Is the MSE-trained U-Net using the same modern U-net architecture as your diffusion-based methods do? - For the (non-universal) amortised models, you train one different model for each unique combination of P | C? How do you explain that the amortised models trained for a specific P | C combination underperform the universal amortised model that might need to distribute capacity between all the P | C combinations that it was trained on? - You could also perform the "universal" training of the amortised models (i.e. train on any combination of P|C) for non-diffusion models, right? That would be very interesting to see (e.g. it may or may not work for a MSE-trained model).

Rating

6

Confidence

3

Soundness

3

Presentation

2

Contribution

3

Limitations

Yes.

Summary

The authors perform an extensive study of generative diffusion models applied to the task of PDE forecasting and data assimilation (DA). Further, they introduce an autoregressive (AR) sampling strategy, and a universal amortised model based with variable context size, based on masking.

Strengths

The authors have done a large number of tests with different configuration, leading to useful insights for the community that go beyond the architectures introduced in this paper. The AR sampling does not suffer from the poor memory scaling of AAO sampling. By leveraging masking during training (similar to how LLMs are trained), their universal amortised model is flexible and robust with respect to the size of the conditioning window.

Weaknesses

AR sampling scales poorly in time as each step costs (1 + c)*p evals, that is, a diffusion sampling process, known to be expensive, has to be performed at every timestep. The inference times for both AR and AAO models are large compared to the baselines in the forecasting literature or to Rozet et al. for DA. The validation metrics presented could be clearer. The RMSD and MSE (appendix) presented are not normalized (in some sense) to account for the magnitude of the data, making the absolute number hard to interpret. Further, insightful metrics are missing, such as the energy spectrum of the predicted trajectories or long term trajectory stability. There are only a few comparison trajectories for the forecasting tests. While qualitative results are always subjective, it is known to be hard to judge the quality of a predicted trajectory just from quantitative summary statistics. Similarly, for the offline DA results, the partial observations are not shown.

Questions

1. What are the training times? How do they relate to model performance, i.e. validation metrics? 2. What is the long-term behaviour of the AR model? I.e. longer than seen during training. 3. How do the Fourier modes of the predicted trajectories look like? Given that PDE-Refiner is a baseline it would be interesting to see how it compares. 4. How well do the predicted trajectories match the physics? This is related to the question above (the energy spectrum, for example). -- End of questions -- A brief explanation of the review and grades given, put here for lack of a better place: Soundness: 3 - extensive benchmarking Presentation: 2 - appropriate but limited validation metrics (see Weaknesses) Contribution: 2 - extensive benchmarking, approach novel in this context, lacking some interesting metrics Rating 6 - A technically solid paper with extensive tests. It is held back by the validations done, the scalability issues, and inference times. Rating after rebuttal 7 - Added validation metrics and spectral studies. Scalability, accuracy in forecasting task and inference times remain an open question for future work.

Rating

7

Confidence

3

Soundness

3

Presentation

2

Contribution

2

Limitations

Some of the limitations are properly acknowledged. The time scaling of the AR sampling, and the lower performance on forecasting tasks compared to SOTA approaches. The fact that while flexible, the universal amortised model has a maximum conditioning size (as opposed to AR models) should be mentioned.

Summary

- The paper tackles of the problem of ML-based PDE modelling. Specifically two sub-problems: (a) *forecasting*: to generate rollouts given initial observations; and (b) *data assimilation*: to refine a trajectory given partial and noisy observations. - The proposed approach extends a score-based diffusion model solution (notably, related to Rozet & Louppe NeurIPS '23, and Lippe et al. NeurIPS '23) that jointly tackles the two sub-problems. - First, by proposing a *joint model* that allows for auto-regressive sampling of the trajectory (in contrast to "all-at-once" of Rozet & Louppe). Second, to consider *amortization* using a conditional diffusion model that no longer requires multiple trained models for choices of prediction and correction steps. - The approach is evaluated on standard benchmarks 1D Kuramoto-Sivashinsky ("KS") and 2D Kolmogorov flow ("Kolmogorov") and compared with standard baselines (e.g., PDERefiner of Lippe et al.). For forecasting, the results are on-par or behind baselines. For data assimilation, the approach outperforms Rozet & Louppe.

Strengths

1. The paper proposes a common framework to tackle multiple PDE simulation related tasks of forecasting, offline and online data assimilation. This appears to constrast existing work that solely focus on one. 2. The paper is overall well-written. I especially appreciated the related works discussions, which also highlight concurrent works. 3. The evaluation is reasonable, and considers multiple relevant baselines (including very recent ones). 4. The approach (although involves more computation) has some practical merits e.g., using "universal amortized" strategy to rely on a single trained model.

Weaknesses

## Major concerns I have no major concerns with the paper. ## Minor concerns **1. Results - Forecasting** - While a common framework for multiple PDE simulation tasks is appealing, I wonder the specific contributions for the forecasting problem given that many baselines outperform the proposed approach (and if I understand right, with lower computation cost). - Although, the paper claims (p8, L299) the approach is appealing for flexibility, I am unclear on scenarios that require this at the price of empirical gains. **2. Slightly unclear - experimental settings/results on offline DA vs. AAO** - Firstly, I have to preface that the AAO paper (Rozet & Louppe, 2023) doesn't appear to have an extensive quantitative analysis. - However, I am curious of the difference of the experimental setup between this and AAO paper: is the focus here sparse-observations, whereas AAO reports dense observations? - Moreover, in the dense observation setting, I see that errors (Fig. 3) are comparable. As a result, I wonder whether the results are in favour of the proposed approach only in the sparse observation regime (where the errors, although better than baselines, are still significantly large).

Questions

1. **Forecasting results**: can the authors comment on gains of the proposed approach for forecasting? 2. **Experimental settings vs. AAO**: can the authors comment on how the setup/evaluation regime differs from that of AAO?

Rating

7

Confidence

2

Soundness

3

Presentation

3

Contribution

3

Limitations

Yes, this is a dedicated and balanced discussion on this.

Reply to the rebuttal

I would like to thank the authors for the efforts and clarification, the majority of the concerns and questions I raised have been addressed and clarified. I think the proposed framework is very flexible and versatile for forecasting/inverse problems, despite the performance is not sota yet but I believe there is potential for a lot of improvements. The authors have also carried out comprehensive empirical study, which is informative will be practically beneficial to many data-driven PDE applications. I've raised my score to 7.