Source-linked AI summary

DYffusion: A Dynamics-informed Diffusion Model for Spatiotemporal Forecasting

Salva Rühling Cachay, Bo Zhao, Hailey Joren, Rose Yu

arXiv:2306.01984v2cs.LGcs.AIstat.ML

TL;DR

Probabilistic spatiotemporal forecasting needs accurate long-range predictions, but standard diffusion models do not explicitly exploit temporal dynamics. DYffusion couples diffusion steps to those dynamics through stochastic interpolation and forecasting, achieving competitive forecasts with improved efficiency while retaining limitations in output-space compatibility and inference speed.

  • Problem

    Accurate long-range probabilistic dynamics forecasts are difficult, and standard diffusion models do not explicitly leverage temporal structure for forecasting.

  • Method

    DYffusion trains a stochastic time-conditioned interpolator and a forecaster that replace standard diffusion’s degradation and restoration while coupling diffusion steps to dynamical time.

  • Results

    DYffusion provides strong probabilistic forecasts and improves computational efficiency over standard Gaussian diffusion, performing best on Navier-Stokes and close second to MCVD on SST CRPS.

  • Takeaways & Limitations

    The dynamics-informed process supports multi-step and long-range probabilistic forecasting with flexible sampling schedules and lower memory needs.

  • Takeaways & Limitations

    DYffusion currently does not support output spaces different from input spaces, and sequential diffusion sampling can be slower than single-pass forecasting.

Abstract

from arXiv · show

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 Introduction

Accurate long-range probabilistic dynamics forecasts are important but difficult, while standard diffusion models do not explicitly exploit temporal structure. DYffusion couples diffusion steps to temporal dynamics to improve forecasting efficiency and supports strong probabilistic forecasts.

  • Accurate long-range probabilistic forecasts matter for policy, risk management, resource optimization, and strategic planning, but operational numerical simulations can require supercomputers and reduced spatial resolution.
  • Standard diffusion models can generate realistic samples but do not explicitly leverage temporal data to produce accurate dynamics forecasts.
  • DYffusion replaces conventional degradation and restoration with temporal interpolation and forecasting, coupling diffusion steps to dynamical time steps without physical-system assumptions.
  • DYffusion targets multi-step and long-range forecasting while reducing diffusion-model memory, data, and training-step requirements through a temporal inductive bias.
  • The empirical study finds strong probabilistic forecasts and improved computational efficiency over standard Gaussian diffusion.

2 Background

The paper formulates probabilistic spatiotemporal forecasting as learning a conditional distribution over a future horizon from observed snapshots. Standard diffusion progressively adds noise and then learns to reverse that degradation to generate data.

  • The forecasting task is to learn P(x_t+1:t+h | x_t-l+1:t), predicting h future snapshots from l past snapshots, with the paper focusing on a single initial condition.
  • The data space may include spatial dimensions such as latitude and longitude together with channels such as velocity, temperature, or humidity.
  • In standard diffusion, a forward process progressively degrades data with Gaussian noise, while a reverse process removes noise to generate samples.
  • A denoising network is trained to restore the original data from a degraded state, with generalized objectives covering standard diffusion training.

3 DYffusion: DYnamics-Informed Diffusion Model

DYffusion replaces Gaussian degradation and denoising with temporal interpolation and forecasting, so diffusion states correspond to dynamical times and sampling proceeds in observation space. Its schedules support continuous-time forecasts, flexible inference, and reduced memory demands.

  • DYffusion: DYnamics-Informed Diffusion Model: DYffusion replaces the degradation operator with a stochastic interpolator and the restoration network with a deterministic forecaster for spatiotemporal sequences.
  • DYffusion: DYnamics-Informed Diffusion Model: DYffusion trains in two stages: it first learns the interpolator, then freezes it while training a forecaster across a configurable diffusion-dynamics schedule.
  • DYffusion: DYnamics-Informed Diffusion Model: The interpolator learns time-conditioned intermediate states, supports continuous time inputs, and uses Monte Carlo dropout to produce stochastic forecasts.
  • DYffusion: DYnamics-Informed Diffusion Model: During sampling, the method alternates forecasting and interpolation, beginning from initial conditions and progressing forward in time until the final horizon forecast.
  • DYffusion: DYnamics-Informed Diffusion Model: DYffusion can forecast beyond the training horizon autoregressively, while its sampling schedule can be changed at inference time because schedules discretize the same learned ODE.
  • DYffusion: DYnamics-Informed Diffusion Model: Its memory footprint remains constant in the forecasting horizon because training requires only endpoint snapshots, unlike methods that store all h + 1 timesteps.

