HydroGAT: Distributed Heterogeneous Graph Attention Transformer for Spatiotemporal Flood Prediction

2025-09-02

AI Summary

Flood forecasting requires joint modeling of time-varying runoff drivers and spatial interactions across river networks; however, existing methods either neglect topological structure or sacrifice spatial resolution to curb GNN training costs, and commonly decouple spatiotemporal dependencies. This paper proposes the Distributed Spatio-Temporal Graph Attention Network (DSTGAN), which constructs a fine-grained heterogeneous watershed graph integrating surface and river pixels—enabling, for the first time, pixel-level heterogeneous graph modeling. DSTGAN jointly learns local hydrological responses and structured inter-subbasin hydrological propagation by unifying dynamic temporal attention with an interpretable upstream influence discovery mechanism. It incorporates physically informed flow-direction constraints and cross-basin relational modeling, trained via distributed data parallelism. In hourly streamflow forecasting across multiple basins, DSTGAN achieves a Nash–Sutcliffe Efficiency of 0.97, Kling–Gupta Efficiency of 0.96, and bias within ±5%; on 64 A100 GPUs, it attains a 15× speedup over baseline implementations.

Technology Category

• Machine Learning: Graph-based Machine Learning
• Planning, Routing, and Scheduling: Optimization of Spatio-temporal Systems
• Computer Vision: Generative Adversarial Networks (GANs) for Vision

Application Category

• Graph Algorithms and Modeling for the Web: Graph neural networks and deep learning approaches for Web-related graphs
• Systems and Infrastructure for Web, Mobile and WoT: Federated Web and WoT systems, including distributed, federated and edge-based data processing
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs

Abstract

Accurate flood forecasting remains a challenge for water-resource management, as it demands modeling of local, time-varying runoff drivers (e.g., rainfall-induced peaks, baseflow trends) and complex spatial interactions across a river network. Traditional data-driven approaches, such as convolutional networks and sequence-based models, ignore topological information about the region. Graph Neural Networks (GNNs) propagate information exactly along the river network, which is ideal for learning hydrological routing. However, state-of-the-art GNN-based flood prediction models collapse pixels to coarse catchment polygons as the cost of training explodes with graph size and higher resolution. Furthermore, most existing methods treat spatial and temporal dependencies separately, either applying GNNs solely on spatial graphs or transformers purely on temporal sequences, thus failing to simultaneously capture spatiotemporal interactions critical for accurate flood prediction. We introduce a heterogenous basin graph where every land and river pixel is a node connected by physical hydrological flow directions and inter-catchment relationships. We propose HydroGAT, a spatiotemporal network that adaptively learns local temporal importance and the most influential upstream locations. Evaluated in two Midwestern US basins and across five baseline architectures, our model achieves higher NSE (up to 0.97), improved KGE (up to 0.96), and low bias (PBIAS within $pm$5%) in hourly discharge prediction, while offering interpretable attention maps that reveal sparse, structured intercatchment influences. To support high-resolution basin-scale training, we develop a distributed data-parallel pipeline that scales efficiently up to 64 NVIDIA A100 GPUs on NERSC Perlmutter supercomputer, demonstrating up to 15x speedup across machines. Our code is available at https://github.com/swapp-lab/HydroGAT.

Problem

• Modeling local time-varying runoff drivers and complex spatial interactions
• Capturing spatiotemporal dependencies simultaneously for flood prediction
• Overcoming computational costs of high-resolution basin-scale training

Innovation

• Distributed heterogeneous graph attention transformer for spatiotemporal prediction
• Pixel-level nodes connected by hydrological flow directions
• Adaptively learns temporal importance and upstream influences