An enhanced multi-head attention-based LSTM model for forecasting the surface water quality index Open Access

Surface water plays a crucial role in supporting environmental sustainability and the inhibition of diverse forms of life that mostly rely on it (Young et al. 2021; Suleman & Shridevi 2022; Mishra 2023). In recent decades, water quality has experienced a notable decline due to both natural and human-induced factors (Georgescu et al. 2023). Natural influences, such as climate change, flooding, and alterations in atmospheric and hydrological conditions, directly affect the water quality. Additionally, anthropogenic activities, including industrial effluent, municipal sewage, agricultural practices, and soil erosion contribute significantly to the deterioration of water quality (Vlad et al. 2012; Calmuc et al. 2021; Uddin et al. 2021). This decline in water quality has made it imperative to assess and check surface water regularly with precision as a result. Organizations like the WHO and various governments are prioritizing water quality evaluation in a regular interval as well (Tavakoly Sany et al. 2014).

The Water Quality Index (WQI) models are used widely and a well-known method for water quality monitoring and forecasting which is defined as a numerical tool that combines multiple water quality parameters into a unified value classification (Uddin et al. 2022, 2023a, b; Syeed et al. 2023). The WQI provides an overall assessment of water quality by classifying this into various categories. However, the method is susceptible to limitations such as eclipses and ambiguity classification calculations (Syeed et al. 2023). Eclipsing occurs when a sudden change of a water quality parameter disproportionately affects the WQI, leading to misclassification by overshadowing other parameters (Ding et al. 2023). Ambiguity arises when the WQI classification is worse than expected and even when the underlying water quality conditions fail to undertake such a surmise (Georgescu et al. 2023; Uddin et al. 2024). These issues, which are primarily driven by fluctuations in water quality parameters due to natural and anthropogenic influences hinder the overall accuracy of water quality assessments (Georgescu et al. 2023; Uddin et al. 2024).

In order to address the eclipses and ambiguities in the WQI prediction, contemporary literature employs a variety of data-driven forecasting models and techniques. Forecasting, in general, involves predicting future trends based on historical data. Unlike traditional WQI models, this approach uses historical time-series WQI data to project future values by analyzing past trends and patterns, rather than focusing on individual parameter values (Rouf et al. 2021; Alsharef et al. 2022; Sajib et al. 2023). Forecasting is prevalent in many prediction domains, including weather forecasting (Chen et al. 2023) and air quality forecasting (Méndez et al. 2023). Recent studies have used both conventional machine learning (ML) and sophisticated deep learning models for surface water quality forecasting. ML models, such as the automatic exponential smoothing model (AESM), are used to evaluate temporal changes in groundwater quality over six years in Xianyang City (Méndez et al. 2023). A fuzzy expert system is proposed to predict the WQI at various locations in Solapur City, India, using triangular and trapezoidal membership functions (Patki et al. 2015). Although these models perform well in WQI forecasting, they often do not capture the intricate non-linear relationships and patterns inherent in the water quality data. Moreover, they often depend on features that require extensive human expertise to prepare and, therefore, potentially overlook data complexities. Advanced deep learning algorithms are used for the automated extraction of complex features from raw water quality datasets to improve predictive and forecasting capabilities (Gambín et al. 2021; Kulisz et al. 2021; Zamili et al. 2023). Another study improves the prediction performance by integrating an artificial neural network (ANN) with the constraint Coefficient-based Particle Swarm Optimization and a Chaotic Gravitational Search Algorithm (CPSOCGSA), achieving an R² 0.965, a mean absolute error (MAE) 0.01627, and an RMSE 0.0187 (Zamili et al. 2023). Further advancements are observed with the use of a cascade-forward network (CFN) model combined with a radial basis function (RBF) network, showing mean squared error (MSE) values ranging from 0.083 to 0.319 and R-value between 0.940 and 0.911 across different quarters (Georgescu et al. 2023). However, these models show limitations in encapsulating the temporal complexities of the water quality data, raising concerns about the reliability of their reported performances (Somlyódy & Varis 2006; Wang et al. 2016; Kulisz et al. 2021). Another study proposes a Variational Autoencoder (VAE) with a self-attention mechanism for improved short-term wind power forecasting (Harrou et al. 2024). The SA–VAE model outperforms eight established deep learning methods, achieving superior accuracy (average R² = 0.992) using real-world data from France and Turkey. Besides that, they also proposed another approach for predicting energy consumption in wastewater treatment plants (WWTPs) using data augmentation, feature selection, and deep learning (Harrou et al. 2023). LSTM and BiGRU models, enhanced with augmented data and lagged features, achieve high accuracy with MAPE values of 1.36 and 1.436%, respectively.

