Sustainable Sewage Treatment Prediction Using Integrated KAN-LSTM with Multi-Head Attention

Abstract

The accurate prediction of sewage treatment indicators is crucial for optimizing management and supporting sustainable water use. This study proposes the KAN-LSTM model, a hybrid deep learning model combining Long short-term memory (LSTM) networks, Kolmogorov-Arnold Network (KAN) layers, and multi-head attention. The model effectively captures complex temporal dynamics and nonlinear relationships in sewage data, outperforming conventional methods. We applied correlation analysis with time-lag consideration to select key indicators. The KAN-LSTM model then processes them through LSTM layers for sequential dependencies, KAN layers for enhanced nonlinear modeling via learnable B-spline transformations, and multi-head attention for dynamic weighting of temporal features. This combination handles short-term patterns and long-range dependencies effectively. Experiments showed the model’s superior performance, achieving 95.13% R-squared score for FOss (final sedimentation basin outflow suspended solid, one indicator of our research predictions)and significantly improving prediction accuracy. These advancements in intelligent sewage treatment prediction modeling not only enhance water sustainability but also demonstrate the transformative potential of hybrid deep learning approaches. This methodology could be extended to optimize predictive tasks in sustainable aquaponic systems and other smart aquaculture applications.

1. Introduction

Accurate time-series prediction has emerged as a critical tool for optimizing sewage treatment systems [1], enabling both proactive process control and sustainable resource management. While deep learning approaches show considerable promise in this domain, current methods frequently fail to address two fundamental challenges inherent to sewage treatment data: the complex nonlinear relationships between treatment indicators [2] and the optimal selection of predictive features from highly correlated temporal measurements [3].

Traditional prediction methodologies, including Autoregressive integrated moving average (ARIMA) [4] and Support Vector Regression (SVR) [5], exhibit well-documented limitations when applied to sewage treatment data. These approaches typically demonstrate poor performance when handling multivariate, non-stationary time series [6], struggle to capture long-range dependencies in treatment processes, and show sensitivity to measurement noise and missing data [7]. While recent advances in deep learning have offered potential solutions, these too present significant challenges. Long short-term memory (LSTM) networks [8], despite their effectiveness in sequence modeling, often encounter difficulties with gradient flow in long sequences [9]. We acknowledge that residual connections can mitigate LSTM’s gradient issues [10]. However, for sewage treatment forecasting with moderate sequence lengths, our tests showed that residual LSTM provided limited accuracy gains while complicating temporal interpretation—a key requirement for operational use. This trade-off motivated our baseline LSTM design, though residual variants may benefit longer-sequence extensions. Similarly, Transformer architectures, while excelling at capturing long-distance dependencies, frequently fail to recognize local temporal patterns that prove critical for accurate treatment monitoring. Importantly, neither approach adequately addresses the systematic feature selection required for handling correlated sewage indicators [11].

Our study makes two distinct yet complementary contributions to address these challenges. First, we develop a rigorous time-lag correlation analysis [12] framework specifically designed for identifying the most predictive sewage treatment indicators. This preprocessing methodology quantitatively evaluates delayed relationships between parameters (such as the 3-day lag observed between pH and FOss levels). It selects features based on statistically significant correlations, while deliberately avoiding any modification of the subsequent modeling approach. Second, we introduce an innovative KAN-LSTM hybrid architecture that effectively combines the sequential pattern recognition capabilities of LSTM layers with the complex nonlinear modeling capacity of Kolmogorov–Arnold network (KAN) [13] transformations through learnable B-splines, further enhanced by multi-head attention mechanisms that dynamically weigh important temporal features [14].

This deliberate separation of feature selection from modeling architecture yields several significant advantages. The framework provides enhanced interpretability through explicit indicator screening while achieving superior predictive accuracy via its specialized architecture and maintaining practical deployment advantages for treatment plant operations. Our experimental results demonstrate consistent improvements over conventional methods without compromising computational efficiency [15].

The remainder of this paper is systematically organized: Section 2 introduces our methodology, beginning with the time-lag correlation analysis for feature selection, followed by detailed architectural explanations of the integrated KAN-LSTM model with multi-head attention. Section 3 presents comprehensive experimental results, including comparative analyses with state-of-the-art models and ablation studies validating each component’s contribution. Section 4 discusses the practical implications of our findings, acknowledges limitations, and outlines future research directions for sustainable sewage management. Finally, we conclude with the broader impacts of this work.

2. Methods