4 Related Work

Prior work applies diffusion models mainly to images, text, and video, while high-dimensional spatiotemporal forecasting remains less explored. DYffusion differs by coupling diffusion schedules to temporal dynamics and targeting full physical-system evolution.

  • Diffusion models have succeeded in images and text, but high-dimensional spatiotemporal forecasting has remained comparatively unexplored.
  • Non-Gaussian diffusion approaches have extended diffusion beyond Gaussian noise, but reported results for cold diffusion and related corruptions are usually inferior to Gaussian diffusion.
  • Video diffusion models are closely related, but this paper focuses on full dynamics and underlying physics rather than object-centric prediction, and vision techniques may not transfer well.
  • Autoregressive dynamics models can suffer unstable multi-step rollouts from compounding errors, while multi-step training reduces this issue at substantial memory and efficiency costs.
  • DYffusion couples diffusion steps to continuous dynamical times, extending schedules beyond the discrete temporal resolution used by naive approaches.

5 Experiments

DYffusion is evaluated for probabilistic multi-step and long-range forecasting across sea-surface temperatures, Navier–Stokes flows, and spring-mesh systems, with comparisons spanning diffusion and ensemble baselines. It achieves strong forecasting quality, efficient sampling, stable long rollouts, and improved visual fidelity, while retaining flexible sampling behavior.

  • Experimental Focus: Long-horizon evaluation shows DYffusion and the baselines excel in the long-range regime while remaining competitive with deterministic short-range forecasting methods.The study evaluates probabilistic multi-step rollouts rather than optimizing only next-step prediction because deterministic single-step models can become unstable over long forecasts.
  • Quantitative Results: DYffusion performs best on Navier–Stokes CRPS, ranks close second on SST after MCVD, and beats the barebone baseline on long-horizon Spring Mesh forecasting.The evaluation covers full trajectories of 64 Navier–Stokes timesteps and 804 Spring Mesh steps; Table 1 reports SST and Navier–Stokes comparisons, while Table 2 covers Spring Mesh.
  • Quantitative Results: DYffusion samples with fewer than 50 diffusion steps, whereas MCVD uses 1000, creating a documented inference-speed advantage over that baseline.DDPM is efficient on SST with five steps but trails in performance, illustrating the reported accuracy–efficiency trade-off.
  • Qualitative Results: DYffusion reproduces Navier–Stokes fine-scale flow patterns better than the Dropout baseline across qualitative forecasts from early to final timesteps.The comparison uses five sampled trajectories per method and visualizes the sample with the best trajectory-average MSE.
  • Sampling Analysis: DYffusion supports accelerated sampling by skipping intermediate diffusion steps, with a reported trade-off between forecast accuracy and sampling speed.Performance generally improves as the reverse process approaches the target timestep, and additional interpolation steps can improve results.
  • Data Ablation: Forecast quality degrades more gracefully with limited SST training data, especially when training uses only two years or approximately 8,000 data points.The authors compare CRPS across SST training-set sizes in Table 3.

6 Conclusion

DYffusion couples diffusion processes to dynamical evolution to support long-range probabilistic forecasting, while retaining a broader vision for dynamics-informed conditional generation. Its current limitations include mismatched input-output spaces and slower inference than single-pass methods.

  • 6 Conclusion: DYffusion couples diffusion steps with the dynamical nature of data to support long-range forecasts without relying solely on autoregressive iteration.The authors present this as a dynamics-informed alternative for probabilistic spatiotemporal forecasting.
  • 6 Conclusion: DYffusion currently does not support forecasting problems whose output space differs from the input space.
  • 6 Conclusion: Single-forward-pass methods remain faster than sequential diffusion sampling when inference speed is prioritized.Additional interpolator passes reduce some of DYffusion’s efficiency advantage despite requiring fewer diffusion steps than Gaussian diffusion.

B Methodology

