Using Hybrid Machine Learning to Predict Wastewater Effluent Quality and Ensure Treatment Plant Stability

Abstract

The accurate prediction of wastewater quality parameters is pivotal for evaluating the treatment stability of processes and for ensuring regulatory compliance in wastewater treatment plants. A singular machine learning model often faces challenges in fully capturing and extracting the complex nonlinear relationships inherent in multivariate time series data. To overcome this limitation, this study proposes a dual hybrid modeling framework that effectively integrates LSTM and XGBoost models, leveraging their complementary strengths. The first hybrid model refines the residues to utilize the information, whereas the second hybrid model enhances the input features by extracting temporal dependencies. A comparative analysis against three standalone models reveals that the proposed hybrid framework consistently outperforms them in both predictive accuracy and generalization ability across four key effluent indicators—chemical oxygen demand, ammonia nitrogen, total nitrogen, and total phosphorus. These results demonstrate that the proposed hybrid machine learning framework has great potential to be used to evaluate process stability in wastewater treatment plants, paving a way for smarter, more resilient, and more sustainable wastewater management, which will improve ecological integrity and regulatory compliance.

1. Introduction

Water is essential for sustaining life, maintaining ecosystems, and supporting economic activities [1]. Water quality is critical for sustainable development, public health, and environmental balance. However, the pollution of water resources caused by urbanization and industrialization has emerged as a major global challenge for environmental conservation and sustainability [2]. The direct discharge of industrial wastewater, agricultural runoff, and urban sewage has exacerbated water pollution, posing significant threats to human health and socio-economic development. Water pollutants encompass a wide range of substances, including organic pollutants, heavy metals, nutrients, pathogens, and microplastics. As water treatment technologies advance and the scale of sewage treatment expands, the complexity of water treatment processes has also increased. Wastewater treatment systems are highly nonlinear industrial control processes characterized by substantial time lags and variability [3]. The efficient operation of wastewater treatment plants (WWTPs) is therefore critical in safeguarding water resources and promoting environmental protection through the removal of various pollutants [4]. Key water quality indicators (WQIs) affecting WWTPs’ performance across different stages include chemical oxygen demand (COD), biological oxygen demand (BOD), total nitrogen (TN), total phosphorus (TP), ammonia nitrogen (NH₃-N), and influent flow (IF). For instance, COD represents organic pollution, while NH₃-N reflects eutrophication potential [5]. The accurate prediction of these indicators is essential to ensure the stability of effective wastewater treatment and provides a foundation for optimizing treatment programs. Additionally, it enables the early detection of potential problems, serving as a warning to prevent environmental hazards from non-compliant discharges [6].

Existing water quality prediction models primarily fall into two categories: mechanistic models based on biochemical reactions and data-driven machine learning models. Mechanistic models (such as the IWA ASM family [7,8,9]) employ mathematical frameworks grounded in the principles of wastewater treatment processes, incorporating multiple biochemical and chemical reactions to enable the real-time monitoring and prediction of water quality parameters [10]. Bolyard and Reinhart [11] developed a total nitrogen model based on biochemical reaction mechanisms to describe the relationship between TN and process variables. Similarly, El Shorbagy et al. [12] proposed an energy consumption prediction model rooted in nitrification reaction mechanisms, establishing correlations between energy consumption and variables such as dissolved oxygen and nitrate concentrations. While these models demonstrate utility, their prediction accuracy remains limited due to the complexity and variability of wastewater indicators and reactions [13]. With the development of artificial intelligence [14] in recent years, machine learning (ML) offers new approaches to better capture the complex relationships inherent in wastewater data [15]. Machine learning algorithms excel in uncovering complex nonlinear relationships between variables from real-world data and have achieved notable success in diverse fields such as weather forecasting [16], intelligent transportation [17], disease diagnosis [18], and financial risk control [19]. Over the past decade, machine learning models such as Artificial Neural Networks (ANNs) [20,21], support vector Machines (SVMs) [22,23], extreme learning machines (ELMs) [24,25], Random Forests (RFs) [26,27], and Reinforcement Learning [28,29,30] have demonstrated their potential in influent parameter prediction by effectively capturing nonlinear relationships and temporal dependencies.