With the exponential growth of time-series data across financial forecasting, environmental monitoring, and industrial operations, we face increasing demands for advanced forecasting models that can effectively capture linear and nonlinear patterns while remaining computationally efficient [16]. Current deep learning approaches often struggle with the complex characteristics of time-series data [17], particularly when confronting non-stationary patterns and dependencies that span longer periods [18]. To address these challenges, our research introduces a time-lag algorithm utilizing Pearson correlation matrices [19]. This time-lag screening method helps to select more meaningful indicators for prediction, and the timeliness of sewage treatment is more consistent with the actual process. Building on this foundation, we developed an integrated KAN-LSTM model, which combines the sequential learning capabilities of LSTM networks with the expressive power of KAN and the selective focus of multi-head attention mechanisms. We tested our hybrid model on a practical real-world challenge: predicting critical parameters in sewage treatment processes. This application domain is an ideal testing ground due to its complex, multifaceted nature and the significant operational benefits that accurate predictions can provide for treatment facility management and environmental compliance.

2.1. Time-Series Lag Analysis of Indicators in Sewage Treatment

This section examines the interrelations among key indicators in sewage treatment and develops predictive models for forecasting future trends of these indicators, focusing on the importance of time-series lag. We focus on pH, final sedimentation basin outflow (FOss), treated water transparency (TWT), and treated water total nitrogen (TWTN). We analyze extensive data and establish accurate predictive models with a deep learning network using advanced statistical analysis and deep learning algorithms. Our results reveal the intrinsic connections between these indicators and confirm the high accuracy of our model in predictions [11].

Data elements are often standardized internationally because the sewage treatment system and equipment are based on the activated sludge method, which is an international standard. Therefore, we perform a correlation analysis of the indicators on the sewage treatment data to select the four indicators with the strongest correlations. To address the issue of timeliness in the sewage treatment process, we adopt a different approach from the traditional method of comparing correlations. In preliminary experiments, we found that the correlation between the data is the highest when the lag is 3 days, after testing lags from 1 to 5 days. Therefore, we calculate the correlations based on a 3-day lag between the data points. The specific formula is shown in Equation (1). In this study, k is three. This means analyzing the correlation of time-series data with a 3-day difference. The number of lag days depends on the size of the sewage treatment area, and this experiment uses Fukugawa city:

$$\begin{array}{c}\hfill r_{\mathrm{lag}}^k={\displaystyle \frac{{\sum }_{i=k+1}^n(X_i-\overline{X})(Y_{i-k}-\overline{Y_{\mathrm{lag}}})}{\sqrt{{\sum }_{i=k+1}^n{(X_i-\overline{X})}^2{\sum }_{i=k+1}^n{(Y_i-\overline{Y_{\mathrm{lag}}})}^2}}},\end{array}$$

(1)

where $r_{\mathrm{lag}}^k$ is the Pearson correlation coefficient with a lag of k. $X_i$ and $Y_{i-k}$ are the individual data points for the two variables, with the Y variable data lagged by k periods. $\overline{Y_{\mathrm{lag}}}$ is the mean of the Y data points after applying a lag of k periods. n is the total number of data points available after applying the lag. The summation starts from $i=k+1$ because the first k data points of the Y variable do not have corresponding X values due to the lag. In the correlation matrix of Figure 1, the rows and columns correspond to a variety of metrics that are essential in sewage treatment processes, including IF-BOD (inflow water biochemical oxygen demand), MLSS (mixed liquor suspended solids), and DO (dissolved oxygen), among others. The correlation coefficient displayed by the heat map identifies pH, FOss, TWT, and TWTN as exhibiting the most significant relationships (blue frames). Using the cross-correlation with a lag of 3 days, we only display data in the figure with absolute values of correlations greater than 0.5, including positive and negative correlations. This method explores the relationship between sewage indicators before and after treatment. This analytical approach allows us to examine the interconnections between indicators from various perspectives, uncovering a broader spectrum of relationships and deepening our understanding of the dynamics within sewage treatment systems. The threshold used in this preliminary experiment is greater than 0.5 and more stringent than in related research [20]. Compared to their threshold of 0.3, 0.5 can find stronger correlations. The results are more convincing.

Figure 2 illustrates the relationships among pH, FOss, TWT, and TWTN using pink scatter plots and blue linear regression lines. Diagonal graphs display the distribution of each variable, while off-diagonal graphs depict relationships between variables through data points and linear regression fits. These visualizations allow us to observe how the linear regression model captures the relationships between each pair of variables, providing an intuitive display of the relationships and model fit. Due to computing resource limitations, predictions were made by selecting the four most correlated indicators. Figure 2 verifies that the selected four indicators are reasonable [11].

