Title: DYffusion: A Dynamics-informed Diffusion Model for Spatiotemporal Forecasting
Venue: NeurIPS 2023
Accepted: yes
Keywords: AI for science

diffusion models

scientific machine learning

probabilistic forecasting
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

Abstract

While diffusion models can successfully generate data and make predictions, they are predominantly designed for static images. We propose an approach for efficiently 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 in the model. We train a stochastic, time-conditioned interpolator and a forecaster network that mimic the forward and reverse processes of standard diffusion models, respectively. DYffusion naturally facilitates multi-step and long-range forecasting, 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 in DYffusion imposes a strong inductive bias and significantly improves computational efficiency compared to traditional Gaussian noise-based diffusion models. Our approach performs competitively on probabilistic forecasting of complex dynamics in sea surface temperatures, Navier-Stokes flows, and spring mesh systems. 1

1 Introduction

Dynamics forecasting refers to the task of predicting the future behavior of a dynamic system, involving learning the underlying dynamics governing the system's evolution to make accurate predictions about its future states. Obtaining accurate and reliable probabilistic forecasts is an important component 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 on 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 a 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 are prohibitive, requiring a sequential sampling process over hundreds of diffusion 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

Noiseinterpolation

D

Dend

Reverseprocess

Fore

Xt+i1

Xt+iN-1

Temporal interpolation ZΦ

(N)

S

(0)

S

ise

ast

Xt+h

Initial conditions Xt

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-informed diffusion model for multistep probabilistic forecasting. Inspired by recent findings that highlight the potential of non-Gaussian diffusion processes [2, 28, 10], we introduce a new forward process. This process relies on temporal interpolation and is implemented through a 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 of diffusion models, including their applications to complex physical systems with a large number of dimensions and low data availability.

We introduce DYffusion, a flexible framework for multi-step forecasting and long-range horizons that leverages a temporal inductive bias to accelerate training and lower memory needs. We explore the theoretical implications of our method, proving that DYffusion is 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-art probabilistic methods such as 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 t } T t =1 snapshots with x t ∈ X . 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 a conditional distribution P ( x t +1: t + h | x t -l +1: t ) that uses l snapshots from the past to forecast a subsequent horizon of h snapshots. Here, we focus on the task of forecasting from a single initial condition, where we aim to learn P ( x t +1: t + h | x t ) . This covers many realistic settings [50, 44] while minimizing the computational requirements for accurate forecasts [65].

Diffusion processes. In a standard diffusion model, the forward diffusion 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 t } T t =1 . This process can be generalized to consider a degradation operator D that takes as input the data point s (0) and outputs s ( n ) = D ( s (0) , n ) for varying degrees of degradation proportional to n ∈ { 1 , . . . , N } [2]. Oftentimes, D adds Gaussian noise with increasing levels of variance so that s ( N ) ∼ N ( 0 , I ) . A denoising network parameterized by θ , R θ is trained to restore s (0) , i.e. such that R θ ( s ( n ) , n ) ≈ s (0) . The diffusion model can be conditioned on the input dynamics by considering R θ ( s ( n ) , x t , n ) . In the case of dynamics forecasting, the diffusion model can be trained to minimize the objective

where X is the data distribution, || · || is a norm, and s (0) = x t +1: t + h is the prediction target. U [ [ a, b ] ] describes the uniform distribution over the set { a, a +1 , · · · , b } . In practice, R θ 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 a generalized version of the standard diffusion models, see Appendix B.2 for a more comprehensive analysis.

3 DYffusion: DYnamics-Informed Diffusion Model

The key innovation of our method, DYffusion, is a reimagining of the diffusion processes to more naturally model spatiotemporal sequences, x t : t + h (see Fig. 1). Specifically, we replace the degradation operator, D , with a stochastic interpolator network I ϕ , and the restoration network, R θ , with a deterministic forecaster network, F θ .

At a 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 is well aligned with the task at hand. See Table 4 for a full glossary.

Temporal interpolation as a forward process. To impose a temporal bias, we train a timeconditioned network I ϕ to interpolate between snapshots of data. Specifically, given a horizon h , we train I ϕ so that I ϕ ( x t , x t + h , i ) ≈ x t + i for i ∈ { 1 , . . . , h -1 } using the objective:

Algorithm 1 DYffusion, Two-stage Training

Input: networks F θ , I ϕ , norm || · || , horizon h , schedule [ i n ] N -1 i =0 Stage 1: Train interpolator network, I ϕ

Sample i ∼ Uniform ( { 1 , . . . , h -1 } )

Sample x t , x t + i , x t + h ∼ X

Optimize min ϕ ||I ϕ ( x t , x t + h , i ) -x t + i || 2

Stage 2: Train forecaster network (diffusion model backbone), F θ

Freeze I ϕ and enable inference stochasticity (e.g. dropout)

Sample n ∼ Uniform ( { 0 , . . . , N -1 } ) and x t , x t + h ∼ X

Optimize min θ || F θ ( I ϕ ( x t , x t + h , i n ) , i n ) -x t + h ||

2

Interpolation is an easier task than forecasting, and we can use the resulting interpolator I ϕ 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, I ϕ , to produce stochastic outputs within the diffusion model and during inference time. We enable this using Monte Carlo dropout [17] at inference time.

Forecasting as a reverse process. In the second stage, we train a forecaster network F θ to forecast x t + h such that F θ ( I ϕ ( x t , x t + h , i | ξ ) , i ) ≈ x t + h for i ∈ S = [ i n ] N -1 n =0 , where S denotes a 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

To include the setting where F θ learns to forecast the initial conditions, we define i 0 := 0 and I ϕ ( x t , · , i 0 ) := x t . 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 = [ j ] h -1 j =0 . Generally, the interpolation timesteps should satisfy 0 = i 0 < i n < i m < h for 0 < n < m ≤ N -1 . 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 i n instead of n , we can choose any diffusion-dynamics schedule during training or inference and even use F θ for unseen timesteps (see Fig. 3). The resulting two-stage training algorithm is shown in Alg. 1.

Because the interpolator I ϕ is frozen in the second stage, the imperfect forecasts ˆ x t + h = F θ ( I ϕ ( x t , x t + h , i n ) , i n ) may degrade accuracy when used during sequential sampling. To handle this, we introduce a one-step look-ahead loss term || F θ ( I ϕ ( x t , ˆ x t + h , i n +1 ) , i n +1 ) -x t + h || 2 whenever n +1 < N and weight the two loss terms equally. Additionally, providing a clean or noised form of the initial conditions x t as an additional input to the forecaster net can improve performance. These additional tricks are discussed in Appendix B.

Sampling. Taken together, we can write the generative process of DYffusion as follows:

where s (0) = x t and s ( n ) ≈ x t + i n 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 = 0 refers to the start of the reverse process (i.e. x t ), while n = N refers to the final output of the reverse process (here, x t + h ). Our reverse process steps forward in time, in contrast to the mapping from noise to data in standard diffusion models. As a result, DYffusion should require fewer diffusion steps and data.

