DYffusion: A Dynamics-informed Diffusion Model for Spatiotemporal Forecasting
Authors: Salva Rühling Cachay
Bo Zhao
Hailey Joren
Rose Yu
Venue: NeurIPS 2023 poster (2023)
Keywords: AI for science
diffusion models
scientific machine learning
probabilistic forecasting
TL;DR: We introduce a diffusion model tailored for probabilistic spatiotemporal forecasting by designing temporal forward and reverse processes.

Abstract: While diffusion models can successfully generate data and make predictions, they are predominantly designed for static images. We propose an approach for training diffusion models for dynamics forecasting that leverages the temporal dynamics encoded in the data, directly coupling it with the diffusion steps in the network. We train a stochastic, time-conditioned interpolator and a backbone forecaster network
that mimic the forward and reverse processes of conventional diffusion models, respectively. This design choice naturally encodes multi-step and long-range forecasting capabilities, allowing for highly flexible, continuous-time sampling trajectories and the ability to trade-off performance with accelerated sampling at inference time. In addition, the dynamics-informed diffusion process imposes a strong inductive bias, allowing for improved computational efficiency compared to traditional Gaussian noise-based diffusion models. Our approach performs competitively on probabilistic skill score metrics in complex dynamics forecasting of sea surface temperatures, Navier-Stokes flows, and spring mesh systems.

DYffusion: A Dynamics-informed Diffusion Model for Spatiotemporal Forecasting
Salva Rühling Cachay
Bo Zhao
Hailey Joren
Rose Yu
University of California, San Diego
(sruhlingcachay, bozhao, hjoren, roseyu]@ucsd.com edu
Abstract
While diffusion models can successfully generate data and make predictions, they are predominantly designed for static images. We propose an approach for ef- ficiently training diffusion models for probabilistic spatiotemporal forecasting, where generating stable and accurate rollout forecasts remains challenging, Our method, DYffusion, leverages the temporal dynamics in the data, directly coupling it with the diffusion steps the model. We train stochastic, time-conditioned terpolator and a forecaster network that mimic the forward and reverse processes of standard diffusion models, respectively. DY DYffusion naturally facilitates multi-step and rong-range forecasting, allowing for highly flexible, continuous- time sampling trajectories and the ability to trade-off performance with accelerated sampling diffusion process in DYffusion
inference time. In addition, the dyynamic-informen imposes a strong inductive bias and significantly improves computational efficiency compared to traditional Gaussian boise-based diffusion models. Our approach performs competitively on probabilistic forecasting of complex dynamics sea surface temperatures, Navier-Stokes flows, and spring mesh systems."
Introduction
Dynamics forecasting refers to the task of predicting the future behavior of dynamic system, involv- ing learning the underlying dynamics governing the system's evolution to make accurate predictions about its future states. Obtaining accurate and reliable probabilistic forecasts an important compo- nent of policy formulation [39, 4], risk management [38, 20], resource optimization [7], and strategic planning [46, 72]. In many applications, accurate long-range probabilistic forecasts are particularly challenging to obtain [38, 20, 4]. In operational settings, methods typically hinge sophisticated numerical simulations that require supercomputers to perform calculations within manageable time frames, often compromising the spatial resolution of the grid for efficiency [3].
Generative modeling presents a promising avenue for probabilistic dynamics forecasting. Diffusion models, in particular, can successfully model natural image and video distributions [47, 54, 59]. The standard approach, Gaussian diffusion, corrupts the data with Gaussian noise to varying degrees through the "forward process" " and sequentially denoises random input at inference time to obtain highly realistic samples through the "reverse process" " [60, 25, 32). However, learning to map from noise to realistic data is challenging in high dimensions, especially with limited data availability. Consequently, the computational costs associated with training and inference from diffusion models prohibitive, requiring sequential sampling process over hundreds of diffusioon steps. For example, a denoising diffusion probabilistic model (DDPM) takes around 20 hours to sample 50k images of size 32 32 [61]. In addition, few methods apply diffusion models beyond static images. Video
Code available at: https: stttps:/ggihhann wtweshhb..comoeeeeee www.coo//////eeeee Raaem/osee-aaee htwp//Roeee-TL--aaaaggffssoooooo
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