2.2. Integrated KAN-LSTM Model

2.2.1. Time-Series Prediction Framework

The proposed hybrid model processes 10 historical time steps (look_back = 10) to predict subsequent indicator values $(\mathrm{T}+1)$, addressing two key sewage treatment challenges:

(1)

Long-term dependency capture via LSTM’s gated memory mechanism.

(2)

Dynamic feature weighting through multi-head attention.

This design (Figure 3) serves as the computational backbone for subsequent KAN-enhanced transformations (Section 2.2.2), with empirical results in Section 3 demonstrating their superiority over only LSTM or attention-only architectures for sewage data forecasting. The model structure is shown in Figure 3.

Figure 3 illustrates the architecture of the proposed Integrated KAN-LSTM Model, designed for accurate prediction of key sewage treatment indicators—pH, FOss, TWT, and TWTN. The process begins on the left side, where the selected indicators, determined through a 3-day time-lagged Pearson correlation analysis, are input as multivariate time series features. The light blue arrows indicate the data flow direction throughout the model.

First, the input sequence passes through the LSTM layer (depicted in light blue), which captures temporal dependencies using a gated memory mechanism. Next, the LSTM output is forwarded to the KAN layer (green block), which applies nonlinear transformations using learnable B-spline basis functions. Unlike traditional neural networks that use fixed activation functions on nodes, KAN assigns learnable functions on edges, improving the model’s capacity to capture complex nonlinearities in sewage data. Subsequently, the transformed features are passed into the Multi-Head Attention mechanism (orange block), which computes attention weights across different time steps. Each attention head processes the Query (Q), Key (K), and Value (V) vectors, calculated via linear transformations, and applies scaled dot-product attention. The outputs from all heads are concatenated and linearly transformed, allowing the model to dynamically weight temporal features. Finally, the aggregated features are passed to a fully connected output layer, which produces the final predictions for the next time step. A comparative graph in the top right corner visualizes the alignment between predicted and true values, demonstrating the model’s high forecasting accuracy. Overall, this integrated framework leverages the strengths of LSTM in sequence modeling, the nonlinear expressiveness of KAN, and the dynamic temporal weighting capability of attention mechanisms, all while maintaining interpretability and predictive performance.

2.2.2. KAN Layer (Kolmogorov–Arnold Network Layer)

Traditional neural networks (MLPs) [21] fixed activation functions (e.g., ReLU, tanh) at nodes, while KAN innovates by employing learnable activation functions on edges. Each weight parameter in KAN is replaced by a univariate function (typically B-splines [22]), enabling dynamic adaptation to the sewage treatment data’s complex nonlinearities.The KAN model excels in capturing complex nonlinear relationships in data while maintaining high interpretability. By leveraging B-spline functions for nonlinear transformations, it enhances expressiveness when handling highly nonlinear data. Additionally, the KAN model has relatively low computational complexity [13], allowing it to achieve high prediction accuracy with fewer parameters, making it well suited for efficient time-series modeling tasks. Unlike traditional neural networks that use fixed activation functions, KAN uses a learnable activation function at the network’s edge. This design allows each weight parameter in KAN to be replaced by a single variable function, which is usually parameterized in the form of a spline function, providing extremely high flexibility and the ability to simulate complex functions with fewer parameters, enhancing the interpretability of the model.

The KAN layer is a nonlinear transformation layer based on B-spline functions, capable of capturing complex nonlinear relationships in data. The mathematical representation of the KAN layer is as follows:

$$y=W_{\mathrm{base}}x+\sum _{i=1}^G{B}_i\left(x\right)·W_{\mathrm{spline},i},$$

(2)

where $W_{\mathrm{base}}$ is the base weight matrix, $B_i\left(x\right)$ represents the B-spline basis function, $W_{\mathrm{spline},i}$ is the spline weight, and G is the number of grid points. The B-spline basis function is defined as follows:

$$B_i\left(x\right)=max\left(0,1-\left|x-g_i\right|\right),$$

(3)

where $g_i$ is the grid points uniformly distributed in the $[0,1]$ interval. The KAN layer computes both the linear output from the base weight and the nonlinear output from the B-spline transformation, which are then summed to obtain the final output.

2.2.3. LSTM

In this study, LSTM plays a central role in modeling time-series data by capturing long-term dependencies in the sewage treatment dataset. LSTM effectively mitigates the vanishing gradient problem as a recurrent neural network variant, making it well suited for learning sequential patterns. The LSTM layer extracts temporal features from the input sequences [23], which are subsequently refined through a KAN layer to enhance nonlinear representation. Additionally, a multi-head attention mechanism is applied to further emphasize key temporal dependencies, improving predictive accuracy. Integrated within a Transformer-based framework, LSTM serves as the foundational feature extraction component, providing essential temporal information for subsequent processing. Its ability to model sequential dependencies ensures that the predictive model effectively learns patterns in the time-series data, contributing to robust and interpretable forecasting performance.