Similarly to the forward process in [61, 2], our interpolation stage ceases to be a diffusion process. Our forecasting stage as detailed in Eq. (3), follows the (generalized) diffusion model objectives. This similarity allows us to use many existing 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. F θ always predicts the last timestep, x t + h , but iteratively improves those forecasts as the reverse process comes closer in time to t + h . This is analogous to the iterative denoising of the 'clean' data in standard diffusion models. This motivates line 6

of Alg. 2, where the final forecast of x t + h can be used to fine-tune intermediate predictions or to increase the temporal resolution of the forecast. DYffusion can be applied autoregressively to forecast even longer rollouts beyond the training horizon, as demonstrated by the Navier-Stokes and spring mesh experiments in section 5.

x

0

n

= 0

i

n

= 1

ˆ

x

1

ˆ x

(0)

h

n

= 1

i

n

= 2

ˆ

x

2

ˆ x

(1)

h

n

= 2

i

n

= 3

ˆ

x

3

ˆ x

(2)

h

n

= 3

ˆ x

4

:=

ˆ x

(3)

h

Memory footprint. DYffusion requires only x t and x t + h (plus x t + i during the first stage) to train, resulting in a constant memory footprint as a function of h . In contrast, direct multi-step prediction models including video diffusion models or (autoregressive) multi-step loss approaches require x t : t + h 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 on a maximum of 5 video frames due to GPU memory constraints [67].

Reverse process as ODE. DYffusion can be interpreted as modeling the dynamics by a 'diffusion process', similar to an implicit model [13]. This view helps explain DYffusion's superior forecasting ability, as it learns to model the physical dynamical system x ( t ) . The neural networks I ϕ and F θ are trained to construct a differential equation that implicitly models the solution to x ( t ) .

Let s be the time variable. DYffusion models the dynamics as

The initial condition is given by x ( t ) = x t . During prediction, we seek to obtain

The prediction process is equivalent to evaluating (6) at discrete points in the forecasting window ( t, t + h ] . Sampling schedules i n can be thought of as step sizes in 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 an approximation of the Euler method (derivation in Appendix C.1).

Viewing DYffusion as an implicit model emphasizes its dynamics-informed nature and reveals the relation between DYffusion and DDIM [61]. Specifically, DDIM can be viewed as an implicit model that models the solution to the noise-removing dynamics from a random image to a clean image. In comparison, DYffusion is an implicit model that learns the solutions of a dynamical system using data from the current step x t to the future step x t + h .

4 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 diffusion for multivariate time-series imputation [1] and forecasting [51], high-dimensional spatiotemporal

forecasting has remained unexplored. Methods such as combining blurring and noising techniques [28, 10, 2] and adopting a 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 corruption.

Diffusion Models for Video Prediction. Most closely related to our work are diffusion models for videos [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 is 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, 57, 9, 8, 34, 5, 33, 43]. The issue can 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 is memory-intensive 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 a fixed temporal resolution. For probabilistic dynamics forecasting [74], most methods focus on autoregressive nextstep 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 to generate an 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.

x

0

ˆ x

1

k

+1

. . .

ˆ x

k

k

+1

ˆ x

1

ˆ x

2

. . .

ˆ x

h

p

θ

5 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 × 60 longitude resolution in the eastern tropical Pacific Ocean). Unlike the data based on the NEMO dataset in [11, 71], we choose OISSTv2

as our SST dataset because it contains more data (although it has a lower spatial resolution of 1 / 4 ◦ compared to 1 / 12 ◦ of NEMO). We train, validate, and test all models for the years 1982-2019, 2020, and 2021, respectively. The preprocessing details are in the Appendix D.1.1.

Navier-Stokes flow: benchmark dataset from [44], which consists of a 221 × 42 grid. Each trajectory contains four randomly generated circular obstacles that block the flow. The viscosity is 1 e -3 . The channels consist of the x and y velocities as well as a pressure field. Boundary conditions and obstacle masks are given as additional inputs to all models.

Spring Mesh: benchmark dataset from [44]. It represents a 10 × 10 grid of particles connected by springs, each with mass 1 . The channels consist of two position and momentum fields each.

5.2 Experimental Setup

We focus on probabilistic multi-step forecasting , and especially long rollouts, as opposed to attaining the best next-step or short-term forecasts. This is an important distinction because, while deterministic single-step forecasting models such as those of [44] are appropriate for short-term predictive skill comparison, such models tend to lack stability for long-term forecasts. We demonstrate this in the Appendix D.4.2 (see Fig. 6), where our method and baselines excel in the long-range forecasting regime while being competitive in the short-range regime compared to the baselines from [44].

Neural network architectures. For a given dataset, we use the same backbone architecture for all baselines as well as for both the interpolation and forecaster networks in DYffusion. For the SST dataset, we use a popular UNet architecture designed for diffusion models 2 . For the Navier-Stokes and spring mesh datasets, we use the UNet and CNN from the original paper [44], respectively. The UNet and CNN models from [44] are extended by the sine/cosine-based featurization module of the SST UNet to embed the diffusion step or dynamical timestep.

Baselines. We compare our method against both direct applications of standard diffusion models to dynamics forecasting and methods to ensemble the 'barebone' backbone network of each dataset. The network operating in 'barebone' form means that there is no involvement of diffusion.

DDPM [25] : both next-step (image-like problem) as well as multi-step (video-like problem) prediction mode (we report the best model version only in the results).

MCVD[67]: conditional video diffusion model for video prediction (in 'concat' mode).

Dropout [17]: Ensemble multi-step forecasting based on enabling dropout of the barebone backbone network at inference time.

Perturbation [45]: Ensemble multi-step forecasting based on random perturbations of the initial conditions/inputs with a fixed variance of the barebone backbone network.

MCVD and the multi-step DDPM variant predict the timesteps x t +1: t + h based on x t , which is the standard approach in video (prediction) diffusion models. The barebone backbone network baselines are time-conditioned forecasters (similarly to the DYffusion forecaster) trained on the multi-step objective E i ∼U [ [1 ,h ] ] , x t,t + i ∼X || F θ ( x t , i ) -x t + i || 2 from scratch. We found it to perform very similarly to predicting all h horizon timesteps at once in a single forward pass (not shown). In all experiments, we observed that the single-step forecasting ( h = 1 ) version of the barebone network yielded significantly lower performance compared to any of the multi-step training approaches (see Appendix D.4.2). More details of the implementation are provided in Appendix D.2.

