On conditional diffusion models for PDE simulations

TL;DR

This work addresses the shortcomings of existing models by proposing an autoregressive sampling approach that significantly improves performance in forecasting, a new training strategy for conditional score-based models that achieves stable performance over a range of history lengths, and a hybrid model which employs flexible pre-training conditioning on initial conditions and flexible post-training conditioning to handle data assimilation.

Abstract

Modelling partial differential equations (PDEs) is of crucial importance in science and engineering, and it includes tasks ranging from forecasting to inverse problems, such as data assimilation. However, most previous numerical and machine learning approaches that target forecasting cannot be applied out-of-the-box for data assimilation. Recently, diffusion models have emerged as a powerful tool for conditional generation, being able to flexibly incorporate observations without retraining. In this work, we perform a comparative study of score-based diffusion models for forecasting and assimilation of sparse observations. In particular, we focus on diffusion models that are either trained in a conditional manner, or conditioned after unconditional training. We address the shortcomings of existing models by proposing 1) an autoregressive sampling approach that significantly improves performance in forecasting, 2) a new training strategy for conditional score-based models that achieves stable performance over a range of history lengths, and 3) a hybrid model which employs flexible pre-training conditioning on initial conditions and flexible post-training conditioning to handle data assimilation. We empirically show that these modifications are crucial for successfully tackling the combination of forecasting and data assimilation, a task commonly encountered in real-world scenarios.

Figures (44)

• Figure 1: High correlation time ($\uparrow$) on the KS (left) and Kolmogorov (right) datasets for different $P \mid C$ conditioning scenarios of the joint, amortised and universal amortised models, where $P$ indicates the number of generated states and $C$ the number of states conditioned upon. We show mean $\pm$$3$ standard errors, as computed on the test trajectories.
• Figure 2: The best joint and amortised models compared against standard ML-based benchmarks for forecasting on the KS (left) and Kolmogorov (right) datasets. High correlation time ($\uparrow$) is reported, showing the mean $\pm$$3$ standard errors, as computed on the test trajectories. We show in \ref{['tab:additional_architectures']} the configurations for the best models.
• Figure 3: RMSD (mean $\pm$$3$ standard errors) for KS (left) and Kolmogorov (right) on the offline DA setting for varying sparsity levels (top), and the computational cost associated with each setting (bottom). The latter is the same for all sparsity settings for AAO, but differs for AR since it depends on the $P \mid C$ setting that was used. The $c$ in AAO ($c$) refers to the number of corrector steps used.
• Figure 4: True test trajectories from solving the Burgers' equation following the setup of wang2021learning. In this case, both training, validation, and test trajectories have the same length. We used $\Delta \tau = 0.01s$, so the trajectory contains states from time $\tau=0s$ to $\tau=1s$. The spatial dimension is discretized in 128 evenly distributed states
• Figure 5: Examples from the Kuramoto-Sivashinsky dataset. Training trajectories are shorter than those used to evaluate the models. Training trajectories have $140$ time steps generated every $\Delta \tau = 0.2s$, i.e. training trajectories contain examples from $0s$ to $28s$. Validation and test trajectories, instead, show the evolution of the equation for $640$ steps, i.e. from $0s$ to $128s$. The spatial dimension is discretized in 256 evenly distributed states.