LSTM consists of four main components: forget gate, input gate, cell candidate, and output gate. The forward computation is as follows:

(1)

Forget gate: Determines how much past information should be forgotten:

$$f_t=\sigma \left(W_f{x}_t+U_f{h}_{t-1}+b_f\right).$$

(4)

where $f_t$ is the forget gate output (ranges between 0 and 1). $x_t$ is the input at the current time step. $h_{t-1}$ is the hidden state from the previous time step. $W_f,U_f$, and $b_f$ are learnable parameters. $\sigma$ represents the sigmoid activation function.

(2)

Input gate: Determines the importance of the current input:

$$i_t=\sigma \left(W_i{x}_t+U_i{h}_{t-1}+b_i\right),$$

(5)

$${\tilde{C}}_t=tanh\left(W_c{x}_t+U_c{h}_{t-1}+b_c\right).$$

(6)

where $i_t$ is the input gate output (ranges between 0 and 1). ${\tilde{C}}_t$ is the candidate cell state, with tanh limiting its range to $[-1,1]$.

(3)

Cell-state update: Combines information from the forget and input gates:

$$C_t=f_t\odot C_{t-1}+i_t\odot {\tilde{C}}_t.$$

(7)

where $C_t$ is the updated cell state at time t. ⊙ denotes the Hadamard (element-wise) product.

(4)

Output gate: Determines which information should be passed to the hidden state:

$$o_t=\sigma \left(W_o{x}_t+U_o{h}_{t-1}+b_o\right),$$

(8)

$$h_t=o_t\odot tanh\left(C_t\right).$$

(9)

where $o_t$ is the output gate activation (ranges between 0 and 1). $h_t$ is the hidden state at the current time step. These equations collectively enable LSTM to capture long-term dependencies in sequential data while effectively mitigating the vanishing gradient problem.

2.2.4. Multi-Head Attention Mechanism

As an extension of the attention mechanism, the multi-head attention mechanism further enhances the model’s attention capture ability by processing multiple attention distributions in parallel [24], improves the model’s expression ability and learning efficiency, and becomes one of the core components of the Transformer architecture. After LSTM processes the input data, the KAN layer extracts nonlinear features, and then, the multi-head attention mechanism allocates different attention weights globally, thereby resulting in the following:

(1)

Enhancing long-term dependencies: Compared with LSTM, the attention mechanism can capture long-range dependencies and avoid the gradient vanishing problem.

(2)

Different feature weight allocation: Multiple attention heads can focus on different time steps, improving the model’s representation ability.

(3)

Improving nonlinear modeling capabilities: Combined with the KAN layer, the multi-head attention mechanism further enhances nonlinear mapping and improves prediction capabilities.

Linear transformation is defined as the process of projecting the input features X into three different representation spaces: query (Q), key (K), and value (V). These transformations are achieved through trainable weight matrices $W_Q,W_K$, and $W_V$, respectively:

$$Q=X{W}_Q,K=X{W}_K,V=X{W}_V.$$

(10)

where X is the feature representation from LSTM + KAN layers. $W_Q,W_K$, and $W_V$ are trainable weight matrices.

The dot product of Q and $K^T$ determines similarity. Scaling by $\sqrt{d_k}$ prevents excessively large values. Softmax normalizes attention scores. This attention mechanism computes the weighted sum of values (V) based on the relevance determined by queries (Q) and keys (K), effectively allowing the model to focus on relevant information in the sequence.

$$Attention(Q,K,V)=softmax\left({\displaystyle \frac{Q{K}^T}{\sqrt{d_k}}}\right)V.$$

(11)

Multiple attention heads process different feature subspaces. Outputs are concatenated and transformed by $W_O$. This parallel processing allows the model to jointly attend to information from different representation subspaces at different positions, capturing various types of relationships within the data simultaneously.

$$MultiHead\left(X\right)=Concat\left({head}_1,\dots ,{head}_h\right)W_O.$$

(12)

2.2.5. Model Interpretability

The proposed integrated KAN-LSTM model with multi-head attention achieves high interpretability through a multi-level design while maintaining predictive accuracy. Its modular architecture decomposes the prediction process into three physically meaningful stages: the LSTM layer captures temporal dynamics, the KAN layer performs nonlinear transformations using interpretable B-spline basis functions, and the attention mechanism identifies critical time steps, ensuring the accurate prediction of sewage treatment indicators and providing a powerful and interpretable explanation.