Evaluation. We evaluate the models on the best validation Continuous Ranked Probability Score (CRPS) [42], which is a proper scoring rule and a popular metric in the probabilistic forecasting literature [19, 12, 51, 48, 58]. The CRPS is computed by generating a 20-member ensemble (i.e. 20 samples are drawn per batch element), while we generate a 50-member ensemble for final model selection between different hyperparameter runs. We also use a 50-member ensemble for evaluation on the test datasets to compute the CRPS, mean squared error (MSE), and spread-skill ratio (SSR). The MSE is computed on the ensemble mean prediction. The SSR is defined as the ratio of the square root of the ensemble variance to the corresponding ensemble RMSE. It serves as a measure of the reliability of the ensemble, where values smaller than 1 indicate underdispersion (i.e. the probabilistic forecast is overconfident in its forecasts), and larger values overdispersion [16, 18]. On the Navier-Stokes and spring mesh datasets, models are evaluated by autogressively forecasting the full test trajectories of length 64 and 804, respectively. For the SST dataset, all models are evaluated on forecasts of up to 7 days. We do not explore more long-term SST forecasts because the chaotic nature of the system, and the fact that we only use regional patches, inherently limits predictability.

5.3 Results

Quantitative Results. We present the time-averaged metrics for the SST and Navier-Stokes dataset in Table 1 and for the Spring Mesh dataset in Table 2. DYffusion performs best on the Navier-Stokes dataset, while coming in a close second on the SST dataset after MCVD, in terms of CRPS. Since MCVD uses 1000 diffusion steps 3 , it is slower to sample from at inference time than from our DYffusion which is trained with less than 50 diffusion steps. The DDPM model for the SST dataset is fairly efficient because it only uses 5 diffusion steps but lags in terms of performance. Thanks to the time-conditioned nature of DYffusion's interpolator and forecaster nets, memory is not an issue when scaling our framework to long horizons. In the spring mesh dataset, we train with a horizon of 134 and evaluate long trajectories of 804 steps. Our method beats the barebone network baseline, with a larger margin on the out-of-distribution test dataset. It is worth noting that our reported MSE scores are significantly better than the ones reported in [44], likely due to multi-step training being superior to the single-step forecasting approach of [44] (see Appendix D.4.2 for more details). For the out-of-distribution test set of the Navier-Stokes benchmark, the results are almost identical to the one in Tab. 1, so we show them in Appendix D.4.1.

Qualitative Results. In dynamics forecasting, long-range forecasts of ML models often suffer from blurriness (or might even diverge when using autoregressive models). In Fig. 4 we show exemplary

samples for the best baseline (Dropout) and DYffusion as well as the corresponding ground truth at five different timesteps from a complete Navier-Stokes trajectory forecast. Our method can reproduce the true dynamics over the full trajectory and does so better than the baseline, especially for fine-scale patterns such as the tails of the flow after the right-most obstacle. The corresponding full video can be found at this URL.

Increasing the forecasted resolution. Motivated by the continuous-time nature of DYffusion, we aim to study in this experiment whether it is possible to forecast skillfully beyond the resolution given by the data. Here, we forecast the same Navier-Stokes trajectory shown in Fig. 4 but at 8 × resolution (i.e. 512 timesteps instead of 64 are forecasted in total). This behavior can be achieved by either changing the sampling trajectory [ i n ] N -1 n =0 or by including additional output timesteps, J , for the refinement step of Alg. 2. In this case, we choose to do the latter and find the resulting forecast to be visibly pleasing and temporally consistent; see the full video at this URL.

Note that we hope that our probabilistic forecasting model can capture any of the possible, uncertain futures instead of forecasting their mean, as a deterministic model would do. As a result, some long-term rollout samples are expected to deviate from the ground truth. For example, see the velocity at t = 3 . 70 in the video. It is reassuring that our samples show sufficient variation, but also cover the ground truth quite well (sample 1).

Ablations. In Appendix D.5, we ablate various components of our method. This includes showing that inference stochasticity in the interpolation network is crucial for performance. Our experiments in Appendix D.5.1 confirm the findings of [2] regarding the superiority of cold sampling (see Alg. 2) compared to the naive sampling (Alg. 4) algorithm. We provide a theoretical justification in Appendix C.2, where we show that cold sampling has first-order discretization error, whereas naive sampling does not.

We also study the choice of the interpolation schedule:

For the Navier-Stokes and spring mesh datasets, it was sufficient to use the simplest possible schedule S = [ i n ] N -1 n =0 = [ j ] h -1 j =0 , where the number of diffusion steps and the horizon are equal. For the SST dataset, using k additional diffusion steps corresponding to floating-point i k , especially such that i k < 1 for all k , improved performance significantly. Similarly to the finding of DDIM [61] for Gaussian diffusion, we show that the sampling trajectories of DYffusion can be accelerated by skipping intermediate diffusion steps, resulting in a clear trade-off between accuracy and speed. Finally, we verify that the forecaster's prediction of x t + h for a given diffusion step n and associated input ˆ x t + i n usually improves as n increases, i.e. as we approach the end of the reverse process and x t + i n gets closer in time to x t + h .

The variance in the performance of the standard diffusion models across datasets may be intriguing. One possible explanation could be the size of the training dataset, since the SST dataset is more than ten times larger than the Navier-Stokes and spring mesh datasets, respectively. Sufficient training data may be a key factor that enables Gaussian diffusion models to effectively learn the reverse process

from noise to the data space. We explore this hypothesis in Table 3, where we retrain MCVD as well as DYffusion's F θ and I ϕ on small subsets of the SST training set. We find that DYffusion's performance degrades more gracefully under limited training data, especially when using only two years of data (or ≈ 8 , 000 training data points).

We found it very challenging to train video diffusion models on long horizons. A reason could be that the channel dimension of the video diffusion backbone is scaled by the training horizon, which increases the problem dimensionality and complexity considerably (see Appendix D.6). Unlike common practice in natural image or video problems, we do not threshold the predictions of any diffusion model, as the data range is unbounded. However, this thresholding has been shown to be an important implementation detail that stabilizes diffusion model sampling for complex problems in the domain of natural images [40]. This could be the reason why the studied diffusion model baselines produce poor predictions for the Navier-Stokes and spring mesh datasets.

6 Conclusion

Wehave introduced DYffusion, a novel dynamics-informed diffusion model for improved probabilistic spatiotemporal forecasting. In contrast to most prior work on probabilistic forecasting that often generates long-range forecasts via iterating through an autoregressive model, we tailor diffusion models to naturally support long-range forecasts during inference time by coupling their diffusion process with the dynamical nature of the data. Our study presents the first comprehensive evaluation of diffusion models for the task of spatiotemporal forecasting. Finally, we believe that our dynamicsinformed diffusion model can serve as a more general source of inspiration for tailoring diffusion models to other conditional generation problems, such as super-resolution, where the data itself can natively inform an inductive bias for the diffusion model.

Limitations. While our proposed framework can be applied to any dynamics forecasting problem as an alternative to autoregressive models, it does not support, in its current form, problems where the output space is different from the input space. Additionally, although it requires fewer diffusion steps compared to Gaussian diffusion models to achieve similar performance, the need for additional forward passes through the interpolator network diminishes some of these advantages. When prioritizing inference speed, approaches that only require a single forward pass to forecast outperform sequential sampling of diffusion models, including DYffusion.

