Daily runoff prediction based on lightweight Mamba with partial normalization

Runoff forecasting plays a critical role in water resources planning and management (Liu et al. 2024), and high-precision runoff forecasts are essential for water supply, flood control, and power generation (Williams et al. 2021). Influenced by natural meteorological conditions, watershed characteristics, and human activities, the runoff sequence is characterized by nonlinearity, stochasticity, and periodicity (Sharafati et al. 2020). The search for more accurate runoff forecasting methods has been a major concern of scholars.

Existing runoff prediction methods generally fall into two broad categories: physically driven models and data-driven models (Nourani et al. 2021). Physically driven models help understand the underlying physics of hydrologic processes. These models construct complex mathematical models, such as high-dimensional partial differential equations, using data on hydrology, meteorology, topography, soil properties, and vegetation cover index, to simulate or predict river runoff (Farmer et al. 2003; Fenicia et al. 2008; Bai et al. 2009). Unfortunately, the degree of difficulty and the complexity of parameter estimation are high.

Unlike physically driven models, data-driven models focus only on the relationship between inputs and outputs (Kumar et al. 2022), and it is easy to conduct. As data-driven models continue to evolve, they can be further categorized into statistical models, dynamic and stochastic model, machine learning methods and deep learning models. Statistical models use statistical correlation methods to predict future runoff from historical runoff observations, such as Autoregressive Integrated Moving Average Model (ARIMA) (Farajzadeh et al. 2014). Statistical models can accurately capture the linear relationship between historical and future data, but they cannot account for the effects of external factors related on considered series, and the ability of catching nonlinear relationship between time series is weaker. Stochastic methods model time series based on stochastic processes (Dimitriadis et al. 2021). Complex nonlinear relationships in systems can be described. However, it is difficult to calculate various parameters of stochastic processes.

Machine learning models are capable of fitting nonlinear relationships between inputs and outputs (Mosavi et al. 2018) and can therefore be used for short-term runoff prediction. A variety of machine learning methods have been proposed for runoff prediction, including k-nearest neighbor (KNN) regression (Yang et al. 2020), support vector regression (SVR) (Bigdeli et al. 2023), eXtreme Gradient Boosting (XGBoost) (Szczepanek 2022), and others. However, these machine learning methods have limited capabilities in extracting deep information and lack the ability to obtain prior information about the sequence order. Therefore, temporal information must be added to the model through feature engineering (Dwivedi et al. 2022). However, the predictive performance of the models decreases rapidly as the dimensionality of the input features and the prediction steps increase.

With the rapid development of artificial intelligence and computing capability (e.g., GPUs), deep learning models are widely used in various tasks due to their ability to process high-dimensional input information. Owing to its self-ordering advantage, Recurrent Neural Network (RNN) excels at sequence tasks (Nosouhian et al. 2021). Long Short-Term Memory (LSTM) alleviates the problem of vanishing or exploding gradients in RNNs, and Gated Recurrent Unit (GRU) further improves training efficiency. Therefore, these two variants are often implemented as base models for runoff prediction (Xiang et al. 2020; Man et al. 2023; Yao et al. 2023). However, based on the structure of RNN, these models can only be trained in recursive mode. When the input sequence is long, the training of the neural network becomes inefficient.

Transformer (Vaswani et al. 2017) is widely used in the fields of natural language processing and computer vision, and it is also applied to the runoff prediction task (Yin et al. 2022; Fang et al. 2024). The core of transformer is the self-attention mechanism, which eliminates the sequential dependency problem in traditional recurrent neural networks by introducing an attention mechanism. Informer model (Zhou et al. 2021), which improves the efficiency of the transformer in the task of time series prediction, was used to runoff prediction and achieved good results (Du et al. 2022). However, the attention mechanism suffers from the quadratic complexity of the input sequence length and low inductive bias (Neyshabur 2020).

Recently, state space models (SSMs) have attracted a lot of attention in the NLP and CV communities due to their efficacy in long-range sequence modeling. Mamba introduces a selection mechanism into SSMs, giving them the ability to perform context-aware selection and capture long-range correlations (Gu & Dao 2023). It employs scanning operations for parallel computation during training while still maintaining linear complexity, allowing Mamba-based models to outperform transformers. As a sequence model, Mamba also has great potential for time series prediction tasks, which triggers our interest in exploring Mamba for runoff prediction.

The main contributions of this study can be categorized into three areas:

• The Mamba model was used to predict runoff data to explore the potential application of Mamba as a backbone.

• A new deep learning prediction model LightMamba is proposed. The Multi-Path-Mamba (MPM) module is utilized to reduce the number of parameters of the Mamba model, and the input sequences are locally regularized to improve the prediction accuracy of non-stationary time series such as runoff data.