Noise interpolation D
Gaussian Diffusion
s
6(N-1)
(1)
(
Initial conditions X
Denoise
Reverse process
Forecast
DYffusion (ours)
Xt-in
Xi+iN-1
X
Temporal interpolation Isg
Figure 1: Our proposed framework, DYffusion, reimagines the noise-denoise -ackurd- pro- cesses of conventional diffusion models as interplay of temporal interpolation and forecasting, On the top row, we illustrate the direct application of video diffusion model to dynamics forecasting for probabilistic forecasts
horizon of h 3. On the bottom row, DY ffusion generates continuous for X given the initial conditions, X₃. During sampling, the reverse process iteratively steps forward in time by forecasting X (which plays the role of the "clean data", s(O), in conventional diffusion models) and interpolating to one of N intermediate timesteps, X4+i,,- As a result, our approach operates the data space all times and does not need to model high- diteemsinniii videos at each diffusion state.
diffusion models [27, 24, 77, 67] can generate realistic samples but do not explicitly leverage the temporal nature of the data to generate accurate forecasts.
In this work, we introduce a novel framework to train a dynamics-informe informed diffusion model for multi- step probabilistic forecasting, Inspired by recent findings that highlight the potential of non-Gaune non-Gaussian diffusion processes [2, 28, 10). we introduce a new forward process. This process relies on temporal interpolation and is implemented through time-conditioned neural network. Without requiring assumptions about the physical system, our approach imposes an inductive bias by coupling the diffusion process steps with the time steps in the dynamical system. This reduces the computational complexity of our diffusion model in terms of memory footprint, data efficiency, and the number of diffusion steps required during training. Our resulting diffusion model-based framework, which we call DYffusion, naturally captures long-range dependencies and generates accurate probabilistic ensemble forecasts for high-dimensional spatiotemporal data.
Our contributions can be summarized as follows:
We investigate probabilistic spatiotemporal forecasting from the perspective diffusion models, including their applications to complex physical systems with a large number of dimensions and low data availability.
We introduce DYffusion, flexible framework for multi-step forecasting and long-range horizons that leverages temporal inductive bias to accelerate training and lower memory needs. We explore the theoretical implications our method, proving that DY Dffssiion an implicit model that learns the solutions to a dynamical system, and cold sampling [2] can be interpreted as its Euler's method solution.
We conduct an empirical study comparing performance and computational requirements in dynamics forecasting, including state-of-the-an probabilistic methods such conditional video diffusion models. We find that the proposed approach achieves strong probabilistic forecasts and improves computational efficiency over standard Gaussian diffusion.

---

2 Background
Problem setup. We study the problem of probabilistic spatiotemporal forecasting. We start with a dataset of [x)] snapshots with X; E A. Here, X' represents the space in which the data lies, which may consist of spatial dimensions (e.g., latitude, longitude, atmospheric height) and a channel dimension (e.g., velocities, temperature, humidity). The task of probabilistic forecasting is to learn conditional distribution 1(x++112+h x:-)+1:r) that uses / snapshots from the past forecast subsequent horizon of / snapshots. Here, we focus on the task of forecasting from single initial condition, where we aim learn P(x++1++h xx). This covers many realistic settings [50, 44] while minimizing the computational requirements for accurate forecasts [65].
Diffusion processes. In standard diffusion model, the forward diffiusion process iteratively degrades the data with increasing levels of (Gaussian) noise. The reverse diffusion process then gradually removes the noise to generate data. We denote these diffusion step states, s(n), with a superscript n to clearly distinguish them from the time steps of the data {x,)²21 This process can be generalized to consider degradation operator D that takes input the data point (0)) and outputs s ((n) (ss(0), n) for varying degrees of degradation proportional to n f1...., .N [2]. Oftentimes, D adds Gaussian noise with increasing levels of variance so that s(N) ~ (0, A denoising network parameterized by 9, Ro is trained to restore s("), i.e. such that R (s .n) R (())
The diffusion model can be conditioned on the input dynamics by considering Rg(s("), x2, In the case of dynamics forecasting, the diffusion model can be trained minimize the objective
min Bn-N
[|| [R x. s(0)
n)), xx, - 13
(1)
where A' is the data distribution, || || is norm, and s X x is the prediction target. 2/Ja, b]
describes the uniform distribution over the set {a. a + ... b}. In practice, Rg may be trained to predict the Gaussian noise that has been added to the data point using score matching objective [25].
The objective in Eq. (1) can be viewed as generalized version of the standard diffusion models, see
Appendix B.2 for more comprehensive analysis.
3 DYffusion: DYnamics-Informed Diffusion Model
The key innovation of our method, DYffusion, is reimagining of the diffu- sion processes to more natu- rally model spatiotemporal se- quences, Xe.+A (see Fig. I). Specifically, we replace the degradation operator, D, with a stochastic interpolator net- work Ig. and the restoration network, Ro, with a determin- istic forecaster network, Fo.
Algorithm 1 DYffusion, Two-stage Training
Input: networks Ig. norm ||. horizon h, schedule [i Stage 1: Train interpolator network, Ip
1. Sample i~ ~ Uniform
1)) 1-1))
2. Sample X,, Xt+i; X
X
3. Optimize min I (x,, X i) x++il?
Stage 2: Train foreeaster network (diffusion model backbone), Fe
1. Freeze Ia and enable inference stochasticity (e.g. dropout)
2. Sample ~ Uniform (f0.... .O....... N - 11) and x,, X4th X'
3. Optimize minş (o (Xa, Xith, in) in) x
At high level, our forward and reverse processes emulate the temporal dynamics in the data. Thus. intermediate steps in the diffusion process can be reused as forecasts for actual timesteps in multi-step forecasting. Another benefit of our approach is that the reverse process is initialized with the initial conditions of the dynamics and operates in observation space at all times. In contrast, a standard diffusion model is designed for unconditional generation, and reversing from white noise requires more diffusion steps. For conditional prediction tasks such as forecasting, DYffusion emerges as a much more natural method that well aligned with the task hand. See Table 4 for full glossary.
Temporal interpolation as a forward process.
impose a temporal bias, we train a time- conditioned network Ig to interpolate between snapshots of data. Specifically, given a horizon h, we train I,j so that I&p (x,, X(th, i) Re Xt+i for i E fl....,h h - 1} using the objective:
min Ei 1] 14.,aaa [[ (Xe, Xt i) - x
(2)
3

---

Interpolation is an easier task than forecasting, and we can use the resulting interpolator Ig during inference to interpolate beyond the temporal resolution of the data. That is, the time input can be continuous, with i (0, h 1), where we note that the range (0, 1) is outside the training regime, as discussed in Appendix D.5.6. To generate probabilistic forecasts, it is crucial for the interpolator; Ip. to produce stochastic outputs within the diffusion model and during inference time. We enable this using Monte Carlo dropout [17] at inference time.
Forecasting reverse process. In the second stage, we train a forecaster network Fè to forecast X such that F& (Ig (xe, Xt+h; i|E) ,) s X for i S = [in] where S denotes schedule mapping the diffusion step to the interpolation timestep. The interpolator network, Iø, is frozen with inference stochasticity enabled, represented by the random variable ę. Here, & stands for the randomly dropped out weights of the neural network and is omitted henceforth for clarity. Specifically, we seek to optimize the objective
min -.. [ (Ios (Xe, Xe+hi &n ix) X(+h
(3)
To include the setting where Fg learns forecast the initial conditions, we define ig :- 0 and Is (x) == Xg. In the simplest case, the forecaster network is supervised by all possible timesteps given by the temporal resolution of the training data. That is, N = h and S -i Generally, the interpolation timesteps should satisfy 0 = in < in < im < h for 0 < n < m S Given the equivalent roles between our forecaster and the denoising net in diffusion models, we also refer to them as the diffusion backbones. As the time condition to our diffusion backbone is in instead of n, we can choose any ceffiusion-dynaninn schedule during training or inference and even use Fe for unseen timesteps (see Fig. 3). The resulting two-stage training algorithm is shown in Alg. 1.
Because the interpolator Ig is frozen in the second stage, the imperfect forecasts &14h FI (Xe, X1+h, (x) in) may degrade accuracy when used during sequential sampling. To handle this, we introduce a one-step look-ahead loss term || (Ip (X,, & 1n+1) (2+1) x when- ever n n++1 < N and weight the two loss terms equally. Additionally, providing a clean noised form of the initial conditions X2 as an additional input to the forecaster net can improve performance. These additional tricks discussed in Appendix B.
Sampling.
Taken together, we can write the generative process of DYffusion as follows:
po(s(n-1)ssh ((( xx) ( (Fo(s(n),))) in) (X,, Fo(s(m), (x+) otherwise,
if N 1
(X,,
where s(0) x, and (!)) X(+i, correspond to the initial conditions and predictions of intermediate steps, respectively. In our formulations, we reverse the diffusion step indexing to align with the temporal indexing of the data. That is, n = refers to the start of the reverse process (i.e. x(). while n = N refers to the final output of the reverse process (here, X1+1). Our reverse process steps forward in time, in contrast to the mapping from noise data in standard diffusion models. As result, DYffusion should require fewer diffusion steps and data.
Similarly the forward pro- cess in [61, 2], our interpola- tion stage ceases to be a diffu- sion process. Our forecasting stage as detailed in Eq. (3), fol- lows the (generalized) diffusion model objectives. This similar- ity allows us to use many ex- isting diffusion model sampling methods for inference, such as the cold sampling algorithm from [2] (see Alg. 2). In Appendix C.2, we also discuss a simpler but less performant sampling algorithm. During the sampling process, our method essentially alternates between forecasting and interpolation, as illustrated in Fig. 2. Fo always predicts the last timestep, Xith, but iteratively improves those forecasts as the reverse process comes closer in time to + This is analogous to the iterative denoising of the "clean" data in standard diffusion models. This motivates line 6
Algorithm 2 Adapted Cold Sampling [2] for DYffusion
Input: Initial conditions x; := Xx, schedule [6,] output
timesteps J (by default JJ = {l.....A .h. 1})
for л = 0, .......N NN
do
3:
81th t Fo (8 in)
8iti.t1 = Is (X,, 81th, (n+1) - Is (x,, &eth, in) + &eti,
5: end for
X,4j + Io (x,, (1th.) Vj E
Optional refinement
7: Return: {84t) 1 E J} (844)

---

of Alg. 2, where the final forecast of X can be used to fine-tune intermediate predictions or to increase the temporal resolution of the forecast. DYffusion can be applied autoregressively forecast even longer rollouts beyond the training horizon, as demonstrated by the Navier-Stokes and spring mesh experiments in section 5.
0)
(1)
(2)
&
x
= 0
2
3
1n
1 =
2
Xg
X2
X3
Figure 2: During sampling, DYffusion alternates between forecasting and interpolation, following Alg. 2. In this example, the sampling trajectory follows simple schedule of going through all integer timesteps that precede the horizon of h 4, with the number of diffusion steps N h. The output of the last diffusion step is used as the final forecast for & The black lines represent forecasts by the forecaster network, Fg. The first forecast is based on the initial conditions, Xy. The blue lines represent the subsequent temporal interpolations performed by the interpolator network, Ig-
Memory footprint. DYffusion requires only x, and X++hh (plus X during the first stage) to train, resulting in constant memory footprint function of h. In contrast, direct multi-step prediction models including video diffusion models or (autoregressive) multi-step loss approaches require X to compute the loss. This means that these models must fit h + 1 timesteps of data into memory. which scales poorly with the forecasting horizon h. Therefore, many are limited to predicting a small number of frames. MCVD, for example, trains maximum of 5 video frames due to GPU memory constraints [67].
Reverse process as ODE. DYffusion can be interpreted modeling the dynamics by a "diffusion process", similar to an implicit model [13]. This view helps explain DY ffusion's 's superior forecasting ability, as it learns model the physical dynamical system x(t). The neural networks Ig and Fg are trained to construct a differential equation that implicitly models the solution to x(t).
Let be the time variable. DYffusion models the dynamics
dx(s)
dIg (x,, Fx s),
(5)
ds
ds
The initial condition is given by x(1) = Xx. During prediction, we seek to obtain
((s) = x(1) + J." dIs (x,, Fo(x, ds for E (t, + hj.
(6)
ds
The prediction process is equivalent to evaluating (6) at discrete points in the forecasting window (t, + h.. Sampling schedules i, can be thought of as step sizes ODE solvers. Therefore, different sampling schedules can be flexibly chosen at inference time, since they are discretizations of the same ODE. Different sampling algorithms can be obtained from different ODE solvers and we prove that cold sampling is approximation of the Euler method (derivation in Appendix C.I).
Viewing DYffusion as an implicit model emphasizes its dynamics-informen nature and reveals the relation between DYffusion and DDIM [61]. Specifically, DDIM can be viewed as an implicit model that models the solution the noise-remove removing dynamics from a random image to a clean image. In comparison, DYffusion is an implicit model that learns the solutions of dynamical system using data from the current step x, the future step X
Related Work
Diffusion Models. Diffusion models [25, 60, 62, 63] have demonstrated significant success in
diverse domains, notably images and text [76]. While interesting lines of work explore Gaussian diffu- sion for multivariate time-series imputation [1] and forecasting [51], nigh-dimensional spatiotemporal
5

---

forecasting has remained unexplored. Methods such as combining blurring and noising techniques [28, 10, 2] and adopting unified perspective with Poisson Flow Generative Models [75] have shown promise as alternative approaches to extend diffusion processes beyond Gaussian noise. Although intellectually intriguing, the results of cold diffusion [2] and other recent attempts that use different corruptions [53] are usually inferior to Gaussian diffusion. Our work departs from these methods in that we propose a dynamics-informed process for the conditional generation of high-dimensional forecasts that is not based on multiple levels of data cornuption.
Diffusion Models for Video Prediction. Most closely related to our work are diffusion models for vidcos [67, 27, 77, 59, 26, 24]. The task of conditional video prediction can be placed under the multi-step notation of Eq. (1). Notably, we are interested in modeling the full dynamics and the underlying physics rather than object-centric tasks. There also evidence that techniques tailored to computer vision and NLP do not transfer well to spatiotemporal problems [66]. These video diffusion models also rely internally on the standard Gaussian noise diffusion process inherited from their image-generation counterparts, As such, these models miss the opportunity to incorporate the temporal dynamics of videos into the diffusion process.
Dynamics Forecasting. Deterministic next-step forecasting models can be unrolled autoregressively
to achieve multi-step predictions. However, this may result in poor or unstable rollouts due to compounding prediction errors [11, 9, 8, 34, 5, 33, 43]. The issue be alleviated by unrolling the model at training time to introduce a loss term on the multi-step predictions [45, 34, 36, 23, 5]. However, this approach memory-inten -eemory-nteensive and comes at the expense of efficiency since the model needs to be called sequentially for each multi-step loss term. It can also introduce new instabilities because the gradients have to be computed recursively through the unrolled steps. Even then, this approach may fall behind physics- based baselines for long-range horizons [45]. Alternatively, forecasting multiple steps all at once as in video diffusion models [73, 52, 5] or conditioned on the time [37, 49, 15, 43], training separate models for each required horizon [56. 22, 68], noising the training data [55], and encoding physics knowledge into the ML model [11, 14, 70, 41, 9, 35, 8], can often produce better long-range forecasts, though they still require fixed temporal resolution. For probabilistic dynamics forecasting [74], most methods focus on autoregressive next- step prediction models [58, 18, 29]. Some notable works target precipitation nowcasting [52], or short-term forecasting [15]. Alternatively to using intrinsically probabilistic models, one can use a deterministic model generate ensemble forecast by perturbing its initial conditions [45, 21]. Our framework addresses primary challenges in the field [64], specifically realistic long-range forecasts and efficient modeling of multi-step dynamics.
Pe
X0
81
&2
E1.
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. Xo, and earliest output timestep, & 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).
Experiments
5.1 Datasets
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 x60 longitude resolution the eastern tropical Pacific Ocean). Unlike the data based on the NEMO dataset in [11, 71J, we choose OISSTv2

---

Table Results for sea surface temperature forecasting of 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 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 DY are trained
horizon of h 16. MCVD and DDPM trained on h 4 and I, 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,180
0,411
0.4241
0.090
-001
0.028
0.448 0.00
Dropout
0.267 0.033
0,164 1 0.04
0,406
0002
0,4241
0.078
0.00
0.027
-000
0.715
1.006
DDPM
0.246 0.005
0.177
0.674 0.00
0.3054
0.180
0.044
0.105
0.573 1.000
MCVD
0.216
0.161
0.926
79.167
0.154
0.43
0.070
00083
0.524 0.004
DYffusion 0.224 0.001
0.173 0:011
1,033 0,005
4.6722
0.067 0.003
0.022 0.002
0.877 0.806
as our SST dataset because it contains more data (although has a lower spatial resolution of 1 1/40 compared to 1 1/120 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.I.I.
Navier-Stokes flow: benchmark dataset from [44], which consists of a 221 x 42 grid. Each trajectory contains four randomly generated circular obstacles that block the flow. The viscosity is 1e-3. The channels consist of the æ and y velocities as well as a pressure field. Boundary conditions and obstacle masks are given as additional inputs all models.
Spring Mesh: benchmark dataset from [44]. It represents a 10 x 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 probabilistic multi-step forecasting, and especially long rollouts, as opposed to attaining the best next-step or short-term forecasts. This an important distinction because, while deterministic single-step forecasting models such 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 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 popular UNet architecture designed for diffusion models For the Ravier-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-basee fcaturization module of the SST UNet to embed the diffusion step or dynamical timestep.
We compare our method against both direct applications of standard diffusion models to dynamies 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.
Baselines.
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).
DDPM [25]:
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 back- bone network at inference time.
Perturbation [45]: Ensemble multi-step forecasting based on random perturbations of the initial conditionssinputs with fixed variance of the barebone backbone network.
wtwpaanniiiiiinifffffuieeeeeee diffusion- httpa///ttttooeoooo//aaeeeeeeeee huffuusooo.yoorrcc///ooooinnnddd http://ssiiinniifssssoooooorrorr
7

---

Table 2: Spring Mesh results. Both methods are trained on a horizon of 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 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.
Test
Out of Distribution
Method
CRPS
MSE
SSR
CRPS
MSE
SSR
Dropout
0.0138 + 0.0066
7.27c-4 1.865
1.01 0,02
0.0448
0.0007
7.08c-3
LT7-4
0,70 = 0.01
DYffusion 0.0103 + 0.0022
4.20c-4 + 2.144
1.13
0.0292 0.0009
3.54e-3 2.0-4
0.82 = 0.00
MCVD and the multi-step DDPM variant predict the timesteps X++1:t+h based on x,, which is the standard approach in video (prediction) diffusion models. The barebone backbone network baselines are time-conditioned forecasters (similarly to the DYffusiin forecaster) trained on the multi-step objective E,-H([IA) 1,) x.+++x+) [o, i) x++ill? from scratch. We found it to perform very similarly to predicting all 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 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 RMSF. It serves as a measure of the reliability of the ensemble, where values smaller than indicate underdispersion (i.e. the probabilistic forecast is overconfident in its forecasts), and larger values overdispersion [16, 18). On the Havier-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.
5.3 Results
Quantitative Results. We present the time-averaged metrics for the SST and Navier-Stoker dataset in Table and for the Spring Mesh dataset in Table 2. DYffussion 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',
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 cine-conditioned nature of DYffusion's interpolator and forecaster nets, memory is not 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 distriibiiii test dataset. It is worth noting that our reported MSE scores 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-distribene -istrriiitttoon test of the Navier-Stokes benchmark, the results almost identical 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, were not able successfully train MCVD models with fewer diffusion steps.
8

---

[=2
64
Tut
ongan
omte
Figure 4: Qualitative forecasts for timesteps 2, 6. 24, 46, and 64 (last timestep) of the veloc- ity norm of an example Navier-Stokes test trajectory. Here, we generate five sample trajecto- ries for the best baseline (Dropout) and our method DYffusion, both with h 16, and visu- alize the one with the best trajectory-averaee MSE for each of the methods. Our method (bot- tom row) can reproduce fine-scale details visibly better than the baseline (see c.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 hwtiddiiveegoggl...o google. cowccooo/leeeee fon/// SXKVV422111818SVT ZmAKR159qiHG hwtp://vseeeessssssseee
samples for the best baseline (Dropout) and DYffusion as well as the corresponding ground truth at five different timesteps from complete Havier-Stoker 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 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 DYffusion, we aim study in this experiment whether it is possible forecast skillfully beyond the resolution given by the data. Here, we forecast the same Navier-Stokes trajectory shown in Fig. 4 but at 8x resolution (i.e. 512 timesteps instead of 64 are forecasted in total). This behavior be achieved by either changing the sampling trajectory [ 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 this URL.
Note that we hope that our probabilistic forecasting model can capture any of the possible, uncertain futures instead of forecasting their mean, deterministic model would do. As 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 samples show sufficient variation, but also cover the ground truth quite well (sample l).
Ablations. In Appendix D.5, we ablate various com- ponents of our method. This includes showing that inference stochasticity in the interpolation network is crucial for performance. Our experiments in Ap- pendix D.5.1 1 confirm the findings [2] regarding the superiority of cold sampling (see Alg. 2) compared the naive sampling (Alg. 4) algorithm. We provide theoretical justification 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 training from 2017, 2015, 1982 until the end of 2018, respectively. The stan- dard deviations are around 0.01 or lower.
# years
2
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 the simplest possible schedule S N=
1 where the number of diffusion steps and the horizon equal. For the SST dataset, using & additional diffusion steps corresponding to floating-poin ij., especially such that is < 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 difffusion steps, resulting in a clear trade-off between accuracy and speed. Finally, we verify that the forecaster's 's prediction of X for given diffusion step n and associated input &tti, usually improves as n increases, i.e. as we approach the end of the reverse process and Xi+i, gets closer in time to Xith-
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 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 Fy and I on small subsets of the SST training set. We find that DYffusion's 's performance degrades more gracefully under limited training data, especially when using only two years of data (or 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 the domain of natural images [40]. This could be the reason why the studied diffusion model baselines produce poor predictions for the Navier-Stoker and spring mesh datasets.
Conclusion
We have introduced DYffusion, a novel dynamics dynamics-inforeed 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 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 models to other conditional generation problems, such as super-resolution, where the data itself can natively inform inductive bias for the diffusion model.
more general source of inspiration for tailoring diffusion
Limitations. While our proposed framework can be applied to any dynamics forecasting problem as alternative 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 forecast outperform sequential sampling of diffusion models, including Dfffusion.
Future work. Cold Diffusion [2] indicates that deterministic data degradations cannot meet the performance of stochastic (i.e. Gaussian noise- noise-based) degradations. In our work, we have exper- imented with using an interpolator network with Monte Carlo dropout enabled during inference time to achieve 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-bane 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 in- teresting 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.
Acknowledgements
This work was supported in part by the U.S. Army Research Office under Army-ECASE award W91 INT-07-R-0003-03, the U.S. Department Of Energy, Office of Science, IARPA HAYSTAC Program, NSF Grants #2205093, #2 #2146343, and #2134274.
NOAA OI SST V2 High Resolution Dataset data provided by the NOAA PSL, Boulder, Colorado, USA, from their website at httpt:://// //psll......... noaa. gov..
10

---

References
[1] Juan Lopez. Alcaraz and Nils Strodthoff, Diffusion-base time series imputation and forecasting with structured state space models. Transactions Machine Learning Research, 2023. ISSN 2835-8856. URL /ttp:///////eeeeeereeeeeeee net/ homm?dddhiiibbe
Tbk7ApW.
[2] Arpit Bansal, Eitan Borgnia, Hong-Min Chu, Jie S. Li, Hamid Kazemi, Furong Huang. Micah Goldblum, Jonas Geiping, and Tom Goldstein. Cold diffusion: Inverting arbitrary image transforms without noise. Advances in Neural Information Processing Systems, 2023.
[3] Peter Bauer, Alan Thorpe, and Gilbert Brunet. The quiet revolution numerical weather prediction. Nature, 525(7567):47-55, Sep 2015. ISSN 1476-4687. doi: 10. 10.1038/nature149555
URL https:///ooiii org/ http:///....rr/11111/////////ee/ 10w.//.rrre//////eeee.e 14956.
[4] Emanuele Bevacqua, Laura Suarez-Gutierrez, Aglaé Jézéquel, Flavio Lehner, Mathieu Vrac, Pascal Yiou, and Jakob Zscheischler. Advancing research on compound weather and climate events via large ensemble model simulations. Nature Communications, 14(1):2145, April 2023.
[5] Johannes Brandstetter, Daniel Worrall, and Max Welling. Message passing neural PDE solvers. In International Conference Learning Representations, 2022. URL https:// openreview.. openreview.net/forumn www.:rrrmmiiiiiiiiniiisi.a....e
[6] Johannes Brandstetter, Rianne van den Berg, Max Welling, and Jayesh K Gupta. Clifford neural layers for PDE modeling. In The Eleventh International Conference on Learning Representations, 2023. URL https::////.. httpt::s/ppppeeeeeeeeeeeee
[7] Casey M Brown, Jay R Lund, Ximing Cai, Patrick M Reed, Edith A Zagona, Avi Ostfeld, Jim Hall, Gregory W Characklis, Winston Yu, and Levi Brekke. The future of water resources systems analysis: Toward a scientific framework for sustainable water management. Water resources research, (1)8:6110-6111 2015.
[8] Ashesh Chattopadhyay and Pedram Hassanzadeh. Why are deep learning-based models of geophysical turbulence long-term unstable? NeurIPS, Machine Learning for Physical Sciences, 2022.
[9] Ashesh Chattopadhyay, Mustafa Mustafa, Pedram Hassanzadeh, and Karthik Kashinath. Deep spatial transformers for autoregressive data-driven forecasting of geophysical turbulence. In Proceedings of the 10th International Conference on Climate Informatics, pages 106-112, 2020.
[10] Giannis Daras, Mauricio Delbracio, Hossein Talebi, Alexandros G. Dimakis, and Peyman Milanfar. Soft diffusion: Score matching for general corruptions. 2022. URL httttt/////aiii org/ arg/99bs22099.. 055422.
[11]
Emmanuel de Bézenac, Arthur Pajot, and Patrick Gallinari. Deep learning for physical pro- cesses: Incorporating prior scientific knowledge. In International Conference on Learning Representations, 2018. URL.
net/fformmtt fiiimmidd-BBoa Bid-B4HsfWAZZ.
[12] Emmanuel de Bézenac, Syama Sundar Rangapuram, Konstantinos Benidis, Michael Bohlke-Schneide Richard Kurle, Lorenzo Stella, Hilaf Hasson, Patrick Gallinari, and Tim Januschowski. Normalizing kalman filters for multivariate time series analy- sis. In Advances in Neural Information Processing Systems, volume 33, pages 2995- 3007, 2020. URL https://proceeddiin neurips /tt////fee6666002222666612222220 Paper. pdf.
[13] Laurent El Ghaoui, Fangda Gu, Bertrand Travacca, Armin Askari, and Alicia Tsai. Implicit deep learning. SIAM Journal on Mathematics of Data Science, 3(3):930-953) 2021.
[14] N. Benjamin Erichson, Michael Muchlebach, and Michael W. Mahoney. Physics-informed autoencoders for lyapunov-stable fluid flow prediction. 2019. URL /tttpps:/arxxv.orr. org/ 0bs/1905.10866
[15] Lasse Espeholt, Shreya Agrawal, Casper Sønderby, Manoj Kumar, Jonathan Heek, Carla Bromberg, Cenk Gazen, Rob Carver, Marcin Andrychowicz, Jason Hickey, Aaron Bell, and Nal Kalchbrenner. Deep learning for twelve hour precipitation forecasts. Nature Communications, 13(1):5145, Sep 2022. ISSN 2041-1723 doi: 1111038/44147--0222244555 //doi. /ww.o.///........... 0008//41467-022 222214111-22-222--------
URL https:
[16) V. Fortin, M. Abaza, F. Anetil, and R. Turcotte. Why should ensemble spread match the mmse of the ensemble mean? Journal of Hydrometeorology, 15(4):1708 -1713, 2014. doi:
11

---

D- 10----14-0000 14-0008.1. URL
ametsoc org/view/
Cournals/hydr/155 httm/////////////////// d- 14- 0008-1001111111 xmm1.
[17] Yarin Gal and Zoubin Ghahramani. Dropout bayesian approximation: Representing model uncertainty in deep learning. In international conference on machine learning, pages 1050-1059.
PMLR, 2016.
[18] Sagar Garg, Stephan Rasp, and Nils Thuerey. Weatherbench probability: A benchmark dataset for probabilistic medium-range range weather forecasting along with deep learning baseline models. arXiv preprint arXiv:2205. 00065, 2022.
[19] Tilmann Gneiting and Matthias Katzfuss. Probabilistic forecasting. Annual Review of Statistics and Its Application, 1(1):125-151, 2014.
[20]
Tilmann Gneiting and Adrian E. Raftery. Weather forecasting with ensemble methods. Science,
3 31065746) 2005. doi: 10.1126/science] 111225............ URL https://www //tpp//////ssssieneee... science. org/doi/abs 10. 1...1.1...1 www1/112csiieni ence...cc..... 1115255..
[21] Andre Graubner, Kamyar Kamyar Azizzadenesheli, Jaideep Pathak, Morteza Mardani, Mike Pritchard, Karthik Kashinath, and Anima Anandkumar. Calibration of large neural weather models. In NewrlPS 2022 Workshop on Tackling Climate Change with Machine Learning, 2022.
[22] Yoo-Geur Ham, Jeong- -wan Kim, and Jing-Jia Luo. Deep learning for multi-year enso forecasts. Nature, 573:568- -7288857777998 9 2019.
[23] Xu Han, Han Gao, Tobias Pfaff, Jian-Xun Wang, and Liping Liu. Predicting physics in mesh- reduced space with temporal attention. In International Conference on Learning Representations, 2022. URL https: / /ttp//ppeeeeeieeeeeeeeeeeeee wwwtffoummmiiieeeeeeee werrumiiddxttLddFFE.
[24] William Harvey, Saeid Naderiparizi, Vaden Masrani, Christian Weilbach, and Frank Wood. Flexible diffusion modeling of long videos. Advances in Neural Information Processing Systems, 2022.
[25] Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. In Advances in Neural Information Processing Systems, 2020. URL httpt:////pocceedinnng .dd. [26] Jonathan Ho, William Chan, Chitwan Saharia, Jay Whang, Ruiqi Gao, Alexey Gritsenko, Diederik P. Kingma, Ben Poole, Mohammad Norouzi, David J. Fleet, and Tim Salimans. Imagen video: High definition video generation with diffusion models. 2022. URL https:
neurips. cc ctcccc/pppp/.c
//arxiv. org/ abs/ /2210 2ttp//221...0......2
[27] Jonathan Ho, Tim Salimans, Alexey Gritsenko, William Chan, Mohammad Norouzi, and David Fleet. Video diffusion models. Advances in Neural Information Processing Systems, 2022.
[28] Emiel Hoogeboom and Tim Salimans. Blurring diffusion models. In The Eleventh International Conference on Learning Representations, 2023.
[29] Yuan Hu, Lei Chen, Zhibin Wang, and Hao Li. Swinvrnn: A data-driven ensem- ble forecasting model via learned distribution perturbation. Journal of Advances in Modeling Earth Systems, 15(2):22022MS003211, 2023. doi: https://doi. http://iiior/100.102.00000 1029/ 2022MS003211. URL https: onlinelibrary 0029/2022H5003211. e2022MS003211 2022MS003211.
caerrrrrwwleey.. wiley. cow/ doi http///oo//b//...//// abs/10.
[30] Boyin Huang, Chunying Liu, Viva Banzon, Eric Freeman, Garrett Graham, Bill Hankins, Tom Smith, and Huai-Min Zhang. Improvements of the daily optimum interpolation sea surface temperature (doisst) version 2.1. Journal of Climate, 34(8):2923 - -9939, 2021. doi:
http:://doi.org/10.11111 1ttp///011.1/1111111 20-0666.111111 URL
ametsoc. wwtesoo..org//// org/view/
journals/clim/34/ -CLI-D-200 --------2-0------ 2066166... 1 xml.
[31] Steeven Janny, Aurélien Bénéteau, Madiha Nadri, Julie Digne, Nicolas Thome, and Christian Wolf. EAGLE: Large-scale learning of turbulent fluid dynamics with mesh transformers. In The Eleventh International Conference on Learning Representations, 2023. URL https:
/ /openreview nwt/ffoummiddmiiii IXQQPAS..
[32] Tero Karras, Miia Aittala, Timo Aila, and Samuli Laine. Elucidating the design space of diffusion-based generative models. In Proc. NeurIPS, 2022.
[33] K. Kashinath, M. Mustafa, A. Albert, J-L. Wu, C. Jiang, S. Esmaeilzadeh, K. Azizzade- nesheli, R. Wang, A. Chattopadhyay, A. Singh, A. Manepalli, D. Chirila, R. Yu, R. Wal- ters, B. White, H. Xiao, H. A. Tchelepi, P. Marcus, A. Anandkumar, P. Hassanzadeh, and
12

---

null Prabhat. Physics-informed machine learning: case studies for weather and climate modelling. Philosophical Transactions of the Royal Society A: Mathematical, Physical and Engineering Sciences, 379:2194);202200099, 2021. doi: (098//sta.2020000095 URL
https:/ http:::/rooalsscceeeppllliiiiiie http://puub..iiiinnooogooooooooo http://///////o///////////////// att///.aa////////////../.../...a 1098000.001001111 0000////aaa...... 2020. 0093.
[34] Ryan Keisler. Forecasting global weather with graph neural networks. 2022. URL. https:
//arxiv. hts//bb//2000 07575.
[35] Dmitrii Kochkov, Jamie A. Smith, Ayya Alieva, Qing Wang, Michael P. Brenner, and Stephan
Hoyer. Machine seeming-acceleraan computational fluid dynamics. Proceedings of the National Academy of Sciences, 118(21))e1001784118 2021. doi: (0.1073/pnas.201173
URL htthht:s//wwwg wttp:ppp..oooooop///ooo//ossssss 10t.///aa......ppaa.aa 2101784118.
[36] Remi Lam, Alvaro Sanchez- Matthew Willson, Peter Wirnsberger, Meire Fortunato, Alexander Pritzel, Suman Ravuri, Timo Ewalds, Ferran Alet, Zach Rosen, Weihua Hu, Alexander Merose, Stephan Hoyer, George Holland, Jacklynn Stott, Oriol Vinyals, Shakir Mohamed, and Peter Battaglia. GraphCast: Learning skillful medium-range global weather 2212. 12794.
forecasting. 2022. URL
[37] Nathan Lambert, Albert Wilcox, Howard Zhang, Kristofer Pister, and Roberto Calandra. Learning accurate long-term dynamics for model-based reinforcement learning. In IEEE Conference on Decision and Control (CDC), 2021.
[38] Jeffrey K. Lazo, Rebecca E. Morss, and Julie L.. Demuth. 300 billion served: Sources, perceptions, uses, and values of weather forecasts. Bulletin of the American Meteoro- logical Society, 90(6):785 - 798, 2009. doi: 10. 10111//2008BAM222611. URL https: //journals. ametsoc. org/ooo///.. jouraal//ooas////aaaaaa /baass//////////a/ 20100///2000bm22604.4.2 xml.
[39] Jane A Leggett. The united nations framework convention climate change, the kyoto protocol, and the paris agreement: a summary. UNFCC: New York, NY, USA, 2, 2020.
[40] Aaron Lou and Stefano Ermon. Reflected diffusion models. In International Conference on Machine Learning, 2023.
[41] Giorgos Mamakoukas, Ian Abraham, and Todd Murphey. Learning stable models for prediction and control. IEEE Transactions on Robotics, 2020.
[42] James E. Matheson and Robert L. Winkler. Scoring rules for continuous probability distributions. Management Science, 22(10):0087-109, 1976.
[43] Tung Nguyen, Johannes Brandstetter, Ashish Kapoor, Jayesh K Gupta, and Aditya Grover. ClimaX: A foundation model for weather and climate. International Conference on Machine Learning, 2023.
[44] Karl Otness, Arvi Gjoka, Joan Bruna, Daniele Panozzo, Benjamin Peherstorfer, Teseo Schneider, and Denis Zorin. An extensible benchmark suite for learning to simulate physical systems. Advances in Neural Information Processing Systems (NeurIPS 2021) Track on Datasets and Benchmarks, 2021. URL httttt:t////stteeeee org/ abs///bs/2100 07799.
[45] Jaideep Pathak, Shashank Subramanian, Peter Harrington, Sanjeev Raja, Ashesh Chattopadhyay, Morteza Mardani, Thorsten Kurth, David Hall, Zongyi Li, Kamyar Azizzadeneshelli Pedram Hassanzadeh, Karthik Kashinath, and Animashree Anandkumar. FourCastNet: global data- driven high-resolution weather model using adaptive fourier neural operators. 2022. URL
https: / /ttp://aaxvvvv/// wwwooo/g//ooooooo 2202. 11214.
[46] Haiganoush K Preisler and Anthony L Westerling. Statistical model for forecasting monthly large wildfire events in western united states. Journal of Applied Meteorology and Climatology, 46(7):1020-1030, 2007.
[47] Aditya Ramesh, Prafulla Dhariwal, Alex Nichol, Casey Chu, and Mark Chen. Hierarchical text-conditional image generation with clip latents. arXiv preprint 2XXx:22044006225 2022.
[48] Stephan Rasp and Sebastian Lerch. Neural networks for postprocessing ensemble weather forecasts. Monthly Weather Review, (46((1)):388) - 3900, 2018
[49] Stephan Rasp and Nils Thuerey. Data-driven medium-range weather prediction with a resnet pretrained on climate simulations: A new model for weatherbench. Journal of Advances in Modeling Earth Systems, 13(2), 2021. doi: http://ddo.urg///...oo/2222220s0 URI. https:/ aggpubs. aulinnlibbarry wiley. http:///s/abb//10000//22220222//
13

---

[50] Stephan Rasp, Peter D. Dueben, Sebastian Scher. Jonathan A. Weyn, Soukayna Mouatadid, and Nils Thuerey. WeatherBench: A benchmark data set for data-driven weather forecasting. Journal of Advances in Modeling Earth Systems, 2020. https:/ org/ http//..199//// 2F2020ms002203.
[51] Kashif Rasul, Calvin Seward, Ingmar Schuster, and Roland Vollgraf. Autoregressive denoising diffusion models for multivariate probabilistic time series forecasting. In ICMI., 2021.
[52] Suman Ravuri, Karel Lenc, Matthew Willson, Dmitry Kangin, Remi Lam, Piotr Mirowski, Megan Fitzsimons, Maria Athanassiadou, Sheleem Kashem, Sam Madge, et al. Skilful precipi- tation nowcasting using deep generative models of radar. Nature, 597(7878):672-677, 2021. URL htt.hhttp///oo org/ 10. 1038/ 00088//445500 s11///441885611 038-022...
[53] Severi Rissanen, Markus Heinonen, and Arno Solin. Generative modelling with inverse heat dissipation. In The Eleventh International Conference on Learning Representations, 2023.
[54] Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Bjöm Ommer. High- resolution image synthesis with latent diffusion models. In JEEE Conference on Computer Vision and Puttern Recognition, pages 10684-10695, 2022.
[55] Alvaro
Jonathan Godwin, Tobias Pfaff, Rex Ying, Jure Leskovec, and Peter W. Battaglia. Learning to simulate complex physics with graph networks. In International Conference Machine Leurning. 2020.
[56] S. Scher. Toward data-driven weather and climate forecasting: Approximating a sim- ple general circulation model with deep learning. Geophysical Research Letters, 45(22): 12,616-12,62,, 2018. doi: https///bb.org/11000099/12210099 8GL080704 URL URR.hhttps///gguubb onlinelibrary wiley. wwt//oo/aws.///occct.oo..toota 1029 11w.11029/2200118800000000000
[57] S. Scheer and G. Messori. Generalization properties of feed-forward neural networks trained on lorenz systems. Nonlinear Processes in Geophysics, 26(4):381
2019. doi: 10.5194/ httpss:::: //pg. coperni cus. org/ aataarrticcll /26 /tt//////////////////
nng-26-381-200
[58] Sebastian Scher and Gabriele Messori. Ensemble methods for neural network-baid -baed weather forecasts. Journal of Advances in Modeling Earth Systems, 13(2), 2021. doi: https://doi.org
10.1029/0020MSS03333. URL https: http:///gupuub.uee www.:uppww....neeeeene....e...ee colineeibbarry wtw.wbrrrrrr..r.........rrreeeee www/ hon//do// /bbs/ 10. 20.1122220220S0023311
[59] Uriel Singer, Adam Polyak, Thomas Hayes, Xi Yin, Jie An, Songyang Zhang, Qiyuan Hu, Harry Yang. Oron Ashual, Oran Gafni, Devi Parikh, Sonal Gupta, and Yaniv Taigman. Make-a-video: Text-to-video generation without text-video data. 2022. URL https: org/ /bs// 1209.14799
[60] Jascha Sohl- Sick-Dicsttei, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsuper- vised learning using nonequilibrium thermodynamics. In Proceedings of the 32nd International Conference on Machine Learning, 2015. URL
mlr www..presssveee
sohl-dicksteinll html.
[61] Jiaming Song. Chenlin Meng, and Stefano Ermon. Denoising diffusion implicit models. In International Conference on Learning Representations, 2021.
[62] Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution. In Advances in Neural Information Processing Systems, volume 32, 2019.
[63] Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. In International Conference on Learning Representations, 2020.
[64] Jannik Thuemmel, Matthias Karlbauer, Sebastian Otte, Christiane Zarfl, Georg Martius, Nicole Ludwig, Thomas Scholten, Ulrich Friedrich, Volker Wulfmeyer, Bedartha Goswami, and Martin V. Butz. Inductive biases in deep learning models for weather prediction. 2023.
[65] Alasdair Tran, Alexander Mathews, Lexing Xie, and Cheng Soon Ong. Factorized fourier neural operators. In The Eleventh International Conference on Learning Representations, 2023. URL https:///penreviier
[66] Renbo Tu, Nicholas Roberts, Vishak Prasad, Sibasis Nayak, Paarth Jain, Frederic Sala, Ganesh Ramakrishnan, Ameet Talwalkar, Willie Neiswanger, and Colin White. Automl for climate change: A call to action. 2022. URL https: //txxv/r/vvv///////////ee////av org abs 2td.b//211 2210..221002 03324.
14

---

[67] Vikram Voleti, Alexia folicoeur-Martinede and Christopher Pal. MCVD: Masked conditional video diffusion for prediction, generation, and interpolation. Advances in Neural Information Processing Systems, 2022.
hattps://arxae org/ ab/aaa//s//22/ 098533.
[68] Etienne E Vos, Ashley Gritzman, Sibusisiwe Makhanya, Thabang Mashinini, and Camphell D) Watson. Long-range seasonal forecasting of 2m-temperatuur with machine learning. NeurIPS, Tackling Climate Change with Machine Learning, 2020.
[69] Zhong Yi Wan, Leonardo Zepeda-Nunin 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 tett/forumi1d-747333JJ
[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 Brown, Daniel R Cayan, and Michael D) Dettinger. Climate and wildfire in the western united states. Bulletin of the American Meteoro- logical 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, 118):2680-2699 2019. doi: https: //doi.or//00. 10to:.or//00.002//0001 10011000000000000000010000000000 705. URL https: //gga///////..////// onl onlineiibrary wiley com/ doi/abs/ 10. 1029/ /2229/22122 22//2001100000 22199M00000
[74] Dongxia Wu, Liyao Gao, Matteo Chinazzi, Xinyue Xiong, Alessandro Vespignani, Yi-An Ma, and Rose Yu. Quantifying uncertainty deep spatiotemporal forecasting. In Proceedings of the 27th ACM SIGKDD Conference 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-inspire generative models. 2023. URL https:
//arxiv. serxiv.org// org/ prg//bs//ae /aaaab//// 22022 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:///aaaiiii org/ /a//ab//2//0// 00796.
[77] Ruihan Yang, Prakhar Srivastava, and Stephan Mandt. Diffusion probabilistic modeling for vidco generation. arXiv preprint 2Et.:2203.0948,, 2022.
15

---

Contents
Introduction
2 Background
3 DYffussio: DYnamics-Informed Diffusion Model
4 Related Work
5 Experiments
5.1
Datasets
5.2
Experimental Setup
5.3
Results
6
Conclusion
10
Glossary
17
Methodology
17
B.1 DYffusion
17
B.2
Connection standard diffusion models
18
C Diffusion process ODE
18
C.I
Cold Sampling from the Euler method
C.2 Why cold sampling better than naive sampling?
19
D Experiments
20
D.I
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
Huvier-Stokes out-of 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 DYffusion and Baselines
27
D.7
Example Sampling Trajectories Visualized.
28
16

---

Appendix
A Glossary
The glossary is given in Table below.
Symbol
Used for
L
Indexer 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 #f 1-times during evaluation Total number diffusion steps Indexer the diffusion step, where S n S Distribution from which the data snapshots are sampled from Data point at timestep & Predicted data point timestep f Shorthand for sequence data points xx, x X
N
n
X
X
Xr
Initial conditions, i.c. input snapshot for the forecasting model Sequence of target snapshots that the forecasting model predicts given Xt Data point at diffusion step n Stochastic temporal interpolation neural network; Trained in the first-stage of DYffusion Forecasting neural network; Trained in the second-stage of DYffusion
g(n)
I
Fe
Interpolation schedule for DYffuusio that couples every diffusion step, Maps diffuusion step to timestep i.. where
timestep n < m S N
ia bn < in h Number of non-integer auxiliary diffusion steps for DYffusion. If l: = 0, then S = [in] bb Forward diffusion process operator, e.g. Gaussian noise, of conventional diffusion model Denoising neural network of conventional diffusion model
&
Ne
U[a,b]] b]
Descrihes the uniform distribution over the set {a, a 11.....b
Table Glossary variables and symbols used in this paper. Note that for DYffusion, the reverse process starts at 0 and progresses forward (together with the corresponding interpolation timestep, in) as + N ending with the final output s(N) For conventional diffusion models, this notation is reversed and the final output of the reverse diffusion process corresponds to s(
B Methodology
B.1 DYfflusion
Algorithm 3 DYffusion, Two-stage Training, all details
networks Fe, To, norm ||-I], horizon h, schedule [', conditioning
Input:
N-1 loss coefficients λ1 Ag. forecaster
Stage 1: Train interpolator network, I'
1. Sample ~ Uniform (fI,...
13)
2.
Sample x,, X X X'
3. Optimize min I (x,, X1+h, 1) x
Stage 2: Train forecaster network (diffusion model backbone), Fe
1. Freeze Ig and enable inference stochasticity (e.g. dropout)
2. Sample
Uniform (LO, .., N 1}) and Xr, Xt+h X' C4th + Fg (Io (Xe, Xith, in), in, c(x,,
3. & (1)
4. & (+h f- Fo (xe, R2 in+1) int1 c(xi, 1)) if n N - else Xt+h
(2)
5. Optimize ming Yo [1] x1+al12 + X2 22112(?7) -x+tAl2
In Alg. we give the full training algorithm for DYffusion. The only differences to the truncated training algorithm in the main text (Alg. 1) are that also include:
17

---

One-step look-ahead loss term trick, (Z (X,, & in+1), ( (n+1)) -x whenever n n+ N. Essentially, the look-ahead loss term simulates one step of the sampling process
(Alg. 2) and backpropagates through it, so that the network is trained with an objective that is closer to and partially mimics the iterative sampling process. In all experiments, we use \1 X2 2. We believe that it would be interesting to explore alternative approaches to this loss term, such as fine-tuning the interpolator network in the recond-stage too.
Optional conditioning of the forecaster network on a noised or clean version of the initial conditions, Xz- We explore three different approaches: 1) No conditioning at all, i.e. c(:-)) mull; 2) Clean initial conditions, i.e. c(x,,,) = x,; 3) Noised initial conditions, where we sample random noise (00, I) of the same shape of X,, and then define c(x,, n)) - x= , + (1 x= The second variant is analogous to the conditioning conventional diffusion models, which DYffusion does not necessarily require since the initial conditions indirectly influence the main input of the forecaster net (i.c. &tti, = Lg (x,, Xt+h: ). The motivation for the third variant stems from trying to give the forecaster net more explicit conditioning on the initial conditions when the reverse process goes farther away from (i.e. as >> N and &eti, - X1+h). An ablation for the choice of с is reported in D.5.4.
B.2 Connection to standard diffusion models
Our work builds upon cold diffusion [2], which "paves the way for generalized diffusion models that invert arbitrary processes". The cold sampling algorithm, which we use to sample from DYffusion (see Alg. 2), generalization of DDIM sampling for "generalized diffusion models" (see Appendix A.6 of [2]) and a key ingredient to make DYffusion work (see D.5.I). Our proposed forward and reverse processes are specifically designed to fall under this framework. Given this context, it becomes clear that DYffusion is generative model for forecasting that falls in the category of "generalized diffusion models""
Training objectives. To see the similarity between standard diffusion models and DYffusson in terms training objectives, we write out the corresponding equations Eq. and Eq. 3. For a standard "generalized" diffusion model,
rewrite the main objective inside Eq. 1 as follows:
Ro (D(s n), xx, n) - 122 1| (s(m)) xx, n) 10
(7)
where s. 1+ and s(n) is noisy version of (0) (its level of corruption increases with m))
In the context of DYffusion's forecasting network, we can express its training objective (Eq, 3) in the following manner:
[Io X4th, X4th |IFo (N)) kn), - 112 = =I(o(s!m) in) - 1(
(8)
(9)
where ((N) Xx+h and ((1) s Xe+i, is now stepped backward in time version of g(N), Note that the diffusion step indexing (superscript n) for DYffusion is reversed so that it aligns with the temporal indexing (subscript (), such that e.g. ((m) Xt+i,, temporally precedes (mn+1) Xt+ 11" Accounting for the reversed indexing of the diffusion steps, it is easy to see the equivalence of the denoising network, Rg, in a conventional diffusion model, and its forecasting network counterpart, Fe. in DYffusion. The reason why DYffusion 's forecaster network does not necessarily require the initial conditions, x,, as additional input is because the diffusion state, s("), depends on them. The central difference the design and semantic meaning of the diffusion processes accompanying them. We visualize this in the diagrams in Fig. 1.
C Diffusion process as ODE
In this section, we treat x as a continuous function of time, i.e. x R + R", + x(s). Here s denotes the time variable (since is already used to define the index of temporal data snapshots). During prediction, x, given, @ fixed, and I, only depends on & and Therefore, we simplify notations by writing Ig (x,, Seth, i) Ix,x (Reth, 2). We will further omit the subscripts & and x
18

---

C.1 Cold Sampling from the Euler method
In this section, we show that Cold Sampling is an approximation of the Euler method for (5).
The Euler method for integrating x is
dI (xx(; 8),
X, As
(10)
ds
for small As. We do not have access to dF( approximate (I by its first-order Taylor expansion around s:
However, since we know Fø and Iz. we can
da
de dI (x), [(Fo(Xx+..: s + As), + As) (XF) As
(11)
ds
This step can also be interpreted evaluating the integral in (6) using the fundamental theorem of calculus. Then the Euler method becomes
X As X, + I(Fo(Xx+ As), s + As) ((F 8)
(12)
Note that X As on the right-hand side is unknown because it is the quantity we want to approximate in this step. A reasonable way approximate Fg (X++Ae) 8 + As) is to replace it by (x,,, 8) because they both predict x(h) and use nearby points (assuming As is small and x behaves nicely around s). The resulting update,
X3- x, I(Fo( s), + As) - [(F. s).
(13)
is exactly the Cold Sampling algorithm (Alg. 2). By formulating the diffusion process as an ODE, we have provided a new theoretical explanation for Cold Sampling.
Note that we made the approximation that Fe s + As) = (o s) to obtain the update rule (previous equation). The error introduced by this approximation is expected be larger when is small. The intuition as follows. When 8 is small, the distance between s and + is large, and Fx s) has make prediction farther ahead. Predicting far into the future is generally harder than predicting the near future. Therefore, the prediction error Fx(x,, X is expected to be larger when s small. The larger uncertainty may lead larger difference between Fe (x*+As, As) and (x)
To reduce the error brought by this approximation, it makes sense sample more densely around the early part of the prediction window.
C.2 Why is cold sampling better than naive sampling?
The cold sampling algorithm (adapted for DYffusion in Alg. 2) was proposed as an improved version over so-called "naive" sampling in the original Cold Sampling paper [2] for generalized diffusion model Both
Algorithm 4 Adapted Naive Sampling [2]
1: Input: Initial conditions X; := xt, schedule [im.]/= for n = 0, 0, .1,.....N........... N do
3: 81th +- Fo Fo in)
S = Is (Xi, 8eths in+1) # difference w.r.t Alg.
5: end for
forward processes.
sampling algorithms can be directly used to sample from DYffusion. For completeness, we adapt the simpler, naive sampling algorithm our notation Alg. 4, where the only difference to cold sampling (Alg. 2) lies in the missing error-correction terms in line 4 of the algorithms. In our ablation experiments, we find that cold sampling indeed outperforms naive sampling by large margin (see Naive sampling row in Table 7), especially for the SST dataset where we use auxiliary diffusion steps. We explain by analyzing the discretization errors in the two sampling algorithms. In cold sampling, the discretization error per step is bounded by term proportional to the step size As. Naive sampling does not have this property.
The true value of x at 8 As according to Equation (6) is
As
dI (Fe (Fo s)
x(s + As)
x(s)
ds
= =x(s) + I + As), s + As), As) - I(F((x(s), s).
(14)
19

---

Recall from Alg. 2 that given x(s), cold sampling predicts x(s + As)
* As) = x(s) + I (x Fo(x(s), 8 + As) I (x (o(x(s)) 8).
(15)
The discretization error e(x) of one step of cold sampling is the difference between the exact and predicted x(s + As):
e(x, As) = x(5 + As) 8(+ + As)
= (x(( + As), # AS), s + As) - I(Fo((((())()))))))))))())())))) As).
(16)
The following proposition states that e(x: is bounded by term proportional to the step size As. Proposition C.I. Assume that (x) 8) is Lipschitz in s. Assume also that I& (Xeth, 8) is Lipschitz in Xithe The norm of the cold sampling discretization error, ||e(x, As) 112. is bounded by O(As).
Proof. The proof relied on applying definitions of Lipschitz functions twice. Let L1 be the Lipschitz constant for Fx(x(s), 8) in s. Let L2 be the Lipschitz constant for Z(x, s) x. Since F6 (x((s,,, s) is Lipschitz in we have [x(ss + As), s + As) Fa(x(s), s) I LıAs. Since I(x, s) is Lipschitz in x, we have ( + As), s As). s As) (x s). As)ll2 S LsL
Therefore ||e(x)]l & L2L,As, As, which means the discretization error is bounded by first-order term of the step size.
Under the same Lipschitz assumptions, the discretization error of the naive sampling is not guaranteed to be in the first order step size. naive sampling, the predicted x at time + As
(s + As)
8), As).
(17)
The discretization error of one step of naive sampling is x(s + As):
c(x, As) xx(s + As) x(5 + As)
(x(((x(5 + As). s As), s + As) - I(F(((((()))))))) As) + x(s) I[ (xx s),
(18)
Note that the first two terms are the same as the discretization error in cold sampling. However, the last two terms not bounded by first-order terms of As. Hence, naive sampling have larger discretization errors.
D Experiments
D.1 Datasets
D.1.1 SST data preprocessing
We create a new sea surface temperatures (SST) dataset based on NOAA OI SST V2 [30], which comes daily time scale. These data are available from 1982 to the present at resolution of 1/4 degrees. For training,
156
157
132
133
134
135
139
the years 1982-2018, for validation 2019, and for testing 2020. We have preprocessed the NOAA OI SST V2 dataset as follows:
108
109
111
113
84
85
87
88
89
1. First, the globe is divided into 60 x 60 latitude x longitude grid tiles,
60
61
63
65
66
Figure 5: Visualization the SST dataset that created. It divides the globe into 60 x 60 latitude x longitude grid tiles. We only the subset delineated in red, i.e. boxes 84-89 and 108-112
2. all tiles with less than 95% of ocean out,
filtered
3. standardize the raw SSTs using daily means and stan- dard deviations (computed on the training set only, i.e. 1982-2018).
4. replace continental NaNs with zeroes (after standard- ization), and
5. we subsample 11 grid tiles (covering mostly the east- emn tropical Pacific, shown in Fig. 5).
20

---

Table 5: The hyperparameters used for each dataset. For the learning rates, we sweep over each value and report the best of runs based their validation CRPS computed on 50 samples. DYffusion refers to the number of artificial diffusion steps used, see Fig. 3. For architectural details, see D.3. Hyperparameters for each dataset
Hyperparameter
SST
Navier-Stokes
Spring Mesh
Batch size
64
32
64
Accumulate gradient batches
4
1
Max. Epochs
50
200
300
Gradient clipping (norm)
1.0
1.0
1.0
Learning rate(s)
7e-4, 3e-4, 5e-5, 1e-5
7e-4, 3e-4
4e-4
Weight decay
le-5
le-4
le-4
AdamW 31
0.9
0.9
0.9
AdamW 32
0.99
0.99
0.99
DYffusion k
25
0
0
D.1.2 Navier-Stokes and spring mesh datasets
We refer to the physical systems benchmark paper for more details regarding the Navier-Stokes and spring mesh benchmark datasets [44]. We follow the same evaluation splits and procedure (only complemented by the probabilistic skill metries that we employ too). We always use the full training set provided by [44]. We use the Navier-Stoke dataset with four obstacles.
D.2 Implementation Details
The set of hyperparameters that we use for each dataset, such as the learning rate and maximum number of epochs, can be found in Table 5. For all experiments, we use floating point precision of 16 and do not use learning rate scheduler. All diffusion models, including DYffusion, are trained with the L.1 loss, while all barebone UNet/CNN networks are trained on the I.2 loss. We use three different dropout rates for the SST UNet: 1) before the query-key-value projection of each attention layer, dr... :r: 2) After the first sub-block of each ResNet block, dra,; 3) After the second sub-block of each ResNet block, drw,, where the first ResNet sub-block consists of convolution + normalization + time-embedding scale-shift - activation function, and the second sub-block is the same but without the time-embedding scale-shift.
Perturbation baseline We perturb the initial conditions, x,, with small amounts of Gaussian noise 6 (0, II.). We found that o- 0.05 gave the lowest CRPS scores among all variances that we tried, E (0.01, 0.03, 0.05, 0.07, 0. 0.15, 0.2). We note that choosing larger variances results in better SSR scores, but significantly lower CRPS and MSE scores. Inference dropout was disabled for this bascline variant.
Dropout baseline For this baseline, we enable the barebone model's dropout during both training and inference. For the SST dataset, similarly to the interpolator network of DY ffusion, we found that using high dropout rates results better performance. An explanation could be that the SST UNet has more capacity than the other backbone architectures. Concretely, the following SST UNet dropout rates resulted in the best performance: drat = drw,, 0.6, drM, = 0.3. For Navier-Stoke and spring mesh, there is only one dropout hyperparameter, and the corresponding best model uses 0.2 and 0.05 as dropout rate, respectively (selected from a sweep over (0.05, 0.1, 0. 0.2, 0.25, 0.3)).
We found that the cosine (linear) noise schedule gives better results for the SST (Navier- Stokes) dataset, and always uses the "predict noise" objective. For the SST dataset, the best performing DDPM is trained with 5 diffusion steps, while for Navier-Stokes it is trained with 500 steps. For Navier-Stokes, we found that while a DDPM with 5 or 10 diffusion steps can give good validation scores (or even better ones than the 500-step DDPM), it ends up diverging at test time after a few autoregressive iterations when used to forecast full trajectories.
DDPM
MCVD [67] We train MCVD with 1000 diffusion steps for all datasets, as we were not able to successfully train with fewer diffusion steps. We use a linear noise schedule (we found the cosine schedule produce inferior results) using the "predict noise" objective. Due to the inference runtime
21

---

complexity of using 1000 diffusion steps, we only report one MCVD run in our main SST results. Note that in none of the experiments we use the UNet-based diffusion model backbone originally used by MCVD. In preliminary experiments on the SST dataset, however, we found that wrapping MCVD around the SST UNet resulted in slightly improved scores over wrapping it around the original MCVD UNet.
DYffusion For the SST dataset, we use 25 artificial diffusion steps (analogous to the schedule in green in Fig. 3), while for the Navier-Stokes and spring mesh datasets, we do not use any, i.e. S = j=n Furthermore, for the SST dataset so we did not use it there, whereas it did improve performance for Navier-Stoker and spring mesh (see D.5.2). As for the choice of the interpolator network, Ig. we conduct sweep over the dropout rates for each dataset (analogous to the "Dropout" baseline). This is an important hyperparameter since we found that stochasticity in the interpolator is crucial for the overall performance of DYffusion. The interpolator network is selected based on the lowest validation CRPS. For the SST dataset, the selected Ig uses dr drw, = 0.6, drw, = O. The dropout rates for Navier-Stokes and spring mesh are 0.15 and 0.05, respectively. Generally, we found that the optimal amount of dropout for any given dataset strongly correlates between the "Dropout" multi-step forecasting baseline and DYffusion's 's interpolator network, Ig. Thus, for a new dataset or problem, it a valid strategy to sweep over the dropout rate for just one of the two model types. Motivated by the intuition that temporal interpolation is a simpler task than forecasting, the channel dimensionality of I& is only 32 (instead of 64) for the first downsampling block of the SST UNet. To feed the two inputs x, and X into the interpolator, we simply concatenate these snapshots along the channel dimensions, effectively doubling the number of input channels.
found that the refinement step of Alg. 2 did not improve performance
Evaluation All experiments are usually run over three random seeds when compute constraints allow. We use a 50-member ensemble for all reported results unless noted otherwise 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 I indicate underdispersion (i.e. the probabilistic forecast is overconfident in its forecasts), and larger values overdispersion [16, 18). We use the implementation in the xskillscore Python package to compute the CRPS of the ensemble forecasts. The CRPS penalizes both the absolute skill as well as the sharpness of the ensemble of forecasts. It rewards small spread (i.e. sharpness) if the forecast is accurate, and measures the difference between the forecasted and observed cumulative distribution function (CDF). For a deterministic forecast, the CRPS reduces to the mean absolute error. For CRPS and MSE lower is better, for SSR closer to 1 is better.
D.3 Neural Architecture Details
SST UNet For the SST dataset, we use a UNet implementation commonly used as backbone architecture of diffusion models". The UNet for the SST dataset consists of three downsampling and three upsampling blocks. Each block consists of two convolutional residual blocks (ResNet blocks), followed by an attention layer, and a downsampling (or upsampling) module. Each ResNet block can be further divided into two sub-blocks. The first one consists of convolution - normalization - time-embedding scale-shift - activation function, and the second sub-block is the same but without the embedding scale-shift. The downsampling module a 2D convolution that halves the spatial size of the input (with a 4 x 4 kernel and stride= 2). The upsampling module doubles the spatial size via nearest neighbor upsampling followed by 2D convolution (with end of each downsampling (upsampling) block the spatial size is halved (doubled) and the channel dimension is doubled (halved). We use 64 channels for the initial downsampling block, which means that the channel dimensionalities are 64 + 128 - 256 and the spatial dimensions (60, 60) - (30, 30) > (15, 15) in the corresponding downsampling blocks (reversed for the upsampling blocks). We use three different dropout rates for the SST UNet: 1) before the query-key-vare projection of each attention layer, drai; 2) Aferr the first sub-block of each ResNet block, dro,; 3) After the
kernel). At the
"nctpss//xskiillscoree www.:ooorerrrrttcceo.s......e 10/
diffuusioo_pyooor//ddoiiiia wtwiinngdddffffii www.:iiwn.ddffusssiooooooooooooo
22

---

Table Navier-Stokes results, including on the cattof-distributtoo (OOD) test set. Left columns are same as in Table 1. As observed by Fig. 12 in [44], the differences across main vs OOD test set are minor.
Test
Out of Distribution
Method
CRPS
MSE
SSR
CRPS
MSE
SSR
Dropout
0.078 1: 0.001
0.027 +: 0.001
0.715
0.005
0.078 1 0.001
0.027 = 0.001
0.715 1: 0.006
DDPM
0.180 + 0.004
0,105 = 0.010
0.573
0.001
0.189
+ 0.012
0,113 0.013
0.555
0.023
MCVD
0.154 0.043
0.070
0.033
0.524
0.064
0.170
0.016
0.082
0.036
0.489
0.061
DYffusion
0.067 +: 0.003
0.022 + 0.002
0.877
0.006
0.069
1 0.002
0.023 0.002
0.882 0.002
3:004
Diituaion
CW (Dropoul-mma mww.oummm...... step)
OVN (43)
3.000
Horten
(a) Navier-Stokes
(b) Spring Mesh
Figure 6: Comparison against single-step deterministic baselines from [44]. We plot the MSE as function of the rollout step (predicted horizon). It is clear that both our stochastic multi-step baseline, propout-multi-see Navier-Stoker (spring mesh) datasets of [44]. As reported in [44], the single-step CNN diverges after a few autoregressive steps on the spring mesh dataset. In addition, DY ffusion performs especially well long-range forecasts relative to the baselines. For spring mesh plot the three forecasted rollout samples of each model separately due to higher variance between samples.
and DYffusion outperform the UNet (CNN) models from [44] on the
second sub-block of each ResNet block, dratg. For models except the Dropout baseline and Lg in DYffusion, which use higher dropout rates, we use dr 0.1, dror, 0, dr 0.3.
Navier-Stokes UNet and spring mesh CNN For the Navier-Stokes and spring mesh benchmark datasets from [44], we simply re-use their proposed UNet and CNN architecture for the respective dataset. The only modification is the integration of the SST UNet time embedding module into the architectures of [44], described below.
Time embedding module The time-emmedee embeddinn scale-shift, taken from the SST UNet, is a key component of all architectures since enables them to condition the diffusion step (for DDPM and MCVD) or the dynamical timestep (for the time-conditioned barebone models as well as for both Fg and I, in DYffusion). it implemented by sine and cosine-based featurization of the scalar diffusion step/dynamical timestep. These features are projected by linear layer, followed by a GeLU activation, and another linear layer, which results in "time embedding". Then, separately for each convolutional (or ResNet) block of the neural architecture the "time embedding" is further processed by a SiLU activation and another linear layer whose output is interpreted as two vectors that are used to scale and shift the block's 's inputs. In all architectures, the scale-shift operation is performed after the convolution and normalization layers, but before the activation function and dropout layer.
D.4 Additional results
D.4.1 Navier-Stokes -ustt--disttibuttin
In Table 6 we show the out-of- distriibutiion test results for the Navier- Stokes datasets (analogously Table 2 for the spring mesh dataset). As mentioned in the main text and observed by [44], the differences compared to the main test set are marginal.
D.4.2 Benchmark against deterministic models
In this set of experiments, we directly compare our method against the deterministic single-step forecasting baselines presented in the benchmark dataset for the Navier-Stokes and spring mesh exper-
23

---

Table 7: DYffusion ablation. We change one component at a time in DYffusion, starting with all components enabled (first row). For SST, we perform the ablation only on a subset of the test dataset (box 88 in Fig. 5). For the Navier-Stokes and spring mesh datasets, we use the full test sets. The second to fifth rows are ablations at sampling/inference time: No refinement refers to not using the refinement step in line of Alg. No Ia dropout refers to disabling the inference dropout of the interpolator network, Ig. No Ig dropout & refers to disabling the inference dropout of Ig and perturbing the inputs by a, = 0.05 (as for the Perturbation baseline). Naive sampling refers to swapping cold sampling (Alg. 2) with the naive sampling algorithm from [2]. The bottom three rows apply to both training and inference and ablate the choice of the forecaster conditioning c(x,, n)) where for the last row Og N-1
Here, the best choice varies across datasets.
SST
Navier-Stokes
Spring Mesh
CRPS
MSE
SSR
CRPS
MSB
SSK
CRPS
MSE
SSR
Base
0,181 =:9
0.111 2:30
2:11
2:00
2:031
0.0103 cozD
4.2e-4
refiemmen
0.181 a.002
0.111 a.22
a:000
1.12 4:0.2
0.0246
08012
6.8c-4 2
No Ia dropout
0.320 a:39
0.206
2:2222
2:00
0 103
0.028 2.0
0.00 2:0.0
0.0337
+ 2.08
2.5e:3 a
0.01 a
No 22 dropout & o,
0.308 3:300
0.197
1.0312
0.40:00
1:00
0.095
0.028 2.0
0.25 1:0
0.0282
à 20031
2.76-3 a6
0.98 ont
Naive sampling
0.681 0.002
0.945 8.1117
0.52 0:033
0.088
0.44
0.029
0.0.02
0.54 8:0.1
0,0115
0.035
4.5e-4 2004
0.76 a:
mull
0,182 a:00
0.111
1::00
0.94 2:02
0.067
0.022
0.88 4:901
104 111
9.9c47 1
1.75 slte
0.196
0.128
0.77 4:09
0.0103
4.2e-4
c 0. (1 o.
0.181 0:00
0.111
9:30
1.04 2:01
8:0
0.000
0.83 0:0.2
2.2044
0: 1889
3,6c+2 3.ed
iments [44] in terms of MSE. It is important to emphasize that our method is specifically designed to generate ensemble-based probabilistic forecasts, which why our baselines and evaluation procedure focus on the probabilistic forecasting setting. Besides, MSE can only be part of the evaluation of a probabilistic forecasting model, since it does not capture the skill of the forecasted distribution like the CRPS SSR metrics do. In Fig 6 it is visible that the baselines from [44] either degrade or diverge during the inference rollouts for the Navier-Stokes and spring mesh test evaluations, respectively. Both method as well our multi-step Dropout baseline outperform [44], especially in the long-range forecasting regime. This is likely because our models are trained to forecast multiple timesteps, while the models from [44] only learn to forecast the next timestep. As a result, the training objective significantly deviates from the evaluation procedure. The struggle with long-term stability and long-range forecasts by the benchmark datasets baselines was already noted by the dataset paper itself [44). Note that architecturally all these models are almost identical, using the neural backbones proposed in [44].
D.5
Ablations
D.5.1
Cold sampling versus naive sampling
As mentioned in section C.2, we find clear evidence that the DDIM-like cold sampling (Alg. 2) proposed by [2] for "generalized diffusion models" very important to attain good samples from DYffusion. In Table 7, cold sampling corresponds to the first row, Base, and the Naive sampling row refers to swapping cold sampling (Alg. 2) with the naive sampling algorithm from [2] (see Alg. 4). The finding that naive sampling yields poor results for Dyffusion and, more generally, "generalized diffusion models" is consistent with the findings from [2].
D.5.2 Inference dropout in the interpolator network
In Table 7 show that disabling the inference dropout in the interpolator network, Ig, results in considerably worse scores. This is to be expected, since without stochasticity in Izg our current framework collapses generating deterministic forecasts (since the sampling algorithm and forecaster network are deterministic, and we assume that the given initial conditions are fixed). In such a case, computing the CRPS collapses the mean absolute error, and the SSR becomes 0 since there is no spread in the predictions. To attain an ensemble of forecasts, but keeping the interpolator dropout disabled, we also include an ablation row where we perturb the initial conditions with small random noise E ~ (O II), where we use a. = 0.05.
D.5.3 Refining the interpolation forecasts after sampling
We find in Table (No ref. row), that the addition of line 6 to the cold sampling algorithm (see Alg. 2) can sometimes improve performance. However, this is not consistent across datasets: While for the
24

---

SST dataset, we hardly observe any difference, the scores degrade considerably when not refining the forecasts for the spring mesh dataset. A reason could be the relatively long training horizon used for spring mesh (134 for spring mesh versus 16 or 7 for Navier-Stokes or SST). In practice, we recommend that practitioners train DYffusion with the refinement step disabled to accelerate inference time (since the refinement step requires an additional forward pass per output timestep). Then, during evaluation it is encouraged to perform inference with DYffusion with both disabled and enabled refinement step to analyze whether enabling the refinement step can meaningfully improve the forecasts.
D.5.4 Forecaster network conditioning
In the bottom three rows of Table 7 we ablate the three forecaster conditioning options (see B for detailed definitions), c(x,, ni). The optimal choice of varies across datasets. In particular, for the spring mesh experiments, it turns out to be crucial to use the clean initial conditions (variant 2, where c(x,, :) = X2) as conditioning (the two other variants diverge when forecasting the full test trajectories). A likely explanation for this is that the training horizon of h = 134 is relatively long for result, the forecaster network can greatly benefit from being explicitly conditioned
spring mesh. on the initial conditions during all diffusion states (as opposed to only heing indirectly influenced by x, through the interpolator network). practice and if compute allows, may be worthwhile to train DYffusion with at least the first two variants for a few iterations to analyze if any conditioning variant strongly outperforms the others.
D.5.5 Choosing the training horizon
In any multi-step forecasting model, the training horizon, h, is a key hyperpa- rameter choice. Usually, its choice is constrained by the number of timesteps that fit into GPU memory, and it is ex- pected that larger training horizons will improve performance when evaluated on long (autoregressive) rollouts. How-
Table 8: Navier-Stokes ablation of the training horizon, h. All methods are evaluated on the 64-step test trajectories. For example, for h = 1 (h = 16) the corresponding methods are unrolled autoregressively 64 (4) times.
Method
CRPS
MSE
SSR
Dropout (h 1)
0.132 0.066
0.046 + 0,006
0.002 0,00
Dropout (h = 8)
0.086 0,012
0.026 1 0,002
0,416 1 0.29
ever, for continuous-time models includ- ing ours, where the number of timesteps needed in GPU memory does not change as a function of h (see Table 10), the choice of h is flexible. In Table 8, we ex- plore using three different training hori- zons (h E {8, 16, 32}) for the Stvier-Stokes dataset for both the barebone time-conditioned Dropout model as well as DYffusion. For ffffsiion, this means that we train both an interpolator network as well as corresponding forecaster net with the same training horizon. Note that h = 16 corresponds to the main results, and is a sweet spot for DYffusion for this task. However, any of the used horizons results in better scores than the best baseline (for any baseline training horizon). We also include next-step forecasting baseline (h- = 1), which results in significantly worse results than any multi-step forecasting approach. This is likely due the compounding effect of autoregressive errors.
Dropout (ht = 16)
0.078 0.001
0.027 1 0,001
0.715 1 0.01
Dropout (h = 32)
0.078 0.001
0.025 + 0.001
0.651 = 0.01
DY Dffuusion (h
0.076
0,002
0.027
0.701
DYffusion (h 16)
0.067 0.003
0.022
0.877
DYffusion (h 32)
0.075 0.000
0.028
0.862
D.5.6 Using auxiliary diffusion steps can improve performance
In Table we ablate the number of artificial diffusion steps, k used to train/evaluate DYffusion on the SST dataset. As in Table 7 we evaluate on the test box 88 only. It is clear, that & > 0 is important for optimal CRPS and SSR performance. Here, the auxiliary diffusion steps correspond to the floating- point interpolation timesteps { 2 1 as visualized for a toy example in Fig. 9b. Interestingly, we did not observe this phenomenon for the favier-Stoker or spring mesh datasets, where using k > did not improve performance over using the simplest schedule S N-1 1i)=00- where the number of diffusion steps and the horizon are equal (k- 0).
D.5.7 Forecaster network iteratively refines its forecasts
In this set of experiments, we verify that the forecaster's prediction of x for a given diffusion step and associated input & usually improves as increases, ie. approach the end of
25

---

Table 9: For the SST dataset, using k extra artificial diffusion steps corresponding to floating-polts is- especially such that i for all (similarly Fig. 9b), improved performance significantly. The simplest possible schedule with k 0 corresponds diffusion steps and the horizon are equal. The CRPS does not vary significandly when k 0, while the SSR increases steadily with larger k. This can be explained by larger & creating longer sampling sequence with extra diffusion steps that may lead more diverse samples. We evaluate only
[ - ر where the number of
subset of the test dataset (box 88 in Fig. 5).
k
CRPS
MSE
SSR
1 0
0.208 0.002
0.107 + 0.001
0.49 +: 0.00
E 10
0.184
0.113
0.84
k 25
0.181
0.001
0.1 = 0.001
0.99 + 0.02
K 40
0.184
0.002
0.1 0.002
1.06 0.03
k 45
0.189 0.004
0. 19 = 0.005
1.09 + 0.03
2.4
0.90
336
3.8
0.29
0227
(a) OISSTV2
1.55
0.3
0.770
0.03
32
(b) Navier-Stokes
1.5
3.777
1.30
143000
0.0
2.75
2.52
113004
1.25
11300
117 25g
LiPuain StRp
O.Piston Stg
(c) Spring Mesh
Figure 7: We verify that the forecasts & Fo in) of the forecaster network (i.e. the diffusion model backbone) usually improve as the reverse process comes closer in time to f. f.+ h (i.e. and increase) in terms of CRPS and spread-skill ratio (SSR). Interestingly, the MSE scores tend to degrade for the SST and spring mesh datasets (7a and 7c). For the SST dataset (7a), we also visualize dotted vertical, orange line at diffusion step 36, which corresponds auxiliary diffusion steps (i.e. with floating-point in < 1), which are followed by the diffusion steps corresponding to the data resolution (i.e. in after that diffusion step. To achieve clearer visualization of the spring mesh dataset, where the variance across random seeds is relatively large, we plot each run separately. Interestingly, for spring mesh the CRPS starts degrading after around 30 to 70 diffusion steps, which may indicate that shorter training horizon could be more amenable for performance.
the end of the 35
The SSR score improves significantly
26

---

0.38
1.00
0.55
0.40
3250
0.59
0.382
2345
0.30
0240
0.25
0.378
0.99
Figure 8: There is a clear trade-off between inference speed (X-axis) and performance (y-axis) as a function of the number of diffusion steps used for inference by DYffusion. Here, we show the SST test scores for one run of DYffusion, which was trained with 35 auxiliary diffusion steps (on top of the given by the data). Each dot from left to right represents performing inference with an increasing number of diffusion steps. Each dot uses the base schedule Sbase {0, 1, 1........ . 1}. where h = 7. plus Nauz additional diffusion steps drawn from the ones used for training. Naux = (1, 3, 4, 5, 6, 8,,9 8, 10, 12, 14, 16, 18, 20, 23, 26, 29, 35), and each such additional diffusion step corresponds to an implicit dynamical timestep in (0, 1). Interestingly, almost equivalent (for CRPS) slightly better (for MSE) scores be sometimes obtained by using fewer diffusion steps than used for training (right-most dots of each subplot), which immediately benefits inference speed.
the reverse process and Xr+i, gets closer in time to Xr+-h (see Fig. 7). This analysis is conducted on a subset of the SST test set (box 112 in Fig. 5) and validation sets of Navier-Stokes and spring mesh, for the respective datasets, using 20-member ensemble. We use the validation physical systems benchmark because the scope of the analysis only extends up to the training horizon (as opposed the full test set trajectories).
of the
D.5.8 Accelerated DYffusion sampling
In Fig. 8 we study how sampling from DYffusion can be accelerated in a similar way to how DDIM [61] can accelerate sampling from Gaussian diffusion models. The continuous-tin nature of the backbone networks in DYffusion invites the use of arbitrary dynamical timesteps as diffusion states during inference. Thus, to accelcrate sampling, we can skip some of the N diffusion steps used for training, Strain {in ) === which automatically results in fewer neural network forward passes. For example, we can only use the base schedule Sbuse {0, 1. h 1}. where the diffusion states correspond to the temporal resolution of the dynamical data a one-to-one mapping (see black arrows in Fig. 3). In Fig. 8, we start with Shose as inference schedule (left-most dots in cach subplot) for the SST dataset, and then incrementally add more diffusion steps from Strain Sbase to it, until reaching the full training schedule (right-most dots). We find that sampling can be significantly accelerated with marginal drops in CRPS and MSE performance. Note that the dynamical timesteps needed as outputs of DY or for downstream applications pose lower bound (here, Stase) on how much we can accelerate our method since any such output timestep needs be included in the sampling schedule or in the set of output timesteps, J (line 6 in Alg. 2). For that reason, it is not straightforward to accelerate the Stoier-Stokes and spring mesh DYffusion runs, since their training schedules already correspond to Spase-
D.6 Modeling Complexity of DYffusion and Baselines
In Table 10 we enumerate the different modeling and computing requirements needed for each of the baselines and our method. Dropout (multi-step) refers to the barebone backbone network that forecasts all timesteps X in a single forward pass. Dropout (continuous) refers to the barebone backbone network that forecasts one timestep X.++ for & E (1,h) in single forward pass, conditioned on the time, k. Both methods perform similarly our exploratory experiments (not shown), and in our experiments, we always report the scores of the time-conditioned (i.e. continuous) variant. The multi-step approach corresponds to the way the backbone model of a video diffusion model operates, while the continuous variant is similar to how the forecaster network in DYffusion
27

---

Table 10: We report the requirements needed to train a method to forecast up to h steps into the future, where refers to the number of input/output channels (e.g. for SST data), and w to the window size (here, wo = 1 for all experiments). In the second and third columns, we report the input and output channel dimensions, respectively, that the (backbone) neural network needs have (assuming that the window dimension is concatenated to the channel dimension). [Mem(xx) refers to the number of timesteps that need to be present in (GPU) memory to compute the training objective. The last column denotes how many network forward passes are needed to get the forecasts for all timesteps. Here, N1. N₂ refers to the number of (sampling) diffusion steps used by DDPM/MCVD and DY ffusion, respectively. Usually, Nı > N₂ 2 h, since Gaussian noise diffusion models will require more diffusion steps to attain comparable predictive skill. For Navier-Stokes and spring mesh N2 = h, for SST N₂ 35. For MCVD, N 1000. The factor of 3 for DYffusion is result of the two extra interpolator network forward passes needed line of Alg. 2. For large horizons, the model size and memory requirements of multi-step models and conventional diffusion models can be prohibitive. It is clear that (video) diffusion models do not scale well for long horizons.
Modeling complexity
Method
Cin
Cout
Mem(x.)]
#Forward
Dropout (continuous)
w C w c
w + 1
h
Dropout (multi-step)
* c
w + h
DDPM / MCVD
(h w: C
h * c
w + h
N1
DYffusion
w C
w +
N2
operates. Video diffusion models have higher modeling complexity because they need to model the full "videos". Xt+1-t+hs at each diffusion state (or corrupted versions of it). Especially for long horizons, h, and high-dimensional data with several channels, this complicates the learning task for the neural network. Meanwhile, our method is only slightly impacted by the choice of
D.7 Example Sampling Trajectories Visualized
Example sampling trajectories of our approach, as a function of the schedule i,, are visualized in Fig. 9, where the top one corresponds the simplest case where we use a one-to-one mapping between diffusion steps, n. and interpolation/dynamiie, timesteps, i,. Each of the intermediate &; can be used as a forecast for timestep i. The forecaster network, Fe, repeatedly forecasts X,, but does with increasing levels of skill (analogously to how conventional diffusion models iteratively denoise/refine the predictions of the "clean" data, s0)); see Fig. 7).
28

---

(0)
X (2)
(3)
h
84 =
%/
11 =
in
in
in
11
1
X0
&
X2
X3
(a) Basic schedule, 1-to-1 diffusion step dynamical timestep mapping
8
%
(0)
h
(1)
8 (2)
g (3)
h
h
(4)
84
8%
(6)
5
in
0.25 / +
0.5 1
4
in
0.75
Vp
1
in
bn
3
11
1/
1
1
Xg
0025
X0.5
Xg.75
X1
X2
X3
(b) Auxiliary, additional diffusion steps (here, 3) map uniformly to (0, 1)
Figure 9: Exemplary sampling trajectories of two different schedules for mapping diffusion steps, n. to interpolation timesteps. i, For convenience, Fig. 9a is a copy of Fig. 2. The schedules are illustrated using horizon of h 4. The black lines represent forecasts performed by the forecaster network, Fg. The first forecast is performed based on the initial conditions, Xo- The blue lines represent the subsequent temporal interpolations performed by the interpolator network, Ig.
29


Reviews:
[{"review_id": "gih4V43BGF", "reviewer": "Reviewer_WiMo", "summary": "In this paper, authors tackle the long-term forecasting problem applied to dynamics system. To solve this problem, they propose to use diffusion principle along with interpolating and forecaster mechanisms. The former interpolates timestep data in between lookback and target windows (therefore, at a lower resolution). \nAuthors evaluate their proposal on three datasets, analyze and interpret the results as well as further discuss directions for improvements.", "questions": "The proposal is very interesting and sound promising. However, in my opinion, we are missing a representation of the proposed architecture to better understand the paper and how each module interacts with each other.\n\n## Reproducibility\nThe details for training /validation and testing for Navier-Stokes (NSFlow) and Spring-Mesh (SM) are not provided in the dataset description, but in the experiments caption. And yet we can wonder if they are complete, which could make more difficult to reproduce results.\n\n## Performance comparison\nI think we are missing the time column in table 2 and the one for NSFlow experiment in Table 1.\n\n`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 do not show them.` Why not include them in the appendix and let readers be the judge and make their own opinion?\n\n## Comparison with reference [43]\nAs authors build up on ref. [43] models for NSFlow and SM, why are they not comparing with the model (unet and CNN) from such a reference. Especially looking at results from [43], for NSFlow from Figure 5, unet seems to achieve MSE between 0.05 and 0.007. These variations depend on the number of obstacles and there might be some difference in the set-up, but it would be good to see how DYffusion do compared to these baselines. As is it does not look like to me that the proposed model is doing better contrary to what authors are claiming. Other readers might feel the same and wonder what is the advantages of DYffusion. As a consequence, it would be good to define precisely the condition of NSFlow and SM and compare with original model.\n\n`It is worth noting that our reported MSE scores are significantly better than the ones reported in [43]` I have difficulty agreeing with this claim. First, results in [43] are not presented in a table format so it is more difficult to compare. Authors should either provide their results with the same representation (stepMSE vs. OoD MSE) for instance in the appendix, for reader to better judge and compare both solutions. Or auhtors should include [43] models as mentioned above as a baseline in their table 2 (or both). While clearly setting the configuration of the forecast to make sure that we are dealing with similar set-up.\n\n`our MSE results significantly surpass the ones reported in Fig. 8 of [43]`: Again, if authors wish to compare their results with model from [43], why not running forecast with CNN from reference [43] in the same condition? As Authors, in this paper, are re-using CNN from ref. [43], it should then be possible and prevent readers to go check paper [43] to make sure the settings are the same and try to grasp the similitude between papers as each paper use different representations of the results. In addition, if CNN performed so poorly in [43], why choosing it as the base model for SM dataset? Why not choose an MLP or nn kernel that seems to perform better on the figure mentioned by authors?\n\n## Model evaluation\n`resulting forecast to be visibly pleasing and temporally consistent` Not really a scientific judgement in my opinion…But it is indeed quite similar. Nevertheless, at t=3.70, the velocity of sample 2 to 5 are not really “visibly pleasing”, how should we interpret this?\n\n## Ablation study\nIn my understanding, the proposal relies on the forecaster and interpolator. Therefore, I would have expected to see an ablation study, without these modules to judge of their importance. Is it impossible to do so?\nIs it not possible also to do ablation study with different frame number of the interpolator (k-1)? BTW, unless I missed it, we don’t know what the value of k for each dataset is. Does it have an impact on the performance?\n\n## Proof-read\nI found some points to be corrected and authors should pay some more attention on possible other typos or issues:\n * Line 25. and and require -> and require\n * Line 92, should not it be x_{t+1}:{t+h}\n * Line 196, Sentence is a bit long, consider breaking it\n * Line 271, In the spring mesh dataset -> In the Spring Mesh dataset\n * Line 288, whether it is actually able -> whether it is actually possible\n * Line 302/303, For the Navier Stokes and spring mesh dataset it was sufficient -> For the Navier Stokes and Spring Mesh dataset, it was sufficient\n * Line 327, could be the reason for why -> could be the reason why\n", "limitations": "Authors discuss their results and limitations as well as future directions.", "rating": 7, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": " * Interesting approach\n * Extensive ablation study, with summary in the main paper section\n * Interesting discussion of the results", "weaknesses": " * Reproducibility (no code, Navier-Stokes and Spring-Mesh experiment set-ups are presented in the table caption, but I wonder if they are complete for anyone who would like to re-run the same experiments)\n * Comparison with the main reference [43] is in my opinion very subjective and should be improved (as authors build-up on [43] for two datasets, it seems to me that [43] should be included as a baseline)\n * Require some proof-reading\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We would first like to thank you for the positive comments and valuable feedback. We respond to your comments and questions below.\n\n**Q1:** _Reproducibility_\n\n**A1:**\n1. **Code:** We have shared with the AC an anonymous link to our code for reproducibility. We will open-source our code and data when the manuscript is published (noted in the submission questions), \n2. **Navier-Stokes and Spring Mesh set-ups:** We include instructions on how to reproduce our experimental results in the code README. We note that our evaluation procedure is *identical* to the benchmark dataset paper [43] (i.e. the same validation/test sets and the same metrics), except that we adapt it to probabilistic models (e.g. sample multiple forecasts per initial condition, compute CRPS and SSR). We will make it clear that we always train on the full training datasets of [43] and that we use the Navier-Stokes with 4 obstacles. We will note all these details in the revised appendix.\n\n**Q2:** _Performance comparison_\n\n**A2:**\n> Why not include Navier-Stokes out-of-distribution test set results in the appendix and let readers be the judge and make their own opinion?\n\nThis is a good point. We have added these results to the revised appendix, and also show them in Table 10 of the joint rebuttal PDF.\n\n**Q3:** _Comparison with reference [43]_\n\n**A3:** We appreciate this feedback and have revised the manuscript to describe the experimental setup for Navier-Stokes and spring mesh in more detail (see response 2. of **A1** above).\n\nHowever, we would like to point out that [43] is a dataset paper. All baselines therein are ***deterministic*** models, and as such is not an adequate baseline for the ***probabilistic*** forecasting setup that we study in this paper. \nNonetheless, **please see our Figure 7 in the joint rebuttal PDF, where we reproduce the relevant models from [43] and show that we substantially outperform them with all our baselines in terms of MSE**.\n\nOur Navier-Stokes and Spring Mesh ensemble mean MSE scores can be best compared to Fig. 10 (bottom; step prediction Unet with 4 obstacles) and Fig. 8 (bottom; step prediction CNN), respectively, of [43]. While the authors of [43] have not provided tabular results, comparing these figures reveals that 1) our reported MSE scores for Navier-Stokes ($0.022$) are competitive to the ones from Fig. 10 in [43], and 2) our spring mesh MSE ($4.74e-4$) is much lower than any reported MSE in Fig. 8 of [43], where none of the step prediction models achieve MSE less than $10^{-2}$. We have reached out to the authors from [43] to get exact values for these results so that we can report them in our paper as supplementary results.\nWhile helpful, this comparison is insufficient on its own, as MSE can only be part of the evaluation of a probabilistic forecasting model, since it does not capture the skill of the forecasted distribution like the metrics CRPS and SSR. \n\nAdditionally, please note that our Perturbation baseline can be used by the reader to assess how direct applications of the models from [43] perform on our probabilistic skill evaluation. Indeed, except for the baselines being trained to forecast multiple steps, our Perturbation baseline, which forecasts deterministically (see lines 649-653), for Navier-Stokes is identical to the UNet used in [43].\n\n>  In addition, if CNN performed so poorly in [43], why choosing it as the base model for SM dataset? Why not choose an MLP or nn kernel that seems to perform better on the figure mentioned by authors?\n\nDiffusion models usually use UNet backbones, so to stay as close as possible to that, we chose the CNN model. In addition, our results show that the spring mesh CNN can actually perform very competitively when employing stochastic multi-step training as done for our baselines.\n\n**Q4:** _Model evaluation_\n\n**A4:**\n> _\"resulting forecast to be visibly pleasing and temporally consistent\"_. Not really a scientific judgement in my opinion\n\nWe believe that a qualitative evaluation of the forecasted videos is important to complement the reported metrics.\n\n> Nevertheless, at t=3.70, the velocity of sample 2 to 5 are not really “visibly pleasing”, how should we interpret this?\n\nSuch behavior is to be (sometimes) expected for long-rollouts of a probabilistic forecasting model. This is because we hope that the probabilistic model can capture any of the possible, uncertain futures (instead of forecasting their mean, as a deterministic model would do). As such it is reassuring that our samples show sufficient variation but also covers the ground truth (sample 1).\n\n**Q5:** _Ablation study_\n\n**A5:**\n> In my understanding, the proposal relies on the forecaster and interpolator. Therefore, I would have expected to see an ablation study, without these modules to judge of their importance\n\nThe first sentence is correct. However, it is not possible to use our framework without both modules, as the sampling procedure relies on both. The Dropout baseline, i.e. a pure forecasting model, is the closest we can get to this.\n\n>  Is it not possible also to do ablation study with different frame number of the interpolator (k-1)? BTW, unless I missed it, we don’t know what the value of k for each dataset is. \n\nWe have added a table ablating the choice of k for the SST dataset to the appendix and in Table 12 of our joint rebuttal PDF. We report the value of k for each dataset in lines 671-673, but we will also add it to a new row in the hyperparameters table 4.\n\n**Q6:** _Proof-read_\n\n**A6:**\nThank you for helping pointing out several typos – we have fixed these. To clarify, We chose to write spring mesh in lowercase in the main text, because this is what the original dataset paper [43] does (while Navier-Stokes is capitalized).\n> Line 92, should not it be x_{t+1}:{t+h}\n\nNo. We write $\\mathbf{x_{t:{t+h}}}$ because we want to include the initial conditions $\\mathbf{x_t}$ in this description (in addition to the to-be-forecasted snapshots $\\mathbf{x_{t+1:{t+h}}}$).\n"}, {"review_id": "TgjpgDT25S", "reviewer": "Reviewer_7msc", "summary": "This paper proposes a new forecasting model for spatiotemporal data. The idea is based on separately training an interpolator and a forecaster network and applying them in an alternating fashion at inference time to iteratively refine the forward prediction. The inference procedure loosely resembles the denoising process in a diffusion model, where one computes a new denoising target and moves slightly towards it at every step. Numerical experiments are conducted on several datasets to validate its skills.", "questions": "* Equation (4), the second line is not a probability. Do you mean a $\\delta$ distribution at the interpolator output?\n* Also equation (4), adding Gaussian noise only at the last step looks a bit cryptic to me. What is this trying to achieve besides adding randomness to the prediction and how is the noise level $\\sigma$ determined?", "limitations": "The authors mention that input and output space must be the same and that inference costs are higher than predicting directly.", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "* The idea of iteratively refining forward predictions is interesting and quite novel to my knowledge.\n* The paper is generally well written and easy to follow.\n* The numerical experiments cover multiple non-trivial tasks and show moderate improvement in the metrics against benchmarked diffusion-based and ensemble prediction models.", "weaknesses": "* Although advertised as a diffusion model, the actual connection (at least to diffusion-based generative models) is handwavy at best. It is not clear that the underlying process arising from the defined noising and denoising operations form a well-defined generative model (for example see how DDPM is derived with clearly defined conditional and marginal distributions at each step). Only the variational posterior is provided in equation (4). \n* The model is mostly deterministic (except for the last step that injects some Gaussian noise) so I think it’s important to show comparison against deterministic forecast models. The high memory footprint for multi-step training is somewhat fair, but there exist ways to improve single-step models (e.g. noise injection, curriculum training, etc.) which should be considered when validating the proposed model. This is necessary to justify doing the additional steps of interpolation.\n\nOverall I feel in its current state given the strengths and weaknesses I vote for borderline reject tentatively.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We would like to thank you for the positive comments regarding the novelty, clarity, and experiments of our work, as well as your valuable feedback to which we respond below. \n\n_**Potential misunderstanding:**_ We would like to point out that there seems to be a key misunderstanding regarding our method being deterministic while it is actually probabilistic . As  noted in our problem statement section (see Setup, lines 65-67) we specifically study the problem of _probabilistic_ forecasting. Consequently,, our interpolator network,  $\\mathcal{I_\\phi}$, is designed to produce stochastic outputs (see the abstract, Algorithm 1 (Stage 2, line 1), and lines 106 & 111). We achieve this by enabling dropout at inference time. We show that this is a key component of our framework in the last two ablation rows (No. Dr, and No Dr. & $\\sigma_\\epsilon$) of Table 5 in the Appendix.\nBased on this misunderstanding, we will strive to make our generative formulation more clear in the revised methodology section as mentioned in answer A3 below.\n\n\n**Q1:** _Connection to conventional diffusion models_\n\n**A1:** \n> Although advertised as a diffusion model, the actual connection (at least to diffusion-based generative models) is handwavy at best. It is not clear that the underlying process arising from the defined noising and denoising operations form a well-defined generative model\n\nOur work builds upon cold diffusion [2], which \"_paves the way for generalized diffusion models that invert arbitrary processes._\" The cold sampling algorithm, which we use to sample from DYffusion (see Alg. 2), is a generalization of DDIM sampling for \"generalized diffusion models\" (see Appendix A.6 of [2]) and a key ingredient to make DYffusion work (see Table 10 of the joint rebuttal PDF). Our proposed forward and reverse processes are specifically designed to fall under this framework.\n\n*Given this context, it becomes clear that DYffusion is a generative model for forecasting that falls in the category of \"generalized diffusion models.\"* We can see how this relationship may not have been clear in the text, and have added a new paragraph discussing these connections to our appendix.\n\n>  DDPM is derived with clearly defined conditional and marginal distributions at each step. Only the variational posterior is provided in equation (4).\n\nYou are correct that  DYffusion loses some of the benefits of using a simple Gaussian distribution as the forward process. This is because DYffusion's forward process is based on a stochastic interpolator network, $\\mathcal{I_\\phi}$. We believe that this is a fair trade-off to make, especially because ultimately it is the posterior that we care about for forecasting. Future work can explore regaining some of the benefits from DDPM/DDIM.\n\n\n**Q2:** _Comparison against deterministic models_\n\n**A2:**\n> The model is mostly deterministic (except for the last step that injects some Gaussian noise) so I think it’s important to show comparison against deterministic forecast models\n\nAs mentioned in the \"potential misunderstanding\" paragraph above, DYffusion is NOT a deterministic model. Its main source of stochasticity comes from the forward process/interpolator network. As such, deterministic models are not an adequate baseline and metrics such as CRPS and SSR are necessary for properly evaluating the probabilistic forecasts. In addition, we would like to highlight that our Perturbation baseline forecasts deterministically (see lines 649-653) and allows for the assessment of an ensembled deterministic model you might be looking for. **Also, please see the discussion of Fig. 7 in the joint rebuttal and PDF where we compare against a deterministic single-step model.**\n\n> The high memory footprint for multi-step training is somewhat fair, but there exist ways to improve single-step models (e.g. noise injection, curriculum training, etc.) which should be considered when validating the proposed model.\n\nThanks for the suggestion. Noise injection and curriculum training have been only studied in the context of deterministic forecasting, not probabilistic forecasting as our paper. These techniques are orthogonal to the contributions proposed by this work to address multi-step forecasting challenges. For example, noise injection is a complementary method that can be added to further improve DYffusion too. \n\nWhile we would appreciate a reference for \"curriculum training\", we believe that you are thinking of approaches like the one used by [35]. It is one particular kind of performing multi-step training. In this paper, we choose the more common approach of directly predicting any of the multiple steps in a single forward pass (i.e. without needing to backpropagate gradients recursively as in [35]). An extensive benchmarking on how to best perform multi-step training is beyond the scope of this paper.\n\n**Q3:** _Equation (4)_\n\n**A3:**\n> Equation (4), the second line is not a probability. Do you mean a \\delta distribution at the interpolator output?\n\nAs mentioned in the \"potential misunderstanding\" paragraph above, the interpolator network is designed to be stochastic by enabling dropout at inference time. We will make this more clear in our writing.  For example, we have updated the text to make it explicit that $\\mathcal{I_\\phi}$ depends on a random variable $\\xi$ (here, the randomly dropped out weights of the neural network) by writing $\\mathcal{I_\\phi}(\\mathbf{x_t}, \\mathbf{x_{t+h}}, i | \\xi)$.\n\n> Also equation (4), adding Gaussian noise only at the last step looks a bit cryptic to me. What is this trying to achieve besides adding randomness to the prediction and how is the noise level  determined?\n\nWe added Gaussian noise to make our Eq. (4) be consistent with Eq. (10) of DDIM [60]. However, in practice (see our Alg. 2) we do not add any Gaussian noise in the last step (as in DDIM). This is just of notational nature, so wee will remove this part in our revised version.\n"}, {"review_id": "hacwogc7DC", "reviewer": "Reviewer_XpdM", "summary": "In this paper, the authors propose to build on diffusion model for modelling spatio-temporal data.\nIn particular, they first train a time-dependent interpolation network which learn to interpolate the temporal dynamics given a frame at the horizon time $x_{t+h}$, a frame $x_t$, and an index $i$ interpolating between the two distributions.\nThen, the idea is to frame a forecast model as a denoising model between $x_t$ and $x_{t+h}$, that is given an interpoled frame to be able to recover the original 'denoised' $x_{t+h}$ frame.\nThis approach allows for memory efficient multi-step predictions, as opposed to existing diffusion models for videos.\nThey apply the developed model to the forecasting of sea surface temperatures, Navier-Stokes flows, and spring mesh systems.", "questions": "- Equation 3: How is the architecture handling the two inputs $x_t$ and $x_{t+h}$? By doubling the number of input channels?\n- line 119: Why do we need this? This deserves more explanation.\n- Equation 4: Why do we need the interpolator network here? Can't we directly predict $s^{(n+1)}$ with $F_\\theta(s^{(n)}, i_{n+1})$? Aren't we actually extrapolating with the interpolator here?\n- Algorithm 2: What is happening line 4? How does this iteratively refine the prediction of $x_{t+h}$? This is likely worthy of some more explaination.\n- Equation: Isn't this ODE implying that the dynamics is Markovian?\n- line 249: What is the Continuous Ranked Probability Score (CRPS)? Worth at least giving a high level idea I think.\n- line 268-269: Why not training the DDPM with more diffusion steps? Since the number of steps is an hyperparameter for all methods it would be easier to compare them with the same number of steps, albeit could be varied to see how performance vs computation evolves.\n- Table 1: Would you know why the performance of the models is widely different betweent the SST and the Navier-Stokes datasets?", "limitations": "N/A", "rating": 7, "confidence": 3, "soundness": 2, "presentation": 1, "contribution": 3, "strengths": "- The main advantage of the proposed approach seems to be the small memory footprint as it is constant w.r.t. the horizon, in contrast with standard dynamics forecasting (if I understand correctly).\n- It also allows for 'upscaling' the resolution, or for making prediction at any time value, thus is not constrained to a regular discretisation of the time axis.\n- The experiment on the sea surface temperatures dataset is quite promising, as the proposed method is able to accurately forecast not only the predicted mean but also the uncertainty with a lower computational cost than the video diffusion model MCVD. It would be valuable to explore whether one could obtain better performance at the expense of more computational cost.", "weaknesses": "- I found the writing to be confusing at times, I think that Section 3 could be better conveyed, in particular the choice of notations, but also the relation with diffusion models, and the 'cold posterior' sampling scheme.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the positive comments regarding the memory efficiency, continuous-time forecasts, and SST experiments of our work, as well as your valuable feedback to which we respond below. \n\n**Q1:** _Confusing section 3_\n\n**A1:**\n> I found the writing to be confusing at times, I think that Section 3 could be better conveyed, in particular the choice of notations, but also the relation with diffusion models, and the 'cold posterior' sampling scheme.\n\nThank you for bringing this to our attention. We will add a glossary to the appendix that contains all our choice of notations in a single place. Could you let us know which specific parts were confusing? We would like to address these. \nAdditionally, we will write out the objective for a conventional diffusion model (Eq. (1)) alongside our corresponding objective out (Eq. (3)) to make the connection in terms of analogous objectives more clear. \n\nConcretely, for a \"generalized diffusion model\" [2] (equation 1):\n$$||R_\\theta(D(\\mathbf{s}^{(0)}, n), \\mathbf{x_t}, n) - \\mathbf{s}^{(0)}||^2 = ||R_\\theta(\\mathbf{s}^{(n)}, \\mathbf{x_t}, n) - \\mathbf{s}^{(0)}||^2,$$\nwhere $\\mathbf{s}^{(0)} = \\mathbf{x_{t+1:t+h}}$, and $\\mathbf{s}^{(n)}$ is a noisy version of $\\mathbf{s}^{(0)}$ (the level of noise increases with $n$).\n\nAnd for DYffusion (equation 3):\n$$||F_\\theta(\\mathcal{I_\\phi}(\\mathbf{x_t}, \\mathbf{x_{t+h}}, {i_n}), {i_n}) - \\mathbf{x_{t+h}}||^2 \n= ||F_\\theta(\\mathcal{I_\\phi}(\\mathbf{x_t}, \\mathbf{s}^{(N)}, {i_n}), {i_n}) - \\mathbf{s}^{(N)}||^2\n= ||F_\\theta(\\mathbf{s}^{(n)}, {i_n}) - \\mathbf{s}^{(N)}||^2,$$\nwhere $\\mathbf{s}^{(N)} = \\mathbf{x_{t+h}}$ and $\\mathbf{s}^{(n)} \\approx \\mathbf{x_{t+i_n}}$ is now a stepped backward in time version of $\\mathbf{s}^{(N)}$. Note that the diffusion step indexing (superscript $n$) for DYffusion is reversed so that it aligns with the temporal indexing (subscript $t$), such that e.g. $\\mathbf{s}^{(n)}  \\approx \\mathbf{x_{t+i_n}}$ temporally precedes $\\mathbf{s}^{(n+1)}  \\approx \\mathbf{x_{t+i_{n+1}}}$. Accounting for the opposite order of indexing, the similarity between both approaches becomes clear.\n\n>  the 'cold posterior' sampling scheme\n\nThe cold sampling algorithm is directly taken from [2] (their Alg.2; only the notation is adapted), and thus we refer the reader to [2] for intuition on how/why it works. For example, an alternative to sample from DYffusion (or any \"generalized diffusion model\") would be to replace line 4 in Alg. 2 with simply: $\\mathbf{x_{t+i_{n+1}}} = \\mathcal{I_\\phi}(\\mathbf{x_t}, \\hat{\\mathbf{x_{t+h}}}, i_{n+1})$. This corresponds to the _naive sampling_ algorithm of [2] (their Alg. 1) which performs worse than cold sampling as shown in [2] and in our new Table 10 in the joint rebuttal PDF. We will make this more clear in our revised version.\n\n**Q2:** _More questions_\n\n**A2:**\n> Equation 3: How is the architecture handling the two inputs $\\mathbf{x_t}$ and $\\mathbf{x_{t+h}}$? By doubling the number of input channels?\n\nYes, exactly. We have updated the text to mention this explicitly.\n\n> line 119: Why do we need this? This deserves more explanation.\n\nEssentially, the look-ahead loss term simulates one step of the sampling process (Alg. 2) and backpropagates through it, so that the network is trained with an objective that is closer to (partially mimics) the sequential sampling process.\n\n> Equation 4: Why do we need the interpolator network here? Can't we directly predict  $\\mathbf{s}^{(n+1)}$ with $F_\\theta(\\mathbf{s}^{(n)}, i_{n+1})$? Aren't we actually extrapolating with the interpolator here?\n\n$F_\\theta$ always forecasts $\\mathbf{x_{t+h}}=\\mathbf{s}^{(N)}$ (see line 109), so we cannot use it to predict $\\mathbf{s}^{(n+1)}$. As noted in line 132, $\\mathbf{s}^{(n+1)}$ corresponds to $\\mathbf{x_{t+i_{n+1}}}$. We will make clear that for all $n \\in \\{1, ..., N-1\\}$ it holds that $0 < i_n < h$, so that $\\mathbf{s}^{(n+1)}=\\mathbf{x_{t+i_{n+1}}}$ is always between $\\mathbf{x_{t}}$ and $\\mathbf{x_{t+h}} \\approx F_\\theta(\\mathbf{s}^{(n)}, i_{n+1})$, and thus the interpolator can be used to interpolate the timestep $t+i_{n+1}$.\n\n> Algorithm 2: What is happening line 4? How does this iteratively refine the prediction of x_t+ℎ? This is likely worthy of some more explaination\n\nSee the last part of our response **A1** (regarding _\"the 'cold posterior' sampling scheme\"_). Essentially, it allows for more robust sampling over naive sampling defined above.\n\n> Equation: Isn't this ODE implying that the dynamics is Markovian?\n\nNo. We assume you are referring to equation (5) or (6): The dynamics is not Markovian because each step in the prediction is dependent on $\\mathbf{x_t}$ (initial state), in addition to the previous state $\\mathbf{x_s}$.\n\n> line 249: What is the Continuous Ranked Probability Score (CRPS)? Worth at least giving a high level idea I think.\n\nThe CRPS is a common metric for probabilistic forecasts [12, 19, 48, 50, 57]. It penalizes both absolute skill as well as sharpness of the ensemble of forecasts. It rewards small spread (i.e. sharpness) if the forecast is accurate, and measures the difference between the forecasted and observed cumulative distribution function (CDF). For a deterministic forecast, the CRPS reduces to the mean absolute error. We will add a more detailed explanation to our revised paper. \n\n> line 268-269: Why not training the DDPM with more diffusion steps?\n\nWe tried 1000 diffusion steps for the DDPM as well, but found it decreased the performance  on the SST dataset.\n\n> Table 1: Would you know why the performance of the models is widely different between the SST and the Navier-Stokes datasets?\n\nWe are also intrigued by this observation. Our current hypothesis is that the limited dataset size of the Navier Stokes/spring mesh dataset could be an important factor. Indeed, conventional diffusion models are data hungry since they need to learn how to map a noise distribution to the high-dimensional data distribution. We hope to investigate this more in the future.\n"}, {"review_id": "Q6mAoGmVuA", "reviewer": "Reviewer_4TEk", "summary": "This paper proposes DYffusion, an approach that mimics the diffusion models for spatiotemporal forecasting. \nIt treats the noising process as interpolation (parameterized by $I_\\phi$) and the denoising process as forecasting (parameterized by $F_\\theta$), i.e., it reimages the noising step $T$ in the original diffusion models as the temporal step in forecasting. \nDYffusion parameterizes the probabilistic transitions in forecasting continuous-time trajectories, which is similar to parameterizing the diffusion of SDEs. \nIt performs competitively on spatiotemporal forecasting tasks, including sea surface temperatures, Navier-Stokes flows, and spring mesh systems, in terms of probabilistic skill score metrics.", "questions": "Why does Dropout perform so poorly at $t=2$ in Figure 3? Does this observation indicate the mentioned Weakness 3?", "limitations": "This paper includes a separate section for discussing limitations. There is no potential negative societal impact.", "rating": 4, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "1. DYffusion reimagines the continuous-time probabilistic forecasting problem as a diffusion process. It benefits from existing diffusion algorithms and methods to accelerate the inference sampling.\n2. DYffusion achieves the best or competitive scores while significantly reducing the time cost compared to diffusion approaches.", "weaknesses": "1. Eq. (6) is incorrect. There should be a differential instead of a derivative in the integral.\n2. DYffusion's method is closely related to neural ODE and SDE. However, empirical studies lack corresponding baselines. For example, one possible baseline could involve extrapolating $dF_\\theta/ds$ using an ODE solver.\n3. It is unfair to use $F_\\theta$ trained with Algorithm 1 in the baseline method Dropout. The forecaster used in Dropout should be trained with an objective such as $||F_\\theta(x_{t+i}, i)-x_{t+h}||^2$ instead.\n4. One of the main contributions mentioned is that DYffusion reduces complexity. However, a complexity analysis is missing.\n5. Neural SDE [1,2,3,4] parameterizes the stochastic dynamics for modeling continuous-time processes and is therefore inherently suitable for probabilistic continuous-time forecasting. However, there is no discussion on a range of related works on neural SDE.\n\n[1] Deng, Ruizhi, et al. \"Modeling continuous stochastic processes with dynamic normalizing flows.\" Advances in Neural Information Processing Systems 33 (2020): 7805-7815.\n\n[2] Liu, Xuanqing, et al. \"Neural sde: Stabilizing neural ode networks with stochastic noise.\" arXiv preprint arXiv:1906.02355 (2019).\n\n[3] Jia, Junteng, and Austin R. Benson. \"Neural jump stochastic differential equations.\" Advances in Neural Information Processing Systems 32 (2019).\n\n[4] Kidger, Patrick, et al. \"Efficient and accurate gradients for neural SDEs.\" Advances in Neural Information Processing Systems 34 (2021): 18747-18761.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you  for your positive comments regarding the novelty, connection to existing diffusion models, and efficiency of our approach, as well as your valuable feedback to which we respond below. \n\n**Q1:** \n> Eq. (6) is incorrect. There should be a differential instead of a derivative in the integral.\n\n**A1:**\nYes, there should be a $ds$ at the end of the equation. We apologize for the confusion, and have fixed this for the revised version.\n\n**Q2:** \n> DYffusion's method is closely related to neural ODE and SDE. However, empirical studies lack corresponding baselines. For example, one possible baseline could involve extrapolating dF_\\theta/ds  using an ODE solver\n\n**A2:**\n\nThank you for making the connection of DYffusion to neural ODE/SDE methods.\n\nHowever, **DYffusion is not directly related to neural SDEs**. They are different types of deep generative models (DGM). DYffusion is a diffusion model-based DGM that leverages the theory of SDE to learn high-dimensional distributions via score matching. In contrast, neural ODE/SDE is a type of autoregressive DGM that assumes the hidden states of a neural network to follow a particular ODE/SDE dynamics. Despite the shared use of SDE in these two lines of works, they have very different learning mechanisms and applications.\n\n To the best of our knowledge, existing neural SDE methods (including your references [1-4]) only focus on low-dimensional problems. For example, the maximum dimensionality in [1] and [4] is four, respectively two, using an air quality dataset. Meanwhile, [2, 3] do not study dynamics forecasting: [2] only has image experiments, and [3] studies event prediction of low dimensional datasets. In our work, the spring mesh dataset with 400 = 10 x 10 x 4 has lowest dimensionality among all datasets, and we note that this dimensionality can further increase when performing multi-step forecasting.\n\nThe reason for why these neural SDE papers only experiment with low-dimensional multivariate timeseries is exactly because of the ODE/SDE assumption on the hidden state mentioned above. Due to the need of solving a ODE/SDE through a numerical solver, neural ODE/SDE currently struggle with high-dimensional data such as video data in our work.  We would be more than willing to add any baselines that have been applied to high-dimensional spatiotemporal data similar to the one we study in this paper. We would appreciate any such pointers to potential baselines. \n\n**Q3:** \n\n> It is unfair to use $F_\\theta$ trained with Algorithm 1 in the baseline method Dropout. The forecaster used in Dropout should be trained with an objective such as $||F_\\theta(x_{t+i},i)−x_{t+ℎ}||^2$ instead.\n\n\n**A3:**\nThe Dropout baseline was NOT trained with Alg. 1 (that one was only designed/used for DYffusion). Instead it was trained on the objective $||F_\\theta(x_t, i) - x_{t+i}||^2$ for $1\\leq i \\leq h$. We realize that lines 237-238 and 242-244 may be confusing in this respect, and we will improve their clarity in our revised version. We will add the aforementioned objective that we used to train the Dropout baseline to the appendix (and refer to it in the main text). When we say that the baseline is trained _\"analogously to the DYffusion forecaster\"_ (line 243-44) we meant to only refer to the fact that both methods rely on a time-conditioning mechanism.\n\n**Q4:** \n\n> One of the main contributions mentioned is that DYffusion reduces complexity. However, a complexity analysis is missing.\n\n**A4:**\nSee table 7 in the appendix for a comparison on how DYffusion effectively reduces the neural network input and output dimensionality, memory needs, and efficiency compared to a video diffusion model/MCVD.\n\n**Q5:** \n\n> Neural SDE [1,2,3,4] parameterizes the stochastic dynamics for modeling continuous-time processes and is therefore inherently suitable for probabilistic continuous-time forecasting. However, there is no discussion on a range of related works on neural SDE.\n\n**A5:**\nThank you for the references. We will make sure to discuss and cite them, as well as the neural SDE literature more generally, in our revised version. As mentioned in **A2**, at this moment we don't believe that neural SDE methods are an appropriate baseline for the high-dimensional spatiotemporal forecasting problem studied in our paper, since all Neural SDE papers that we are aware of (incl. [1-4]) only cover low dimensional time series.\n\n**Q6:** _Figure 3_\n \n**A6:**\n> Why does Dropout perform so poorly at t=2 in Figure 3? \n\nThat is a good observation. Generally, the beginning of the Navier-Stokes trajectories look qualitatively different to following timesteps where the fluid is already clearly \"flowing\". As such it is possible that those initial timesteps are harder to qualitatively forecast well. Plus, by definition, there are only relatively few training examples of the very start of the training trajectories compared to mid/end points of them. This might be especially a problem for underdispersive models (SSR < 1) that do not capture the full data distribution well.\n\n> Does this observation indicate the mentioned Weakness 3?\n\nNo, see our answer **A3**."}]

Author rebuttal:
We thank the reviewers for their thoughtful comments. 

We are particularly encouraged by the reviewers finding our work  _”interesting”_ (WiMo and 7msc), _”quite novel”_(7msc) and _”quite promising”_ (XpdM).
We are glad to hear that reviewers found our paper _“well written and easy to follow”_ (7msc) with _“numerical experiments cover multiple non-trivial tasks”_ (7msc) and an _“extensive ablation study”_ (WiMo). 

One common issue with our work was insufficient clarity of our methodology section (XpdM) which may have caused misunderstandings with regard to our method being “deterministic” (7msc) or requiring being benchmarked against deterministic baselines (WiMo).
As reviewer 4TEk correctly notes our work _“reimagines the continuous-time probabilistic forecasting problem as a diffusion process”_, and we would like to re-emphasize that our **method was specifically designed to _generate probabilistic_ forecasts**, which is why our baselines and evaluation procedure are for the probabilistic forecasting setting too.

In our joint rebuttal PDF we have added the following new figures and tables:
- Figure 7 shows that deterministic single-step forecasting as in [43] is substantially outperformed by all our own probabilistic multi-step baselines as well as DYffusion. In addition, you can observe that DYffusion performs especially well on long-range forecasts relative to the baselines. These results were requested by reviewer **WiMo** and should also be interesting to reviewer **7msc** who requested to “show comparison against deterministic forecast models”. To attain Fig. 7, we have reproduced the UNet and CNN baselines from [43]. The only difference between them and our Dropout-multi-step baseline (called “Dropout” in our paper) is that the latter is trained to forecast multiple steps and has inference dropout enabled.
- Table 10 shows that cold sampling (Alg. 2) substantially outperforms naive sampling, especially on the SST dataset.  Cold sampling was proposed by [2], while naive sampling is their generalization of DDPM sampling for “generalized diffusion models”. Naive sampling corresponds to replacing line 4 in Alg. 2 with: $\mathbf{x_{t+i_{n+1}}} = \mathcal{I_\phi}(\mathbf{x_t}, \hat{\mathbf{x_{t+h}}}, i_{n+1})$. This finding is consistent with [2] who find that naive sampling “works well for noise-based diffusion” but “yields poor results” for generalized diffusion models, and should be especially interesting to reviewers **7msc** and **XpdM**. 
- Table 11 provides the ablation study requested by reviewer **WiMo** on ablating the value of k, i.e. the number of artificial diffusion steps used by DYffusion, for the SST dataset. 
- Table 12 provides the Navier-Stokes out of distribution results requested by reviewer **WiMo**.

We are of the opinion that the valuable feedback provided by the reviewers has significantly enhanced the quality and clarity of our paper. As a result, the core contributions of our work, centered around the introduction of a pioneering framework for probabilistic dynamics forecasting and complemented by robust empirical evidence and an exhaustive ablative study, have been elevated.


Meta-review: This paper proposes a diffusion model for spatiotemporal forecasting, which goes beyond generating static images to dynamics forecasting by leveraging temporal dynamics learned from the data.

The proposed diffusion model is interesting. It can effectively reduce the memory footprint and achieve speedup as demonstrated by the experimental results. Unlike deterministic forecasting methods which only predict the mean, prediction uncertainty can also be quantified by the proposed model. The experiments reported in the paper are generally strong. Nevertheless, some concerns were raised by the reviewers in their initial reviews, which include both the technical content and the paper presentation. We thank the authors for responding to them and conducting further experiments to provide more details. Consequently, three of the four reviewers decided to increase their overall ratings after some of their concerns were clarified.

While we are inclined to recommending this paper for acceptance, we strongly recommend the authors to take all our comments and suggestions into consideration when revising the paper for a broader audience.
Decision: Accept (poster)