However, most ML-based models are not time-dependent models. The rise in deep learning has further enhanced complex system modeling through its strengths in automatic feature extraction when modeling temporal dependencies and its recognition, particularly for large and intricate datasets, thereby significantly mitigating the limitations of traditional machine learning. In the context of wastewater quality prediction, deep learning models have demonstrated remarkable success. Gated Recurrent Units (GRUs), Long Short-Term Memory (LSTM), and Transformer models have shown promise in predicting water quality parameters [31,32,33]. Furthermore, a series of hybrid approaches that integrate signal decomposition techniques, hyperparameter optimization algorithms, and machine learning algorithms have been proposed to enhance prediction accuracy and stability. For instance, Zhu et al. [34] integrated an adaptive particle swarm optimization algorithm with least squares support vector machines (LSSVMs) to predict COD. Xie et al. [35] combined a temporal convolutional network with LSTM to simulate hourly TN, achieving a 33% accuracy improvement over single models. Yao et al. [36] developed a hybrid model leveraging multi-source spatiotemporal data for water quality prediction, classification, and regulation, achieving a high forecasting accuracy across six key water quality indicators. Similarly, Zhang et al. [37] proposed an integrated model that combined Empirical Mode Decomposition (EMD) with LSTM networks. Their approach demonstrated notable performance improvements in predicting critical water quality parameters when compared to conventional data-driven prediction algorithms. Chen et al. [38] developed a deep learning hybrid prediction model that combined Ensemble Empirical Mode Decomposition (EEMD) and a dynamic feature selection mechanism to predict the influent parameters of wastewater treatment plants. Yin et al. [33] developed a deep learning model that incorporated encoder–decoder LSTM networks, significantly improving the robustness of real-time water quality predictions, particularly during shock load events. More recently, Lv et al. [39] introduced a mechanistically enhanced hybrid model that combined mechanistic modeling with data-driven approaches.

Despite these advancements, two significant challenges persist in practical applications and warrant further investigation. One is that the lack of interpretability has always hindered the engineering applications of deep learning models, and the other is that the reasons for the combination of different models and strategies are still lacking in detailed comparison and analysis. To address this, this study enhances both prediction reliability and interpretability while offering practical applications for wastewater management by adopting the hybrid ML framework. Two types of hybrid models combining LSTM and eXtreme Gradient Boosting (XGBoost) [40] are proposed for wastewater quality prediction. The comparisons with singular machine learning techniques were drawn to verify the advantages of the proposed models. The findings inform the optimization of wastewater management strategies and advance sustainable water management, which will eventually benefit the health of the water environment. The necessary computations are performed through programming in Python (Version 3.13), utilizing essential libraries such as pandas, scikit-learn, numpy, SHAP, and matplotlib.

2. Materials and Methods

This section introduces all the data and methods involved in this work, including data collection, baseline models, interpretable machine learning methods, the proposed hybrid framework, the hyperparameter optimization technique, and the performance evaluation criteria.

2.1. Data Collection

The dataset for this study comes from a sewage treatment plant in Jiujiang, China. The dataset is a record of hourly influent water quality parameters from 00:00 on 12 December 2024, to 08:00 on 19 February 2025. It includes five input parameters: influent chemical oxygen demand (COD), influent ammonia nitrogen (NH₃-N), influent total phosphorus (TP), influent total nitrogen (TN), and influent flowrate (IF), as well as four output parameters: effluent COD, NH₃-N, TP, and TN, with 1665 samples for each parameter.

