Urban Runoff Pollution Forecasting in the Yangtze River Basin: A Physics-Informed Data-Driven Framework Enhanced with Cluster-Based Transfer Learning

Abstract

Accurate forecasting of urban rainfall-runoff pollution across large river basins is essential for urban water management. However, this task faces formidable challenges due to the scarcity of locally monitored data and the heterogeneity in hydrological and pollution processes. To address these challenges, we proposed a novel three-tiered framework comprising (1) functional area clustering using 16-dimensional features to identify zones with shared pollution mechanisms and establish a physical parameter library; (2) a hybrid physics-informed data-driven model integrating SWMM with a Residual-BiLSTM-Multi-Head Attention (RLA) model; and (3) cluster-based transfer learning enabling predictions in data-scarce zones. The framework’s efficacy was demonstrated through a multi-tiered dataset for the Yangtze River Basin. First, a knowledge base comprising 2390 reported rainfall events across 57 functional areas was synthesized to inform the functional clustering and establish a shared physical parameter library. Subsequently, intensive field monitoring from two representative residential areas was used to train and validate the hybrid model. In data-rich zones within a cluster, the model achieved high accuracy (R² > 0.82). For data-scarce zones within the same functional cluster, the model maintained a promising performance (R² > 0.5). This study presents a novel basin-scale framework, with its initial application and preliminary validation in the Yangtze River Basin.

1. Introduction

Urban rainfall-runoff pollution transports contaminants from impervious surfaces into receiving waters through drainage systems. This process is characterized by multiple sources, spatio-temporal heterogeneity, and nonlinear dynamics. Consequently, accurate basin-scale forecasting is essential for urban water management [1,2,3]. Traditional process-based models rely on intricate parameterization to describe runoff generation and pollutant transport mechanisms [4,5]. However, in large river basins such as the Yangtze, rapid urbanization has markedly increased the complexity of impervious surfaces and land-use patterns, intensified spatial mechanistic heterogeneity across functional areas, and resulted in severe data scarcity in numerous regions [6,7,8]. This greatly limits the wide application of traditional physical models at the basin scale [9,10]. Recent advances in data-driven models, particularly deep learning, which demonstrate powerful nonlinear mapping capabilities, can accurately capture the dynamic evolution of rainfall-runoff pollution processes [11,12,13]. Coupling physical and data-driven models preserves the mechanistic constraints of runoff pollution processes while exploiting the data-driven advantages in nonlinear fitting, thereby offering the potential for basin-scale prediction [14,15,16].

The physical model and data-driven model are coupled through observation bias, inductive bias and learning bias, which have application advantages in complex scenes with scarce data and known physical parts [17]. Wang et al. proposed a physics-encoded deep learning framework for distributed hydrological modeling over the Amazon Basin, which integrates physical constraints into neural networks and achieves high accuracy in streamflow prediction under data-rich conditions [18]. Zubelzu et al. proposed a hybrid methodology that integrates the Green-Ampt infiltration model with deep neural networks, employing a physics-guided approach to incorporate both static soil characteristics and dynamic pre-event meteorological data, thereby improving the prediction of runoff occurrence and volume in a small basin [19]. Although hybrid models demonstrate strong advantages in various complex environmental scenarios, their application to the large basin-scale runoff pollution prediction remains severely constrained. This limitation primarily arises from significant disparities in data availability across cities, especially in small and medium-sized urban areas [20]. Moreover, mechanistic heterogeneity results from variations in dominant pollution pathways across different functional land-use types and rainfall regimes [21,22]. To effectively generalize models under such complex conditions, it is crucial to identify shared regional characteristics in hydrological processes, pollutant transport dynamics, and land-use functions [23,24,25].

Multi-source data clustering offers a promising approach to uncover these common patterns. For instance, Lapo et al. demonstrated unsupervised extraction of spatio-temporal patterns, and Wang et al. used clustering to elucidate physical mechanisms in climate phenomena [26,27]. Similarly, for urban runoff pollution, studies reveal significant pattern clustering tendencies based on macro-dimensions such as climate zoning and land use type [28,29,30]. However, despite these advances, most existing clustering approaches rely heavily on superficial indicators that fail to capture the underlying physical mechanisms of runoff and pollutant generation. This results in groupings that may reflect apparent similarities but lack the runoff pollution generation mechanism. As a consequence, the potential for shared modeling strategies, particularly in constructing transferable physical parameter libraries, is limited. A critical gap remains in the development of physics-informed clustering frameworks that can group heterogeneous urban functional areas according to their intrinsic runoff pollution generation mechanisms. Constructing feature spaces that capture the core physical processes is essential for enabling meaningful parameter sharing across zones with similar hydrological responses and pollution behaviors. This approach helps reduce data requirements and enhances model generalization [31,32,33].