The model revealed in this study aims to overcome all limitations of traditional ML and deep learning methods by developing a novel approach integrating long short-term memory (LSTM) networks with a multi-head attention mechanism for the very first time. LSTM networks are well-suited for handling sequential data and capturing long-term dependencies which is crucial for right water quality forecasting data (Huang & Wu 2024; Li et al. 2024a, b). The multi-head attention mechanism enhances this capability by allowing the model to focus on multiple aspects of the temporal relationships presented in the data (Sahoo et al. 2019). This combination holds immense potential for improving the reliability and accuracy of water quality predictions. The demonstration in other fields such as greenhouse temperature and solar irradiance forecasting (Li et al. 2024a, b; Sakib et al. 2024) further solidifies the potential. The results reveal that proposed model achieves a mean square error (MSE) of 3,987.56 and 4,356.39, a root mean square error (RMSE) of 63.14 and 66.00, an MAE of 62.49 and 65.43, and a coefficient of determination (R²) of 0.91 and 0.88 for training and testing datasets, respectively. These metrics show the model's strong predictive accuracy and ability to generalize well across large datasets. A comparative analysis with existing ML and deep learning models for WQI forecasting highlights a significant performance improvement (around 5–8%) achieved by our approach. The enhanced accuracy and robustness underscore the effectiveness of integrating LSTM networks with multi-head attention mechanisms in capturing intricate temporal patterns in water quality data. The integration of these advanced techniques offers a significant step forward in addressing the complexities of water quality forecasting promptly.

The main goal of this study is to develop a new forecasting model that enhances the reliability and accuracy of surface water quality predictions with the help of the forefront multi-head attention-based network. The key contributions of this research are as follows:

• Development of a multi-head attention-based LSTM model for surface water quality forecasting, a novel approach in WQI prediction.

• Improvement of the resilience with outliers shown by lower RMSE and MAE values.

• Enhancement of the sensitivity to data variability, reflected in higher R² values.

• Reduction of forecasting errors, proven by improved MSE values.

These contributions aim at advancing water quality forecasting methodologies, providing policymakers and environmental managers with a more correct, data-driven tool for decision-making and sustainable water management. In the rest of the paper, the method is described in section ‘Materials and methods’. Section ‘Results and discussion' focuses on the result and description and finally, section ‘Conclusion’ concludes the paper with the new room for research.

The following sections detail each of these steps.

To train, test, and validate proposed model, the WQI is calculated using the dataset. WQI models are the widely adopted approach for calculating the WQI (Syeed et al. 2023; Uddin et al. 2024). These are mathematical predictors that take several complex water quality parameter values as input and produce a simple unit less number as the indicator of water quality, which is easy-to-comprehend and act upon (Uddin et al. 2021; Syeed et al. 2023). Worldwide, there are around 23 WQI models currently in practice that can be classified into two categories, weighted and unweighted models. Weighted models' execution is carried out in four steps to calculate the WQI, namely:

• (a) selecting water quality parameters

• (b) producing the parameter sub-indices

• (c) selecting weights for the parameters, and finally

• (d) applying an aggregation function to calculate the WQI. For unweighted WQI models, the parameter weighting (step c) is ignored.

For this study, the WQI values is calculated using the Canadian Council of Ministers of the Environment (CCME) model. This is an unweighted model, developed by the British Columbia Ministry of Environment, Canada. CCME is a widely adopted WQI model applied to a wide range of surface water bodies, namely, river, lake, marine and coastal water (Uddin et al. 2021; Syeed et al. 2023). This model is highly stable and consistent to produce reliable WQI classification under diverse environmental conditions and water quality parameter values (Banda & Kumarasamy 2020). The detail computation of the WQI using the CCME model is presented in Equations (4)–(9).