To train and test the model, data from 00:00 on 12 December 2024 to 13:00 on 29 January 2025 (the first 70%) were used as the training set, while data from 14:00 on 29 January 2025 to 08:00 on 19 February 2025 (the remaining 30%) were used as the test set. The influent and effluent parameters for the designated period are illustrated in Figure 1. Both input and output variables were normalized and preprocessed to address missing values before establishing different types of models.

2.2. Baseline Models

2.2.1. XGBoost

XGBoost builds upon the Gradient Boosting Decision Tree (GBDT) [41] ensemble learning method by enhancing both its loss function formulation and corresponding optimization process. The main objective of XGBoost is to minimize an objective function that combines two primary terms: the loss term (L) and the regularization term (Ω). Mathematically, the objective function (obj) can be formulated as follows:

$$obj={\sum }_{i=1}^{n}L\left(y_i,y_i^{\mathrm{*}}\right)+{\sum }_{j=1}^K\Omega \left(f_j\right),f_j\in f$$

(1)

During training, XGBoost employs additive learning, where new trees are added to minimize the loss function iteratively. This is different from singular models, which tend to have a pre-defined structure and are optimized in a Euclidean parameter space [42].

2.2.2. Long Short-Term Memory

LSTM [43] is an improved version of the Recurrent Neural Networks (RNNs) created by replacing conventional recurrent nodes with memory cells. As illustrated in Figure 2, each memory cell in LSTM contains three types of gates (i.e., input gate, forget gate, and output gate), implemented as fully connected layers with sigmoid activations to control data flow within the cell. In addition to the gating mechanisms, LSTM networks incorporate an input node, denoted as c_t, which serves as the candidate cell state for updating the memory cell. Unlike the gates that typically use the sigmoid activation function, the candidate state employs the hyperbolic tangent (tanh) activation function, mapping values to the range (−1, 1). This design enables the introduction of both positive and negative contributions to the cell’s internal state. The forget gate regulates the extent to which the previous cell state h_t−1 is retained; the input gate determines how much of the new candidate information c_t is incorporated into the current cell state; and the output gate determines how much of the updated internal state (c_t) contributes to the output (h_t) at the current timestep. These three gating structures work together to effectively control the selective storage and updating of memory content within the LSTM unit. The internal computational process of LSTM can be described using Equations (2)–(8).

$$f_t=\sigma \left(W_f·\left[h_{t-1},x_t\right]+b_f\right)$$

(2)

$$i_t=\sigma \left(W_f·\left[h_{t-1},x_t\right]+b_i\right)$$

(3)

$$c_t=\sigma \left(W_c·\left[h_{t-1},x_t\right]+b_c\right)$$

(4)

$${\tilde{c}}_t=\tanh \left(W_c·\left[h_{t-1},x_t\right]+b_c\right)$$

(5)

$$o_t=\sigma \left(W_o·\left[h_{t-1},x_t\right]+b_o\right)$$

(6)

$$c_t=f_t{c}_{t-1}+i_t{\tilde{c}}_t$$

(7)

$$h_t=o_t·\tanh \left(c_t\right)$$

(8)

where W_f, W_i, W_c, and W_o are weight matrices; b_f, b_i, b_c, and b_o are bias terms; and σ is the sigmoid activation function.

2.3. Interpretable Machine Learning

The XGBoost algorithm offers three importance metrics based on gain, coverage, and weight. In recent years, a model-agnostic method called SHapley Additive exPlanations (SHAP) has become the preferred approach to interpret XGBoost outputs. The following discussion briefly introduces the use of SHAP with XGBoost and the attention mechanism to explain the LSTM model.

2.3.1. SHAP

SHAP was initially proposed by Shapley [44] and computes Shapley values from cooperative game theory; its values explain the contribution of each feature to the prediction. More recently SHAP values have been used as a tool [45] for interpreting ensemble tree models. The specific mathematical expressions of SHAP are shown in Equations (9)–(11) [46]:

$$f\left(x\right)=g\left(z\right)+{\phi }_0+{\sum }_{i=1}^M{\phi }_i{z}_i$$

(9)

where z = {0,1}^M, z_i = 1 represents a present feature and z_i = 0 represents an absent feature; M is the total number of input feature variables; φ₀ = f (h_x(0)); and φ_i is the SHAP value:

$${\phi }_i={\sum }_{S\subseteq F\backslash \left\{i\right\}}{\displaystyle \frac{\left|S\right|!\left(M-\left|S\right|!-1\right)!}{\left|M\right|!}}\left(f_x\left(S\cup \left\{i\right\}\right)-f_x\left(S\right)\right)$$

(10)

where F is the set of nonzero inputs in z, S is a subset of the features incorporated into the model, and f_x(S) is the prediction for feature values in set S.

2.3.2. Attention Mechanism

The attention mechanism [47] was inspired by the human visual system’s ability to selectively focus on important regions of a scene. Unlike traditional neural network architectures that apply fixed weights to all inputs, the attention mechanism computes a probability distribution over input elements to amplify the impact of the most critical features on the output. To implement the attention mechanism, the raw input data are first transformed into a set of ⟨Key, Value⟩ pairs. The attention module then computes a similarity score between each key and a given Query, which serves as a measure of relevance. These similarity scores are subsequently used to produce a weighted combination of the corresponding values, enabling the model to selectively focus on the most pertinent information in the input sequence. Denoting Query, Key, and Value, respectively, as Q, K, and V, the formulas for calculating weight coefficient W and output a can be expressed as follows:

$$W=softmax\left(Q{K}^T\right),K^{T}isthetransposeofK$$

(11)

$$a=attention\left(Q,K,V\right)=W·V$$

(12)

2.4. Combination of LSTM and XGBoost

LSTM networks excel in capturing both long- and short-term temporal dependencies within time series data. Meanwhile, XGBoost algorithms effectively model nonlinear relationships and mitigate overfitting through regularized optimization. However, neither architecture alone achieves an optimal balance between predictive accuracy and computational efficiency. This underscores the necessity to develop a hybrid framework that synergistically integrates LSTM and XGBoost, thereby establishing an innovative integration methodology to enhance both prediction precision and operational stability in water quality prediction. This paper proposes two approaches combining the LSTM model with an attention mechanism (i.e., Attention–LSTM model) and the XGBoost model. Each integration approach aims to optimize the application of water quality management in wastewater treatment plants. Figure 3 illustrates the two distinct types of integration processes.

The first approach (Attention–LSTM₁–XGBoost in the left-hand red-dashed box in Figure 3) leverages the XGBoost model to capture the component of water quality variation not explained by the LSTM model. Specifically, the Attention–LSTM generates the initial predicted values y_{Attention-LSTM} and the residual ε_{Attention-LSTM} between the predicted values and the actual values is computed. Then the XGBoost model is trained on these residuals to obtain the predicted residuals ε_XGBoost. The final prediction is obtained by summing y_{Attention-LSTM} and ε_XGBoost. This hybrid configuration is referred to as the Attention–LSTM₁–XGBoost model.

The second approach (Attention–LSTM₂–XGBoost in the right-hand red-dashed box in Figure 3) utilizes the temporal features extracted by Attention–LSTM to assist the XBoost model in water quality prediction. First, the Attention–LSTM model is trained on time series data to obtain predicted value y_{Attention-LSTM}. The obtained prediction is concatenated with the original independent variables to form a new input. The XGBoost model is trained on this combined input to generate the final predicted value. This hybrid model is referred to as the Attention–LSTM₂–XGBoost model.

2.5. Hyperparameter Optimization