This study proposes a novel physics-informed data-driven framework. It aims to address key challenges in basin-scale urban rainfall-runoff pollution forecasting, particularly under data scarcity and spatial heterogeneity. The specific objectives of this framework are: (i) to identify urban functional areas with shared pollution generation mechanisms by clustering 16-dimensional features, thereby establishing a transferable physical parameter library; (ii) to construct a hybrid prediction model that integrates physical simulations with deep learning, leveraging physical simulations as inputs to a deep learning architecture; and (iii) to enable cluster-based transfer learning for predictions in data-scarce zones. The remainder of this paper is organized as follows. Section 2 describes the study area, the acquisition of rainfall-runoff pollution event data across 57 functional areas within the Yangtze River Basin, and the detailed methodology of the proposed framework. Section 3 presents the results and discussion, including clustering outcomes, comparative evaluation of the hybrid model against benchmark models, ablation studies quantifying the contribution of each model component, and the basin-scale prediction performance after transfer learning. Section 4 summarizes the main conclusions of this study and outlines recommendations for future research.

2. Materials and Methods

2.1. Study Area

The Yangtze River Basin spans the transition zone between China’s eastern monsoon region and the Qinghai–Tibet Plateau region (Figure A1). Characterized by a typical subtropical to warm temperate climate, the basin receives an average annual precipitation ranging from 500 to 2500 mm, predominantly concentrated in summer [34]. Encompassing a vast area of 1.8 million square kilometers, the Yangtze River Basin is not only the largest river basin in China but also ranks as the third-largest globally. It serves as a critical economic zone experiencing rapid urbanization, hosting over 30 large and medium-sized cities [35,36]. This high urbanization rate, coupled with a dense population and intensive human activities across its extensive territory, renders the basin both a representative and challenging region for studying and managing large-scale urban non-point source pollution.

To evaluate the model’s performance specifically for basin migration applications (i.e., transferring knowledge from data-rich to data-scarce areas), detailed field monitoring of rainfall-runoff pollution events was conducted in two representative residential areas within the basin. One is located in Wujiang District, Suzhou, China (Area A), and the other is located in Lianxi District, Jiujiang, China (Area B). In Area A, a rain-pollution diversion drainage system was implemented, covering a drainage area of 55,677 square meters. The land cover can be categorized into four types: green space, roads, roofs, and squares. The basin exhibits a north subtropical monsoon maritime climate, with an average annual temperature of 18.2 °C. The average annual rainfall is 1490 mm, of which summer precipitation accounts for nearly 50%. In Area B, a rain-pollution diversion drainage system was implemented, covering a drainage area of 410,160 square meters. The land cover can be categorized into six types: green space, roofs, concrete pavements, tarred pavements, permeable pavements, and artificial turf playgrounds. The basin features a subtropical humid monsoon climate, with an average annual temperature ranging from 15 to 21 °C and an average annual rainfall between 1583 and 1766 mm. Over 40% of the rainfall is concentrated during the plum rain season from April to June, with monthly precipitation in May and June reaching approximately 200 mm.

2.2. Data Acquisition

This study takes the Yangtze River Basin as the research object. Through compilation of public literature, urban rainfall-runoff pollution event data for 17 cities and 57 functional areas within the Yangtze River Basin were acquired (

Table A1. Classification of functional areas in the research area.

Functional Urban Area	Cities
cultural and educational areas	Chongqing, Wuhan, Anqing, Wuhu, Nanjing, Zhenjiang, Yangzhou, Changzhou, Suzhou, Shanghai, Wuxi,
storage areas	Chongqing, Panzhihua, Wuhan, Jiujiang, Chizhou, Wuhu, Nanjing, Zhenjiang, Yangzhou, Changzhou, Suzhou, Shanghai, Wuxi
commercial areas	Chongqing, Wuhan, Wuhu, Zhenjiang, Changzhou, Nantong, Suzhou, Shanghai, Wuxi
residential areas	Chongqing, Wuhan, Huangshi, Jiujiang, Wuhu, Nanjing, Zhenjiang, Yangzhou, Changzhou, Nantong, Suzhou, Shanghai, Wuxi
scenic areas	Chongqing, Yangzhou
industrial areas	Jingzhou, Anqing, Wuhu, Zhenjiang, Yangzhou, Changzhou, Suzhou, Shanghai, Wuxi

). These literature-derived data were used exclusively for constructing the 16-dimensional feature vectors (

Table A3. Feature Selection for Clustering.

Type	Feature
rainfall characteristics	Rainfall, Duration, Max intensity, Antecedent dry days, Temperature, Annual rainfall
pollution characteristics	Pollution concentration (COD, TN, TP SS)
social characteristics	Slope, City area, Built-up area, Green-coverage rate, Population, Per capita GDP

) to perform functional area clustering (Section 2.2). The clustering results subsequently guided the selection of target domains in the transfer learning framework (Section 2.3.3). Following the Standard for Basic Terminology of Urban Planning [37], urban functional area types were classified. According to this standard, the urban functional areas in the study area were categorized into six types: cultural and educational areas, storage areas, commercial areas, residential areas, scenic areas, and industrial areas.