• We tested our model and compared it with several benchmark models in four evaluation metrics, and the results show that our model has better performance.

In the time series forecasting task, the input to the model is ⁠, where ⁠. Predict future sequence using input X. L and T are the lengths of the time windows for inputting the past and predicting the future, respectively, and are referred to as the retrospective window and the predictive horizon. p_i is the variable and D is the dimension of the variable.

Algorithm 1 Mamba Block with Selective SSM

Input:

Output:

1:

2:

3:

4:

5:

6:

7:

8:

9:

10: returnY

The Mamba’s number of parameters is primarily determined by its internal linear projection layer. When the input and output dimensions are both d, the linear layer contains parameters, which is the square level of the input variable’s dimension. Thus, reducing the input variable’s dimension to or less will result in a smaller number of parameters in the linear layer.

The MPM Algorithm 2 is proposed based on this idea. First, the input sequence X is split into different subspaces along the variable dimensions to obtain ⁠. Then, the smaller Mamba model is used to extract the sequence features under the subspace. Finally, the features from each subspace are merged to create the final output. MPM can help the model to learn time series features representations under different subspaces while using fewer parameters than the original model.

Algorithm 2 Multi-Path-Mamba

Input:

Output:

1:

2:

3: for each do

4:

5:

6: end for

7:

8: returnY

The models utilized in this paper are as follows:

• Benchmark: BM.

• Statistical and ML methods: ARIMA, SVR, KNN, XGBoost.

• Recurrent neural network: RNN, LSTM, BiLSTM, GRU.

• Attention-based models: Transformer, Informer.

• Mamba-based models: Mamba, LightMamba.

⁠

The data is divided into training, validation, and testing sets in a 7:1:2 ratio along the time axis. Normalization is performed on the training set, and the validation and test sets are normalized using the scale parameters of the training set. After making a prediction using the model, the results are denormalized using the same scale, to avoid information leakage (Li et al. 2023). The selected training model is the best performer from the validation set, ensuring the model’s generalization ability. The daily runoff for the next 7 days was predicted by all models using runoff data from the past 30 days, that is the in this study. The parameter settings for each model are shown in Table 2, in which the deep learning model shares some hyperparameters. All models will be trained three times, and the average of the evaluation metrics will be taken as the final result.

Table 2

Parameter settings for each model

Model .	Parameter .
ARIMA	⁠, ⁠,
KNN	Number of neighbors
SVR	Kernel type is linear
XGBoost	Number of estimators 100, max depth 6
RNNLSTM	Hidden layer size 64, layers of RNNs 1 (BiLSTM 2), attention heads of transformers 4, dropout rate 0.1, loss MSE, optimizer AdamW, learning rate 0.001, learning rate decay factor 0.9, maximum training period 30, early-stop patience 5
BiLSTM
GRU
Transformer
Informer
Mamba
LightMamba

Model .	Parameter .
ARIMA	⁠, ⁠,
KNN	Number of neighbors
SVR	Kernel type is linear
XGBoost	Number of estimators 100, max depth 6
RNNLSTM	Hidden layer size 64, layers of RNNs 1 (BiLSTM 2), attention heads of transformers 4, dropout rate 0.1, loss MSE, optimizer AdamW, learning rate 0.001, learning rate decay factor 0.9, maximum training period 30, early-stop patience 5
BiLSTM
GRU
Transformer
Informer
Mamba
LightMamba

View Large

From Tables 3–5 we can see that the ARIMA model has excellent forecasting performance on the 1-day forecasting and also performs better over the whole forecasting period compared to listed ML methods. This indicates that the ARIMA model can capture short-term linear relationships by considering the historical behavior and trend of the series, and the prediction is more robust in terms of the long-term trend (Fard & Akbari-Zadeh 2014). However, the ARIMA model has poorer prediction performance at relatively long step sizes. ARIMA is not suitable for predicting runoff, as it cannot handle time series with complex patterns or external factors influencing the series. Additionally, selecting the appropriate ARIMA model requires complex data preprocessing, making it difficult to apply in multivariate forecasting and practical applications.

The predictive power of machine learning methods in this study is inadequate. These models lack sequential structure and only learn the relationship between input features and outputs, without the ability to infer temporal order. Although XGBoost achieved an impressive accuracy of 0.991 NSE on the first day, its prediction error increased significantly with each subsequent step. As a result, the model’s overall performance was poor, with a score of only 0.788 NSE for the entire prediction period. This suggests that machine learning methods struggle to capture the temporal evolution of runoff sequences.