Future work. Cold Diffusion [2] indicates that deterministic data degradations cannot meet the performance of stochastic (i.e. Gaussian noise-based) degradations. In our work, we have experimented with using an interpolator network with Monte Carlo dropout enabled during inference time to achieve a stochastic forward process, which we have shown to be important for optimal performance. Exploring more advanced approaches for introducing stochasticity into our general framework presents an interesting avenue for future work. A more advanced approach could draw insights from the hypernetwork-based parameterization of the latent space from [69]. Furthermore, given our finding of the equivalence of cold sampling with the Euler solver (C.1), it could be interesting to consider more advanced ODE solvers for improved sampling from DYffusion. Recent advances in neural architectures for physical systems [6, 31] are complementary to our work. Using such alternative neural architectures as interpolator or forecaster networks within our framework could be a promising approach to further improve performance.

References

Juan Lopez Alcaraz and Nils Strodthoff. Diffusion-based time series imputation and forecasting with structured state space models. Transactions on Machine Learning Research , 2023. ISSN 2835-8856. URL https://openreview.net/forum?id=hHiIbk7ApW .

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.

Peter Bauer, Alan Thorpe, and Gilbert Brunet. The quiet revolution of numerical weather prediction. Nature , 525(7567):47-55, Sep 2015. ISSN 1476-4687. doi: 10.1038/nature14956. URL https://doi.org/10.1038/nature14956 .

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.

Johannes Brandstetter, Daniel E. Worrall, and Max Welling. Message passing neural PDE solvers. In International Conference on Learning Representations , 2022. URL https:// openreview.net/forum?id=vSix3HPYKSU .

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://openreview.net/forum?id=okwxL_c4x84 .

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 , 51(8):6110-6124, 2015.

Ashesh Chattopadhyay and Pedram Hassanzadeh. Why are deep learning-based models of geophysical turbulence long-term unstable? NeurIPS, Machine Learning for Physical Sciences , 2022.

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.

Giannis Daras, Mauricio Delbracio, Hossein Talebi, Alexandros G. Dimakis, and Peyman Milanfar. Soft diffusion: Score matching for general corruptions. 2022. URL https://arxiv. org/abs/2209.05442 .

Emmanuel de Bézenac, Arthur Pajot, and Patrick Gallinari. Deep learning for physical processes: Incorporating prior scientific knowledge. In International Conference on Learning Representations , 2018. URL https://openreview.net/forum?id=By4HsfWAZ .

Emmanuel de Bézenac, Syama Sundar Rangapuram, Konstantinos Benidis, Michael Bohlke-Schneider, Richard Kurle, Lorenzo Stella, Hilaf Hasson, Patrick Gallinari, and Tim Januschowski. Normalizing kalman filters for multivariate time series analysis. In Advances in Neural Information Processing Systems , volume 33, pages 29953007, 2020. URL https://proceedings.neurips.cc/paper_files/paper/2020/ file/1f47cef5e38c952f94c5d61726027439-Paper.pdf .

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-958, 2021.

N. Benjamin Erichson, Michael Muehlebach, and Michael W. Mahoney. Physics-informed autoencoders for lyapunov-stable fluid flow prediction. 2019. URL https://arxiv.org/ abs/1905.10866 .

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: 10.1038/s41467-022-32483-x. URL https: //doi.org/10.1038/s41467-022-32483-x .

V. Fortin, M. Abaza, F. Anctil, and R. Turcotte. Why should ensemble spread match the rmse of the ensemble mean? Journal of Hydrometeorology , 15(4):1708 - 1713, 2014. doi:

https://doi.org/10.1175/JHM-D-14-0008.1. URL https://journals.ametsoc.org/view/ journals/hydr/15/4/jhm-d-14-0008_1.xml .

Yarin Gal and Zoubin Ghahramani. Dropout as a bayesian approximation: Representing model uncertainty in deep learning. In international conference on machine learning , pages 1050-1059. PMLR, 2016.

Sagar Garg, Stephan Rasp, and Nils Thuerey. Weatherbench probability: A benchmark dataset for probabilistic medium-range weather forecasting along with deep learning baseline models. arXiv preprint arXiv:2205.00865 , 2022.

Tilmann Gneiting and Matthias Katzfuss. Probabilistic forecasting. Annual Review of Statistics and Its Application , 1(1):125-151, 2014.

Tilmann Gneiting and Adrian E. Raftery. Weather forecasting with ensemble methods. Science , 310(5746):248-249, 2005. doi: 10.1126/science.1115255. URL https://www.science. org/doi/abs/10.1126/science.1115255 .

Andre Graubner, Kamyar Kamyar Azizzadenesheli, Jaideep Pathak, Morteza Mardani, Mike Pritchard, Karthik Kashinath, and Anima Anandkumar. Calibration of large neural weather models. In NeurIPS 2022 Workshop on Tackling Climate Change with Machine Learning , 2022.

Yoo-Geun Ham, Jeong-Hwan Kim, and Jing-Jia Luo. Deep learning for multi-year enso forecasts. Nature , 573:568-572, 9 2019.

Xu Han, Han Gao, Tobias Pfaff, Jian-Xun Wang, and Liping Liu. Predicting physics in meshreduced space with temporal attention. In International Conference on Learning Representations , 2022. URL https://openreview.net/forum?id=XctLdNfCmP .

William Harvey, Saeid Naderiparizi, Vaden Masrani, Christian Weilbach, and Frank Wood. Flexible diffusion modeling of long videos. Advances in Neural Information Processing Systems , 2022.

Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. In Advances in Neural Information Processing Systems , 2020. URL https://proceedings. neurips.cc/paper/2020/file/4c5bcfec8584af0d967f1ab10179ca4b-Paper.pdf .

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: //arxiv.org/abs/2210.02303 .

Jonathan Ho, Tim Salimans, Alexey Gritsenko, William Chan, Mohammad Norouzi, and David J Fleet. Video diffusion models. Advances in Neural Information Processing Systems , 2022.

Emiel Hoogeboom and Tim Salimans. Blurring diffusion models. In The Eleventh International Conference on Learning Representations , 2023.

Yuan Hu, Lei Chen, Zhibin Wang, and Hao Li. Swinvrnn: A data-driven ensemble forecasting model via learned distribution perturbation. Journal of Advances in Modeling Earth Systems , 15(2):e2022MS003211, 2023. doi: https://doi.org/10.1029/ 2022MS003211. URL https://agupubs.onlinelibrary.wiley.com/doi/abs/10. 1029/2022MS003211 . e2022MS003211 2022MS003211.

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 - 2939, 2021. doi: https://doi.org/10.1175/JCLI-D-20-0166.1. URL https://journals.ametsoc.org/view/ journals/clim/34/8/JCLI-D-20-0166.1.xml .

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.net/forum?id=mfIX4QpsARJ .

Tero Karras, Miika Aittala, Timo Aila, and Samuli Laine. Elucidating the design space of diffusion-based generative models. In Proc. NeurIPS , 2022.

