SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Salva Ruhling Cachay
Bo Zhao
Hailey Joren
{sruhlingcachay, bozhao, hjoren, roseyu}@ucsd.edu
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 mimics 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}$
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 \times 32$  [61]. In addition, few methods apply diffusion models beyond static images. Video
$^{1}$ Code is available at: https://github.com/Rose-STL-Lab/dyffusion
Figure 1: Our proposed framework, DYffusion, reimagines the noise-denoise forward-backward processes of conventional diffusion models as an interplay of temporal interpolation and forecasting. On the top row, we illustrate the direct application of a video diffusion model to dynamics forecasting for a horizon of  $h = 3$ . On the bottom row, DYffusion generates continuous-time probabilistic forecasts for  $\mathbf{x}_{t + 1:t + h}$ , given the initial conditions,  $\mathbf{x}_t$ . During sampling, the reverse process iteratively steps forward in time by forecasting  $\mathbf{x}_{t + h}$  (which plays the role of the "clean data",  $\mathbf{s}^{(0)}$ , in conventional diffusion models) and interpolating to one of  $N$  intermediate timesteps,  $\mathbf{x}_{t + i_n}$ . As a result, our approach operates in the data space at all times and does not need to model high-dimensional videos at each diffusion state.
diffusion models [27, 24, 77, 67] can generate realistic samples but do not explicitly leverage the temporal nature of the data to generate accurate forecasts.
In this work, we introduce a novel framework to train a dynamics-informed diffusion model for multi-step 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.
- 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.
Problem setup. We study the problem of probabilistic spatiotemporal forecasting. We start with a dataset of  $\{\mathbf{x}_t\}_{t=1}^T$  snapshots with  $\mathbf{x}_t \in \mathcal{X}$ . Here,  $\mathcal{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(\mathbf{x}_{t+1:t+h} \mid \mathbf{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(\mathbf{x}_{t+1:t+h} \mid \mathbf{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,  $\mathbf{s}^{(n)}$ , with a superscript  $n$  to clearly distinguish them from the time steps of the data  $\{\mathbf{x}_t\}_{t=1}^T$ . This process can be generalized to consider a degradation operator  $D$  that takes as input the data point  $\mathbf{s}^{(0)}$  and outputs  $\mathbf{s}^{(n)} = D(\mathbf{s}^{(0)}, n)$  for varying degrees of degradation proportional to  $n \in \{1, \dots, N\}$  [2]. Oftentimes,  $D$  adds Gaussian noise with increasing levels of variance so that  $\mathbf{s}^{(N)} \sim \mathcal{N}(\mathbf{0}, \mathbf{I})$ . A denoising network parameterized by  $\theta$ ,  $R_{\theta}$  is trained to restore  $\mathbf{s}^{(0)}$ , i.e., such that  $R_{\theta}(\mathbf{s}^{(n)}, n) \approx \mathbf{s}^{(0)}$ . The diffusion model can be conditioned on the input dynamics by considering  $R_{\theta}(\mathbf{s}^{(n)}, \mathbf{x}_t, n)$ . In the case of dynamics forecasting, the diffusion model can be trained to minimize the objective
$$
\min  _ {\theta} \mathbb {E} _ {n \sim \mathcal {U} [ 1, N ], \mathbf {x} _ {t}, \mathbf {s} ^ {(0)} \sim \mathcal {X}} \left[ | | R _ {\theta} (D (\mathbf {s} ^ {(0)}, n), \mathbf {x} _ {t}, n) - \mathbf {s} ^ {(0)} | | ^ {2} \right], \tag {1}
$$
where  $\mathcal{X}$  is the data distribution,  $||\cdot ||$  is a norm, and  $\mathbf{s}^{(0)} = \mathbf{x}_{t + 1:t + h}$  is the prediction target.  $\mathcal{U}[[a,b]]$  describes the uniform distribution over the set  $\{a,a + 1,\dots ,b\}$ . In practice,  $R_{\theta}$  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,  $\mathbf{x}_{t:t + h}$  (see Fig. 1). Specifically, we replace the degradation operator,  $D$ , with a stochastic interpolator network  $\mathcal{I}_{\phi}$ , and the restoration network,  $R_{\theta}$ , with a deterministic forecaster network,  $F_{\theta}$ .
Algorithm 1 DYffusion, Two-stage Training
Input: networks Fθ, Πφ, norm || · ||, horizon h, schedule [in]N-1i=0
Stage 1: Train interpolator network, Πφ
1. Sample i ~ Uniform ( {1, . . ., h-1} )
2. Sample xt, xt+i, xt+h ~ X
3. Optimize minφ || Πφ (xt, xt+h, i) - xt+i||2
Stage 2: Train forecaster network (diffusion model backbone), Fθ
1. Freeze Πφ and enable inference stochasticity (e.g. dropout)
2. Sample n ~ Uniform ( {0, . . ., N-1} ) and xt, xt+h ~ X
3. Optimize minθ || Fθ( Πφ (xt, xt+h, in), in) - xt+h||2
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 time-conditioned network  $\mathcal{I}_{\phi}$  to interpolate between snapshots of data. Specifically, given a horizon  $h$ , we train  $\mathcal{I}_{\phi}$  so that  $\mathcal{I}_{\phi}(\mathbf{x}_t,\mathbf{x}_{t + h},i)\approx \mathbf{x}_{t + i}$  for  $i\in \{1,\dots ,h - 1\}$  using the objective:
$$
\operatorname * {m i n} _ {\phi} \mathbb {E} _ {i \sim \mathcal {U} [ 1, h - 1 ]}, \mathbf {x} _ {t, t + i, t + h} \sim \chi \left[ | | \mathcal {I} _ {\phi} (\mathbf {x} _ {t}, \mathbf {x} _ {t + h}, i) - \mathbf {x} _ {t + i} | | ^ {2} \right]. \tag {2}
$$
Interpolation is an easier task than forecasting, and we can use the resulting interpolator  $\mathcal{I}_{\phi}$  during inference to interpolate beyond the temporal resolution of the data. That is, the time input can be continuous, with  $i \in (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,  $\mathcal{I}_{\phi}$ , 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_{\theta}$  to forecast  $\mathbf{x}_{t + h}$  such that  $F_{\theta}(\mathcal{I}_{\phi}(\mathbf{x}_t,\mathbf{x}_{t + h},i|\xi),i)\approx \mathbf{x}_{t + h}$  for  $i\in S = [i_n]_{n = 0}^{N - 1}$ , where  $S$  denotes a schedule mapping the diffusion step to the interpolation timestep. The interpolator network,  $\mathcal{I}_{\phi}$ , is frozen with inference stochasticity enabled, represented by the random variable  $\xi$ . Here,  $\xi$  stands for the randomly dropped out weights of the neural network and is omitted henceforth for clarity. Specifically, we seek to optimize the objective
$$
\min  _ {\theta} \mathbb {E} _ {n \sim \mathcal {U} [ [ 0, N - 1 ] ], \mathbf {x} _ {t, t + h} \sim \mathcal {X}} \left[ | | F _ {\theta} \left(\mathcal {I} _ {\phi} (\mathbf {x} _ {t}, \mathbf {x} _ {t + h}, i _ {n} | \xi), i _ {n}\right) - \mathbf {x} _ {t + h} \| ^ {2} \right]. \tag {3}
$$
To include the setting where  $F_{\theta}$  learns to forecast the initial conditions, we define  $i_0 \coloneqq 0$  and  $\mathcal{I}_{\phi}(\mathbf{x}_t,\cdot ,i_0) \coloneqq \mathbf{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]_{j = 0}^{h - 1}$ . Generally, the interpolation timesteps should satisfy  $0 = i_0 < i_n < i_m < h$  for  $0 < n < m \leq 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_{\theta}$  for unseen timesteps (see Fig. 3). The resulting two-stage training algorithm is shown in Alg. 1.
Because the interpolator  $\mathcal{I}_{\phi}$  is frozen in the second stage, the imperfect forecasts  $\hat{\mathbf{x}}_{t + h} = F_{\theta}(\mathcal{I}_{\phi}(\mathbf{x}_t,\mathbf{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_{\theta}(\mathcal{I}_{\phi}(\mathbf{x}_t,\hat{\mathbf{x}}_{t + h},i_{n + 1}),i_{n + 1}) - \mathbf{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  $\mathbf{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:
$$
p _ {\theta} \left(\mathbf {s} ^ {(n + 1)} \mid \mathbf {s} ^ {(n)}, \mathbf {x} _ {t}\right) = \left\{ \begin{array}{l l} F _ {\theta} \left(\mathbf {s} ^ {(n)}, i _ {n}\right) & \text {i f} n = N - 1 \\ \mathcal {I} _ {\phi} \left(\mathbf {x} _ {t}, F _ {\theta} \left(\mathbf {s} ^ {(n)}, i _ {n}\right), i _ {n + 1}\right) & \text {o t h e r w i s e ,} \end{array} \right. \tag {4}
$$
where  $\mathbf{s}^{(0)} = \mathbf{x}_t$  and  $\mathbf{s}^{(n)}\approx \mathbf{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.  $\mathbf{x}_t$ ), while  $n = N$  refers to the final output of the reverse process (here,  $\mathbf{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
Algorithm 2 Adapted Cold Sampling [2] for DYffusion
1:	Input: Initial conditions x̂t := xt, schedule [in]N-1i=0, output timesteps J (by default J = {1,...,h-1})
2:	for n = 0,1,...,N-1 do
3:	x̂t+h ← Fθ(x̂t+in, i_n)
4:	x̂t+in+1 = Iφ(x_t, x̂t+h, in+1) - Iφ(x_t, x̂t+h, in) + x̂t+in
5:	end for
6:	x̂t+j ← Iφ(x_t, x̂t+h, j), ∀j ∈ J # Optional refinement
7:	Return: {x̂t+j | j ∈ J} ∪ {x̂t+h}
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_{\theta}$  always predicts the last timestep,  $\mathbf{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  $\mathbf{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.
Figure 2: During sampling, DYffusion alternates between forecasting and interpolation, following Alg. 2. In this example, the sampling trajectory follows a simple schedule of going through all integer timesteps that precede the horizon of  $h = 4$ , with the number of diffusion steps  $N = h$ . The output of the last diffusion step is used as the final forecast for  $\hat{\mathbf{x}}_4$ . The black lines represent forecasts by the forecaster network,  $F_{\theta}$ . The first forecast is based on the initial conditions,  $\mathbf{x}_0$ . The blue lines represent the subsequent temporal interpolations performed by the interpolator network,  $\mathcal{I}_{\phi}$ .
Memory footprint. DYffusion requires only  $\mathbf{x}_t$  and  $\mathbf{x}_{t + h}$  (plus  $\mathbf{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  $\mathbf{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  $\mathbf{x}(t)$ . The neural networks  $\mathcal{I}_{\phi}$  and  $F_{\theta}$  are trained to construct a differential equation that implicitly models the solution to  $\mathbf{x}(t)$ .
Let  $s$  be the time variable. DYffusion models the dynamics as
$$
\frac {d \mathbf {x} (s)}{d s} = \frac {d \mathcal {I} _ {\phi} (\mathbf {x} _ {t} , F _ {\theta} (\mathbf {x} , s) , s)}{d s}. \tag {5}
$$
The initial condition is given by  $\mathbf{x}(t) = \mathbf{x}_t$ . During prediction, we seek to obtain
$$
\mathbf {x} (s) = \mathbf {x} (t) + \int_ {t} ^ {s} \frac {d \mathcal {I} _ {\phi} (\mathbf {x} _ {t} , F _ {\theta} (\mathbf {x} , s) , s)}{d s} d s \quad \text {f o r} s \in (t, t + h ]. \tag {6}
$$
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  $\mathbf{x}_t$  to the future step  $\mathbf{x}_{t + h}$ .
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 next-step prediction models [58, 18, 29]. Some notable works target precipitation nowcasting [52], or short-term forecasting [15]. Alternatively to using intrinsically probabilistic models, one can use a deterministic model 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.