• (a) F₁: Stands for the percentage of parameters not meeting specified objectives:

  

  (4)
• (b) F₂: Stands for the percentage of individual test values not meeting objectives:

  

  (5)
• (c) F₃: Measures the deviation of test values from goals using an asymptotic function:

  

  (6)

The divisor of 1.732 in the aggregation equation normalizes the WQI to a range of 0–100, where 0 denotes the worst water quality and 100 is the best, considering the maximum values of the individual index factors (Canadian Water 2001). The specific classification scheme for CCME WQI (Syeed et al. 2023; Bouteldjaoui & Taupin 2024) is demonstrated in Table 2.

This is a required step as ML models show greater robustness and effectiveness when trained on scaled data rather than unscaled ones.

As shown in Figure 6, the proposed model has two LSTM layers with 128 and 64 units each. The output of the first LSTM layer (having 128 units) is passed as input to the second LSTM layer (having 64 units). The second LSTM layer produces two outputs, as ATSO and LTSO, where ATSO is a sequence of predictions for all time steps, represented as ⁠, and LTSO is the output for the final time step ⁠, which produces a single prediction for the final time point in the sequence.

To prevent overfitting, a dropout layer is introduced after the attention mechanism, with a dropout rate of 0.1. This introduces regularization by randomly dropping out connections during training, ensuring the model generalizes better to unseen data. The dropout layer effectively reduces reliance on specific neurons, forcing the model to spread its learning across the entire network. Figure 9 summarizes the working principle of the multi-head attention mechanism. The output produced in this layer is fed into a fully connected neural network layer, where integration and transformation are carried out to forecast the WQI value.

The neural network layer that precedes the Multi-Head Attention mechanism consists of several neural network layers to refine the extracted features to forecast the WQI value. In the first layer it applies two pooling operations, namely, Global Average Pooling and Global Max Pooling. These pooling layers summarize the temporal features extracted by the attention mechanism. The Global Average Pooling layer computes the average value across ATSO, highlighting the overall trend in the data. On the other hand, the Global Max Pooling layer finds the most prominent feature at each time step. These two outputs provide complementary perspectives on the temporal patterns of the water quality data. These outputs are concatenated into a single feature vector, which is later passed through a fully connected layer (i.e., dense layer) with 64 units and a sigmoid activation function. This process allows the model learning more complex relationships by reducing the feature vector dimensionality. Finally, the output from this dense layer is fed into another single unit fully connected layer and a sigmoid activation function for producing the final forecasting.

To train and evaluate the proposed model, the prepared dataset is split into training (80%) and test (20%) sets. The training set is used to fit the model, while the test set evaluates the generalization performance. Hyperparameter tuning plays a pivotal role to improve performance of the model. Table 3 provides a summary of all the candidate hyperparameters tuned for this purpose. By fine-tuning these hyperparameters and training the model under best conditions, effective forecasting performance is achieved.

All possible combinations of hyperparameters in Table 3 are evaluated, and the combination that produces the best results on the test dataset are selected. As highlighted in Table 3, a batch size of 64, a learning rate of 0.001, a dropout rate of 0.5, an input sequence length of 50, and 100 epochs are selected. Alongside, MSE is chosen as the loss function, and the Adam optimizer is selected for optimization. These selections are based on the high performance of the model in terms of MSE, RMSE, MAE, and R² using Equations (20)–(23). Python scripts are written to automate these calculations.

Performance and reliability of the proposed multi-head attention-based LSTM model is evaluated in three axes, (a) using the conventional performance assessment metrics, e.g., MSE, RMSE, MAE, R², and DM test, (b) benchmark the reported performance against existing ML WQI forecasting models using the listed metrics, and finally, (c) provide a time-series assessment of the forecasting relative to the WQI measurements in the recent past. It is worth mentioning that all the candidate models (both traditional and deep learning) that are selected for performance benchmarking are trained and tested with the same dataset (Rahman et al. 2025).