Grid search performs exhaustive optimization by discretizing all possible parameter combinations into a grid within the parameter space and evaluating each combination. For each parameter combination, cross-validation is used to calculate the error, and the combination with the minimum error is selected as the global optimal solution. Grid search is effective in finding the optimal solution from multiple parameters. Therefore, in this study, we employ grid search to determine the optimal hyperparameters for both the XGBoost and LSTM models. For the XGBoost model, the search space included n_estimators, which determines the number of weak learners; max_depth, which limits tree depth to prevent overfitting; learning_rate, which adjusts the contribution of each tree; subsample, which controls the proportion of data sampled for training each tree; and child_weight, which controls the weight of splits at child nodes. For the LSTM model, grid search optimizes epochs representing the number of full training passes; batch size defines the number of samples used in each training step; units control the number of LSTM units and thus the model’s expressive capacity; and timestep refers to the number of timesteps in each input sequence and thus influences the model’s temporal context.

Figure 4 illustrates the example of a grid search procedure applied to the model’s parameters. The selection ranges of all hyperparameters are listed in

Table 1. Hyperparameters for different prediction models.

Model	Range	COD	NH3-N	TP	TN
XGBoost	n_estimators = [50, 500]	n_estimators = 482	n_estimators = 385	n_estimators = 442	n_estimators = 352
	max_depth = [2, 7]	max_depth = 3	max_depth = 5	max_depth = 3	max_depth = 3
	learning_rate = [0.01, 0.3]	learning_rate = 0.3	learning_rate = 0.16	learning_rate = 0.23	learning_rate = 0.24
	subsample = [0.5, 1]	subsample = 0.5	subsample = 0.65	subsample = 0.38	subsample = 0.6
	child_weight = [0.5, 1]	child_weight = 0.9	child_weight = 0.85	child_weight = 0.84	child_weight = 0.5
LSTM	epoch = [50, 350]	epoch = 200	epoch = 300	epoch = 200	epoch = 250
	batch size = [20, 80]	batch size = 50	batch size = 45	batch size = 65	batch size = 60
	units = [120, 250]	units = 190	units = 185	units = 210	units = 205
	timesteps = [3, 7]	timesteps = 5	timesteps = 5	timesteps = 5	timesteps = 4
Attention–LSTM	epoch = [50, 350]	epoch = 200	epoch = 250	epoch = 200	epoch = 200
	batch size = [20, 80]	batch size = 45	batch size = 50	batch size = 40	batch size = 55
	units = [120, 250]	units = 210	units = 210	units = 225	units = 190
	timesteps = [3, 7]	timesteps = 4	timesteps = 4	timesteps = 4	timesteps = 3
Attention–LSTM2–XGBoost	n_estimators = [50, 500]	n_estimators = 376	n_estimators = 475	n_estimators = 435	n_estimators = 386
	max_depth = [2, 7]	max_depth = 4	max_depth = 4	max_depth = 4	max_depth = 4
	learning_rate = [0.01, 0.3]	learning_rate = 0.29	learning_rate = 0.22	learning_rate = 0.24	learning_rate = 0.2
	subsample = [0.5, 1]	subsample = 0.3	subsample = 0.48	subsample = 0.7	subsample = 0.5
	child_weight = [0.5, 1]	child_weight = 0.7	child_weight = 0.86	child_weight = 0.8	child_weight = 1
Attention–LSTM1–XGBoost	n_estimators = [50, 500]	n_estimators = 400	n_estimators = 444	n_estimators = 500	n_estimators = 435
	max_depth = [2, 7]	max_depth = 5	max_depth = 4	max_depth = 4	max_depth = 5
	learning_rate = [0.01, 0.3]	learning_rate = 0.25	learning_rate = 0.29	learning_rate = 0.26	learning_rate = 0.26
	subsample = [0.5, 1]	subsample = 0.4	subsample = 0.56	subsample = 0.82	subsample = 0.35
	child_weight = [0.5, 1]	child_weight = 0.8	child_weight = 0.76	child_weight = 0.9	child_weight = 0.97

. Hyperparameter selection ranges for different models were determined based on the characteristics of the time series dataset collected in this work.