Drainage network data for Areas A and B were sourced from residential property management departments and the National Engineering Research Center of Eco-environment Protection for the Yangtze River Economic Belt. Additionally, pollutant concentrations were monitored at the outlets of storm-sewer pipelines from June 2023 to July 2024. Rainfall data were recorded manually using JQR-1 rain gauges (Jinshui Huayu Information Technology Co., Ltd., Weifang, China) and computed using the local rainstorm intensity formula. During this period, five rainfall events in Area A and fourteen rainfall events in Area B were monitored (

Table A2. Basic information on the monitored rainfall events.

Event ID	Date	Rainfall (mm)	Duration (min)	Antecedent Dry Days (d)	Region
1	19 June 2024	8.9	180	1	A
2	20 June 2024	3.1	190	0	A
3	21 June 2024	3.5	180	0	A
4	27 June 2024	4.3	180	1	A
5	12 July 2024	23.3	180	1	A
6	30 June 2023	13	60	0	B
7	4 July 2023	66.5	60	3	B
8	17 July 2023	13.5	65	0	B
9	19 July 2023	43.6	57	0	B
10	30 July 2023	12.8	55	0	B
11	6 August 2023	20.4	100	7	B
12	13 August 2023	9.8	49	6	B
13	27 August 2023	34	120	0	B
14	9 November 2023	22.8	76	0	B
15	25 March 2024	13.2	65	6	B
16	19 April 2024	13.2	60	2	B
17	26 May 2024	40.4	209	1	B
18	22 June 2024	24.3	60	0	B
19	29 June 2024	24.2	145	0	B

). Water sampling was conducted at 5 min intervals during the initial storm phase and approximately 30 min intervals thereafter until runoff cessation. According to the standard methods outlined in the Water and Wastewater Monitoring and Analysis Method [38], the chemical oxygen demand (COD), total nitrogen (TN), suspended solids (SS), and total phosphorus (TP) of the samples were analyzed using the potassium dichromate method, potassium persulfate oxidation-UV spectrophotometry method, filter drying method, and ammonium molybdate spectrophotometry method, respectively. These field monitoring data from Areas A and B constituted the core dataset for training, validating, and testing the hybrid model and served as the source domain for the subsequent transfer learning experiments.

Based on the influencing factors of urban rainfall runoff pollution and the representativeness and availability of related data, a 16-dimensional feature vector was constructed, incorporating rainfall characteristics, pollution characteristics, and human and social characteristics (Table A3) [39,40,41]. The hierarchical clustering algorithm with the Canberra distance was employed to conduct unsupervised classification based on the similarity of pollution mechanisms in functional areas, resulting in the output of multiple classes of typical functional areas [42,43,44].

2.3. Hybrid Model Construction

In response to the limitations of traditional models in basin-scale urban rain-runoff pollution load accounting, a computational framework integrating the physical model and the data-driven model was proposed (Figure 1). The systematic errors in physically based models, often stemming from initial parameterization and leading to reduced prediction accuracy, can be learned and corrected using data-driven models, as established in prior research on similar modeling frameworks [45]. Therefore, by integrating the preliminary simulation results of the process model, time-series monitoring data, and deep learning approaches, high-precision and transferable rainfall-runoff pollution prediction can be achieved.

2.3.1. SWMM

Storm Water Management Model (SWMM) is a dynamic rainfall-runoff simulation model developed by the U.S. Environmental Protection Agency (EPA) [46]. It is primarily utilized to simulate stormwater runoff, pipeline network transportation, water quality, and the performance of low-impact development facilities in urban drainage systems. In the SWMM, the study area is divided into several sub-catchment areas, and surface runoff production and confluence calculations are performed separately in each sub-catchment area. The sub-catchment area is divided into pervious and impervious surfaces [47]. The Horton method was used to simulate the process of surface runoff generation, and the Horton equation is an empirical model proposed by Horton. The principle of the model is to describe long-duration precipitation events, where the infiltration attenuation index decreases exponentially from the maximum infiltration rate to a minimum value over time, as shown in Equation (1) [48,49]. The model employs the nonlinear reservoir method to estimate the surface confluence process [50,51]. An exponential function model was used to simulate the accumulation process of surface pollutants in Equation (2), while an exponential function was employed to simulate the scouring process of pollutants in Equation (3) [52], and the dynamic wave method was utilized to simulate water flow movement in pipes [53].

$$f=f_t+\left(f_0-f_t\right)e^{-kt}$$

(1)

$$B=C_1\left(1-e^{{-C}_{2}t}\right)$$

(2)

$$W=C_3{q}^{C_4}B$$

(3)

where $f$ is the infiltration rate (mm/h); $f_t$ is the steady infiltration rate (mm/h); $f_0$ is the initial infiltration rate (mm/h); $t$ is precipitation duration (h); $k$ is the attenuation coefficient of infiltration (1/h); $B$ is the accumulation of pollutants (kg); $C_1$ is the maximum accumulation of pollutants (kg); $C_2$ is the half-saturation constant; $W$ is the surface pollutant washout (kg); $C_3$ is the wash-off coefficient; $C_4$ is the wash-off exponent; $q$ is the runoff rate (mm/h).

