HydroGAT: Graph Attention for Flood Prediction

• HydroGAT is a distributed heterogeneous graph attention transformer that represents drainage basins as high-resolution graphs encoding both land and river pixels for spatiotemporal flood prediction.
• It combines a temporal transformer encoder with graph-attentional GRU-based spatial propagation to capture short-term event responses and long-horizon hydrologic memory.
• Empirical results on Midwestern U.S. basins show HydroGAT achieves high performance (NSE up to 0.97, KGE up to 0.96) while efficiently scaling via distributed training.

HydroGAT is a distributed heterogeneous graph attention transformer for spatiotemporal flood prediction that represents a drainage basin as a high-resolution graph whose nodes are land and river pixels and whose edges encode hydrological flow directions and inter-catchment relationships. In its primary formulation, HydroGAT predicts future river discharge at gauge stations from the past $T$ hours of meteorological and hydrologic variables together with a basin graph $\mathcal{G}=(\mathcal{V},\mathcal{E},\mathbf{X})$, via

$$f(\mathbf X^{t}, \mathcal G;\theta) \to \hat{\mathbf Y}^{\,t+\Delta t}_{\mathcal V_\rho}\in\mathbb R^{|\mathcal V_\rho|\times\Delta t},$$

where $\mathcal V_\rho\subseteq\mathcal V$ denotes gauge nodes (Sarkar et al., 2 Sep 2025). More broadly, the term has been used in adjacent work to denote hydrology-focused graph-attention or graph-transformer methods and, in a more abstract systems-engineering sense, graph-based hydrologic representations that preserve modularity, topology, and heterogeneous process structure (Vemula et al., 31 Mar 2025, Harris et al., 31 May 2025).

1. Problem setting and conceptual scope

HydroGAT addresses hourly flood prediction in settings where local, time-varying runoff drivers and complex spatial interactions across a river network must be modeled simultaneously. The motivating difficulties are explicitly hydrologic: rainfall intensity, antecedent soil moisture, snowmelt, baseflow, and evapotranspiration vary strongly in time; water moves along a directed river network and hillslopes; land-use changes, seasonal crop phenology, winter freeze–thaw, and complex terrain introduce heterogeneity and non-stationarity; and pixel-scale modeling greatly increases graph size relative to catchment-based abstractions (Sarkar et al., 2 Sep 2025).

The model is positioned against three limitations in prior data-driven flood modeling. First, convolutional and sequence models typically treat catchments as independent time series and ignore explicit river-network topology. Second, many GNN-based flood models use only gauges, reaches, or catchment outlets as nodes and therefore encode land processes only implicitly or ignore them altogether. Third, transformer-style global attention on large graphs is computationally prohibitive, especially when thousands of pixels and multiday sequences are retained at full spatial resolution. HydroGAT therefore combines a temporal transformer encoder with graph-attentional spatial propagation on a heterogeneous basin graph, while using distributed data-parallel training to avoid the coarse polygon aggregation that earlier graph models often required (Sarkar et al., 2 Sep 2025).

A common misconception is that HydroGAT is simply a graph neural network over river reaches. In the reported formulation, every land and river pixel is a node, and heterogeneity arises primarily through edge types and node roles rather than through a reach-only river graph. Another misconception is that the model is a fully global transformer over all nodes and all times. The architecture instead applies temporal self-attention per node along time and restricts spatial propagation to graph neighborhoods defined by hydrologically constructed edge sets (Sarkar et al., 2 Sep 2025).

2. Basin graph construction and data representation

HydroGAT represents a basin as a heterogeneous directed graph with three edge categories. The first is a set of flow-direction edges $\mathcal E_F$, derived from a DEM with the D8 algorithm, where each pixel points to its steepest downhill neighbor. The second is a set of catchment edges $\mathcal E_C$, defined only between target nodes and intended to encode upstream–downstream relationships among gauge stations. The third is a self-loop $(v,v)$ for every node, ensuring retention of local hidden state during updates. Nodes are 4 km $\times$ 4 km raster cells and include both land pixels and river pixels, while gauge stations form the target-node subset $\mathcal V_\rho$ (Sarkar et al., 2 Sep 2025).

