Mass Conserving LSTM with Dual States for Improved Streamflow Prediction through Quickflow and Slow Storage Separation

Long-Short Term Memory (LSTM) shows exceptional performance for rainfall-runoff modelling, but lacks physical realism. Efforts to integrate mass conserving into the model architecture didn’t translate to significant improvement in predictive accuracy. We build on the Mass-Conserving LSTM (MC-LSTM) by proposing a novel architecture that incorporates two complementary cell states to represent distinct fast and slow hydrologic memory components while maintaining strict mass balance. We introduce a new partition gate for segregating the mass input for long- and short-term memory, and made required architectural changes to incorporate additional cell state. We benchmarked our model against LSTM and MC-LSTM on CAMELS-IND (158 basins) and CAMELS-US (531 basins) using NSE, KGE, Pearson-r, FHV, FLV, and peak timing/magnitude. For the Indian dataset, MC-LSTM-DS surpasses both LSTM and MC-LSTM across all metrics except Pearson-r and FLV, where it exceeds the performance of LSTM but falls short of MC-LSTM. In the low flow regime (FLV), our model decreases the overestimation of LSTM significantly, while MC-LSTM shows severe underestimation. Analysis of the spatial distribution revealed it to be aligned with hydroclimate, where all the models performed better in humid/tropical climates, while performance lacked in arid regions. Investigation of the cell states revealed that the added cell state represents the long-term processes effectively, while the original cell state captures short-term processes. The change in their relative contributions according to the climate characteristic is observed, thus confirming our hypothesis and also providing an interpretable decomposition of the simulated flows. On the CAMELS-US dataset, MC-LSTM-DS demonstrates equal performance as MC-LSTM and LSTM on NSE, and outperforms all the models in KGE, FHV, and Pearson-r. In FLV, it outperforms all the mass-conserving models by a huge margin and is just short of LSTM. This study proposes a novel mass-conserving model that provides an interpretable prediction. We claim MC-LSTM-DS to be the current state-of-the-art for large sample rainfall runoff modelling, as it showed superior performance across two diverse regions. To the best of our knowledge, this study is the first to investigate the effects of strict mass conservation on the diverse Indian region.