Deep Learning the River Network: Message-Passing LSTMs for Robust Stream Water Temperature Prediction

Predicting stream water temperature (SWT) under non‑stationary hydroclimatic conditions is essential for ecosystem management yet remains challenging for deep learning applications in hydrology due to spatially structured network processes and disturbance‑driven variability. We present a graph‑informed deep learning framework that combines Long Short‑Term Memory (LSTM) networks with message passing and multi‑head attention to jointly capture temporal dynamics and upstream connectivity for daily SWT forecasting. 

Sagehen Creek, a snowmelt‑dominated montane watershed in the northern Sierra Nevada (California, United States, U.S.), served as a benchmark for evaluating robustness in climate‑sensitive mountain systems. Its pronounced seasonality, groundwater influence, and sensitivity to climate variability provide an ideal setting to assess model robustness in underrepresented montane systems and demonstrate practical scalability to larger river networks. The architecture integrates shared LSTM layers for temporal feature extraction with a graph‑based message‑passing module that weights upstream contributions via multi‑head attention. Inputs include meteorological drivers (air temperature, precipitation, solar radiation), land cover, elevation, and seasonality (day of year), derived from long‑term observations and national datasets. Hyperparameters were tuned using Bayesian methods to improve model accuracy and reliability. Applied to Sagehen Creek and thousands of gages across the U.S., the model achieves strong performance in gaged settings (RMSE ≈ 1.32°C) and maintains comparable skill in ungaged scenarios (RMSE ≈ 1.35 °C), demonstrating generalization across heterogeneous basins. Explicit representation of seasonality improves predictions of extremes, and attention weights provide insight into upstream influence. 

Overall, this work advances deep learning in hydrology by introducing a scalable, network‑aware architecture suited to non‑stationary conditions, employing structured training methods to improve reliability, and enabling ungaged predictions with minimal reliance on local observations. These results demonstrate the potential for network‑aware deep learning approaches to support more flexible and transferable hydrologic prediction strategies as environmental conditions evolve. Future work aims to include systematic comparisons with traditional statistical models to better contextualize performance gains and clarify where deep learning provides distinct advantages for SWT forecasting.