Each node carries a multivariate time series $\mathbf X_v \in \mathbb R^{T\times F}$. In the reported experiments, precipitation is available at all nodes, whereas discharge is available only at target nodes and is used both as an input history and as the prediction target. Static information such as terrain routing is encoded primarily through graph structure rather than through a long list of static node features. Missing discharge values are handled by linear interpolation for downstream stations with sparse missing data and by linear regression against downstream stations for upstream sites,

$\mathcal{G}=(\mathcal{V},\mathcal{E},\mathbf{X})$0

with $\mathcal{G}=(\mathcal{V},\mathcal{E},\mathbf{X})$1 and $\mathcal{G}=(\mathcal{V},\mathcal{E},\mathbf{X})$2 estimated from periods with concurrent observations. Both precipitation and discharge are transformed by

$\mathcal{G}=(\mathcal{V},\mathcal{E},\mathbf{X})$3

followed by min–max scaling to $\mathcal{G}=(\mathcal{V},\mathcal{E},\mathbf{X})$4, a preprocessing step intended to reduce skewness and stabilize training (Sarkar et al., 2 Sep 2025).

The published experiments use two Midwestern U.S. basins and preserve full pixel-level connectivity rather than collapsing the domain to coarse hydrologic units.

Basin	Graph statistics	Reported HydroGAT performance
Cedar River Basin	1,288 nodes; 1,247 flow edges; 17 catchment edges; 18 gauges	NSE $\mathcal{G}=(\mathcal{V},\mathcal{E},\mathbf{X})$5, KGE $\mathcal{G}=(\mathcal{V},\mathcal{E},\mathbf{X})$6, PBIAS $\mathcal{G}=(\mathcal{V},\mathcal{E},\mathbf{X})$7
Des Moines River Basin	2,226 nodes; 2,157 flow edges; 32 catchment edges; 33 gauges	NSE $\mathcal{G}=(\mathcal{V},\mathcal{E},\mathbf{X})$8, KGE $\mathcal{G}=(\mathcal{V},\mathcal{E},\mathbf{X})$9, PBIAS $$f(\mathbf X^{t}, \mathcal G;\theta) \to \hat{\mathbf Y}^{\,t+\Delta t}_{\mathcal V_\rho}\in\mathbb R^{|\mathcal V_\rho|\times\Delta t},$$0

These graph statistics are significant because they show that HydroGAT’s heterogeneity is not merely semantic. The graph mixes dense pixel-scale routing with sparse gauge-to-gauge dependencies, allowing local hillslope-to-channel routing and larger-scale inter-catchment effects to be learned within the same architecture (Sarkar et al., 2 Sep 2025).

3. Architecture and learning dynamics

HydroGAT consists of four principal components: a temporal transformer encoder applied per node, two GAT-GRU spatial branches, a learnable fusion mechanism, and a predictor that incorporates forecasted rainfall. The temporal encoder maps each node’s recent history into a sequence of temporal embeddings $$f(\mathbf X^{t}, \mathcal G;\theta) \to \hat{\mathbf Y}^{\,t+\Delta t}_{\mathcal V_\rho}\in\mathbb R^{|\mathcal V_\rho|\times\Delta t},$$1. For each time step $$f(\mathbf X^{t}, \mathcal G;\theta) \to \hat{\mathbf Y}^{\,t+\Delta t}_{\mathcal V_\rho}\in\mathbb R^{|\mathcal V_\rho|\times\Delta t},$$2, raw features $$f(\mathbf X^{t}, \mathcal G;\theta) \to \hat{\mathbf Y}^{\,t+\Delta t}_{\mathcal V_\rho}\in\mathbb R^{|\mathcal V_\rho|\times\Delta t},$$3 are projected to dimension $$f(\mathbf X^{t}, \mathcal G;\theta) \to \hat{\mathbf Y}^{\,t+\Delta t}_{\mathcal V_\rho}\in\mathbb R^{|\mathcal V_\rho|\times\Delta t},$$4 and combined with sinusoidal positional encoding: $$f(\mathbf X^{t}, \mathcal G;\theta) \to \hat{\mathbf Y}^{\,t+\Delta t}_{\mathcal V_\rho}\in\mathbb R^{|\mathcal V_\rho|\times\Delta t},$$5 Causal multi-head self-attention is then applied with a 24-hour sliding window, so that timestep $$f(\mathbf X^{t}, \mathcal G;\theta) \to \hat{\mathbf Y}^{\,t+\Delta t}_{\mathcal V_\rho}\in\mathbb R^{|\mathcal V_\rho|\times\Delta t},$$6 attends only to $$f(\mathbf X^{t}, \mathcal G;\theta) \to \hat{\mathbf Y}^{\,t+\Delta t}_{\mathcal V_\rho}\in\mathbb R^{|\mathcal V_\rho|\times\Delta t},$$7. With