K. Kashinath, M. Mustafa, A. Albert, J-L. Wu, C. Jiang, S. Esmaeilzadeh, K. Azizzadenesheli, R. Wang, A. Chattopadhyay, A. Singh, A. Manepalli, D. Chirila, R. Yu, R. Walters, B. White, H. Xiao, H. A. Tchelepi, P. Marcus, A. Anandkumar, P. Hassanzadeh, and

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):20200093, 2021. doi: 10.1098/rsta.2020.0093. URL https://royalsocietypublishing.org/doi/abs/10.1098/rsta.2020.0093 .

Ryan Keisler. Forecasting global weather with graph neural networks. 2022. URL https: //arxiv.org/abs/2202.07575 .

Dmitrii Kochkov, Jamie A. Smith, Ayya Alieva, Qing Wang, Michael P. Brenner, and Stephan Hoyer. Machine learning-accelerated computational fluid dynamics. Proceedings of the National Academy of Sciences , 118(21):e2101784118, 2021. doi: 10.1073/pnas.2101784118. URL https://www.pnas.org/doi/abs/10.1073/pnas.2101784118 .

Remi Lam, Alvaro Sanchez-Gonzalez, Matthew Willson, Peter Wirnsberger, Meire Fortunato, Alexander Pritzel, Suman Ravuri, Timo Ewalds, Ferran Alet, Zach Eaton-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 forecasting. 2022. URL https://arxiv.org/abs/2212.12794 .

Nathan Lambert, Albert Wilcox, Howard Zhang, Kristofer S. J. Pister, and Roberto Calandra. Learning accurate long-term dynamics for model-based reinforcement learning. In IEEE Conference on Decision and Control (CDC) , 2021.

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 Meteorological Society , 90(6):785 - 798, 2009. doi: 10.1175/2008BAMS2604.1. URL https: //journals.ametsoc.org/view/journals/bams/90/6/2008bams2604_1.xml .

Jane A Leggett. The united nations framework convention on climate change, the kyoto protocol, and the paris agreement: a summary. UNFCC: New York, NY, USA , 2, 2020.

Aaron Lou and Stefano Ermon. Reflected diffusion models. In International Conference on Machine Learning , 2023.

Giorgos Mamakoukas, Ian Abraham, and Todd Murphey. Learning stable models for prediction and control. IEEE Transactions on Robotics , 2020.

James E. Matheson and Robert L. Winkler. Scoring rules for continuous probability distributions. Management Science , 22(10):1087-1096, 1976.

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.

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 https://arxiv.org/abs/2108.07799 .

Jaideep Pathak, Shashank Subramanian, Peter Harrington, Sanjeev Raja, Ashesh Chattopadhyay, Morteza Mardani, Thorsten Kurth, David Hall, Zongyi Li, Kamyar Azizzadenesheli, Pedram Hassanzadeh, Karthik Kashinath, and Animashree Anandkumar. FourCastNet: A global datadriven high-resolution weather model using adaptive fourier neural operators. 2022. URL https://arxiv.org/abs/2202.11214 .

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.

Aditya Ramesh, Prafulla Dhariwal, Alex Nichol, Casey Chu, and Mark Chen. Hierarchical text-conditional image generation with clip latents. arXiv preprint arXiv:2204.06125 , 2022.

Stephan Rasp and Sebastian Lerch. Neural networks for postprocessing ensemble weather forecasts. Monthly Weather Review , 146(11):3885 - 3900, 2018.

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: https://doi.org/10.1029/2020MS002405. URL https://agupubs.onlinelibrary.wiley.com/doi/abs/10.1029/2020MS002405 .

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. URL https://doi.org/10.1029% 2F2020ms002203 .

Kashif Rasul, Calvin Seward, Ingmar Schuster, and Roland Vollgraf. Autoregressive denoising diffusion models for multivariate probabilistic time series forecasting. In ICML , 2021.

Suman Ravuri, Karel Lenc, Matthew Willson, Dmitry Kangin, Remi Lam, Piotr Mirowski, Megan Fitzsimons, Maria Athanassiadou, Sheleem Kashem, Sam Madge, et al. Skilful precipitation nowcasting using deep generative models of radar. Nature , 597(7878):672-677, 2021. URL https://doi.org/10.1038/s41586-021-03854-z .

Severi Rissanen, Markus Heinonen, and Arno Solin. Generative modelling with inverse heat dissipation. In The Eleventh International Conference on Learning Representations , 2023.

Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Björn Ommer. Highresolution image synthesis with latent diffusion models. In IEEE Conference on Computer Vision and Pattern Recognition , pages 10684-10695, 2022.

Alvaro Sanchez-Gonzalez, Jonathan Godwin, Tobias Pfaff, Rex Ying, Jure Leskovec, and Peter W. Battaglia. Learning to simulate complex physics with graph networks. In International Conference on Machine Learning , 2020.

S. Scher. Toward data-driven weather and climate forecasting: Approximating a simple general circulation model with deep learning. Geophysical Research Letters , 45(22): 12,616-12,622, 2018. doi: https://doi.org/10.1029/2018GL080704. URL https://agupubs. onlinelibrary.wiley.com/doi/abs/10.1029/2018GL080704 .

S. Scher and G. Messori. Generalization properties of feed-forward neural networks trained on lorenz systems. Nonlinear Processes in Geophysics , 26(4):381-399, 2019. doi: 10.5194/ npg-26-381-2019. URL https://npg.copernicus.org/articles/26/381/2019/ .

Sebastian Scher and Gabriele Messori. Ensemble methods for neural network-based weather forecasts. Journal of Advances in Modeling Earth Systems , 13(2), 2021. doi: https://doi.org/ 10.1029/2020MS002331. URL https://agupubs.onlinelibrary.wiley.com/doi/abs/ 10.1029/2020MS002331 .

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://arxiv.org/abs/ 2209.14792 .

Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. In Proceedings of the 32nd International Conference on Machine Learning , 2015. URL https://proceedings.mlr.press/v37/ sohl-dickstein15.html .

Jiaming Song, Chenlin Meng, and Stefano Ermon. Denoising diffusion implicit models. In International Conference on Learning Representations , 2021.

Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution. In Advances in Neural Information Processing Systems , volume 32, 2019.

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.

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.

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://openreview.net/forum?id=tmIiMPl4IPa .

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://arxiv.org/abs/2210.03324 .

Vikram Voleti, Alexia Jolicoeur-Martineau, and Christopher Pal. MCVD: Masked conditional video diffusion for prediction, generation, and interpolation. Advances in Neural Information Processing Systems , 2022. URL https://arxiv.org/abs/2205.09853 .

Etienne E Vos, Ashley Gritzman, Sibusisiwe Makhanya, Thabang Mashinini, and Campbell D Watson. Long-range seasonal forecasting of 2m-temperature with machine learning. NeurIPS, Tackling Climate Change with Machine Learning , 2020.

Zhong Yi Wan, Leonardo Zepeda-Nunez, Anudhyan Boral, and Fei Sha. Evolve smoothly, fit consistently: Learning smooth latent dynamics for advection-dominated systems. In The Eleventh International Conference on Learning Representations , 2023. URL https: //openreview.net/forum?id=Z4s73sJYQM .

