Interpreting LSTM and MC-LSTM internal states with hydrological physics

The use of artificial neural networks (ANNs) in hydrological modelling has gained increasing popularity due to their ability to represent non-linear relationships and complex system dynamics.  In particular, Long Short-Term Memory (LSTM) networks have become the state-of-the-art approach for streamflow simulations, as they incorporate memory cells and gating mechanisms capable of learning both short- and long-term dependencies. However, standard LSTM models have limitations for predicting extreme high-flow events and suffer from limited interpretability due to their lack of explicit physical grounding.

The Mass-Conserving LSTM (MC-LSTM) is a variant of the standard LSTM architecture designed to address the lack of physical consistency by embedding mass conservation directly into the internal model structure. Hence, the information stored in MC-LSTM cell states is expected to correspond more directly to hydrological processes contributing to the basin water balance.

This study analyses and compares the internal processes of standard LSTM and MC-LSTM networks trained on four snowmelt-dominated watersheds located in the Central Spanish Pyrenees. We first evaluate the ability of both models to conserve water volume, showing that the MC-LSTM maintains volumetric consistency due to the imposed physical constraint, whereas the standard LSTM exhibits substantial discrepancies between observed and simulated volumes. We then investigate the learning behaviour of the MC-LSTM using two independent physical datasets not included as model inputs: snow and evapotranspiration (ETO), both of which play a key role in the local water balance. Using a wavelet-based methodology, snow-cells and ETO-cells are identified within the MC-LSTM cell state. Snow-cells exhibit Pearson correlations exceeding 0.5 across all watersheds, while ETO-cells reproduce the observed variability despite low temporal correlation. Furthermore, ETO-cells show a limited contribution to the model output, consistent with their physical role as water losses.

Overall, this analysis highlights the limitations of standard LSTM models in representing volumetric consistency and physical conservation processes, while demonstrating the enhanced physical interpretability of the MC-LSTM architecture, which achieves comparable or superior performance to standard LSTM models while preserving hydrological coherence.

Acknowledgments: This work is funded by the European Research Council (ERC) through the Horizon Europe 2021 Starting Grant program under REA grant agreement number 101039181-SED@HEAD.