$$f(\mathbf X^{t}, \mathcal G;\theta) \to \hat{\mathbf Y}^{\,t+\Delta t}_{\mathcal V_\rho}\in\mathbb R^{|\mathcal V_\rho|\times\Delta t},$$8

the temporal embedding is

$$f(\mathbf X^{t}, \mathcal G;\theta) \to \hat{\mathbf Y}^{\,t+\Delta t}_{\mathcal V_\rho}\in\mathbb R^{|\mathcal V_\rho|\times\Delta t},$$9

This design is meant to capture both short-term event response and longer-lag hydrologic memory without the vanishing-gradient issues associated with purely recurrent temporal modules (Sarkar et al., 2 Sep 2025).

Spatial propagation is performed at each time step through two parallel branches. One branch operates on the flow graph $\mathcal V_\rho\subseteq\mathcal V$0 over all nodes; the other operates on the catchment graph $\mathcal V_\rho\subseteq\mathcal V$1 over gauge nodes only. Each branch is a GRU in which the linear transformations are replaced by graph attention operators. For branch $\mathcal V_\rho\subseteq\mathcal V$2, update and reset gates are

$\mathcal V_\rho\subseteq\mathcal V$3

A candidate hidden state is constructed from

$\mathcal V_\rho\subseteq\mathcal V$4

then updated through

$\mathcal V_\rho\subseteq\mathcal V$5

and finally

$\mathcal V_\rho\subseteq\mathcal V$6

This construction is central to HydroGAT’s formulation: temporal context is computed first, then spatial routing is modulated through graph attention within GRU-style gating (Sarkar et al., 2 Sep 2025).

The two spatial branches are fused only for target nodes through a learnable vector $\mathcal V_\rho\subseteq\mathcal V$7, with $\mathcal V_\rho\subseteq\mathcal V$8: $\mathcal V_\rho\subseteq\mathcal V$9 For non-target nodes, the hidden state remains $\mathcal E_F$0. Future rainfall forecasts over the next $\mathcal E_F$1 hours are processed by a lightweight convolutional layer, concatenated with the fused hidden-state summary, and passed to a predictor that outputs discharge forecasts for gauge nodes (Sarkar et al., 2 Sep 2025).

4. Empirical performance, lead-time behavior, and distributed scaling

HydroGAT is evaluated with standard hydrologic metrics, including Nash–Sutcliffe Efficiency,

$\mathcal E_F$2

Kling–Gupta Efficiency, Percent Bias,

$\mathcal E_F$3

as well as NRMSE, NMAE, and MAPE. On the reported 72-hour input and 72-hour output experiments, HydroGAT achieves NSE $\mathcal E_F$4, KGE $\mathcal E_F$5, NRMSE $\mathcal E_F$6, NMAE $\mathcal E_F$7, MAPE $\mathcal E_F$8, and PBIAS $\mathcal E_F$9 on the Cedar River Basin. On the Des Moines River Basin, it reports NSE $\mathcal E_C$0, KGE $\mathcal E_C$1, NRMSE $\mathcal E_C$2, NMAE $\mathcal E_C$3, MAPE $\mathcal E_C$4, and PBIAS $\mathcal E_C$5. Across the two basins and five baseline architectures, the abstract summarizes these outcomes as higher NSE up to $\mathcal E_C$6, improved KGE up to $\mathcal E_C$7, and low bias with PBIAS within $\mathcal E_C$8 in hourly discharge prediction (Sarkar et al., 2 Sep 2025).