The recurrent neural network contains a priori information on sequence order and has a better prediction effect on time series. The GRU exhibits the best performance on four recurrent neural networks in terms of MAE, RMSE, and NSE, with values of 889.00, 1371.60, and 0.894, respectively, on St. Louis station, and has 25.5K parameters relatively lower than others. The effectiveness of GRU in runoff prediction is verified. Due to its superior prediction capacity and the advantage of the number of parameters, GRU is widely used in various runoff prediction works. However, the recursive structure of RNNs limits their ability to be trained in parallel. This inefficiency becomes more pronounced when the input sequence is long, requiring additional computational resources for neural network training.

However, our results show that the transformers did not achieve the expected accuracy. We analyzed that this is due to the low inductive bias of the transformers and its encoder and decoder structure, which results in a large number of parameters in this type of model. In this experiment, the parameters of the transformer and informer are as high as 93.5K and 106.0K. Although larger model parameters can increase the upper learning limit of the model, it is crucial to train it adequately with a significant amount of data. Therefore, this type of model is more suitable for large-scale datasets (Wen et al. 2022). In contrast, the runoff data is small, making it challenging to realize the potential of transformer models.

The Mamba-based model demonstrates superior performance in comparison to the other models, particularly in terms of the relatively long number of prediction steps and the overall performance. The MAE, RMSE, and NSE of Mamba are 856.00, 1320.39, and 0.900 on St. Louis station, respectively. The number of parameters is only 16.9K, which exceeds the prediction performance of GRU with a lower number of parameters. The results demonstrate that the Mamba-based model is capable of performing well in the runoff prediction task. The performance of LightMamba is particularly noteworthy, with MAE, RMSE, and NSE values of 833.75, 1286.01, and 0.904, respectively. These values illustrate a greater degree of accuracy in prediction than that observed in Mamba. This is due to the incorporation of local normalization and MPM into LightMamba. Local normalization enables the output of the model to be maintained within the same range as the inputs, thereby enhancing the prediction accuracy to a certain extent when confronted with a non-stationary time series such as runoff. Concurrently, our MPM algorithm, analogous to the multi-head attention mechanism, enables the Mamba to model the sequence relationship in disparate subspaces, thereby enhancing the model’s representation capability. Furthermore, MPM can effectively reduce the model parameters. LightMamba has only 10.2K parameters, which constitutes 60% of Mamba.

The average performance of each model across the three stations is summarized in Table 6. The best results are highlighted in bold, while the second-best results are shown in underlined text. The results indicate that the transformers exhibit a lower overall performance compared to the RNNs when confronted with this limited dataset. In contrast, the Mamba-based model achieves the best combined performance across the three stations, with higher predictive accuracy of the model while maintaining a lower number of parameters. In particular, an examination of the comprehensive performance of the three stations, as reflected in the RMSE metrics, reveals that our proposed model, LightMamba, exhibits a reduction of 14.71, 49.14, 7.67, and 16.16% in comparison to BM, XGBoost, GRU, and Informer, respectively.

This paper examines the potential of the Mamba model for runoff prediction. Compared to the statistical model, machine learning methods, recurrent neural networks, and attention-based models, the Mamba model demonstrates superior overall performance. The Mamba model is adept at extracting nonlinear and long-term dependencies in runoff data, and its prediction accuracy is enhanced while the number of parameters is reduced. The scanning algorithm in the Mamba block enables the model to be trained in parallel, resulting in a faster training speed and lower GPU memory usage than other models. This demonstrates the great potential of Mamba as a backbone model for runoff prediction tasks.

In this paper, we also propose a deep learning model for runoff prediction: LightMamba. It utilizes partial normalization and an MPM module, which improves the accuracy of Mamba on runoff prediction tasks and reduces the number of parameters. We validate LightMamba on runoff data from three stations, and the results demonstrate that LightMamba can effectively capture the temporal dependence in runoff data and outperforms previous methods.

However, in this study, only historical runoff data were utilized to predict future runoff, which is a univariate forecasting problem. Mamba is also adept at extracting synergistic relationships between multiple variables. Therefore, other variables related to runoff (e.g., rainfall, temperature, potential evapotranspiration, vegetation cover) can be subsequently collected and utilized. In addition, various decomposition methods are frequently employed in runoff prediction tasks, as they can significantly improve the prediction accuracy. Mamba may be considered as a backbone model, combined with data decomposition methods to improve the precision of runoff prediction, which is also our future research direction.

This work was supported by the National Natural Science Foundation of China (Grant No. 42130113). The numerical calculations in this paper were supported by the Supercomputing Center of Lanzhou University.

J.D. conceived and designed the study; J.D., H.D., and C.S. conducted the workshops; J.D. and H.D. conducted the survey; J.D., H.D., and C.S. performed the inspections; J.D. and H.D. analyzed the data; J.D., H.D., and L.W. wrote the draft paper. All authors have read and agreed to the published version of the manuscript.

All relevant data are available from an online repository: https://grdc.bafg.de/.

The authors declare there is no conflict.