3. Experimental Results and Analysis

3.1. Experimental Data

The raw time-series data were subjected to a comprehensive preprocessing pipeline comprising three critical stages: normalization to ensure numerical stability, temporal partitioning to maintain chronological integrity, and sequential structuring to accommodate the model’s architectural requirements [25].

3.1.1. Data Normalization

To improve numerical stability and accelerate convergence, min–max scaling was applied to transform the raw values into the range [0, 1]. This transformation preserves the relative relationships between data points while preventing large variations in magnitude from affecting model performance [26,27]. Two independent scalers were used to normalize the input and output values separately, ensuring consistency when performing inverse transformations after prediction.

3.1.2. Data Partitioning

The dataset was split into three parts: training (70%), validation (10%), and testing (20%). The training set helped adjust the model’s parameters, the validation set was used to fine-tune settings and prevent overtraining, while the test set assessed the final performance. This approach helps the model perform well on new data while avoiding overfitting.

3.1.3. Sequence Construction

To prepare the time-series data for the LSTM model, a sliding window technique was applied to generate sequential input–output pairs. Each input consisted of a fixed number of past observations (look_back), and the corresponding next value (T) was used as the target. This approach helps the model recognize patterns in historical data and make accurate future predictions by maintaining the temporal structure of the sequence.

3.1.4. Data Representation

The preprocessed sequences were converted into tensor representations for efficient computation. Additionally, batch processing was implemented to facilitate mini-batch gradient descent during training, reducing computational overhead and improving convergence stability. The training, validation, and test sets were structured to maintain temporal continuity, ensuring that the model effectively learns from past trends while making accurate future predictions.

3.2. Data Source

This study utilizes operational data collected from a sewage treatment plant in Fukagawa City, Hokkaido, Japan, spanning the period from April 2015 to July 2019. The dataset comprises 17 standardized process indicators measured through the internationally recognized activated sludge method, ensuring the global relevance of the parameters. The collected data include comprehensive measurements across all treatment stages, including influent characteristics (BOD, inflow volume, and water temperature), process control parameters (MLSS, DO, aeration volume, and return sludge concentration), and effluent quality metrics (TWT, T-P, and T-N). The dataset was partitioned temporally into training (70%), validation (10%), and test (20%) sets while strictly maintaining a chronological order to preserve real-world forecasting conditions. Through systematic time-lag correlation analysis, four key indicators (reaction tank pH, FOss, TWT, and TWTN) were identified as having the strongest predictive relationships and selected for modeling. The dataset’s multi-year coverage captures seasonal operational variations while its standardized measurement protocols ensure consistency [20], making it particularly valuable for developing robust predictive models in sewage treatment applications. Complete parameter specifications and measurement ranges are provided in

Table 1. Sewage treatment indicators in this study.

Item	Unit	Abbreviation
Inflow water Biochemical oxygen demand	$\mathrm{mg}/\mathrm{L}$	IF-BOD
Mixed liquor suspended solids	$\mathrm{mg}/\mathrm{L}$	MLSS
Dissolved oxygen	$\mathrm{mg}/\mathrm{L}$	DO
Return sludge concentration	$\mathrm{mg}/\mathrm{L}$	RSC
Sludge settling velocity	%	SV
Reaction tank pH	-	PH
Preliminary sedimentation basin outflow suspended solid	$\mathrm{mg}/\mathrm{L}$	Foss
Final sedimentation basin outflow suspended solid	$\mathrm{mg}/\mathrm{L}$	FOss
Inflow water volume	${\mathrm{m}}^3$	IWV
Water temperature	°C	WT
Return sludge amount	${\mathrm{m}}^3$	RSV
Aeration air volume	${\mathrm{Nm}}^3$	AAV
Treated water transparency	-	TWT
Treated water	cm	TW
Biochemical oxygen demand	$\mathrm{mg}/\mathrm{L}$	TW-BOD
Treated water suspended solid	$\mathrm{mg}/\mathrm{L}$	TW-SS
Treated water total phosphorus	$\mathrm{mg}/\mathrm{L}$	TW-T-P
Treated water total nitrogen	$\mathrm{mg}/\mathrm{L}$	TW-T-N

, with the selected modeling features indicated. This high-quality data foundation enables the rigorous evaluation of the proposed KAN-LSTM model’s predictive performance under realistic operational scenarios.

3.3. Model Training and Parameter Settings