In this hybrid framework, SWMM is employed to provide a physically consistent baseline simulation rather than a precisely calibrated prediction. Accordingly, the SWMM parameters (

Table A4. SWMM parameter values.

Target Variable	Cluster	N-imperv	N-perv	S-imperv	S-perv	PctZero	MaxRate	MinRate	DryTime	Decay	Manning-N
Runoff	1	0.008	0.223	2	4	75	79.2	3.9	24	4	0.008
	2	0.014	0.2	2	4	80	74.3	3.7	30	4.2	0.008
	3	0.016	0.271	1.5	3	80	68.4	2.5	24	3.2	0.009
Target Variable	Cluster		Kdecay		Coeff-1 (build)		Coeff-2 (build)		Coeff-3 (wash)		Coeff-42 (wash)
SS	1		0.1		140		2		0.024		2
	2		0.1		140		2		0.037		2
	3		0.1		200		2		0.038		1.5
TN	1		0.6		10		1.5		0.029		2
	2		0.1		12		1.5		0.014		2.2
	3		0.1		12		1.5		0.015		1.8
TP	1		0.8		2		1.2		0.005		2
	2		0.8		2		1.1		0.005		2.2
	3		0.8		2		1		0.015		1.5
COD	1		0.5		100		1.65		0.007		2
	2		0.5		100		1.55		0.007		1.5
	3		0.5		100		1.28		0.005		1.5

) are derived from literature-based values. This design preserves the generalizability of the framework across data-scarce regions, as the subsequent RLA model is tasked with learning the systematic bias between the uncalibrated SWMM output and observed concentrations.

2.3.2. Residual-BiLSTM-Multi-Head Attention Model

We employ the Bi-directional Long Short-Term Memory (BiLSTM) model as the fundamental machine learning module and integrate residual linking and the multi-head attention mechanism to construct our model architecture (RLA) (Figure A2). The BiLSTM model is an improvement of the traditional LSTM network [54]. In the traditional LSTM network, information is propagated backward only from past time steps. However, the BiLSTM introduces an additional forward layer, enabling bidirectional information propagation from both past and future time steps. At each time step, the outputs of the forward and backward channels are concatenated to generate the complete bidirectional output. This allows the BiLSTM to capture temporal information in sequence data more comprehensively. The fundamental computational method of the BiLSTM algorithm is presented in Equations (4) and (5) [55]. In this study, the BiLSTM model had an input dimension of 64, a hidden layer dimension of 64, and consisted of 6 hidden layers.