Rui Wang, Karthik Kashinath, Mustafa Mustafa, Adrian Albert, and Rose Yu. Towards physicsinformed deep learning for turbulent flow prediction. Proceedings of the 26th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining , 2020.

Rui Wang, Robin Walters, and Rose Yu. Meta-learning dynamics forecasting using task inference. In Advances in Neural Information Processing Systems (NeurIPS) , 2022.

Anthony L Westerling, Alexander Gershunov, Timothy J Brown, Daniel R Cayan, and Michael D Dettinger. Climate and wildfire in the western united states. Bulletin of the American Meteorological Society , 84(5):595-604, 2003.

Jonathan A. Weyn, Dale R. Durran, and Rich Caruana. Can machines learn to predict weather? using deep learning to predict gridded 500-hpa geopotential height from historical weather data. Journal of Advances in Modeling Earth Systems , 11(8):2680-2693, 2019. doi: https: //doi.org/10.1029/2019MS001705. URL https://agupubs.onlinelibrary.wiley.com/ doi/abs/10.1029/2019MS001705 .

Dongxia Wu, Liyao Gao, Matteo Chinazzi, Xinyue Xiong, Alessandro Vespignani, Yi-An Ma, and Rose Yu. Quantifying uncertainty in deep spatiotemporal forecasting. In Proceedings of the 27th ACM SIGKDD Conference on Knowledge Discovery & Data Mining , pages 1841-1851, 2021.

Yilun Xu, Ziming Liu, Yonglong Tian, Shangyuan Tong, Max Tegmark, and Tommi Jaakkola. Pfgm++: Unlocking the potential of physics-inspired generative models. 2023. URL https: //arxiv.org/abs/2302.04265 .

Ling Yang, Zhilong Zhang, Yang Song, Shenda Hong, Runsheng Xu, Yue Zhao, Yingxia Shao, Wentao Zhang, Bin Cui, and Ming-Hsuan Yang. Diffusion models: A comprehensive survey of methods and applications, 2022. URL https://arxiv.org/abs/2209.00796 .

Ruihan Yang, Prakhar Srivastava, and Stephan Mandt. Diffusion probabilistic modeling for video generation. arXiv preprint arXiv:2203.09481 , 2022.

Contents

Appendix

A Glossary

The glossary is given in Table 4 below.

B Methodology

B.1 DYffusion

Algorithm 3 DYffusion, Two-stage Training, all details

Input: networks F θ , I ϕ , norm ||·|| , horizon h , schedule [ i n ] N -1 i =0 , loss coefficients λ 1 , λ 2 , forecaster conditioning c Stage 1: Train interpolator network, I ϕ 1. Sample i ∼ Uniform ( { 1 , . . . , h -1 } ) 2. Sample x t , x t + i , x t + h ∼ X 3. Optimize min ϕ ||I ϕ ( x t , x t + h , i ) -x t + i || 2 Stage 2: Train forecaster network (diffusion model backbone), F θ 1. Freeze I ϕ and enable inference stochasticity (e.g. dropout) 2. Sample n ∼ Uniform ( { 0 , . . . , N -1 } ) and x t , x t + h ∼ X 3. ˆ x (1) t + h ← F θ ( I ϕ ( x t , x t + h , i n ) , i n , c ( x t , n )) 4. ˆ x (2) t + h ← F θ ( I ϕ ( x t , ˆ x (1) t + h , i n +1 ) , i n +1 , c ( x t , n +1)) if n ≤ N -2 else x t + h 5. Optimize min θ λ 1 || ˆ x (1) t + h -x t + h || 2 + λ 2 || ˆ x (2) t + h -x t + h || 2

In Alg. 3 we give the full training algorithm for DYffusion. The only differences to the truncated training algorithm in the main text (Alg. 1) are that we also include:

One-step look-ahead loss term trick, λ 2 || F θ ( I ϕ ( x t , ˆ x t + h , i n +1 ) , i n +1 ) -x t + h || 2 , whenever n +1 < 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 = λ 2 = 1 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 second-stage too.

Optional conditioning of the forecaster network on a noised or clean version of the initial conditions, x t . We explore three different approaches: 1) No conditioning at all, i.e. c ( · , · ) = null ; 2) Clean initial conditions, i.e. c ( x t , · ) = x t ; 3) Noised initial conditions, where we sample random noise ϵ ∼ N ( 0 , I ) of the same shape of x t , and then define c ( x t , n ) = n N -1 x t + (1 -n N -1 ) ϵ . The second variant is analogous to the conditioning of conventional diffusion models, which DYffusion does not necessarily require since the initial conditions indirectly influence the main input of the forecaster net (i.e. ˆ x t + i n = I ϕ ( x t , x t + h , i n ) ). The motivation for the third variant stems from trying to give the forecaster net a more explicit conditioning on the initial conditions when the reverse process goes farther away from t (i.e. as n → N and ˆ x t + i n → x t + h ) . An ablation for the choice of c 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), 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 D.5.1). Our proposed forward and reverse processes are specifically designed to fall under this framework. Given this context, it becomes clear that DYffusion is a generative model for forecasting that falls in the category of 'generalized diffusion models'.

Training objectives. To see the similarity between standard diffusion models and DYffusion in terms of training objectives, we write out the corresponding equations Eq. 1 and Eq. 3. For a standard 'generalized' diffusion model, we can rewrite the main objective inside Eq. 1 as follows:

where s (0) = x t +1: t + h , and s ( n ) is a noisy version of s (0) (its level of corruption increases with n ).

In the context of DYffusion's forecasting network, we can express its training objective (Eq. 3) in the following manner:

where s ( N ) = x t + h and s ( n ) ≈ x t + i n is now a stepped backward in time version of 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. s ( n ) ≈ x t + i n temporally precedes s ( n +1) ≈ x t + i n +1 . Accounting for the reversed indexing of the diffusion steps, it is easy to see the equivalence of the denoising network, R θ , in a conventional diffusion model, and its forecasting network counterpart, F θ , in DYffusion. The reason why DYffusion's forecaster network does not necessarily require the initial conditions, x t , as additional input is because the diffusion state, s ( n ) , depends on them. The central difference is 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 n , s ↦→ x ( s ) . Here s denotes the time variable (since t is already used to define the index of temporal data snapshots).

During prediction, x t is given, ϕ is fixed, and I ϕ only depends on ˆ x t + h and i . Therefore, we simplify notations by writing I ϕ ( x t , ˆ x t + h , i ) = I ϕ, x t ( ˆ x t + h , i ) . We will further omit the subscripts ϕ and x t .

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

for a small ∆ s . We do not have access to d I ( F θ ( x s ,s ) ,s ) ds . However, since we know F θ and I ϕ , we can approximate d I ( F θ ( x s ,s ) ,s ) ds by its first-order Taylor expansion around s:

This step can also be interpreted as evaluating the integral in (6) using the fundamental theorem of calculus. Then the Euler method becomes