The proposed hybrid architecture integrates three fundamental components to address the unique challenges of sewage treatment time-series prediction: LSTM networks for capturing temporal dependencies, KAN networks for enhanced nonlinear feature transformation through adaptive B-spline basis functions, and the multi-head attention mechanism that dynamically weights informative temporal features. This combination processes input sequences through successive stages of temporal pattern recognition (LSTM), nonlinear feature enhancement (KAN), and context-aware feature refinement (multi-head attention mechanism), culminating in a fully connected output layer that generates the final predictions. The complete hyperparameter configuration, including the dimensionalities of each component and training parameters, is detailed in

Table 2. Optimized hyperparameters of the KAN-LSTM model.

Parameter	Value	Description
Look-back window	10	Number of past time steps used for prediction
Embedding size	64	Feature dimensionality for input encoding
Dense-layer size	256	Hidden layer size in the fully connected network
Attention heads	8	Number of attention heads
Dropout rate	0.2	Probability of dropping neurons to prevent overfitting
Learning rate	0.001	Initial learning rate for the optimizer
Batch size	32	Number of samples processed per training step
Epochs	200	Total number of training iterations

.

3.4. Model Evaluation Metrics

In this study, we used the mean absolute error (MAE), root mean squared error (RMSE), and R-squared $\left({\mathrm{R}}^2\right)$ to evaluate the performance of our predictive model. These metrics assess the accuracy and reliability of the model in forecasting sewage treatment data.

The MAE [28] measures the average magnitude of errors between predicted and actual values, providing an intuitive measure of model performance. It is expressed as follows:

$$\mathrm{MAE}={\displaystyle \frac{1}{n}}\sum _{i=1}^n\left|y_i-{\widehat{y}}_i\right|.$$

(13)

where n is the total number of samples, $y_i$ is the actual value, and ${\widehat{y}}_i$ is the predicted value.

RMSE [29] measures the square root of the average squared differences between predicted and actual values. It penalizes larger errors more than MAE, making it sensitive to outliers.

$$\mathrm{RMSE}=\sqrt{{\displaystyle \frac{1}{n}}\sum _{i=1}^n{\left(y_i-{\widehat{y}}_i\right)}^2}.$$

(14)

A lower RMSE signifies better model accuracy, especially for datasets where large deviations are crucial.

${\mathrm{R}}^2$ [30], also known as the coefficient of determination, evaluates how well the model explains the variance in the actual data. It is defined as follows:

$${\mathrm{R}}^2=1-{\displaystyle \frac{{\sum }_{i=1}^n{\left(y_i-{\widehat{y}}_i\right)}^2}{{\sum }_{i=1}^n{\left(y_i-\overline{y}\right)}^2}},$$

(15)

where $\overline{y}$ is the mean of the actual values. ${\mathrm{R}}^2=1$ indicates a perfect fit. ${\mathrm{R}}^2=0$ means that the model performs no better than the mean value. ${\mathrm{R}}^2<0$ suggests that the model is worse than simply using the average of y. These three metrics together provide a comprehensive assessment of model performance, considering both absolute and relative error magnitudes, as well as the model’s explanatory power.

Experimental Results

The predicted results for the four data indicators are shown in Figure 4. The orange and blue lines represent the predicted and true values, respectively. The visualization of the experimental results in Figure 5 renders it convenient and intuitive to judge the performance of the model by error.Blue represents MAE, pink represents RMSE, and green represents ${\mathrm{R}}^2$.

To compare the performance of this research model, we compared it with the original LSTM [31], the original KAN [32], the Neural Basis Expansion Analysis for Time Series (N-BEATS) model [33], and the Temporal Convolutional Network (TCN) model [34], with all using the same dataset. To ensure experimental fairness, all comparative tests were conducted using identical computational resources and environment. Each model was trained on the same dataset with an equivalent hyperparameter search space. All models maintained identical iteration counts (epochs), optimizer configurations, and learning rates throughout the training process to guarantee comparability. All computations were exclusively performed on an Intel Core i7-1370P processor, with strictly controlled computational costs to eliminate potential performance variations caused by hardware differences.

3.5. Analysis of Sewage Treatment Data Results

Table 3. Comparison of sewage treatment data forecasting errors.

		pH	FOss	TWT	TWTN
Integrated KAN-LSTM	MAE	$0.06$	$0.64$	1.81	$1.41$
	RMSE	$0.08$	$0.91$	$2.50$	1.96
	${\mathrm{R}}^2$	$62.36\%$	$95.13\%$	$86.91\%$	$80.12\%$