$$\left\{\begin{array}{c}\overrightarrow{f_t}=\sigma \left(W_f·\left[\overrightarrow{h_{t-1}};x_t\right]+b_f\right)\\ \overrightarrow{i_t}=\sigma \left(W_i·\left[\overrightarrow{h_{t-1}};x_t\right]+b_i\right)\\ \overrightarrow{\widetilde{Z_t}}=\tanh \left(W_z·\left[\overrightarrow{h_{t-1}};x_t\right]\right)+b_Z\\ \overrightarrow{Z_t}=\overrightarrow{f_t}\times \overrightarrow{Z_{t-1}}+\overrightarrow{i_t}\times \overrightarrow{\widetilde{Z_t}}\\ \overrightarrow{o_t}=\sigma \left(W_o·\left[\overrightarrow{h_{t-1}};x_t\right]+b_o\right)\\ \overrightarrow{h_t}=\overrightarrow{o_t}\times \tanh \left(\overrightarrow{Z_t}\right)\end{array}\right.\hspace{1em}\left\{\begin{array}{c}\overleftarrow{f_t}=\sigma \left({W^{\prime }}_f·\left[\overleftarrow{h_{t+1}};x_t\right]+{b^{\prime }}_f\right)\\ \overleftarrow{i_t}=\sigma \left({W^{\prime }}_i·\left[\overleftarrow{h_{t+1}};x_t\right]+{b^{\prime }}_i\right)\\ \overleftarrow{\widetilde{Z_t}}=\tanh \left({W^{\prime }}_z·\left[\overleftarrow{h_{t+1}};x_t\right]\right)+{b^{\prime }}_Z\\ \overleftarrow{Z_t}=\overleftarrow{f_t}\times \overleftarrow{Z_{t+1}}+\overleftarrow{i_t}\times \overleftarrow{\widetilde{Z_t}}\\ \overleftarrow{o_t}=\sigma \left({W^{\prime }}_o·\left[\overleftarrow{h_{t+1}};x_t\right]+{b^{\prime }}_o\right)\\ \overleftarrow{h_t}=\overleftarrow{o_t}\times \tanh \left(\overleftarrow{Z_t}\right)\end{array}\right.$$

(4)

$$h_t=\left[\overrightarrow{h_t};\overleftarrow{h_t}\right]$$

(5)

where $\overleftarrow{.}$ is the forward computation and $\overrightarrow{.}$ is the backward computation. $W_f$,$W_i$,$W_z$,$W_o$ are the weight matrices for the forward LSTM, while ${W^{\prime }}_f$,${W^{\prime }}_i$,${W^{\prime }}_z$,${W^{\prime }}_o$ are the weight matrices for the backward LSTM. $b_f$,$b_i$,$b_Z$,$b_o$ are the bias vectors for the forward LSTM, and ${b^{\prime }}_f$,${b^{\prime }}_i$,${b^{\prime }}_Z$,${b^{\prime }}_o$ are the bias vectors for the backward LSTM. $x_t$ is the model input at time step t. $[ ; ]$ is the vector concatenation operation. $\widetilde{Z_t}$ is the new cell state candidate generated by the hyperbolic tangent function (tanh). $Z_t$ is the updated cell state. $\sigma$ is the sigmoid activation function. $o_t$ is the output gate information. $f_t$ is the forget gate information. $i_t$ is the input gate information. $h_t$ is the output of the current sequence.

It is widely acknowledged that the deeper the deep learning layers, the richer the high-level features that can be extracted, thereby enhancing model performance [56]. However, a network that is excessively deep may cause the gradient to either gradually vanish or explode during backpropagation across layers, thereby hindering the model’s ability to learn effectively. Therefore, the concept of residual learning was proposed to address this issue [57]. In this study, the residual block was adopted. This block comprises two mappings: the residual mapping and the identity mapping. It can be expressed as follows:

$$y=x+F\left(x\right)$$

(6)

$$F\left(x\right)=LayerNorm\left(W_1·Dropout\left(SiLU\left(LayerNorm\left(W_1·x+b_1\right)\right),p\right)+b_2\right)$$

(7)

where x is the input feature map. $F\left(x\right)$ is the residual function, representing the “residual” that the network needs to learn. $F\left(x\right)$ is composed of a series of nonlinear transformations. $W_1$, $W_2$ are the weight matrices of the linear layers. $b_1$,$b_2$ are the bias terms. $LayerNorm$ is the layer normalization operation, used to stabilize training. $SiLU$ is the Sigmoid Linear Unit activation function. $Dropout\left(x,p\right)$ is the regularization operation that randomly drops units to prevent overfitting. $p$ denotes the regularization rate and is set to 0.1 in this study.

Residual connections enable the network to learn incremental changes in the input instead of directly learning complex transformations. This structure helps mitigate the vanishing gradient problem and facilitates the training of deeper networks. In this paper, the input linear layer has 312 dimensions and maps the 6 input features to a 312-dimensional space. The activation function layer: The SiLU function is employed to activate two residual blocks. Each residual block comprises two linear layers, each followed by a layer normalization operation. The features are mapped from 312 dimensions to 64 dimensions in the final linear layer.

To further focus on the prediction of urban rainfall-runoff pollution, the attention mechanisms are frequently employed to identify and weight relevant features [58]. The attention mechanism determines the importance level of different features by calculating attention scores, while the multi-head attention mechanism not only focuses on various historical time points in the time series but also attends to different heads. Compared with the single attention mechanism, it exhibits stronger comprehensive efficiency [59]. This mechanism maps the input to each head by randomly initializing the mapping matrix, enabling each head to focus on information from different positions within the input feature. As a result, the model is able to learn a variety of complex features. The corresponding formula is:

$$MultiHead\left(Q,K,V\right)=Concat\left({head}_1,{head}_2,\dots ,{head}_h\right)W^0$$

(8)

$${head}_i=Attention\left(Q_i,K_i,V_i\right)=softmax\left({\displaystyle \frac{Q_i{K_i}^{⊺}}{\sqrt{d_k}}}\right)V_i$$

(9)

where ${head}_i$ is the result of a single attention head, $Q_i$, $K_i$ and $V_i$ are the query, key, and value matrices, respectively. $d_k$ is the dimensionality of the keys. In this paper, the number of heads $h$ is 8.

Owing to the deep network architecture of the model in this study and the distinct characteristics of rainfall data across different time periods, conventional training strategies struggle to achieve ideal outcomes. In this study, a self-paced learning method was employed, where the samples in the training set were categorized into different difficulty levels through the design of a loss function [60]. Initially, the model prioritizes samples that are easier to predict and progressively incorporates more challenging samples as training progresses. The following expression can be used:

$$Loss=\sum _{i=1}^N{\displaystyle \frac{v_i{\left(y_i-{y^{\prime }}_i\right)}^2}{N}}$$

(10)

$$v_i=\left\{\begin{array}{c}0 when{\displaystyle \frac{{\left(y_i-{y^{\prime }}_i\right)}^2}{N}}\ge \lambda \\ 1 when{\displaystyle \frac{{\left(y_i-{y^{\prime }}_i\right)}^2}{N}}<\lambda \end{array}\right.$$

(11)

$$\lambda ={\lambda }^{\prime }+K$$

(12)

where $y_i$ is the model’s predicted value, and ${y^{\prime }}_i$ is the ground truth value. $v_i$ is the weight assignment, where simple samples are assigned 1 and complex samples are assigned 0. $N$ is the number of samples. $\lambda$ is the threshold of the loss function, which is progressively increased by adding $K$ per training epoch. In this study, the initial value of $\lambda$ is 0.1, and $K$ is 0.01.

The RLA input (X_t) includes instantaneous rainfall intensity, the duration of the dry period preceding the rain, and the SWMM simulation values (TN, TP, COD, SS). X_t is first extracted and transformed by the residual encoder, and the resulting encoded features are then input into the bidirectional LSTM to capture the temporal dependency. The output of the BiLSTM is enhanced to focus on key information through the multi-head attention mechanism.

2.3.3. Transfer Learning

Transfer Learning refers to the learning process where the similarity between data, tasks, or models is leveraged to apply a model learned in an old domain to a new domain [61]. Before applying transfer learning, the dataset needs to be divided into a source domain and a target domain. In this study, based on the clustering results of functional areas in the Yangtze River Basin, functional areas with sufficient data (the number of rainfall monitoring data in each event ≥ 4) within the same class were designated as the source domain, while the remaining regions were assigned as the target domain.

This study adopts a model-based fine-tuning strategy for transfer learning. The RLA model pre-trained on the source domain (detailed in Section 2.3.2) provided the initial weights for the target-domain model, expressed as:

$${\theta }^{\left(0\right)}={\theta }_{source}$$

(13)

where θ denotes the comprehensive set of learnable parameters within the RLA network architecture. This parameter set encompasses the input encoder (comprising a linear projection layer and two subsequent residual blocks), the stacked bidirectional LSTM, the multi-head self-attention module, the temporal-compression fully connected layer, and the final output regression layer. During the fine-tuning phase on the target-domain training set, the physical parameters of the SWMM component were strictly frozen, ensuring that only the parameters within the deep-learning backend underwent joint optimization.

The robust feature representations acquired from the source domain were preserved via three implicit regularization mechanisms: (i) initialization with near-optimal source-domain weights, (ii) execution over a shortened training horizon, and (iii) application of a self-paced adaptive loss-threshold weighting scheme inherited from the source-domain training phase, which actively filters out high-error batches during the initial fine-tuning iterations. The parameter update rule is defined as:

$${\theta }^{\left(t+1\right)}={\theta }^{\left(t\right)}-\eta \hspace{0.17em}{\nabla }_{\theta }\hspace{0.17em}{\mathcal{L}}^{\left(t\right)}$$

(14)

where the learning rate is set to η = 5 × 10⁻⁵, consistent with the rate utilized during the initial source-domain pre-training. To maintain methodological consistency and stabilize optimization across both training phases, the identical objective was employed for both source-domain pre-training and target-domain fine-tuning.

The fine-tuning procedure was executed over 50 epochs on the target-domain training set, utilizing a batch size of 312, with optimization driven by the Adam algorithm at a learning rate of η = 5 × 10⁻⁵. In contrast to the initial source-domain training, which required 120 epochs with a batch size of 256, the fine-tuning phase necessitated fewer than half the total epochs. This intentionally abbreviated training horizon constrains the cumulative parameter drift, restricting deviation from the original ${\theta }_{source}$. Model evaluation was conducted on the target-domain test set at defined step intervals, and the model checkpoint yielding the lowest test-set MSE was preserved as the definitive fine-tuned model. Hyperparameters, including the window length, feature dimensions, and LSTM structural configurations (depth and width), were strictly aligned with those established in Section 2.3.2.

3. Results and Discussion

This section may be divided by subheadings. It should provide a concise and precise description of the experimental results, their interpretation, and the experimental conclusions that can be drawn.

3.1. Clustering Results and Analysis of Functional Areas in the Yangtze River Basin

Building on established knowledge of the key drivers influencing urban rainfall-runoff pollution (Section 2.2), a comprehensive 16-dimensional feature space was developed to capture rainfall characteristics, pollution signatures, and socio-environmental factors. These features directly reflect the underlying hydrological and pollutant generation mechanisms, ensuring that the resulting clusters are not merely statistically similar but physically meaningful. Based on 16-dimensional feature vectors incorporating land use types and runoff pollution characteristics, hierarchical clustering categorized typical urban functional areas within the Yangtze River Basin into three distinct groups sharing common pollution mechanisms. The specific classification is presented in Figure 2, where the first type represents traffic-industry-education oriented functional areas, and the second type corresponds to commercial-residence-comprehensive service functional areas. The third category comprises industrial-recreation-mixed ecological functional areas. Based on the clustering results, a shared physical parameter library (Table A4) for these functional areas was established. To quantitatively justify the selection of three clusters, we performed a silhouette coefficient analysis for k values ranging from 2 to 9 (Figure A3). The silhouette score measures how similar an object is to its own cluster compared to other clusters, with higher values indicating better-defined clusters. The results show that k = 3 yields the highest silhouette score.

As shown in Figure A4, the storage-industrial-education dominant areas exhibited characteristics of large rainfall amounts, long rainfall durations, extended intervals between rainfalls, prolonged dry periods prior to rain, as well as long-duration and high-intensity rainfall events. Combined with the dominant characteristics of this cluster, the pollution concentration was closely associated with traffic volume and industrial activities, exhibiting high concentrations of COD, TN, and SS. The second type was the commercial-residential-integrated service area. Its rainfall events were relatively small in amount but large in intensity, with a short dry period before rain. Overall, the rainfall events in this cluster exhibited characteristics of a short pre-rain dry period and medium-to-high intensity rainfall. The third cluster was the industrial-scenic-mixed ecological area, which demonstrated relatively large rainfall events, a short interval between consecutive rainfall events, and a high average rainfall intensity. Overall, this cluster showed characteristics of a short pre-rain dry period and high-intensity rainfall. Critically, this approach explicitly grouped regions according to their shared physical drivers of runoff pollution. Thereby, it identified underlying similarities in hydrological response and contaminant generation across areas within each group, providing a solid basis for parameter sharing. This grouping into homogeneous clusters allowed the development of representative shared physical parameter libraries (Table A4), overcoming the high data demands of traditional basin-scale modeling. Notably, our clustering units, urban functional areas, are consistent with those identified by X. Wang et al. [62]. This deliberate selection prioritized categorizations that meet practical model application requirements, thereby serving basin-scale runoff pollution prediction.

3.2. Prediction Performance of the Hybrid Model

To evaluate the predictive capability and data efficiency of the proposed framework, this study employs the R² to quantify model performance. Given that LSTM has demonstrated remarkable effectiveness in simulating sequential data and is widely applied in runoff pollution prediction, they were selected as the benchmark model to provide a comparative reference for capturing complex temporal dynamics.

Firstly, according to the accuracy evaluation under different data volumes (Figure 3a), the model maintained a low prediction error (R² > 0.5) even in data-scarce scenarios (≤3 events), demonstrating its robustness in handling small-sample data. In data-sufficient regions (>3 events), the model performance was further improved and stabilized (R² > 0.82), indicating its high data utilization efficiency. By integrating residual connections, bidirectional temporal processing, and multi-head attention mechanisms, RLA adaptively optimized initial predictions based on physical processes, thereby significantly mitigating the impact of SWMM parameter uncertainty on the final results.

Second, four rainfall-runoff pollution event datasets were selected to predict individual rainfall events, further validating the model’s time-series predictive capability (Figure 3c–f). In the prediction of TP, the RLA model exhibited a high degree of alignment with the observed values throughout the entire period, particularly during the critical concentration peak phase from 20 to 60 min after rainfall. For COD prediction, the RLA model accurately captured the rising concentration trend from the early stage of rainfall, demonstrating its strong capability to model nonlinear and transient pollution dynamics, which provides direct evidence for the effectiveness of its physics-based error correction mechanism.

As shown in Figure 3b, when training data are limited (≤3 events), the prediction accuracy of LSTM decreased significantly (R² ≈ 0.17). Although the accuracy of LSTM improved with increasing data volume (>3 events) (R² ≈ 0.39), its performance remained consistently inferior to that of RLA. The comparison of time trajectories in Figure A5 further confirmed that LSTM is prone to prediction deviations during the initial stage of rainfall or during fluctuation phases, whereas RLA demonstrates higher accuracy and stability. These results align with the findings of Chen et al. [15] and Xu et al. [63], who also showed that integrating physics-based models with data-driven approaches substantially improves prediction accuracy, particularly when physical simulation outputs are used as input features. The present study further extends this advantage through the innovative RLA architecture, especially in the most challenging small-data scenarios.

3.3. Analysis of Hybrid Model Component Contributions

To validate the contribution of each component while mitigating the risk of single-event overfitting, we performed ablation experiments on the complete test dataset consisting of 10 rainfall events (Figure 4a,b). As illustrated in this figure, the mean R² value of the BiLSTM model was as low as 0.167 in the small-data scenario and improved only marginally to 0.3877 in the large-data scenario. This suggests that traditional recurrent neural networks have difficulty capturing the highly nonlinear features of rain-runoff pollution. After the introduction of residual connections, the performance of the model exhibited a breakthrough improvement: the average R² value in the small-data scenario increased to 0.339, while that in the large-data scenario reached 0.624. The residual structure effectively alleviated the vanishing gradient problem in deep LSTM networks. After incorporating the multi-head attention mechanism, the model achieved R² increments of 0.09 and 0.13 in the small-data and large-data zones, respectively. This dynamic feature focusing mechanism successfully decoupled the timing deviation between SWMM simulation outputs and the corresponding measured data. In addition, the significant performance improvement in the small-data scenario revealed the critical role of the attention mechanism in enhancing information via feature importance recalibration when data are scarce. After integrating the self-paced learning component, the model achieved optimal performance, demonstrating its critical role in dynamically adjusting learning priorities and optimizing overall simulation performance. The prediction examples from different component models further illustrate this improvement. Specifically, the RLA model demonstrates a clear ability to capture the complex rainfall-runoff pollution process (Figure 4c–f).

These ablation results collectively confirm that the complexity of the RLA model is justified. It can be concluded from the above results that the residual connection and feature selection mechanism effectively compensate for the lack of data information density and prevent the model from falling into a local optimum. That is, under realistic conditions with limited monitoring data, more efficient knowledge mining can be achieved through architectural innovation rather than by merely increasing the parameter scale. This model introduces a novel technical paradigm for environmental temporal modeling, serving as a collaborative framework that ensures gradient propagation via the residual structure, decouples features through the attention mechanism, and optimizes the training trajectory using self-paced learning. Furthermore, our ablation results are consistent with the well-established consensus in machine learning that model architecture design and training strategies have a decisive impact on final performance [64,65,66]. We systematically integrated these advanced techniques, residual learning, multi-head attention, and self-paced learning, into an urban hydro-environmental prediction framework, revealing their strong synergistic effects. Crucially, this work represents the first systematic quantification of the individual and joint contributions of these three pivotal components specifically for urban runoff pollution forecasting. This quantification offers clear architectural and training guidelines for constructing efficient and robust physics-informed data-driven models in this domain.

3.4. Pollution Prediction at the Basin Scale After Transfer Learning

To further demonstrate the performance advantages of the proposed framework at the basin scale and under varying data availability conditions, this study implemented the cluster-based transfer learning strategy as outlined in Section 2.3.3. The prediction performance of the LSTM and RLA models in target functional areas within the Yangtze River Basin was then quantitatively analyzed, specifically categorized by the number of monitored rainfall-runoff pollution events in the target domain. The results (Figure 5) show that in the Yangtze River Basin, when sufficient data exists in the target functional area (>3 events), the RLA model achieves a significantly higher R² value compared to the LSTM model. Meanwhile, the results also indicate that the fine-tuning strategy has higher prediction accuracy than the freezing strategy and is a better transfer learning strategy. Cluster-based transfer learning enabled the RLA model to maintain robust performance even in target zones with very limited monitored events (≤3 events), demonstrating its effectiveness under data scarcity. This result validates the critical role of the transfer learning framework in overcoming the data barrier [67]. By leveraging knowledge transfer from shared pollution mechanisms, data-driven models can effectively preserve prediction performance in areas with limited monitoring data.

The robust predictive performance maintained after cluster-based transfer learning (R² > 0.5 with ≤3 events) demonstrates the framework’s effectiveness across data-scarce functional areas within the Yangtze River Basin. This validates the core premise that organizing knowledge transfer around areas with similar pollution mechanisms enables basin-scale prediction. While this study rigorously validated the framework at the functional-area scale, the approach provides a foundation for larger-scale assessment. Future work incorporating finer-resolution spatial data can further enhance the precision of basin-wide pollution pattern analysis.

3.5. Critical Discussion on Model Limitations and Assumptions

Despite the promising performance of the proposed framework, several limitations need to be acknowledged. First, the hybrid RLA model, while accurate, has high structural complexity, which imposes computational demands. More importantly, its current design predicts time-series pollutant concentrations at the outlet but does not yet provide spatially distributed pollution loads across sub-catchments, limiting its utility for spatially explicit management. Second, the cluster-based transfer learning strategy assumes that the hydrological and pollution mechanisms of the target domain remain similar to those of the source domain. This assumption may weaken under rapidly evolving urban forms or extreme climate events that alter the original pollution generation regime. Therefore, the framework is most robust when applied within the same cluster and under climate and land-use conditions comparable to those represented in the training data. Explicitly recognizing these boundary conditions is essential for avoiding overgeneralization.

4. Conclusions

The “functional area clustering-coupled model-transfer learning” framework developed in this study clusters the functional areas of the Yangtze River Basin into three pollution mechanism common groups using 16-dimensional feature vectors. The established physical model parameter library reduces the cost of cross-regional parameter calibration. In data-sufficient regions, the prediction R² of the coupled model reaches 0.82, while in data-scarce regions, transfer learning maintains the R-squared above 0.5, significantly enhancing the prediction accuracy and generalization ability of the basin-scale pollution load model.

This accounting system offers a technical paradigm for the precise prevention and control of non-point source pollution in the Yangtze River Basin and represents a feasible pathway for upscaling water environment modeling from individual cities to basin-wide applications. This approach may offer insights for pollution control in other complex watersheds globally.

While this study demonstrates the effectiveness of the proposed framework in the Yangtze River Basin, future work should extend its validation to other major river basins under diverse climatic and land-use conditions. Further efforts will also focus on quantifying uncertainty propagation from the SWMM parameters, the RLA model, and the clustering process, as well as developing lightweight model versions with online learning capabilities to enhance generalizability.