To offer a comprehensive performance comparison, the proposed model is evaluated against the forecasting performance of several deep learning and traditional ML models. First, the comparison is drawn in terms of MSE, RMSE, MAE, and R² metrics, and then the DM test is carried out to verify that the reported performance improvement of the proposed model is statistically significant. The results for MSE, RMSE, MAE, and R² metrics for each of the model is summarized in Table 4, and the DM test result is presented in Table 5. As can be seen from Table 4, the proposed model achieves the lowest MSE, RMSE, and MAE, along with the highest R² values (0.91 train, 0.88 test), establishing its' ability to capture the complex temporal dependencies of water quality data with minimal prediction error. Overall, 5–8% improvement in accuracy and reliability is observed for the proposed model as compared to the existing ones.

A detailed performance comparison with the deep learning models (from Table 4) shows that the proposed model achieves significant improvements in comparison to Vanilla LSTM, Stacked LSTM, gated recurrent units (GRUs), bi-directional LSTM (Bi-LSTM), and ANN. For example, the Stacked LSTM records higher MSE (4,307.56 train, 4,509.94 test), RMSE (65.63 train, 67.15 test), and lower R² values (0.83 train, 0.81 test), reflecting less accurate predictions. GRU, while slightly better in MSE (4,189.10 train, 4,395.66 test) and RMSE (64.72 train, 66.29 test), still underperforms compared to the proposed model in terms of lower error values and higher R². The multi-head attention mechanism in the proposed model provides a nuanced understanding of temporal dependencies, allowing it to adapt efficiently to fluctuations in data and reducing forecast errors compared to simpler architectures like GRU and Stacked LSTM. Alongside, the advanced models like ANN with Particle Swarm Optimization and CFN also fall short with MSE values exceeding 4,200 and R² scores below 0.85, highlighting their limited sensitivity to the variability in the water quality data.

Furthermore, performance comparison with the traditional ML models, e.g., linear regression (LR), decision tree (DT), random forest (RF), and support vector regression (SVR) (from Table 4) also reported better performance for the proposed model. Even though these models show relatively low error values, with MSE ranging from 4,243.16 to 4,265.65 (train) and RMSE between 65.14 and 65.31, their ability to capture temporal dependencies and non-linear relationships is limited. For instance, while RF achieves a slightly better R² (0.83) value, it still lags in handling the complex pattern compared to the proposed model. These traditional models fail to offer the forecasting accuracy required for handling fluctuations and capturing data variability as demonstrated by their higher errors and lower R² values. In contrast, the proposed model achieves enhanced accuracy and reliability, with reduced forecasting errors and better alignment with observed trends.

The DM test among the models is carried out with a null hypothesis posits that there is no significant difference between the forecasting performances of the two models, and an alternative hypothesis suggests a significant difference (Chen et al. 2014). Table 5 shows the test results comparing the forecasting performance of the proposed Multi-head Attention-based LSTM model with other models on the train and test datasets. The DM-MSE and DM-MAE values are the test statistics based on MSE and MAE respectively, which assess the variance and magnitude of the forecast errors. The p-value indicates the statistical significance of the test results, with a smaller p-value indicating stronger evidence against the null hypothesis. The p-values for different models help to identify whether the differences in their forecasting performance are statistically significant.

The DM test results summarized in the Table 5 show the following findings for the proposed model compared to Stacked LSTM:

• Train set: The DM-MSE value of −2.3145 with a p-value of 0.0206 rejects the null hypothesis at the 5% significance level, indicating a meaningful difference in MSE. However, the DM-MAE value for MAE is −1.8742 with a p-value of 0.0609, which does not reject the null hypothesis at the 5% significance level, suggesting borderline significance.

• Test set: The DM-MSE value of −2.5748 with a p-value of 0.0101 and the DM-MAE value of −2.0147 with a p-value of 0.0439 both reject the null hypothesis at the 5% significance level, confirming that the proposed model outperforms the Stacked LSTM in terms of forecasting accuracy on the test set.