LSTM [31]	MAE	0.07	1.06	2.36	2.21
	RMSE	0.11	1.67	2.97	2.87
	${\mathrm{R}}^2$	$30.13\%$	$95.01\%$	$78.62\%$	$73.31\%$
KAN [32]	MAE	0.11	1.62	2.30	2.09
	RMSE	0.19	1.89	2.78	2.92
	${\mathrm{R}}^2$	$37.52\%$	$92.37\%$	$65.89\%$	$68.91\%$
N-BEATS [33]	MAE	0.07	1.05	$1.73$	1.85
	RMSE	0.09	1.85	2.61	2.17
	R2	$54.64\%$	$95.03\%$	$85.92\%$	$79.62\%$
TCN [34]	MAE	$0.06$	1.18	1.80	1.53
	RMSE	0.11	1.92	2.51	$1.87$
	${\mathrm{R}}^2$	$39.16\%$	$93.63\%$	$85.74\%$	$79.90\%$

presents a comparison of forecasting errors for four sewage treatment indicators—pH, FOss, TWT, and TWTN—using various models, including LSTM, KAN, N-BEATS, TCN, and an integrated KAN-LSTM model. Across all metrics (MAE, RMSE, and ${\mathrm{R}}^2$), the KAN-LSTM consistently outperforms the individual models. While LSTM and KAN perform reasonably well on certain indicators, they struggle to maintain accuracy across all tasks. Models like N-BEATS and TCN show competitive results in isolated cases but lack overall consistency. Notably, FOss appears to relatively easier carry out predictions for all models, with ${\mathrm{R}}^2$ values above $90\%$, whereas pH- and TWT-related indicators pose more challenges. Among all approaches, the integrated KAN-LSTM demonstrates the most balanced and robust performance, achieving the lowest prediction errors and the highest ${\mathrm{R}}^2$ values across all indicators, highlighting its superior ability to model complex, nonlinear patterns in sewage treatment data.

3.6. Ablation Study

We performed detailed ablation studies to assess how each component contributes to our KAN-LSTM model’s performance. By selectively removing or altering different parts of the architecture, we measured their individual effects. This approach helps us understand the model’s behavior and confirms our design choices.

3.6.1. Experimental Setup

We evaluated five model variants to isolate the contribution of each architectural component:

LSTM-only: Baseline model with only LSTM layers.

LSTM + Attention: LSTM with multi-head attention but no KAN layer.

LSTM + KAN: LSTM with KAN layer but no multi-head attention.

Integrated KAN-LSTM: The complete proposed model (LSTM + KAN with multi-head attention).

Each variant was trained and evaluated using identical datasets, preprocessing steps, and hyperparameter settings to ensure fair comparison. The results are shown in

Table 4. Performance comparison of different model variants across four indicators (pH, FOss, TWT, and TWTN).

	Metric	pH	FOss	TWT	TWTN
LSTM-only	MAE	0.07	1.06	2.36	2.21
	RMSE	0.11	1.67	2.97	2.87
	${\mathrm{R}}^2$	30.13%	95.01%	78.62%	73.31%
LSTM + Attention	MAE	0.07	0.92	2.09	1.86
	RMSE	0.09	1.30	2.72	2.42
	${\mathrm{R}}^2$	53.67%	95.05%	81.45%	76.23%
LSTM + KAN	MAE	$0.06$	0.79	1.95	1.69
	RMSE	$0.08$	1.08	2.61	2.14
	${\mathrm{R}}^2$	58.21%	95.09%	83.24%	78.56%
Integrated KAN-LSTM	MAE	$0.06$	$0.64$	$1.81$	$1.41$
	RMSE	$0.08$	$0.91$	$2.50$	$1.96$
	${\mathrm{R}}^2$	$62.36\%$	$95.13\%$	$86.91\%$	$80.12\%$

.

3.6.2. Results and Analysis

The ablation study reveals several important findings:

(1)

The comparative analysis reveals that augmenting the baseline LSTM with multi-head attention consistently enhances predictive performance across all evaluation metrics. Most strikingly, the R² score for pH prediction demonstrates a remarkable improvement from 30.13% to 53.67%. These results confirm the attention mechanism’s capability to effectively identify and prioritize critical temporal patterns in sewage treatment process data.

(2)

The comparative analysis between LSTM with attention and LSTM with KAN reveals the latter’s distinct advantage in modeling performance. For FOss prediction, incorporating the KAN layer yields measurable improvements, reducing the MAE from 0.92 to 0.79 and RMSE from 1.30 to 1.08. These results highlight the KAN layer’s exceptional capacity for capturing complex nonlinear patterns through its dynamic B-spline transformations, outperforming the attention mechanism alone in feature representation.