3. Results and Discussion

3.1. Statistical Analysis

Before training the model, a statistical analysis of the dataset was performed, calculating the correlation between all measurements at the wastewater treatment plant, including both input and output variables. Figure 5 presents the heat map of correlation coefficients in the dataset. As can be observed, there is a notable correlation between NH₃-H and TN among the input variables, and COD and IF demonstrate statistically weak correlations with all other parameters in the feature space. In contrast, the correlation among the output variables is relatively weak.

3.2. Performance of Singular Models

Figure 6 and Figure 7 show the predicted effluent parameters obtained from the trained LSTM and XGBoost models compared to the measured data. The black line indicates the actual measurements, while the red line shows the predicted values. The prediction performance evaluation metrics of the three singular machine learning methods are listed in

Table 2. Performance indicators of effluent COD, NH₃-N, TN, and TP using singular models.

Models	Effluent Parameters	Datasets	Metrics
			R2	RMSE	MAE
LSTM	COD	Training	0.896	0.856	0.705
		Test	0.867	0.968	0.732
	NH3-N	Training	0.873	0.169	0.069
		Test	0.810	0.122	0.047
	TN	Training	0.943	0.543	0.360
		Test	0.869	0.724	0.567
	TP	Training	0.889	0.053	0.046
		Test	0.859	0.049	0.030
XGBoost	COD	Training	0.893	0.871	0.725
		Test	0.841	0.404	0.292
	NH3-N	Training	0.855	0.181	0.082
		Test	0.827	0.117	0.056
	TN	Training	0.922	0.638	0.457
		Test	0.877	0.702	0.504
	TP	Training	0.879	0.056	0.041
		Test	0.840	0.053	0.036
At-LSTM	COD	Training	0.948	0.608	0.431
		Test	0.887	0.339	0.244
	NH3-N	Training	0.923	0.132	0.040
		Test	0.865	0.103	0.038
	TN	Training	0.941	0.554	0.350
		Test	0.899	0.636	0.416
	TP	Training	0.926	0.044	0.030
		Test	0.850	0.040	0.050

. The results indicate that the LSTM model demonstrates a marginally better performance compared to the XGBoost model on the training set. Specifically, the R² values for LSTM are 0.896 and 0.873 for COD and NH₃-N, respectively, while the R² values for XGBoost are 0.893 and 0.855 for COD and NH₃-N, respectively. Additionally, for the three effluent parameters—COD, NH₃-N, and TP—the model fitting results on the training set are consistently below the measured values. On the test set, the predictive performance of the LSTM and XGBoost models is largely comparable. The R² for the two models across different effluent parameters are all above 0.800. However, it is important to note that the prediction performance should not be solely evaluated based on the value of R². For instance, while the LSTM model yields a R² for the COD parameter compared to the XGBoost model, its RMSE and MAE values are also higher. Furthermore, the Attention–LSTM model, which incorporates the attention mechanism, outperforms the previous two models both on the training and test sets, demonstrating its superior generalization capability (Figure 8). In short, the LSTM, XGBoost, and Attention–LSTM models are capable of effectively approximating the irregular periodicity of effluent parameters during the measurement period.

3.3. Interpretability Analysis