Note that x s +∆ s on the right-hand side is unknown because it is the quantity we want to approximate in this step. A reasonable way to approximate F θ ( x s +∆ s , s + ∆ s ) is to replace it by F θ ( x s , s ) because they both predict x ( h ) and use nearby points (assuming ∆ s is small and x behaves nicely around s ). The resulting update,

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 F θ ( x s +∆ s , s +∆ s ) ≈ F θ ( x s , s ) to obtain the update rule (previous equation). The error introduced by this approximation is expected to be larger when s is small. The intuition is as follows. When s is small, the distance between s and t + h is large, and F θ ( x s , s ) has to make a prediction farther ahead. Predicting far into the future is generally harder than predicting the near future. Therefore, the prediction error F θ ( x s , s ) -x t + h is expected to be larger when s is small. The larger uncertainty may lead to larger difference between F θ ( x s +∆ s , s +∆ s ) and F θ ( x s , s ) .

To reduce the error brought by this approximation, it makes sense to 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 forward processes. Both

sampling algorithms can be directly used to sample from DYffusion. For completeness, we adapt the simpler, naive sampling algorithm to our notation in 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 a 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 a term proportional to the step size ∆ s . Naive sampling does not have this property.

The true value of x at s +∆ s according to Equation (6) is

Recall from Alg. 2 that given x ( s ) , cold sampling predicts x ( s +∆ s ) as

The discretization error e ( x ) of one step of cold sampling is the difference between the exact and predicted x ( s +∆ s ) :

The following proposition states that e ( x ) is bounded by a term proportional to the step size ∆ s .

Proposition C.1. Assume that F θ ( x ( s ) , s ) is Lipschitz in s . Assume also that I ϕ ( x t + h , s ) is Lipschitz in x t + h . The norm of the cold sampling discretization error, || e ( x , ∆ s ) || 2 , is bounded by O (∆ s ) .

Proof. The proof relied on applying definitions of Lipschitz functions twice. Let L 1 be the Lipschitz constant for F θ ( x ( s ) , s ) in s . Let L 2 be the Lipschitz constant for I ( x , s ) in x . Since F θ ( x ( s ) , s ) is Lipschitz in s , we have || F θ ( x ( s +∆ s ) , s + ∆ s ) -F θ ( x ( s ) , s ) || 2 ≤ L 1 ∆ s . Since I ( x , s ) is Lipschitz in x , we have ||I ( F θ ( x ( s +∆ s ) , s +∆ s ) , s +∆ s ) -I ( F θ ( x ( s ) , s ) , s +∆ s ) || 2 ≤ L 2 L 1 ∆ s . Therefore || e ( x ) || ≤ L 2 L 1 ∆ s , which means the discretization error is bounded by a 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 of step size. In naive sampling, the predicted x at time s +∆ s is

The discretization error of one step of naive sampling is x ( s +∆ s ) :

Note that the first two terms are the same as the discretization error in cold sampling. However, the last two terms are not bounded by first-order terms of ∆ s . Hence, naive sampling can 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 at a daily time scale. These data are available from 1982 to the present at a resolution of 1 / 4 ◦ degrees. For training, we use the years 1982-2018, for validation 2019, and for testing 2020. We have preprocessed the NOAA OI SST V2 dataset as follows:

First, the globe is divided into 60 × 60 latitude × longitude grid tiles,

all tiles with less than 95% of ocean cover are filtered out,

standardize the raw SSTs using daily means and standard deviations (computed on the training set only, i.e. 1982-2018),

replace continental NaNs with zeroes (after standardization), and

we subsample 11 grid tiles (covering mostly the eastern tropical Pacific, as shown in Fig. 5).

Hyperparameters for each dataset

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 metrics that we employ too). We always use the full training set provided by [44]. We use the Navier-Stokes 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 a floating point precision of 16 and do not use a learning rate scheduler. All diffusion models, including DYffusion, are trained with the L1 loss, while all barebone UNet/CNN networks are trained on the L2 loss. We use three different dropout rates for the SST UNet: 1) before the query-key-value projection of each attention layer, dr at ; 2) After the first sub-block of each ResNet block, dr bl 1 ; 3) After the second sub-block of each ResNet block, dr bl 2 , 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 t , with small amounts of Gaussian noise ϵ ∼ N (0 , σ ϵ I ) . We found that σ ∗ ϵ = 0 . 05 gave the lowest CRPS scores among all variances that we tried, σ ϵ ∈ { 0 . 01 , 0 . 03 , 0 . 05 , 0 . 07 , 0 . 1 , 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 baseline 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 DYffusion, we found that using high dropout rates results in 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: dr at = dr bl 2 = 0 . 6 , dr bl 1 = 0 . 3 . For Navier-Stokes 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 . 15 , 0 . 2 , 0 . 25 , 0 . 3 } ).

DDPM We found that the cosine (linear) noise schedule gives better results for the SST (NavierStokes) 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.

MCVD [67] We train MCVD with 1000 diffusion steps for all datasets, as we were not able to successfully train it with fewer diffusion steps. We use a linear noise schedule (we found the cosine schedule to produce inferior results) using the 'predict noise' objective. Due to the inference runtime

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 MCVDaround 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 ] h -1 j =0 . Furthermore, we found that the refinement step of Alg. 2 did not improve performance for the SST dataset so we did not use it there, whereas it did improve performance for Navier-Stokes and spring mesh (see D.5.2). As for the choice of the interpolator network, I ϕ , we conduct a 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 I ϕ uses dr at = dr bl 2 = 0 . 6 , dr bl 1 = 0 . 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 interpolator network, I ϕ . Thus, for a new dataset or problem, it is 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 t and x t + h into the interpolator, we simply concatenate these snapshots along the channel dimensions, effectively doubling the number of input channels.

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 1 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 4 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 5 . 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 time-embedding scale-shift. The downsampling module is a 2D convolution that halves the spatial size of the input (with a 4 × 4 kernel and stride = 2 ). The upsampling module doubles the spatial size via nearest neighbor upsampling followed by a 2D convolution (with 3 × 3 kernel). At the 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-value projection of each attention layer, dr at ; 2) After the first sub-block of each ResNet block, dr bl 1 ; 3) After the

0

10

20

30

40

50

60

Horizon

0.00

0.02

0.04

0.06

0.08

0.10

0.12

test

MSE

DYffusion

UNet (Dropout-multi-step)

UNet [43]

(a) Navier-Stokes

0

100

200

300

400

500

600

700

800

Horizon

0.0000

0.0005

0.0010

0.0015

0.0020

0.0025

0.0030

0.0035

0.0040

test

MSE

DYffusion

CNN (Dropout-multi-step)

CNN [43]

(b) Spring Mesh

second sub-block of each ResNet block, dr bl 2 . For all models except the Dropout baseline and I ϕ in DYffusion, which use higher dropout rates, we use dr at = 0 . 1 , dr bl 1 = 0 , dr bl 2 = 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], as described below.