DYffusion replaces Gaussian degradation and denoising with a stochastic temporal interpolator and a forecaster that operate along the system’s dynamics. Its training uses two stages, optional conditioning, and a look-ahead objective that partially mimics iterative sampling.

  • B Methodology: A one-step look-ahead loss simulates part of sampling and backpropagates through it, with λ1 = λ2 = 1 in all experiments.
  • B Methodology: The forecaster can receive no conditioning, clean initial conditions, or noised initial conditions as optional inputs.The noised variant interpolates between the initial state and Gaussian noise according to the diffusion-step index.
  • B Methodology: Cold sampling provides the generalized reverse-process procedure used to sample from DYffusion.
  • B Methodology: DYffusion’s forecasting objective is equivalent in structure to conventional denoising when diffusion indexing is reversed to align with temporal indexing.Its intermediate states approximate future system states rather than progressively noisier versions of a target.

C.1 Cold Sampling from the Euler method

Cold sampling can be derived as an Euler-style discretization of the continuous-time DYffusion process. Under Lipschitz assumptions, its one-step error is first-order in the step size, unlike naive sampling’s error guarantee.

  • C.1 Cold Sampling from the Euler method: Cold sampling is exactly the update obtained after replacing an unknown future forecaster evaluation with the current evaluation in the Euler discretization.This derivation provides a theoretical explanation for the algorithm.
  • C.1 Cold Sampling from the Euler method: Sampling more densely near the early part of the prediction window can reduce approximation error because early states require farther-ahead predictions.
  • C.1 Cold Sampling from the Euler method: Cold sampling differs from naive sampling by including error-correction terms, and experiments report a large margin in its favor, especially for SST with auxiliary diffusion steps.
  • C.1 Cold Sampling from the Euler method: The cold-sampling discretization error is bounded by O(∆s) when the forecaster and interpolator satisfy the stated Lipschitz assumptions.The bound follows from applying the Lipschitz conditions to the forecaster and interpolator composition.
  • C.1 Cold Sampling from the Euler method: Naive sampling lacks a guaranteed first-order error bound because its additional terms are not bounded by first-order quantities in ∆s.The resulting discretization errors can therefore be larger than those of cold sampling.

D.1.1 SST data preprocessing

The SST benchmark is constructed from daily NOAA OI SST V2 data using fixed temporal splits, standardized preprocessing, and a focused subset of ocean grid tiles. The resulting dataset emphasizes mostly the eastern tropical Pacific.

  • D.1.1 SST data preprocessing: The dataset uses NOAA OI SST V2 daily observations from 1982–2018 for training, 2019 for validation, and 2020 for testing.The source data have 1/4° spatial resolution.
  • D.1.1 SST data preprocessing: The globe is divided into 60 × 60 latitude × longitude tiles, after which tiles with less than 95% ocean cover are removed.
  • D.1.1 SST data preprocessing: SST values are standardized using daily training-set means and standard deviations, and continental NaNs are replaced with zero after standardization.
  • D.1.1 SST data preprocessing: The experiments subsample 11 grid tiles covering mostly the eastern tropical Pacific, corresponding to the red-bounded subset in Figure 5.
  • D.1.1 SST data preprocessing: The dataset-specific hyperparameters are reported in Table 5, including the number of artificial diffusion steps used by DYffusion.

D.2 Implementation Details

The experiments use dataset-specific diffusion and forecasting configurations, with ensemble evaluation based on CRPS, MSE, and SSR. DYffusion uses auxiliary diffusion steps for SST, while stochastic interpolator dropout and cold sampling are important design choices.

  • Baselines: The baselines use dataset-specific configurations, including five DDPM steps for SST, 500 for Navier-Stokes, and 1000 MCVD steps across datasets.MCVD is reported only once in the main SST results because of its inference runtime.
  • DYffusion: DYffusion uses 25 artificial diffusion steps for SST but none for Navier-Stokes or spring mesh.The refinement step is omitted for SST but retained for Navier-Stokes and spring mesh based on its observed effect.
  • DYffusion: The interpolator dropout rates are selected by validation CRPS, and stochasticity in the interpolator is described as crucial for DYffusion’s performance.Selected rates are dataset-specific: 0.6/0/0.6 for SST, 0.15 for Navier-Stokes, and 0.05 for spring mesh.
  • Evaluation: All reported results use 50-member ensembles to compute CRPS, MSE on the ensemble mean, and SSR, with experiments generally repeated over three seeds when feasible.SSR is the square root of ensemble variance divided by ensemble RMSE and measures ensemble reliability.

D.3 Neural Architecture Details