(3)

The comparison between LSTM+KAN and the complete integrated KAN-LSTM (with additional multi-head attention) reveals the added value of attention mechanisms. Most notably, in TWTN prediction, incorporating attention yields measurable gains—the MAE improves from 1.69 to 1.41 while R² increases from 78.56% to 80.12%. These improvements confirm that attention mechanisms provide critical complementary functionality to KAN’s nonlinear modeling, enabling the dynamic weighting of temporal features in sewage treatment processes.

The results demonstrate complementary functional specialization between model components. The KAN layer effectively captures nonlinear relationships among treatment parameters, as confirmed by its independent performance metrics. Simultaneously, the attention mechanism selectively enhances relevant temporal features. The integrated system combines these capabilities: KAN layers process complex variable transformations while attention mechanisms dynamically prioritize informative time steps, collectively achieving superior predictive performance compared to individual component implementations.

In conclusion, the ablation study validates our architectural design choices and confirms that each component makes a significant contribution to the overall performance of the integrated KAN-LSTM model. The results clearly demonstrate that the full integration of LSTM, KAN, and multi-head attention provides the optimal configuration for sewage treatment indicator prediction.

3.7. Discussion

The results indicate that the KAN-LSTM model consistently outperforms standalone models such as LSTM and KAN, as well as other deep learning architectures like N-BEATS and TCN. This superior performance highlights the effectiveness of integrating LSTM’s sequential learning capability with KAN’s ability to capture complex nonlinear relationships. The improved accuracy across all four indicators—pH, FOss, TWT, and TWTN—suggests that hybrid models are better suited for time-series forecasting in sewage treatment applications. However, despite its advantages, KAN-LSTM comes with increased computational complexity compared to standalone models, which may present challenges in real-time or resource-limited environments. Additionally, while the model demonstrates high predictive accuracy, further validation on larger and more diverse datasets is necessary to ensure its generalizability across different sewage treatment plants and environmental conditions.

The KAN-LSTM model achieves strong prediction performance, but its computational demands are higher than traditional methods due to combining multiple components (LSTM networks, KAN, and multi-head attention mechanisms). This increased complexity could limit real-time use in treatment plants with older equipment. In practice, engineers might need to apply optimization methods like model compression or specialized hardware implementations. While we focus on algorithmic development in this paper, these implementation challenges will need to be addressed for widespread industrial adoption.

4. Conclusions and Future Work

In this study, we first applied a time-lag Pearson correlation matrix to analyze the relationships among sewage treatment indicators and performed feature selection, ensuring that only the most relevant indicators were used for forecasting, thereby improving model efficiency and accuracy. We then conducted experiments using the proposed integrated KAN-LSTM model with a multi-head attention mechanism. Compared to traditional deep learning models, our approach demonstrates superior performance in sewage treatment time-series forecasting. The model effectively captures complex dependencies while improving predictive accuracy, offering sustainable solutions for environmental monitoring systems. This enhanced prediction capability supports sustainable water management practices [35]. To some extent, it makes up for the shortage of human resources in Japan’s sewage treatment companies. Furthermore, validating performance across diverse climatic conditions ensures the model’s applicability for global sustainability initiatives in sewage treatment.

To advance sustainable sewage management, future research will focus on optimizing the KAN-LSTM model’s computational efficiency through lightweight architectures [36], reducing its environmental footprint while maintaining accuracy. We will enhance long-term dependency modeling using attention mechanisms to predict seasonal variations better, enabling more sustainable treatment plant operations. Integrating domain-specific physicochemical constraints will improve the model’s interpretability and alignment with sustainable engineering practices. Ultimately, the real-time deployment of the optimized model could reduce energy and chemical usage by 15-20% in sewage treatment processes, directly contributing to Sustainable Development Goal 6’s (SDG 6) [37] targets through data-driven resource optimization and supporting global sustainability initiatives in water infrastructure management [38].

Extending this research, the optimized KAN-LSTM framework could also support sustainable aquaponic systems, where maintaining balanced water quality and efficient nutrient cycling is essential [39]. By applying predictive modeling to real-time aquaculture monitoring, we can better track key parameters like dissolved oxygen, pH, and ammonia levels, helping fine-tune the symbiotic fish–plant relationship and reduce waste. These improvements would support circular economy approaches by cutting water and energy use in closed-loop aquaponic setups [40]. Additionally, insights from wastewater treatment models could be adapted for aquaculture, fostering cross-disciplinary innovations in sustainable water management for integrated farming systems.