The basin-level comparison is nuanced rather than uniformly dominant. In the Cedar River Basin, HydroGAT improves on the listed DCRNN baseline in NSE, NRMSE, NMAE, and MAPE while remaining close in bias. In the Des Moines River Basin, several baselines match HydroGAT on NSE, and GraphWaveNet matches it on KGE, but HydroGAT provides the best joint trade-off across efficiency and error metrics. This suggests that the model’s advantage lies not in a universal margin on every scalar summary but in a consistent balance between hydrologic skill, bias control, and long-horizon stability (Sarkar et al., 2 Sep 2025).

Lead-time experiments show that in the Cedar River Basin HydroGAT maintains NSE $\mathcal E_C$9 up to 66 hours, while DCRNN drops below $(v,v)$0 after about 50 hours and GraphWaveNet after about 24 hours. In the Des Moines River Basin, several baselines remain competitive at short lead times, but HydroGAT ends with the highest NSE at 72 hours. In a 5-day forecast setting with 72 hours of input and a 120-hour horizon, performance decays more slowly in the Des Moines River Basin, where NSE decreases from about $(v,v)$1 to about $(v,v)$2 by 120 hours, than in the Cedar River Basin, where NSE falls below $(v,v)$3 by 120 hours. The reported interpretation is that DSMRB exhibits a smoother, more persistent hydrologic response, whereas CRB is more variable (Sarkar et al., 2 Sep 2025).

A defining feature of HydroGAT is its distributed implementation. The training pipeline uses PyTorch Distributed Data Parallel, replicates the full basin graph in each process, partitions the temporal dataset into non-overlapping contiguous chunks through a custom sequential distributed sampler, and synchronizes gradients via AllReduce. On NERSC Perlmutter, with PyTorch 2.1, CUDA 11.8, DGL 1.1, mixed precision, AdamW, hidden dimension $(v,v)$4, two attention heads, dropout $(v,v)$5, 100 epochs, and batch size $(v,v)$6, the model scales up to 64 NVIDIA A100 GPUs and demonstrates up to $(v,v)$7 speedup across machines, with near-linear scaling up to 32 GPUs and diminishing returns at 64 GPUs because of communication overhead (Sarkar et al., 2 Sep 2025).

5. Interpretability, ablations, and evaluation criteria

HydroGAT exposes two attention mechanisms with direct hydrologic interpretation. Spatial attention on catchment edges identifies influential upstream gauges. In the Cedar River example, gauge station 613 has two upstream contributors, 553 and 1688; the learned attention strongly favors 1688 and assigns near-zero weight to 553, which is described as a tiny, noisy catchment with poor prediction performance. Temporal attention at the same gauge over a 72-hour input window emphasizes the recent 24–36 hours and also shows peaks around 42–48 hours and 66–72 hours before prediction. These patterns indicate that the model retains both near-event response and multiday memory, and the paper emphasizes that the resulting attention maps are sparse and structured rather than diffuse (Sarkar et al., 2 Sep 2025).

The ablation studies indicate that HydroGAT’s major components are functionally necessary rather than decorative. Replacing learned graph attention with fixed distance-based edge weights produced no gains and sometimes degraded performance. Replacing the full transformer block with a naive multi-head attention module without positional encoding, layer normalization, and feed-forward structure reduced NSE by more than 10 percentage points. Replacing GAT with Chebyshev convolution inside the GRU caused NSE drops of 5–30 points. Removing future rainfall had only a negligible effect in the Cedar River Basin but reduced NSE by about one point and KGE by more than three points in the Des Moines River Basin. Removing the catchment-edge branch reduced basin-level NSE by about 5 points in DSMRB and about 15 points in CRB. Replacing the learnable head-specific fusion $(v,v)$8 with an MLP-based fusion reduced NSE by about 1 point in CRB and about 2 points in DSMRB (Sarkar et al., 2 Sep 2025).

