HydroDiffusion: Diffusion-Based Probabilistic Streamflow Forecasting with a State Space Backbone

Reading time: 5 minute
...

📝 Original Info

  • Title: HydroDiffusion: Diffusion-Based Probabilistic Streamflow Forecasting with a State Space Backbone
  • ArXiv ID: 2512.12183
  • Date: 2025-12-13
  • Authors: Yihan Wang, Annan Yu, Lujun Zhang, Charuleka Varadharajan, N. Benjamin Erichson

📝 Abstract

Recent advances have introduced diffusion models for probabilistic streamflow forecasting, demonstrating strong early flood-warning skill. However, current implementations rely on recurrent Long Short-Term Memory (LSTM) backbones and single-step training objectives, which limit their ability to capture long-range dependencies and produce coherent forecast trajectories across lead times. To address these limitations, we developed HydroDiffusion, a diffusion-based probabilistic forecasting framework with a decoder-only state space model backbone. The proposed framework jointly denoises full multi-day trajectories in a single pass, ensuring temporal coherence and mitigating error accumulation common in autoregressive prediction. HydroDiffusion is evaluated across 531 watersheds in the contiguous United States (CONUS) in the CAMELS dataset. We benchmark HydroDiffusion against two diffusion baselines with LSTM backbones, as well as the recently proposed Diffusion-based Runoff Model (DRUM). Results show that HydroDiffusion achieves strong nowcast accuracy when driven by observed meteorological forcings, and maintains consistent performance across the full simulation horizon. Moreover, HydroDiffusion delivers stronger deterministic and probabilistic forecast skill than DRUM in operational forecasting. These results establish HydroDiffusion as a robust generative modeling framework for medium-range streamflow forecasting, providing both a new modeling benchmark and a foundation for future research on probabilistic hydrologic prediction at continental scales.

💡 Deep Analysis

📄 Full Content

Deep learning has been widely used in hydrology, especially for streamflow prediction. Deep neural networks offer data-driven methods that accurately and efficiently model complex hydrological responses, and have achieved strong performance across diverse spatial and temporal prediction tasks [Kratzert et al., 2019a,b, Wang et al., 2025b,c, Liu et al., 2024, Nearing et al., 2024, Willard et al., 2025]. However, most deep learning models for hydrology are deterministic, producing single point estimates, without quantifying uncertainty. Deterministic predictions are insufficient for operational predictions that must account for risk under different scenarios, such as for flood forecasting and early warning systems [Krzysztofowicz, 2001, Pechlivanidis et al., 2025]. While ensembles of deep learning models can enable probabilistic forecasts, they are computationally intensive, requiring training of models with many different configurations [Willard and Varadharajan, 2025].

Diffusion-based generative models [Ho et al., 2020, Song et al., 2020a] are a promising and efficient alternative approach for probabilistic forecasting. These models learn a generative process that produces forecasts by transforming random noise into samples of the conditional distribution of future streamflow. This generative process is parameterized by a neural network backbone, whose architecture and design determine a diffusion model’s capacity to approximate the conditional distribution and thus influence its overall forecasting performance. Recently, Ou et al. [2025] introduced the Diffusion-based Runoff Model (DRUM), the first diffusion model for probabilistic hydrological forecasting. DRUM combines a denoising diffusion process with an encoder-decoder Long Short-Term Memory (LSTM) backbone [Hochreiter and Schmidhuber, 1997] to produce medium-range (up to seven-day) ensemble flood forecasts across the contiguous United States (CONUS). DRUM demonstrated strong early warning skills for extreme events. Building on this foundation, Yang et al. [2025] proposed a multi-time scale LSTM diffusion model for hourly prediction (denoted h-Diffusion), showing that diffusion approaches can also extend to fine temporal resolutions and data assimilation settings.

These studies highlight the promise of diffusion-based forecasting for hydrological applications. However, there is an opportunity to improve time series forecasts by adopting design choices from recent advances in diffusion modeling. For example, DRUM and h-Diffusion rely on recurrent LSTM-based backbones that model sequences in an autoregressive manner. Therefore, they predict one step at a time and roll forward iteratively, which can lead to error accumulation and loss of temporal coherence over time. Moreover, DRUM uses an encoder-decoder architecture, which can introduce an information bottleneck because the encoder compresses long-range hydrologic information into a fixed representation, potentially discarding details needed for accurate forecasts [Bahdanau et al., 2014]. Finally, both models are trained using the discrete denoising diffusion probabilistic model (DDPM) formulation [Ho et al., 2020], where the network learns to predict noise only at fixed diffusion steps rather than along a continuous diffusion trajectory. This discretization constrains temporal continuity in the learned generative dynamics as the model is tied to a predetermined noise schedule.

In this work, we present HydroDiffusion, a diffusion-based probabilistic streamflow forecasting framework that leverages several recent key innovations. Methodologically, we move from next-step prediction to full-sequence denoising, where the model learns the joint conditional distribution of streamflow trajectories over the entire forecast horizon. This formulation, now widely used in time-series diffusion models, offers advantages over autoregressive forecasting by mitigating error accumulation and enforcing temporal consistency across lead times [Chen et al., 2024, Yuan and Qiao, 2024, Salinas et al., 2020]. Architecturally, we replace the LSTM backbone with a state space model (SSM) [Gu et al., 2022b]. SSMs have recently been shown to capture long-range hydrologic dependencies more effectively and efficiently than LSTMs, and to achieve better performance in rainfall-runoff simulation [Wang et al., 2025c]. Furthermore, we adopt a decoder-only configuration, which avoids the information bottleneck and directly generates the full forecast sequence conditioned on past and future meteorological forcings as well as static catchment attributes. In terms of the diffusion formulation, we adopt the score-based diffusion formulation [Song et al., 2020a] with a velocity parameterization [Karras et al., 2022], which provides a continuous-time perspective on diffusion. This formulation is known to improve both the training stability and the quality of the generated trajectories.

Our experiments target medium-range streamflow forecasti

Reference

This content is AI-processed based on open access ArXiv data.

Start searching

Enter keywords to search articles

↑↓
ESC
⌘K Shortcut