Similar trends are observed when comparing the proposed model with GRU, Bi-LSTM, ANN, and other models. For instance, when compared with GRU, the null hypothesis for MSE on the train set is marginally rejected at the 5% significance level (DM-MSE = −1.9563, p = 0.0502), but it is definitively rejected for the test set (DM-MSE = −2.1203, p = 0.0340). In comparisons with ANN + PSO, the differences are highly significant for both metrics on both datasets, with DM-MSE values less than −3.0 and p-values under 0.01, clearly demonstrating the superiority of the proposed model. Furthermore, when compared with models like Vanilla LSTM and CFN, the results also show consistent improvements in forecasting performance, with significant DM test outcomes across most metrics. For example, against Vanilla LSTM, the DM-MSE for the test set is −2.7829 with a p-value of 0.0054, strongly rejecting the null hypothesis and highlighting the robustness of the proposed model. These findings substantiate the proposed model's effectiveness in handling water quality forecasting tasks by reducing errors and capturing variability more efficiently, as reflected by consistently lower MSE and MAE values, along with significant DM test result.

To address this, the model was refined by tuning hyperparameters, introducing regularization, and expanding the dataset to include more diverse and recent observations. These steps reduced the width of the confidence intervals, enhancing prediction robustness while maintaining accuracy. The remaining intervals, though still wide, serve as an essential cautionary measure, highlighting potential variability and underscoring the importance of short-term predictions for reliable decision-making. For instance, sudden environmental changes can lead to abrupt fluctuations in WQI, necessitating real-time monitoring alongside forecasts for practical applications.

Along with the positives, the model is prone to few issues. The model is limited by its high computational complexity and extended run-time. This poses a threat to the real-time implementation of the model, particularly in resource-constrained environments. Moreover, on occasion of sudden fluctuations in water quality data, the accuracy of the forecast drops abruptly, as observed by the drop in the R² score and an increase in error values. To address these limitations, our future work could focus on integrating external environmental factors, e.g., weather data or pollution sources, which may provide better context for sudden changes in water quality parameters. Additionally, employing more robust data pre-processing methods, including advanced techniques for outlier detection and null value dealing, could enhance the model's performance during sudden fluctuations but keeping its higher accuracy for long-term trends. Hybrid models that combine deep learning with traditional statistical methods, such as ARIMA, could also be explored.

The multi-head attention-based LSTM model can be adopted for diverse domain of environmental and water resource management. This sophisticated data-driven forecasting technique enhances the capabilities of policymakers, environmental agencies, and water resource managers by delivering reliable predictions essential for informed decisions in water quality control, pollution mitigation, and public health protection. Thus, decision-makers can use the forecasting to prioritize water treatment in regions with predicted pollution spikes, distribute resources more effectively during contamination incidents, or plan proactive measures to prevent the deterioration of the water quality. By proactively managing contamination incidents and ecological threats, the research aids in the conservation of aquatic ecosystems and the assurance of safe drinking water for communities. Through prompt interventions and strategic planning, this method aims to advance environmental sustainability and safeguard human health, aligning with broader pollution mitigation efforts and promoting the long-term sustainability of water resources (Georgescu et al. 2023). Also, the application of the model should be explored in other domains characterized by complex temporal dependencies in time-series data, such as weather forecasting, air quality prediction, and energy demand forecasting. By using the attention mechanism to capture crucial temporal patterns, the proposed model can provide valuable insights and predictive capabilities across diverse fields. Thus, the outcome of this study has substantial bearings on environmental science and water resource management, offering a data-driven autonomous solution for water quality prediction that elucidates long-term trends and enhances the predictive reliability.

In this study, a hybrid Multi-Head Attention-Based LSTM model is proposed for water quality prediction, and a detailed performance assessment is drawn using the Irish water quality dataset. The proposed model shows a 5–8% improvement in accuracy while forecasting WQI as compared to the existing ML and DL models. The use of multi-head attention enhances the ability of the model to capture complex temporal dependencies, resulting in improved prediction with reliability and robustness. The key contribution of this research is the development of a data-driven, automated system that can forecast water quality with greater precision, helping the decision-making process for environmental management. This model can aid in the pollution control, early detection of the water contamination, and overall management of the water resource for sustainable public health protection. However, the computational complexity of the model remains a challenge as training the model requires significant processing time and resources. This may pose challenges for improving the model performance with extended data. Future work should focus on this direction to define an efficient process for faster training and deployment of the system.

No funding was available for this research work.

We adhered to all ethical guidelines and obtained proper consent for the preparation of this manuscript.

All relevant data are available from an online repository or repositories: https://doi.org/10.6084/m9.figshare.28184252.v1.

The authors declare there is no conflict.