For the four effluent parameters, Figure 9 presents the feature importance ranking of the five input variables, derived using three interpretability analysis methods. The vertical axis represents the normalized score. These methods include the weight method of XGBoost, the SHAP method, and the attention mechanism method. Overall, the SHAP and attention mechanism methods exhibit a high degree of consistency in the feature importance ranking. Specifically, for effluent TN prediction, both methods rank influent TN, TP, and NH₃-N in the top three, which may be attributed to NH₃-N being an essential component of TN. For effluent NH₃-N prediction, both methods prioritize influent NH₃-N as the most important feature, while ranking influent flow rate as the least significant. These results align with conclusions drawn in existing research [48,49]. In contrast, the feature importance ranking generated by the weight method of XGBoost demonstrates greater variability across the four effluent parameters. For example, influent COD, NH₃-N, and IF are ranked among the top three for effluent TN prediction, and influent COD is ranked first for effluent NH₃-N prediction. This may be due to the fact that when there is multicollinearity among features, XGBoost tends to prioritize splitting one feature, potentially overlooking other correlated features. It is also worth noting that even the lower-ranked features in Figure 9 may demonstrate scores that are not substantially lower than those of higher-ranked features. For example, NH₃-N shows a certain degree of contribution to the predicted COD time series, as indicated by SHAP values. This may be attributed to the indirect relationships between the volume of food scraps, excreta, and washing wastewater generated by residents [50]. Additionally, NH₃-N in water can undergo transformation through microbial processes under specific conditions. Under aerobic conditions, nitrification consumes oxygen, which may result in a decrease in available oxygen, thereby potentially affecting the oxidation of organic matter.

Figure 10 illustrates the significant interactions between input parameters in the prediction of COD using XGBoost. On the right side of each plot, a color bar is provided to represent the interacting input parameters, with blue and red corresponding to lower and higher values, respectively. Taking Figure 10a as an example, the partial dependence relationship between the influent TP and TN is depicted. The figure demonstrates that the effect of influent TN on effluent COD is dependent on the TN range. Notably, when the influent TP concentration is high, the negative correlation between influent TN and effluent COD becomes more pronounced. This observed trend highlights the intricate interaction between influent TN and influent TP.

3.4. Performance of Hybrid Models

Figure 11 and Figure 12 present the results from the two hybrid models, while

Table 3. Performance indicators of effluent COD, NH₃-N, TN, and TP using hybrid models.

Models	Effluent Parameters	Datasets	Metrics
			R2	RMSE	MAE
Attention–LSTM1–XGBoost	COD	Training	0.992	0.234	0.162
		Test	0.972	0.168	0.138
	NH3-N	Training	0.995	0.033	0.013
		Test	0.982	0.037	0.015
	TN	Training	0.980	0.311	0.214
		Test	0.957	0.417	0.326
	TP	Training	0.988	0.018	0.015
		Test	0.952	0.028	0.024
Attention–LSTM2–XGBoost	COD	Training	0.978	0.399	0.283
		Test	0.952	0.223	0.161
	NH3-N	Training	0.992	0.043	0.013
		Test	0.986	0.034	0.012
	TN	Training	0.969	0.404	0.255
		Test	0.946	0.464	0.304
	TP	Training	0.953	0.035	0.024
		Test	0.914	0.038	0.030

provides the specific values of the performance metrics. It is evident that, regardless of the training or test set used, the Attention–LSTM₁–XGBoost and Attention–LSTM₂–XGBoost models outperform the previous three models across all four effluent parameters. A notable observation is that neither of the two hybrid models exhibits a consistent underestimation of certain parameters throughout the entire monitoring period, as was observed in the previous case with a singular model (e.g., as seen in Figure 6). Throughout the entire period, both models displayed relatively large errors during the TP prediction phase (around the 900th hour), as seen in Figure 11d and Figure 12d. However, the error values remained substantially lower compared to those of the singular model. Furthermore, neither of the two hybrid models exhibited the prediction delay typically observed in traditional autoregressive models. The two hybrid models exhibit an outstanding performance in forecasting effluent quality values. In the case of the training dataset, the minimum R² value of 0.953 and the maximum R² value of 0.995 signify a good alignment with the actual values. For the test set, a more detailed quantitative analysis reveals that the Attention–LSTM₁–XGBoost model achieves R² values of 0.972, 0.957, and 0.952 for predicting COD, TN, and TP, respectively. It demonstrates increases in R² of 2.1%, 1.2%, and 4.2% for these three metrics compared to the Attention–LSTM₂–XGBoost model. However, the Attention-LSTM₂-XGBoost model demonstrates slightly superior predictive performance for NH₃-N compared to the Attention–LSTM1–XGBoost model, with an increase in R² of 0.4%, along with reductions in RMSE and MAE of 8.8% and 25.0%, respectively. In summary, LSTM networks are particularly adept at capturing both long-term and short-term temporal dependencies in time series data. In contrast, XGBoost algorithms are highly effective at modeling nonlinear relationships and reducing overfitting through regularization-based optimization techniques. Attention–LSTM₁–XGBoost utilizes the information by adjusting the residuals, while Attention–LSTM₂–XGBoost enhances the input features through the extraction of temporal dependencies. These two hybrid models leverage the respective strengths of the XGBoost and LSTM models from distinct perspectives, resulting in a substantial improvement in both predictive performance and generalization capability.