The architecture details reuse established dataset-specific backbones while adding time conditioning throughout the models. A sine/cosine time embedding supplies scale-and-shift conditioning for diffusion or dynamical timesteps.

  • SST UNet: The SST UNet has three downsampling and three upsampling blocks with residual blocks, attention layers, and spatial resolutions changing from 60 × 60 to 15 × 15.Its channel widths increase from 64 to 256 during downsampling and reverse during upsampling.
  • Benchmark comparison: Figure 6 plots MSE against rollout step and compares DYffusion and Dropout-multi-step with single-step deterministic UNet/CNN baselines.The figure covers the Navier-Stokes and spring mesh datasets.
  • Backbone architectures: All models reuse the same dataset-specific backbone architecture, with DYffusion’s interpolator and forecaster using the corresponding baseline backbone.The Navier-Stokes and spring mesh architectures come from the benchmark, while the SST model uses a diffusion-style UNet.
  • Time embedding: A sine/cosine-based time embedding enables conditioning on diffusion steps or dynamical timesteps across the architectures.The embedding is projected through linear layers and used to scale and shift convolutional or residual-block inputs.
  • OOD evaluation: The Navier-Stokes OOD results differ only marginally from the main test results.Table 6 presents the OOD evaluation analogously to the spring mesh results in Table 2.

D.4.2 Benchmark against deterministic models

Against deterministic single-step benchmarks, DYffusion and the multi-step Dropout baseline remain more stable over long rollouts. The ablations show that stochastic interpolator dropout and cold sampling are important for probabilistic forecast quality.

  • Benchmark comparison: DYffusion and multi-step Dropout outperform the deterministic benchmark models, especially in the long-range forecasting regime.The benchmark models degrade or diverge during Navier-Stokes and spring mesh rollouts, whereas the proposed multi-step models are trained for multiple timesteps.
  • Benchmark comparison: The deterministic benchmark models diverge or degrade during inference rollouts because they learn only next-step prediction while evaluation requires long autoregressive trajectories.The benchmark models are architecturally similar, so the training-objective mismatch is identified as a likely explanation.
  • Ablations: Cold sampling yields much better DYffusion samples than naive sampling, consistent with the role of error-correction terms in generalized diffusion sampling.The comparison is reported in Table 7 and is especially pronounced for SST.
  • Ablations: Disabling inference dropout in the interpolator substantially worsens scores by removing forecast spread and making the forecasts deterministic.Under this setting, CRPS reduces to mean absolute error and SSR becomes 0.
  • Ablations: The refinement step has dataset-dependent effects, improving spring mesh substantially while making little difference on SST.The authors recommend disabling refinement in practice to accelerate inference, except where its benefit is needed.

D.5.5 Choosing the training horizon

The study examines how training horizon, auxiliary diffusion steps, inference schedules, and model requirements affect forecasting. Multi-step horizons outperform next-step training, while continuous-time schedules permit speed–performance trade-offs and flexible sampling.

  • Training horizon: Every tested Navier-Stokes training horizon outperforms the best baseline, while the next-step model performs significantly worse than multi-step approaches.For DYffusion, h = 16 is identified as a performance sweet spot among h ∈ {8, 16, 32}.
  • Auxiliary diffusion steps: For SST, adding auxiliary diffusion steps improves SSR steadily, although CRPS changes little once k > 0.The extra steps correspond to floating-point interpolation times, whereas Navier-Stokes and spring mesh show no improvement from k > 0.
  • Reverse-process behavior: Forecaster predictions usually improve in CRPS and SSR as reverse-process steps approach the target time, while MSE can worsen for SST and spring mesh.The analysis uses a 20-member ensemble on SST and validation sets for the physical-system benchmarks.
  • Inference schedules: DYffusion can accelerate inference by skipping trained diffusion steps, automatically reducing neural-network forward passes.The continuous-time backbone permits arbitrary dynamical timesteps during inference, including the base schedule aligned with data resolution.
  • Computational requirements: Video diffusion models require modeling the full forecast sequence at every diffusion state, whereas DYffusion is only slightly affected by the forecast horizon.Table 10 compares the modeling and computing requirements of the methods.
  • Sampling trajectories: Sampling trajectories map diffusion steps to interpolation timesteps, with intermediate interpolator outputs serving as forecasts for corresponding future times.The forecaster repeatedly predicts the terminal horizon while increasing interpolation time improves forecast skill.
Loading 2306.01984v2…