# HydroGAT: Distributed Heterogeneous Graph Attention Transformer for Spatiotemporal Flood Prediction

arXiv: https://arxiv.org/abs/2509.02481 · HTML: https://mlsanity.com/paper/2509.02481/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Builds on 15 of its 49 references that are in the map.
- First posted on arXiv 2 Sep 2025.

## Builds on (most cited first, 15 of 15)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Graph Attention Networks](https://mlsanity.com/paper/1710.10903/) (2017)
- [Empirical Evaluation of Gated Recurrent Neural Networks on Sequence Modeling](https://mlsanity.com/paper/1412.3555/) (2014)
- [Convolutional Neural Networks on Graphs with Fast Localized Spectral Filtering](https://mlsanity.com/paper/1606.09375/) (2016)
- [Spatio-Temporal Graph Convolutional Networks: A Deep Learning Framework for Traffic Forecasting](https://mlsanity.com/paper/1709.04875/) (2017)
- [Graph WaveNet for Deep Spatial-Temporal Graph Modeling](https://mlsanity.com/paper/1906.00121/) (2019)
- [Spatio-Temporal Graph Transformer Networks for Pedestrian Trajectory Prediction](https://mlsanity.com/paper/2005.08514/) (2020)
- [A Comprehensive Review of Deep Learning Applications in Hydrology and Water Resources](https://mlsanity.com/paper/2007.12269/) (2020)
- [Sequential Aggregation and Rematerialization: Distributed Full-batch Training of Graph Neural Networks on Large Graphs](https://mlsanity.com/paper/2111.06483/) (2021)
- [Short-term Hourly Streamflow Prediction with Graph Convolutional GRU Networks](https://mlsanity.com/paper/2107.07039/) (2021)
- [Comprehensive Evaluation of GNN Training Systems: A Data Management Perspective](https://mlsanity.com/paper/2311.13279/) (2023)
- [Gated Graph Convolutional Recurrent Neural Networks](https://mlsanity.com/paper/1903.01888/) (2019)
- [Time Series Predictions in Unmonitored Sites: A Survey of Machine Learning Techniques in Water Resources](https://mlsanity.com/paper/2308.09766/) (2023)
- [High-resolution rainfall-runoff modeling using graph neural network](https://mlsanity.com/paper/2110.10833/) (2021)
- [Probabilistic modeling of lake surface water temperature using a Bayesian spatio-temporal graph convolutional neural network](https://mlsanity.com/paper/2109.13235/) (2021)

## Built on by (most cited first, 1 of 1)

- [Physics-Refined Spatiotemporal Forecasting on Open-Boundary Hydrologic Graphs](https://mlsanity.com/paper/2610.01765/) (2026)

## Abstract (arXiv)

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 ±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.