Time embedding module The time-embedding scale-shift, taken from the SST UNet, is a key component of all architectures since it enables them to condition on the diffusion step (for DDPM and MCVD) or the dynamical timestep (for the time-conditioned barebone models as well as for both F θ and I ϕ in DYffusion). It is implemented by a sine and cosine-based featurization of the scalar diffusion step/dynamical timestep. These features are projected by a linear layer, followed by a GeLU activation, and another linear layer, which results in a '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 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 out-of-distribution

In Table 6 we show the out-of-distribution test results for the Navier-Stokes datasets (analogously to 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-

iments [44] in terms of MSE. It is important to emphasize that our method is specifically designed to generate ensemble-based probabilistic forecasts, which is 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 or 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 our method as well as 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' is 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 we show that disabling the inference dropout in the interpolator network, I ϕ , results in considerably worse scores. This is to be expected, since without stochasticity in I ϕ , our current framework collapses to 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 to 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 ϵ ∼ N (0 , σ ϵ I ) , where we use σ ϵ = 0 . 05 .

D.5.3 Refining the interpolation forecasts after sampling

We find in Table 7 ( 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

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 t , n ) . The optimal choice of c 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 t , · ) = x t ) 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 spring mesh. As a result, the forecaster network can greatly benefit from being explicitly conditioned on the initial conditions during all diffusion states (as opposed to only being indirectly influenced by x t through the interpolator network). In practice and if compute allows, it 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 hyperparameter choice. Usually, its choice is constrained by the number of timesteps that fit into GPU memory, and it is expected that larger training horizons will improve performance when evaluated on long (autoregressive) rollouts. However, for continuous-time models including 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 explore using three different training hori-

zons ( h ∈ { 8 , 16 , 32 } ) for the Navier-Stokes dataset for both the barebone time-conditioned Dropout model as well as DYffusion. For DYffusion, this means that we train both an interpolator network as well as a 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 a next-step forecasting baseline ( h = 1 ), which results in significantly worse results than any multi-step forecasting approach. This is likely due to the compounding effect of autoregressive errors.

D.5.6 Using auxiliary diffusion steps can improve performance

In Table 9 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 k > 0 is important for optimal CRPS and SSR performance. Here, the k auxiliary diffusion steps correspond to the floatingpoint interpolation timesteps { j k +1 } k j =1 , as visualized for a toy example in Fig. 9b. Interestingly, we did not observe this phenomenon for the Navier-Stokes or spring mesh datasets, where using k > 0 did not improve performance over using the simplest schedule S = [ i n ] N -1 n =0 = [ j ] h -1 j =0 , 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 t + h for a given diffusion step n and associated input ˆ x t + i n usually improves as n increases, i.e. as we approach the end of

(a) OISSTv2

(b) Navier-Stokes

(c) Spring Mesh

the reverse process and x t + i n gets closer in time to x t + 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 a 20-member ensemble. We use the validation sets of the physical systems benchmark because the scope of the analysis only extends up to the training horizon (as opposed to the full test set trajectories).

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-time nature of the backbone networks in DYffusion invites the use of arbitrary dynamical timesteps as diffusion states during inference. Thus, to accelerate sampling, we can skip some of the N diffusion steps used for training, S train = { i n } N -1 n =0 , which automatically results in fewer neural network forward passes. For example, we can only use the base schedule S base = { 0 , 1 , . . . , h -1 } , where the diffusion states correspond to the temporal resolution of the dynamical data in a one-to-one mapping (see black arrows in Fig. 3). In Fig. 8, we start with S base as inference schedule (left-most dots in each subplot) for the SST dataset, and then incrementally add more diffusion steps from S train \ S base 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 DYffusion or for downstream applications pose a lower bound (here, S base ) on how much we can accelerate our method since any such output timestep needs to 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 Navier-Stokes and spring mesh DYffusion runs, since their training schedules already correspond to S base .

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 h timesteps x t +1: t + h in a single forward pass. Dropout (continuous) refers to the barebone backbone network that forecasts one timestep x t + k for k ∈ (1 , h ) in a single forward pass, conditioned on the time, k . Both methods perform similarly in 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

Modeling complexity

operates. Video diffusion models have higher modeling complexity because they need to model the full 'videos', x t +1: t + h , 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 h .

D.7 Example Sampling Trajectories Visualized

Example sampling trajectories of our approach, as a function of the schedule i n , are visualized in Fig. 9, where the top one corresponds to the simplest case where we use a one-to-one mapping between diffusion steps, n , and interpolation/dynamical timesteps, i n . Each of the intermediate ˆ x i can be used as a forecast for timestep i . The forecaster network, F θ , repeatedly forecasts x h , but does so with increasing levels of skill (analogously to how conventional diffusion models iteratively denoise/refine the predictions of the 'clean' data, s (0) ; see Fig. 7).

x

0

n

= 0

i

n

= 1

ˆ

x

1

ˆ x

(0)

h

n

= 1

i

n

= 2

ˆ

x

2

ˆ x

(1)

h

n

= 2

i

n

= 3

ˆ

x

3

ˆ x

(2)

h

n

= 3

ˆ x

4

:=

ˆ x

(3)

h

(a) Basic schedule, 1-to-1 diffusion step to dynamical timestep mapping

x

0

n

= 0

i

n

= 0

.

25

ˆ

x

0

.

25

ˆ x

(0)

h

n

= 1

i

n

= 0

.

5

ˆ

x

0

.

5

ˆ x

(1)

h

n

= 2

i

n

= 0

.

75

ˆ

x

0

.

75

ˆ x

(2)

h

n

= 3

i

n

= 1

ˆ

x

1

ˆ x

(3)

h

n

= 4

i

n

= 2

ˆ

x

2

ˆ x

(4)

h

n

= 5

i

n

= 3

ˆ

x

3

ˆ x

(5)

h

n

= 6

ˆ x

4

:=

ˆ x

(6)

h

(b) Auxiliary, additional diffusion steps (here,

k

= 3

) map uniformly to

(0

,

1)

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n1. 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.\n\nWEAKNESSES:\n1. 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.\n\nQUESTIONS:\nWhy does Dropout perform so poorly at $t=2$ in Figure 3? Does this observation indicate the mentioned Weakness 3?\n\nLIMITATIONS:\nThis paper includes a separate section for discussing limitations. There is no potential negative societal impact.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-06-28"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nIn 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.\n\nSTRENGTHS:\n- 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.\n\nWEAKNESSES:\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\nQUESTIONS:\n- 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?\n\nLIMITATIONS:\nN/A", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-04"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n* 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.\n\nWEAKNESSES:\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 (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.\n\nQUESTIONS:\n* 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?\n\nLIMITATIONS:\nThe authors mention that input and output space must be the same and that inference costs are higher than predicting directly.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nIn 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.\n\nSTRENGTHS:\n* Interesting approach\n * Extensive ablation study, with summary in the main paper section\n * Interesting discussion of the results\n\nWEAKNESSES:\n* 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\nQUESTIONS:\nThe 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\nLIMITATIONS:\nAuthors discuss their results and limitations as well as future directions.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-06"}