4. Conclusions

This paper presents two hybrid models that integrate LSTM and XGBoost to predict four effluent parameters in wastewater treatment plants, and compares their performance against three singular machine learning models. The key findings are summarized as follows:

(1)

The predictive performance of the individual LSTM and XGBoost models is relatively comparable, with R² values for the four effluent parameters ranging from 0.810 to 0.877. However, the predictions for certain parameters exhibit a consistent underestimation of the actual values throughout the monitoring period. In contrast, the LSTM model incorporating the attention mechanism demonstrates improved predictive accuracy, with R² values ranging from 0.850 to 0.899.

(2)

Interpretability analysis reveals that the feature importance rankings from SHAP and the attention mechanism are largely consistent. For a given effluent parameter, in addition to the corresponding influent parameter, other influent indicators also influence the outcome through interactions. However, the rankings from the weight method of XGBoost differ from those of SHAP and the attention mechanism. This discrepancy may result from XGBoost’s tendency to prioritize a single feature for splitting, disregarding other correlated features.

(3)

The prediction accuracy of the two hybrid models is markedly superior to that of the singular model. The first model utilizes the information by adjusting the residuals, achieving R² values ranging from 0.952 to 0.982. The second model enhances the input features by capturing temporal dependencies, resulting in R² values between 0.914 and 0.986. The hybrid framework strategically combines the strengths of XGBoost and LSTM models from complementary perspectives, offering a robust solution for treatment process stability evaluation and real-time effluent quality monitoring. This integration has the potential to enable proactive adjustments in wastewater treatment processes, thereby optimizing operational efficiency and improving system responsiveness and treatment system resilience.

The hydraulic retention time of the sewage biochemical treatment process is usually several hours or may even be more than ten hours. For operational parameter adjustment, using predictions based on the influent quality and flow rate as the feedforward parameters, and hence predicting the future effluent quality, has better sensitivity than methods using the biochemical process parameters or effluent quality feedback. Therefore, based on the effluent quality in the next few hours, operators of wastewater treatment plants can adjust the operational parameters on their initiative to ensure effluent quality is stable and up to regulated discharge standards.

The hybrid model framework has great potential to evaluate process stability in wastewater treatment plants, paving the way for smarter, more resilient, and more sustainable wastewater management with predictive control and real-time process monitoring, which will improve ecological integrity and regulatory compliance. WWTPs are transformed from reactive to predictive sentinels of public health and environmental integrity. By ensuring effluent meets stringent quality standards, the methods mitigate downstream pollution, protect ecosystems, conserve resources, and reduce human exposure to contaminants. As climate change intensifies water stress and regulatory pressures, smart effluent management will become indispensable for sustainable water cycles.

In the future, according to the different working conditions of a wastewater treatment plant, the mechanistic model and data-driven model can be combined to evaluate the response-ability of the biochemical reaction process facing impact loads, and can meanwhile achieve the adjustment and optimization of operating parameters. Future advancements in edge computing and digital twins will also further amplify these benefits, embedding wastewater treatment into the One Health framework.