These results clarify two recurrent misconceptions. First, future exogenous forcings are not uniformly beneficial: in the Cedar River Basin, omitting future rainfall has limited effect, whereas adding noisy rainfall forecasts is worse than omitting them, leading the paper to the practical conclusion that misinformation is worse than missing information. Second, graph attention should not be conflated with unconditional superiority on every benchmark metric. In a related flood-susceptibility study using a graph transformer, AUC-ROC was $(v,v)$9 for the graph transformer and $\times$0 for XGBoost, yet the graph transformer achieved a higher Moran’s $\times$1 of $\times$2 relative to $\times$3 for XGBoost and yielded better clustering and delineation in susceptibility maps. This comparison suggests that graph-attention methods may be most valuable when spatial coherence and topology-sensitive structure are central evaluation criteria, not only scalar discrimination metrics (Vemula et al., 31 Mar 2025).

6. Relation to adjacent graph-hydrology and heterogeneous-graph research

HydroGAT belongs to a broader family of graph-centric hydrologic methods, but neighboring works instantiate that idea in materially different ways. In flood susceptibility mapping, a graph transformer has been applied to a sparse directed $\times$4-NN similarity graph built from 811 flooded and 811 non-flooded watershed points, with cosine-similarity edge weights, PCA-based graph construction, and Laplacian positional encodings defined from the symmetric normalized Laplacian

$\times$5

That model uses binary cross-entropy, achieves Sensitivity $\times$6, Specificity $\times$7, PPV $\times$8, and NPV $\times$9, and is explicitly described as a hydrology-focused graph transformer. It differs from HydroGAT by targeting node classification rather than discharge forecasting, by operating on a similarity graph rather than a pixel-level flow graph, and by emphasizing susceptibility-map clustering under current and projected climates rather than autoregressive discharge prediction (Vemula et al., 31 Mar 2025).

A second neighboring line is ontological and systems-engineering oriented rather than purely neural. In the MBSE+HFGT framework, hydrologic systems are represented through operands $\mathcal V_\rho$0, resources $\mathcal V_\rho$1, processes $\mathcal V_\rho$2, buffers $\mathcal V_\rho$3, and capabilities $\mathcal V_\rho$4, with dynamics governed by an engineering system net whose discrete transition law is

$\mathcal V_\rho$5

This framework explicitly separates form from function in SysML, treats graph edges as executable processes rather than simple links, and is proposed as a formal, ontology-driven basis that could underpin or extend a HydroGAT-style graph abstraction toolbox. A plausible implication is that such a reference architecture could supply a principled backend for defining component semantics and multi-operand coupling in future HydroGAT-like systems (Harris et al., 31 May 2025).

Adjacent heterogeneous-graph work outside hydrology is also relevant. HGATSolver for fluid–structure interaction encodes fluid and solid nodes, uses relation-specific attention over edge types $\mathcal V_\rho$6, introduces a physics-conditioned gating mechanism

$\mathcal V_\rho$7

and balances domain losses through uncertainty-weighted training. The reported state-of-the-art performance on FSI benchmarks suggests a concrete architectural template for HydroGAT extensions to coupled hydro-elastic or other multi-physics settings, though this is an inference rather than a demonstrated hydrologic application (Zhang et al., 14 Jan 2026). At a more abstract topological level, HL-HGAT generalizes heterogeneous graph attention to simplicial complexes and Hodge-Laplacian filtering on $\mathcal V_\rho$8-simplices, which suggests a possible route toward higher-order HydroGAT variants in which nodes, edges, and faces encode different hydrologic entities or fluxes (Huang et al., 2024).

Taken together, these neighboring works show that “HydroGAT” denotes more than a single implementation. In its narrow sense, it is the pixel-level, distributed heterogeneous graph attention transformer for hourly basin-scale discharge forecasting (Sarkar et al., 2 Sep 2025). In a broader research sense, it refers to an emerging hydrologic graph paradigm in which topology, typed interactions, and structured attention are used either for predictive modeling, for susceptibility mapping, or for executable system abstraction across environmental components (Vemula et al., 31 Mar 2025, Harris et al., 31 May 2025).