SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Figure 3: Example schedules for coupling diffusion steps to dynamical time steps. While a naive approach only uses timesteps given by the temporal resolution of the data (i.e. discrete indices), our framework can accommodate continuous indices. The additional  $k$  diffusion steps are highlighted in green and map uniformly between the input timestep,  $\mathbf{x}_0$ , and earliest output timestep,  $\hat{\mathbf{x}}_1$ . Our experiments using the SST dataset in section 5 demonstrate that increasing the number of diffusion steps with implicit intermediate timesteps can improve performance (see Appendix D.7).
We evaluate our method and baselines on three different datasets:
- Sea Surface Temperatures (SST): a new dataset based on NOAA OISSTv2 [30], which comes at a daily time-scale. Similar to [11, 71], we train our models on regional patches which increases the available data (here, we choose 11 boxes of 60 latitude  $\times$  60 longitude resolution in the eastern tropical Pacific Ocean). Unlike the data based on the NEMO dataset in [11, 71], we choose OISSTv2
Table 1: Results for sea surface temperature forecasting of 1 to 7 days ahead, as well Navier-Stokes flow full trajectory forecasting of 64 timesteps. Numbers are averaged out over the evaluation horizon. Bold indicates best, blue second best. For CRPS and MSE, lower is better. For SSR, closer to 1 is better. For SST, all models are trained on forecasting  $h = 7$  timesteps. The time column represents the time needed to forecast all 7 timesteps for a single batch. For Navier-Stokes, Perturbation, Dropout, and DYffusion are trained on a horizon of  $h = 16$ . MCVD and DDPM are trained on  $h = 4$  and  $h = 1$ , respectively, as we were not able to successfully train them using larger horizons.
Method	SST	Navier-Stokes
CRPS	MSE	SSR	Time [s]	CRPS	MSE	SSR
Perturbation	0.281 ± 0.004	0.180 ± 0.011	0.411 ± 0.046	0.4241	0.090 ± 0.001	0.028 ± 0.000	0.448 ± 0.002
Dropout	0.267 ± 0.003	0.164 ± 0.004	0.406 ± 0.042	0.4241	0.078 ± 0.001	0.027 ± 0.001	0.715 ± 0.005
DDPM	0.246 ± 0.005	0.177 ± 0.005	0.674 ± 0.011	0.3054	0.180 ± 0.004	0.105 ± 0.010	0.573 ± 0.001
MCVD	0.216	0.161	0.926	79.167	0.154 ± 0.043	0.070 ± 0.033	0.524 ± 0.064
DYffusion	0.224 ± 0.001	0.173 ± 0.001	1.033 ± 0.005	4.6722	0.067 ± 0.003	0.022 ± 0.002	0.877 ± 0.006
as our SST dataset because it contains more data (although it has a lower spatial resolution of  $1/4^{\circ}$  compared to  $1/12^{\circ}$  of NEMO). We train, validate, and test all models for the years 1982-2019, 2020, and 2021, respectively. The preprocessing details are in the Appendix D.1.1.
- Navier-Stokes flow: benchmark dataset from [44], which consists of a  $221 \times 42$  grid. Each trajectory contains four randomly generated circular obstacles that block the flow. The viscosity is  $1e - 3$ . The channels consist of the  $x$  and  $y$  velocities as well as a pressure field. Boundary conditions and obstacle masks are given as additional inputs to all models.
- Spring Mesh: benchmark dataset from [44]. It represents a  $10 \times 10$  grid of particles connected by springs, each with mass 1. The channels consist of two position and momentum fields each.
5.2 Experimental Setup
We focus on probabilistic multi-step forecasting, and especially long rollouts, as opposed to attaining the best next-step or short-term forecasts. This is an important distinction because, while deterministic single-step forecasting models such as those of [44] are appropriate for short-term predictive skill comparison, such models tend to lack stability for long-term forecasts. We demonstrate this in the Appendix D.4.2 (see Fig. 6), where our method and baselines excel in the long-range forecasting regime while being competitive in the short-range regime compared to the baselines from [44].
Neural network architectures. For a given dataset, we use the same backbone architecture for all baselines as well as for both the interpolation and forecaster networks in DYffusion. For the SST dataset, we use a popular UNet architecture designed for diffusion models<sup>2</sup>. For the Navier-Stokes and spring mesh datasets, we use the UNet and CNN from the original paper [44], respectively. The UNet and CNN models from [44] are extended by the sine/cosine-based featurization module of the SST UNet to embed the diffusion step or dynamical timestep.
Baselines. We compare our method against both direct applications of standard diffusion models to dynamics forecasting and methods to ensemble the "barebone" backbone network of each dataset. The network operating in "barebone" form means that there is no involvement of diffusion.
- DDPM [25]: both next-step (image-like problem) as well as multi-step (video-like problem) prediction mode (we report the best model version only in the results).
- MCVD [67]: conditional video diffusion model for video prediction (in "concat" mode).
- Dropout [17]: Ensemble multi-step forecasting based on enabling dropout of the barebone backbone network at inference time.
- Perturbation [45]: Ensemble multi-step forecasting based on random perturbations of the initial conditions/inputs with a fixed variance of the barebone backbone network.
$^{2}$ https://github.com/lucidrains/denoising-diffusion-pytorch/blob/main/denoising_diffusion_pytorch/denoising_diffusion_pytorch.py
Table 2: Spring Mesh results. Both methods are trained on a horizon of  $h = 134$  timesteps and evaluated how well they forecast the full test trajectories of 804 steps. Despite several attempts with varying configurations over the number of diffusion steps, learning rates, and diffusion schedules, none of the DDPM or MCVD diffusion models converged. Using the same CNN architecture, our MSE results significantly surpass the ones reported in Fig. 8 of [44], where the CNN diverged or attained a very poor MSE. This is likely because of the multi-step training approach that we use, while the single-step prediction approach ( $h = 1$ ) in [44] can generate unstable rollouts.
Method	Test	Out of Distribution
CRPS	MSE	SSR	CRPS	MSE	SSR
Dropout	0.0138 ± 0.0006	7.27e-4 ± 6.8e-5	1.01 ± 0.02	0.0448 ± 0.0007	7.08e-3 ± 1.7e-4	0.70 ± 0.01
DYffusion	0.0103 ± 0.0022	4.20e-4 ± 2.1e-4	1.13 ± 0.08	0.0292 ± 0.0009	3.54e-3 ± 2.0e-4	0.82 ± 0.00
MCVD and the multi-step DDPM variant predict the timesteps  $\mathbf{x}_{t + 1:t + h}$  based on  $\mathbf{x}_t$ , which is the standard approach in video (prediction) diffusion models. The barebone backbone network baselines are time-conditioned forecasters (similarly to the DYffusion forecaster) trained on the multi-step objective  $\mathbb{E}_{i\sim \mathcal{U}[1,h],\mathbf{x}_{t,t + i}\sim \mathcal{X}}||F_{\theta}(\mathbf{x}_t,i) - \mathbf{x}_{t + i}||^2$  from scratch. We found it to perform very similarly to predicting all  $h$  horizon timesteps at once in a single forward pass (not shown). In all experiments, we observed that the single-step forecasting  $(h = 1)$  version of the barebone network yielded significantly lower performance compared to any of the multi-step training approaches (see Appendix D.4.2). More details of the implementation are provided in Appendix D.2.
Evaluation. We evaluate the models on the best validation Continuous Ranked Probability Score (CRPS) [42], which is a proper scoring rule and a popular metric in the probabilistic forecasting literature [19, 12, 51, 48, 58]. The CRPS is computed by generating a 20-member ensemble (i.e. 20 samples are drawn per batch element), while we generate a 50-member ensemble for final model selection between different hyperparameter runs. We also use a 50-member ensemble for evaluation on the test datasets to compute the CRPS, mean squared error (MSE), and spread-skill ratio (SSR). The MSE is computed on the ensemble mean prediction. The SSR is defined as the ratio of the square root of the ensemble variance to the corresponding ensemble RMSE. It serves as a measure of the reliability of the ensemble, where values smaller than 1 indicate underdispersion (i.e. the probabilistic forecast is overconfident in its forecasts), and larger values overdispersion [16, 18]. On the Navier-Stokes and spring mesh datasets, models are evaluated by autogressively forecasting the full test trajectories of length 64 and 804, respectively. For the SST dataset, all models are evaluated on forecasts of up to 7 days. We do not explore more long-term SST forecasts because the chaotic nature of the system, and the fact that we only use regional patches, inherently limits predictability.
Quantitative Results. We present the time-averaged metrics for the SST and Navier-Stokes dataset in Table 1 and for the Spring Mesh dataset in Table 2. DYffusion performs best on the Navier-Stokes dataset, while coming in a close second on the SST dataset after MCVD, in terms of CRPS. Since MCVD uses 1000 diffusion steps $^3$ , it is slower to sample from at inference time than from our DYffusion which is trained with less than 50 diffusion steps. The DDPM model for the SST dataset is fairly efficient because it only uses 5 diffusion steps but lags in terms of performance. Thanks to the time-conditioned nature of DYffusion's interpolator and forecaster nets, memory is not an issue when scaling our framework to long horizons. In the spring mesh dataset, we train with a horizon of 134 and evaluate long trajectories of 804 steps. Our method beats the barebone network baseline, with a larger margin on the out-of-distribution test dataset. It is worth noting that our reported MSE scores are significantly better than the ones reported in [44], likely due to multi-step training being superior to the single-step forecasting approach of [44] (see Appendix D.4.2 for more details). For the out-of-distribution test set of the Navier-Stokes benchmark, the results are almost identical to the one in Tab. 1, so we show them in Appendix D.4.1.
Qualitative Results. In dynamics forecasting, long-range forecasts of ML models often suffer from blurriness (or might even diverge when using autoregressive models). In Fig. 4 we show exemplary
This is the default, we were not able to successfully train MCVD models with fewer diffusion steps.
Figure 4: Qualitative forecasts for timesteps 2, 6, 24, 46, and 64 (last timestep) of the velocity norm of an example Navier-Stokes test trajectory. Here, we generate five sample trajectories for the best baseline (Dropout) and our method DYffusion, both with  $h = 16$ , and visualize the one with the best trajectory-average MSE for each of the methods. Our method (bottom row) can reproduce fine-scale details visibly better than the baseline (see e.g. right sides of the snapshots). The corresponding video of the full trajectory, including the velocity and pressure fields, can be found at this Google Drive URL https://drive.google.com/file/d/1xklVs42Ii18I8SVT0f1ZmAKR159qiHG_/view?usp=share_link.
samples for the best baseline (Dropout) and DYffusion as well as the corresponding ground truth at five different timesteps from a complete Navier-Stokes trajectory forecast. Our method can reproduce the true dynamics over the full trajectory and does so better than the baseline, especially for fine-scale patterns such as the tails of the flow after the right-most obstacle. The corresponding full video can be found at this URL.
Increasing the forecasted resolution. Motivated by the continuous-time nature of DYffusion, we aim to study in this experiment whether it is possible to forecast skillfully beyond the resolution given by the data. Here, we forecast the same Navier-Stokes trajectory shown in Fig. 4 but at  $8\times$  resolution (i.e. 512 timesteps instead of 64 are forecasted in total). This behavior can be achieved by either changing the sampling trajectory  $[i_n]_{n = 0}^{N - 1}$  or by including additional output timesteps,  $J$ , for the refinement step of Alg. 2. In this case, we choose to do the latter and find the resulting forecast to be visibly pleasing and temporally consistent; see the full video at this URL.
Note that we hope that our probabilistic forecasting model can capture any of the possible, uncertain futures instead of forecasting their mean, as a deterministic model would do. As a result, some long-term rollout samples are expected to deviate from the ground truth. For example, see the velocity at  $t = 3.70$  in the video. It is reassuring that our samples show sufficient variation, but also cover the ground truth quite well (sample 1).
Ablations. In Appendix D.5, we ablate various components of our method. This includes showing that inference stochasticity in the interpolation network is crucial for performance. Our experiments in Appendix D.5.1 confirm the findings of [2] regarding the superiority of cold sampling (see Alg. 2) compared to the naive sampling (Alg. 4) algorithm. We provide a theoretical justification in Appendix C.2, where we show that cold sampling has first-order discretization error, whereas naive sampling does not.
Table 3: CRPS test scores for the SST dataset using different number of training years. 2, 4, 37 refers to training from 2017, 2015, 1982 until the end of 2018, respectively. The standard deviations are around 0.01 or lower.
# years	2	4	37
MCVD	0.262	0.233	0.216
DYffusion	0.239	0.234	0.225
We also study the choice of the interpolation schedule:
For the Navier-Stokes and spring mesh datasets, it was sufficient to use the simplest possible schedule  $S = [i_n]_{n=0}^{N-1} = [j]_{j=0}^{h-1}$ , where the number of diffusion steps and the horizon are equal. For the SST dataset, using  $k$  additional diffusion steps corresponding to floating-point  $i_k$ , especially such that  $i_k < 1$  for all  $k$ , improved performance significantly. Similarly to the finding of DDIM [61] for Gaussian diffusion, we show that the sampling trajectories of DYffusion can be accelerated by skipping intermediate diffusion steps, resulting in a clear trade-off between accuracy and speed. Finally, we verify that the forecaster's prediction of  $\mathbf{x}_{t+h}$  for a given diffusion step  $n$  and associated input  $\hat{\mathbf{x}}_{t+i_n}$  usually improves as  $n$  increases, i.e. as we approach the end of the reverse process and  $\mathbf{x}_{t+i_n}$  gets closer in time to  $\mathbf{x}_{t+h}$ .
The variance in the performance of the standard diffusion models across datasets may be intriguing. One possible explanation could be the size of the training dataset, since the SST dataset is more than ten times larger than the Navier-Stokes and spring mesh datasets, respectively. Sufficient training data may be a key factor that enables Gaussian diffusion models to effectively learn the reverse process
from noise to the data space. We explore this hypothesis in Table 3, where we retrain MCVD as well as DYffusion's  $F_{\theta}$  and  $\mathcal{I}_{\phi}$  on small subsets of the SST training set. We find that DYffusion's performance degrades more gracefully under limited training data, especially when using only two years of data (or  $\approx 8,000$  training data points).
We found it very challenging to train video diffusion models on long horizons. A reason could be that the channel dimension of the video diffusion backbone is scaled by the training horizon, which increases the problem dimensionality and complexity considerably (see Appendix D.6). Unlike common practice in natural image or video problems, we do not threshold the predictions of any diffusion model, as the data range is unbounded. However, this thresholding has been shown to be an important implementation detail that stabilizes diffusion model sampling for complex problems in the domain of natural images [40]. This could be the reason why the studied diffusion model baselines produce poor predictions for the Navier-Stokes and spring mesh datasets.
We have introduced DYffusion, a novel dynamics-informed diffusion model for improved probabilistic spatiotemporal forecasting. In contrast to most prior work on probabilistic forecasting that often generates long-range forecasts via iterating through an autoregressive model, we tailor diffusion models to naturally support long-range forecasts during inference time by coupling their diffusion process with the dynamical nature of the data. Our study presents the first comprehensive evaluation of diffusion models for the task of spatiotemporal forecasting. Finally, we believe that our dynamics-informed diffusion model can serve as a more general source of inspiration for tailoring diffusion models to other conditional generation problems, such as super-resolution, where the data itself can natively inform an inductive bias for the diffusion model.
Limitations. While our proposed framework can be applied to any dynamics forecasting problem as an alternative to autoregressive models, it does not support, in its current form, problems where the output space is different from the input space. Additionally, although it requires fewer diffusion steps compared to Gaussian diffusion models to achieve similar performance, the need for additional forward passes through the interpolator network diminishes some of these advantages. When prioritizing inference speed, approaches that only require a single forward pass to forecast outperform sequential sampling of diffusion models, including DYffusion.
Future work. Cold Diffusion [2] indicates that deterministic data degradations cannot meet the performance of stochastic (i.e. Gaussian noise-based) degradations. In our work, we have experimented with using an interpolator network with Monte Carlo dropout enabled during inference time to achieve a stochastic forward process, which we have shown to be important for optimal performance. Exploring more advanced approaches for introducing stochasticity into our general framework presents an interesting avenue for future work. A more advanced approach could draw insights from the hypernetwork-based parameterization of the latent space from [69]. Furthermore, given our finding of the equivalence of cold sampling with the Euler solver (C.1), it could be interesting to consider more advanced ODE solvers for improved sampling from DYffusion. Recent advances in neural architectures for physical systems [6, 31] are complementary to our work. Using such alternative neural architectures as interpolator or forecaster networks within our framework could be a promising approach to further improve performance.
This work was supported in part by the U.S. Army Research Office under Army-ECASE award W911NF-07-R-0003-03, the U.S. Department Of Energy, Office of Science, IARPA HAYSTAC Program, NSF Grants #2205093, #2146343, and #2134274.
NOAA OI SST V2 High Resolution Dataset data provided by the NOAA PSL, Boulder, Colorado, USA, from their website at https://psl.noaa.gov.
[67] Vikram Voleti, Alexia Jolicoeur-Martineau, and Christopher Pal. MCVD: Masked conditional video diffusion for prediction, generation, and interpolation. Advances in Neural Information Processing Systems, 2022. URL https://arxiv.org/abs/2205.09853.
[68] Etienne E Vos, Ashley Gritzman, Sibusisiwe Makhanya, Thabang Mashinini, and Campbell D Watson. Long-range seasonal forecasting of 2m-temperature with machine learning. NeurIPS, Tackling Climate Change with Machine Learning, 2020.
[69] Zhong Yi Wan, Leonardo Zepeda-Nunez, Anudhyan Boral, and Fei Sha. Evolve smoothly, fit consistently: Learning smooth latent dynamics for advection-dominated systems. In The Eleventh International Conference on Learning Representations, 2023. URL https://openreview.net/forum?id=Z4s73sJYQM.
[70] Rui Wang, Karthik Kashinath, Mustafa Mustafa, Adrian Albert, and Rose Yu. Towards physics-informed deep learning for turbulent flow prediction. Proceedings of the 26th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining, 2020.
[71] Rui Wang, Robin Walters, and Rose Yu. Meta-learning dynamics forecasting using task inference. In Advances in Neural Information Processing Systems (NeurIPS), 2022.
[72] Anthony L Westerling, Alexander Gershunov, Timothy J Brown, Daniel R Cayan, and Michael D Dettinger. Climate and wildfire in the western united states. Bulletin of the American Meteorological Society, 84(5):595-604, 2003.
[73] Jonathan A. Weyn, Dale R. Durran, and Rich Caruana. Can machines learn to predict weather? using deep learning to predict gridded 500-hpa geopotential height from historical weather data. Journal of Advances in Modeling Earth Systems, 11(8):2680-2693, 2019. doi: https://doi.org/10.1029/2019MS001705. URL https://agupubs.onlinelibrary.wiley.com/doi/abs/10.1029/2019MS001705.
[74] Dongxia Wu, Liyao Gao, Matteo Chinazzi, Xinyue Xiong, Alessandro Vespignani, Yi-An Ma, and Rose Yu. Quantifying uncertainty in deep spatiotemporal forecasting. In Proceedings of the 27th ACM SIGKDD Conference on Knowledge Discovery & Data Mining, pages 1841-1851, 2021.
[75] Yilun Xu, Ziming Liu, Yonglong Tian, Shangyuan Tong, Max Tegmark, and Tommi Jaakkola. Pfgm++: Unlocking the potential of physics-inspired generative models. 2023. URL https://arxiv.org/abs/2302.04265.
[76] Ling Yang, Zhilong Zhang, Yang Song, Shenda Hong, Runsheng Xu, Yue Zhao, Yingxia Shao, Wentao Zhang, Bin Cui, and Ming-Hsuan Yang. Diffusion models: A comprehensive survey of methods and applications, 2022. URL https://arxiv.org/abs/2209.00796.
[77] Ruihan Yang, Prakhar Srivastava, and Stephan Mandt. Diffusion probabilistic modeling for video generation. arXiv preprint arXiv:2203.09481, 2022.
1 Introduction 1
2 Background 3
3 DYffusion: DYNAMics-Informed Diffusion Model 3
4 Related Work 5
5 Experiments 6
5.1 Datasets 6
5.2 Experimental Setup 7
5.3 Results 8
6 Conclusion 10
A Glossary 17
B Methodology 17
B.1 DYffusion 17
B.2 Connection to standard diffusion models 18
C Diffusion process as ODE 18
C.1 Cold Sampling from the Euler method 19
C.2 Why is cold sampling better than naive sampling? 19
D Experiments 20
D.1 Datasets 20
D.1.1 SST data preprocessing 20
D.1.2 Navier-Stokes and spring mesh datasets 21
D.2 Implementation Details 21
D.3 Neural Architecture Details 22
D.4 Additional results 23
D.4.1 Navier-Stokes out-of-distribution 23
D.4.2 Benchmark against deterministic models 23
D.5 Ablations 24
D.5.1 Cold sampling versus naive sampling 24
D.5.2 Inference dropout in the interpolator network 24
D.5.3 Refining the interpolation forecasts after sampling 24
D.5.4 Forecaster network conditioning 25
D.5.5 Choosing the training horizon 25
D.5.6 Using auxiliary diffusion steps can improve performance 25
D.5.7 Forecaster network iteratively refines its forecasts 25
D.5.8 Accelerated DYffusion sampling 27
D.6 Modeling Complexity of DYffusion and Baselines 27
D.7 Example Sampling Trajectories Visualized 28
The glossary is given in Table 4 below.
Symbol	Used for
t	Indexer of temporal data snapshots
h	Training horizon, i.e. the number of timesteps a method learns to forecast
H	Evaluation horizon. A model with horizon h, is applied autoregressively [H/h]-times during evaluation
N	Total number of diffusion steps
n	Indexer of the diffusion step, where 0 ≤ n ≤ N-1
X	Distribution from which the data snapshots are sampled from
xt+i	Data point at timestep t+i
xt+i	Predicted data point at timestep t+i
xt:t+h	Shorthand for the sequence of data points xt, xt+1, ..., xt+h
xt	Initial conditions, i.e. input snapshot for the forecasting model
xt+1:t+h	Sequence of target snapshots that the forecasting model predicts given xt
s(n)	Data point at diffusion step n
Iφ	Stochastic temporal interpolation neural network; Trained in the first-stage of DYffusion
Fθ	Forecasting neural network; Trained in the second-stage of DYffusion
S	Interpolation schedule for DYffusion that couples every diffusion step, n, to a timestep in
in	Maps diffusion step n to timestep in, where 0 = i0 < in < im < h for 0 < n < m ≤ N-1
k	Number of non-integer auxiliary diffusion steps for DYffusion. If k=0, then S = [in]n=0 N-1 = [j]h=0
D	Forward diffusion process operator, e.g. Gaussian noise, of a conventional diffusion model
Rθ	Denoising neural network of a conventional diffusion model
U[[a,b]]	Describes the uniform distribution over the set {a, a+1, ..., b}