[Paper Review] DYffusion: A Dynamics-informed Diffusion Model for Spatiotemporal Forecasting

Salva Rühling Cachay, Bo Zhao+2 more|arXiv (Cornell University)|Jun 3, 2023

Hydrological Forecasting Using AI21 citations

TL;DR

DYffusion introduces a dynamics-informed diffusion framework for probabilistic spatiotemporal forecasting by coupling a time-conditioned interpolator with a forecaster, enabling efficient multi-step forecasts and continuous-time sampling in high-dimensional dynamics. It achieves competitive accuracy with improved efficiency on SST, Navier-Stokes flows, and spring-mesh datasets, while reducing memory and step requirements compared to standard diffusion models.

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.

Motivation & Objective

• Motivate probabilistic spatiotemporal forecasting with diffusion models for complex dynamics.
• Propose a dynamics-informed diffusion framework that couples temporal interpolation with forecasting.
• Achieve long-horizon forecasting with improved memory efficiency and sampling flexibility.
• Provide theoretical insight by connecting the model to implicit dynamical system solutions.
• Empirically validate on sea surface temperatures, Navier-Stokes flows, and spring-mesh systems.

Proposed method

• Replace the standard diffusion forward-Denoising with a temporal interpolation forward process via a time-conditioned interpolator I_phi.
• Train a forecaster network F_theta to predict x_{t+h} conditioned on interpolations between x_t and x_{t+h}.
• Two-stage training: (1) train the interpolator to predict intermediate timesteps; (2) freeze I_phi, train the forecaster with inference stochasticity to forecast x_{t+h} for multiple i_n steps.
• Use a diffusion-like schedule where n indexes interpolation timesteps i_n, enabling continuous-time sampling.

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 fo

Experimental results

Research questions

• RQ1 Can a dynamics-informed diffusion process improve probabilistic forecasts for high-dimensional spatiotemporal data over standard Gaussian diffusion?
• RQ2 Does coupling temporal interpolation with forecasting reduce memory footprint and sampling steps while maintaining or improving accuracy?
• RQ3 How does DYffusion perform on complex dynamics datasets (SST, Navier-Stokes, Spring Mesh) in terms of probabilistic metrics and efficiency?
• RQ4 Can continuous-time sampling trajectories be generated effectively and beneficially for long-range forecasting?

Key findings

• DYffusion achieves competitive CRPS and MSE on SST and Navier-Stokes relative to baselines, with faster sampling than MCVD due to fewer diffusion steps.
• DYffusion excels in long-range forecasting and outperforms baselines on long trajectories in Spring Mesh, especially out-of-distribution.
• The method uses a constant memory footprint with horizon h, unlike multi-step prediction baselines that require larger memory for many timesteps.
• Sampling can be accelerated via schedules that skip intermediate diffusion steps, trading accuracy for speed.
• Qualitative results show DYffusion better preserves fine-scale dynamics in Navier-Stokes trajectories than a strong baseline.
• Theoretical view connects DYffusion to an implicit dynamical-system solver, analogous to Euler’s method in cold sampling.

Figure 4 : Qualitative forecasts for timesteps 2, 6, 24, 46, and 64 (last timestep) of the velocity norm of an example Navier-Stokes test trajectory. Here, we generate five sample trajectories for the best baseline (Dropout) and our method DYffusion, both with $h=16$ , and visualize the one with the