Enhancing Daily Runoff Prediction via Uniform Design and Meta-Learning Integrated Hyperparameter Optimization Embedded in Transformer

Abstract

Accurate runoff prediction is an essential foundation for water resource management, flood prevention, and drought warning. Despite the superior performance of deep learning models in runoff prediction, the high-dimensional hyperparameter optimization limits their widespread application. To address this challenge, this study proposed a hyperparameter optimization strategy that integrated Uniform Design (UD) and Meta-Learning (ML) within the Transformer framework (UD-ML-Transformer) for daily runoff prediction. Performance of the proposed model was systematically evaluated against five benchmark models, including the UD-Transformer, Particle Swarm Optimization (PSO)-Transformer, and three Receptance Weighted Key Value (RWKV)-based models (PSO-RWKV, UD-RWKV, and UD-ML-RWKV), using hydroclimatic data spanning 1980 to 2014 from the Rio Pueblo de Taos watershed in USA. Results showed that the UD-ML-Transformer model performed the best in both prediction accuracy and peak flow, with the highest Nash-Sutcliffe Efficiency (NSE) of 0.906, and the lowest Mean Squared Error (MSE), Root Mean Squared Error (RMSE), and Mean Absolute Error (MAE) of 0.004, 0.062, and 0.034, respectively. The UD-Transformer ranked second in performance, followed by the PSO-Transformer. The integrated UD-ML hyperparameter optimization strategy also improved the performance of RWKV-based models. Compared with the PSO-RWKV and UD-RWKV models, the UD-ML-RWKV model exhibited an NSE improvement of 0.45–7.65% and an RMSE reduction of 1.47–21.18%, respectively. Moreover, cross-watershed validation conducted in the Ford River watershed, USA, also demonstrated the satisfactory performance of the proposed UD-ML-Transformer model, with the highest NSE of 0.890, and the lowest MSE, RMSE, and MAE of 0.088, 0.296, and 0.141, respectively. These findings highlight the superiority of integrating UD and ML for hyperparameter optimization in runoff forecasting.

1. Introduction

Global climate change has intensified the frequency and severity of extreme hydroclimatic events, posing formidable challenges to conventional hydrological forecasting methods [1]. As a key flux integrating the terrestrial and atmospheric hydrological cycle, runoff serves as a critical link between atmospheric precipitation, soil moisture, and groundwater. It maintains regional ecological balance [2], supplies agricultural irrigation [3], and safeguards urban water security [4]. Thus, accurate runoff prediction is indispensable for ensuring the resilient operation of hydrological systems under a changing climate and for optimizing sustainable utilization of water resources [5].

Over the past few decades, a large number of runoff prediction models have been successfully developed to meet the rising demands for accuracy and applicability. These models can be generally categorized into physical models and data-driven models [6,7]. The conventional physical models can accurately simulate and predict runoff events by leveraging physical laws and hydrological mechanisms governing natural rainfall-runoff processes [8]. Commonly used physical models include the Stanford model [9], Sacramento hydrological model [10], Xin’anjiang model [11], Soil and Water Assessment Tool (SWAT) model [12], Weather Research and Forecasting hydrological modeling system (WRF-Hydro) [13], and ParFlow hydrologic model (ParFlow) [14]. While these hydrological models offer process-based interpretability, their application remains limited by their complex structure, substantial data requirements, and strong dependence on expert knowledge [15,16]. Initially, data-driven models rely on statistical characteristics from historical time series to predict runoff [17], and representative methods include autoregressive integrated moving averages [18], Kalman filter [19], and cyclic regression models [20]. However, these models fail to satisfy their strict assumptions of linearity and stationarity, and often prove inadequate in capturing the highly nonlinear and non-stationary characteristics inherent in natural runoff sequences, leading to unsatisfactory predictive performance.

With the ongoing advancement of artificial intelligence, machine learning has emerged as a primary data-driven method for runoff prediction. Techniques such as the Artificial Neural Networks (ANN), Support Vector Machines (SVM), and Convolutional Neural Networks (CNN), were commonly employed in runoff prediction for their strong nonlinear mapping ability [21]. For example, Nourani [22] achieved reliable runoff forecasting by integrating simulated emotional parameters into backpropagation of ANN. Ren et al. [23] realized monthly runoff forecast by combining the Mann–Kendall test, T-Test, and SVM. Van et al. [24] obtained accurate runoff prediction with proposed 1D CNN model. Subsequently, these techniques have evolved into more sophisticated models such as Recurrent Neural Network (RNN) architectures, most notably Long Short-Term Memory (LSTM) and Gated Recurrent Unit (GRU), which effectively mitigate the long-term dependency limitation in time-series data through their enhanced learning capabilities. For example, Gao et al. [25] found that the superiority of LSTM and GRU over conventional ANN models in runoff prediction of the Shaxi River Basin. Man et al. [26] demonstrated that the enhanced LSTM model outperformed the widely used hydrological models and data-driven models such as ANN, SVR, and GRU in daily runoff forecasting over the upper Huaihe River Basin. Le et al. [27] revealed that the superior performance of four LSTM-based models relative to feedforward neural networks (FFNN) and CNN models in streamflow forecasting across the Red River Basin. However, the advent of the Transformer, first proposed by Vaswani et al. [28], has significantly challenged the predominance of LSTM networks in time-series prediction. Its self-attention mechanism enables efficient capture of long-term dependencies in sequential data via parallel computation, thereby greatly improving training efficiency and reducing inference latency [29,30]. For example, Sun et al. [31] found that the Transformer outperformed both CNN and LSTM in groundwater-level prediction. Wei et al. [32] embedded the self-attention core of Transformer model into the LSTM/GRU models, achieving further performance gains. Yin et al. [33] developed a novel Transformer-based rainfall-runoff model, which surpassed two LSTM-based sequence-to-sequence models. Despite its strong capabilities in handling sequential data, the Transformer is still limited by the heavy computational burden of self-attention mechanism for long sequences and the difficulty of tuning hyperparameters in its large parameter space [34,35,36].

Hyperparameter optimization is a crucial step in deep learning modeling, directly impacting training stability, convergence speed, and final predictive accuracy [37]. Commonly used hyperparameter optimization algorithms include the Grey Wolf Optimizer (GWO), Particle Swarm Optimization (PSO), and Sparrow Search Algorithm (SSA) [38,39,40]. For example, Zhang et al. [41] used PSO to optimize deep learning models for landslide susceptibility assessment. Mabdeh et al. [42] applied GWO in deep learning models to improve flood susceptibility mapping. Zhang et al. [43] employed SSA in deep learning models for monthly runoff forecasting. However, these optimization algorithms are limited by intensive computation, slow convergence, and susceptibility to local optimum [44,45], particularly in high-dimensional search space of Transformer framework.

Uniform Design (UD) offers a cost-effective alternative by selecting representative point sets within high-dimensional factor spaces to minimize uniform distribution deviation, drastically reducing computational costs [46,47]. Its spatial uniformity makes it robust for optimizing complex multi-factor, multi-level systems [48,49]. Nevertheless, the sampling-focused nature limits its cross-task generalization, potentially compromising tuning stability. Meta-Learning (ML), also termed “learning to learn”, compensates for this limitation by leveraging acquired knowledge from existing tasks to rapidly recommend near-optimal hyperparameter for new tasks [50,51]. By exploiting task distribution commonalities, ML significantly reduces cross-task computational costs and improves optimization robustness, even under limited data or computing resources [52,53]. The integration of UD and ML achieved a perfect balance between global parameter space exploration and adaptive local optimal exploitation, addressing the critical research gap in existing hyperparameter optimization. Traditional Bayesian optimization and AutoML tools suffered from extremely high computational costs in large-scale parameter searching; swarm intelligence algorithms were prone to falling into local optima; and UD lacked adaptive local optimization capability and presented low search efficiency [54]. Unlike these conventional methods, our coupled framework combined the global uniform sampling advantage of UD and the adaptive learning optimization capability of ML, which fundamentally mitigated local optimum defects, reduced computational consumption, and achieved a more stable and efficient hyperparameter optimization performance. To date, few studies have combined UD and ML for hydrological model tuning, and none have been reported for Transformer-based runoff prediction models.

Therefore, this study proposed a novel UD-ML-Transformer runoff prediction model by integrating UD and ML into the Transformer architecture (UD-ML-Transformer) to improve the efficiency and robustness of hyperparameter tuning. For rigorous validation, two benchmark Transformer models optimized by PSO and UD were constructed for comparison. Additionally, the Receptance Weighted Key Value (RWKV) framework was also introduced as a supplementary benchmark. This architecture inherits the strengths of recurrent neural networks, supporting parallel training and linear-time inference to achieve efficient modeling [55]. By integrating RWKV with the three aforementioned optimization strategies, three RWKV-based models (PSO-RWKV, UD-RWKV, and UD-ML-RWKV) were constructed for comprehensive comparison. The remainder of this paper was organized as follows. In Section 2, the study area, dataset collection and preprocessing, model methodologies, and experimental design were described, followed by model evaluation metrics. Section 3 and Section 4 presented the modeling results and discussions. Finally, Section 5 summarized the key conclusions.

2. Materials and Methodology

2.1. Study Area

The Rio Pueblo de Taos watershed, situated in the vicinity of Taos, New Mexico, USA, spans latitudes ranging from 36° to 37° N and longitudes from 105° to 106° W, covering a total area of 172.5 km² (Figure 1). As a typical mountain river, the Rio Pueblo de Taos originates from the Sangre de Cristo Mountains, traverses Taos County from the northeast to the southwest, and ultimately converges with the Rio Grande. Located on a high-altitude inland plateau, the watershed is characterized by a steep terrain with an average elevation of 2600 m. The main river channel extends 48 km, with a vertical drop of 1100 m. The water supply in this basin primarily relies on high-mountain snowmelt, leading to distinct seasonal variations in runoff.

2.2. Data Sources and Preprocessing

Meteorological and hydrological data used in this study were obtained from the Catchment Attributes and Meteorology for Large-sample Studies (CAMELS) dataset [56]. The daily hydroclimatic records include precipitation (PRE), snow water equivalent (SWE), air temperature (AT), and potential evapotranspiration (PET), spanning from 1 October 1980, to 31 December 2014 (a total of 12,510 days). For time-series runoff forecasting tasks, the dataset was chronologically divided into three subsets, including training set, validation set, and test set (Figure 2). Specifically, the training set accounted for 80% of the total data, covering 1 October 1980 to 24 February 2008 (10,008 days). The validation set occupied 10% of the total data, ranging from 25 February 2008 to 29 July 2011 (1251 days). The remaining 10% formed the test set, extending from 30 July 2011 to 31 December 2014 (1251 days).

2.3. Constructing Runoff Prediction Model

2.3.1. UD-ML Optimization Algorithm

The large number of hyperparameters and their wide value ranges result in an extremely huge search space, which poses a considerable challenge in determining an appropriate number of trials when using UD to found optimal hyperparameter combinations. UD was selected for this task due to its excellent space-filling capability. It achieved uniform coverage of the entire parameter domain with a small number of trials, making it highly suitable for a large-scale hyperparameter search. Specifically, a limited number of trials (≤100) led to sparse sampling in a vast search space, which may fail to identify the global optimum. Conversely, a large number of trials (≥2000) would substantially increase computational costs. To tackle this trade-off between search sufficiency and computational cost, meta-learning was introduced to improve tuning efficiency by reusing historical optimization experience, and a meta-learner trained on historical experience was adopted to produce simulated outcomes for a large number of trials (≥2000), which greatly reduced the required experimental iterations and improved the optimization accuracy. Accordingly, this study proposed a hyperparameter optimization framework integrating UD and ML, aiming to improve the efficiency and robustness of hyperparameter tuning for Transformer-based runoff prediction models.

In the UD-ML optimization algorithm, a Multilayer Perceptron (MLP) was used as the backbone of the meta-learner, consisting of an input layer, hidden layers, and an output layer [57]. Four hyperparameters were selected for optimization in the MLP meta-learner, yielding 300 trial configurations. The 300 candidate configurations were determined via extensive preliminary trials. Pre-tests demonstrated that this hyperparameter search setting achieved stable optimization performance. Increasing the number of configurations only yielded negligible performance improvement but substantially elevated computational consumption. Each configuration was trained for 500 epochs. The search ranges were defined as follows: the number of hidden layer neurons ([10, 1000], [10, 300], [100, 3000]), the number of hidden layers ([2, 10], [2, 6], [4, 10]), learning rate settings {0.00001, 0.0001, 0.001, 0.01}, and optimizer types {SGD, Adam, RMSprop}. Three different hyperparameter search spaces were constructed by adopting distinct boundary configurations for the number of neurons and hidden layers.

2.3.2. Modified Transformer Model

The Transformer is a deep learning architecture based on the self-attention mechanism to overcome the inherent inefficiency and insufficient long-range dependency modeling of traditional Recurrent Neural Networks (RNNs) in sequential data processing. Its multi-head attention mechanism supports parallel computation across multiple independent attention heads to capture diverse local and global dependencies within sequences across distinct subspaces, thus effectively enhancing model expressiveness and generalization capability [58,59]. The conventional Transformer adopts an encoder–decoder architecture, consisting of embedding layers, positional encoding, linear layers, and the Softmax layers. The encoder extracts contextual features from input sequences via stacked multi-head self-attention and feed-forward network modules, whereas the decoder generates target sequences using masked self-attention to preserve causal dependencies.

To adapt to the requirements of runoff prediction, this study adopted a modified Transformer architecture with key adjustments, detailed as follows: (1) Since runoff forecasting inputs are continuous numerical features rather than discrete tokens, a linear-embedding layer replaced the word-embedding layer, projecting raw inputs into vector representations of the required dimension; (2) To ensure model outputs retain the original numerical characteristics of runoff values, a linear output layer suitable for regression replaced the Softmax layer, which is usually used for classification tasks; (3) To preserve the temporal sequence of runoff data and compensate for the lack of position awareness in the self-attention mechanism, trainable positional embeddings were dynamically generated during forward propagation. These positional embeddings were added to the linearly projected input vectors to form composite proxies, which were then fed into the encoder. All other components of the Transformer architecture remained unchanged.

2.3.3. UD-ML-Transformer Model

Leveraging advantages of Transformer in handling long-term sequential data and addressing numerous hyperparameters inherent in deep learning models, this study developed a UD-ML-Transformer model for runoff prediction by integrating the UD-ML hyperparameter optimization method within the modified Transformer as the core training architecture. In this model, the Transformer model contained 10 hyperparameters, generating 100 trial configurations. Each configuration was trained for 100 epochs with identical hyperparameter search ranges, as follows: data preprocessing methods (standardization or normalization), embedding dimension {64, 80, 96}, number of multi-head attention heads {2, 4, 8}, number of encoder–decoder layers [2, 6], number of neurons in fully connected layer [10, 300], batch size {16, 32, 64, 128, 256}, length of input sequence {30, 60, 90}, learning rate {0.00001, 0.0001, 0.001, 0.01}, optimizer type {SGD, Adam, RMSprop}, and dropout rate {0.1, 0.2, 0.3, 0.4, 0.5}.

Additionally, a multi-step forecasting framework with a 7-day prediction horizon was adopted for short-term streamflow prediction. This forecast timescale matched the routine operational window of short-term hydrological forecasting, providing practical support for daily water resource scheduling and flood early warning [60]. A 1-day sliding window approach was used for continuous prediction updates. The 1-day sliding step was consistent with the daily sampling frequency of the hydroclimatic dataset adopted in this study, enabling effectively utilizing the latest hydrological information and avoiding redundancy in input sequence [61]. All experiments were implemented in Python 3.12, PyTorch 2.4, and Matlab 2022b, with acceleration on Tesla P100 GPUs via the Kaggle platform. The workflow of the proposed UD-ML-Transformer model was presented in Figure 3, and Algorithm 1 provides the pseudocode of the meta-learner with uniform design-based parameter tuning. The number of optimization trials and training epochs were determined based on pre-experimental convergence tests and mainstream configurations in existing hydrological deep learning studies. These settings guaranteed adequate model fitting, while avoiding excessive computational overhead. The key steps were summarized as follows:

(1)

Based on the 10 selected hyperparameters of the modified Transformer, an initial UD matrix was constructed via the good lattice point method. Taking the minimum centered discrepancy as the optimization objective, the optimal generator vector was identified through column-wise replacement and iterative search, ensuring uniform sampling distribution in the high-dimensional parameter space. Accordingly, two UD tables with 10 factors and 100/2000 levels were established. The discrete level indices were linearly mapped to the actual value ranges of each hyperparameter, finally yielding 100 and 2000 valid hyperparameter configuration sets for model training.

(2)

The modified Transformer model was trained and tested for 100 epochs using the 100 hyperparameter trial configurations, and the corresponding model training results were collected. These hyperparameter trial configurations and their associated performance results were combined to form the training set for the meta-learner.

(3)

Four hyperparameters with significant impacts on the MLP meta-learner performance were selected as optimization targets, namely the number of neurons in the hidden layer, the number of hidden layers, learning rate, and optimizer. A UD table with 4 factors and 100 levels was generated using the UD method, and 300 hyperparameter trial configurations were derived by setting three groups of different value ranges for the selected hyperparameters.

(4)

The meta-learner training set was input into the MLP meta-learner, which was then trained for 500 epochs to obtain a well-trained model. Subsequently, the 2000 hyperparameter trial configurations were fed into this trained meta-learner to generate simulated results. Only the best simulated result and its corresponding hyperparameter configuration were retained. By training the meta-learner with 300 hyperparameter trial configurations, one optimal simulated prediction result and its corresponding configuration were obtained per trial, resulting in a total of 300 optimal configurations.

(5)

Duplicate configurations were removed from the 300 optimal outcomes, and the remaining configurations were used as the parameter tuning schemes for the Transformer model. The Transformer model was trained and tested for 100 epochs using each of these configurations, and the hyperparameter scheme yielding the best test set result was identified as the optimal hyperparameter configuration for the Transformer model.

Algorithm 1. Meta-learner based on uniform design
Input: Training dataset Dtrain: 100 hyperparameter trial configurations for Transformer and their corresponding training results (NSE) Hyperparameter candidate table H: 300 hyperparameter trial configurations for meta-learners Candidate scheme set to be predicted Dpred: 2000 hyperparameter trial configurations for Transformer
Output: Optimal hyperparameter set Dpred′ and its predicted accuracy a∗ Best model parameters θ∗ for each evaluated scheme
1:	Split Dtrain into training (80%) and validation (20%) sets
2:	for each hyperparameter scheme hi in H do
3:	Create MLP model with He initialization
4:	Set loss function to MAE (L1Loss) and L2 regularization (λ = 0.001)
5:	Initialize best_loss ← ∞, best_params ← None
6:	Initialize empty lists for training and validation losses
7:	Initialize early stopping flag stop ← False
8:	for epoch e = 1 to 500 do
9:	Perform one training step (forward, backward, update)
10:	Compute validation loss lval
11:	if lval < best_loss then
12:	best_loss ← lval
13:	Save current model parameters as best_params
14:	end if
15:	if e mod 10 = 0 then
16:	Find minimum validation loss among last 10 epochs, denote as lmin at epoch ek
17:	Store (lmin, ek) in list M
18:	if ∣M∣ ≥ 3 then
19:	Let (A, eA), (B, eB), (C, eC) be the last three entries
20:	if B < A and B < C and eA < eB <eC then
21:	Stop ← True, stop_epoch ← eB
22:	break
23:	end if
24:	end if
25:	end if
26:	end for
27:	Load best_params into model
28:	Predict accuracies a∗ on Dpred using the trained model
29:	Identify index j∗ = argmax a∗
30:	Return Dpred′, a∗
31:	end for

Italics represent variables, and bold type indicates reserved keywords of the pseudocode.

2.4. Shapely Additive Explanations (SHAP)

In this study, the Shapley Additive Explanations (SHAP) was employed to explain the proposed model. The SHAP is an explanatory framework for machine learning that primarily uses game theories to quantitatively decompose the contribution of each input feature to individual prediction [62]. The formula was presented as follows:

$$y_i=y_{base}+{\varphi }_{i1}+{\varphi }_{i2}+\dots +{\varphi }_{ij}$$

(1)

where y_base is the average model prediction across a background dataset, and ϕ_ij represents the SHAP value for the j-th feature of the i-th sample. A positive SHAP value (ϕ_ij > 0) indicates that the feature positively contributes to the prediction, increasing the predicted value. Conversely, a negative SHAP value (ϕ_ij < 0) suggests the feature exerts a negative influence, decreasing the predicted value.

2.5. Evaluation Metrics

In this study, four widely used evaluation metrics [63], including the Nash-Sutcliffe Efficiency (NSE), Mean Squared Error (MSE), Mean Absolute Error (MAE), and Root Mean Squared Error (RMSE), were adopted to assess prediction performance of the proposed UD-ML-Transformers model against five benchmark models, namely PSO-Transformer, UD-Transformer, PSO-RWKV, UD-RWKV, and UD-ML-RWKV. The formulas were presented as follows:

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

(2)

$$\mathrm{MSE}={\displaystyle \frac{1}{n}}{\displaystyle \sum _{i=1}^n{(y_i-{\widehat{y}}_i)}^2}$$

(3)

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

(4)

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

(5)

where n denotes the sample size, y_i represents the observed runoff value of the i-th sample, ŷ_i signifies the predicted runoff value of the i-th sample, and ȳ indicates the mean observed runoff value across all samples.

3. Results and Discussion

3.1. Advantages of UD-ML-Transformer Model

To comprehensively evaluate the performance of the proposed UD-ML-Transformer model, its prediction performance was compared against the five benchmark models, namely PSO-Transformer, UD-Transformer, PSO-RWKV, UD-RWKV, and UD-ML-RWKV models on the Rio Pueblo de Taos test dataset. All models were trained and validated under the same experimental settings and hyperparameter value ranges for fair comparison. Evaluation metrics of each model were listed in

Table 1. The optimal runoff prediction performance of the UD-ML-Transformer model and five benchmark models in the Rio Pueblo de Taos watershed.

Models	NSE	MSE	RMSE	MAE
PSO-Transformer	0.867	0.006	0.075	0.039
UD-Transformer	0.903	0.004	0.064	0.035
UD-ML-Transformer	0.906	0.004	0.062	0.034
PSO-RWKV	0.824	0.007	0.085	0.043
UD-RWKV	0.883	0.005	0.068	0.038
UD-ML-RWKV	0.887	0.004	0.067	0.037

. Simulated hydrographs with peak flows were presented in Figure 4, and linear regression plots were presented in Figure 5.

As expected, the proposed UD-ML-Transformer significantly outperformed all five benchmark models. Specifically, it achieved the highest NSE (0.906), and the lowest MSE (0.004), RMSE (0.062), and MAE (0.034) among all six models, demonstrating a substantial improvement in prediction accuracy. The simulated hydrograph (Figure 4a) further confirmed its superior ability to capture temporal dynamics of runoff, with the highest Goodness of Fit (R² = 0.906 in Figure 5a), and it also exhibited better prediction for peak flow events (e.g., Days 240–285 and 990–1050). When compared with the two based-Transformer models (PSO-Transformer and UD-Transformer), the UD-ML-Transformer model maintained its superiority. The UD-ML-Transformer model outperformed the PSO-Transformer with NSE increasing by 4.50%, and MSE, RMSE, and MAE decreasing by 33.33%, 17.33%, and 12.82%, respectively. Compared to the UD-Transformer, the UD-ML-Transformer exhibited comparable performance under low-flow conditions but superior peak flow capture, with RMSE and MAE reduced by 0.002 and 0.001, respectively. This demonstrated that integrating UD and ML modules into Transformer framework led to improvements in the runoff simulations of standalone UD module embedding in Transformer model. Significant improvements were also observed in Transformer-based models over the alternative RWKV-based benchmark models (PSO-RWKV, UD-RWKV and UD-ML-RWKV). The UD-ML-Transformer achieved a 2.14% increase in NSE, along with a 7.46% decrease in RMSE and an 8.11% decrease in MAE compared with UD-ML-RWKV. Among the three RWKV-based models, UD-ML-RWKV outperformed UD-RWKV and PSO-RWKV, with NSE increasing by 0.45% and 7.65%, MSE declining by 20% and 42.86%, RMSE declining by 1.47% and 21.18%, and MAE declining by 2.63% and 13.95%, respectively. This result confirms the strong generalization of UD-ML hyperparameter optimization methods, which can be effectively applied to various model architectures.

In summary, Transformer-based models performed better than the RWKV-based models in sequential prediction. Within time series prediction tasks, the self-attention mechanism weighs various segments of the input sequence, enabling flexible information organization and ensuring that the model can focus on historical data most relevant to the current prediction. Additionally, the encoder–decoder structure of the Transformer architecture not only enhances understanding of the relationships between inputs and outputs but also allows for more flexible and efficient handling of different types of time series data [64,65]. Moreover, the proposed UD-ML-Transformer runoff prediction model outperformed the five benchmark models by integrating the predictive capability of Transformer architecture in complex long-sequence data, the efficient exploration capability of UD in high-dimensional parameter spaces, and the experience transfer capability of the ML module. These synergistic advantages enabled the model to deliver excellent performance in multi-step runoff prediction.

3.2. Robustness of Peak Flow Prediction

Accurate peak flow forecasting is critical for flood disaster prevention and mitigation, as it directly contributes to reducing potential casualties and economic losses associated with flood events [66]. To further validate the robustness of the proposed UD-ML-Transformer model in predicting extreme runoff events, the top 10–20% runoff data from the Rio Pueblo de Taos test set were extracted and evaluated. The values of evaluation metrics (MSE, RMSE, and MAE) for each model are listed in

Table 2. Performance of peak flow prediction across six models in the Rio Pueblo de Taos watershed.

Models	MSE	RMSE	MAE
PSO-Transformer	0.025	0.159	0.095
UD-Transformer	0.014	0.118	0.074
UD-ML-Transformer	0.013	0.117	0.063
PSO-RWKV	0.016	0.128	0.078
UD-REKV	0.015	0.120	0.073
UD-ML-REKV	0.014	0.118	0.066

.

The proposed UD-ML-Transformer model maintained its superior performance for peak flow prediction (Figure 4a–f), achieving the lowest MSE (0.013), RMSE (0.117), and MAE (0.063) among all six models. Compared to PSO-Transformer, the UD-ML-Transformer reduced the MSE by 48.00%, the RMSE by 26.42%, and the MAE by 33.68%. Even when compared to the closest UD-Transformer model, the UD-ML-Transformer still achieved marginal improvements in peak prediction (MSE: 0.013 vs. 0.014; RMSE: 0.117 vs. 0.118; MAE: 0.063 vs. 0.074), further confirming the advantage of integrating the ML module for capturing extreme hydrological events. Despite modest improvements in peak flow accuracy, minor refinements in simulating peak timing and magnitude can yield substantial practical benefits for flood early warning and watershed risk management, which effectively reduce the under- or overestimation of extreme runoff events. In addition, UD-ML-Transformer also outperformed all RWKV-based models in peak runoff prediction. This is attributed to the inherent robustness of the Transformer architecture in modeling nonlinear and non-stationary data relationships, as well as the synergistic effects of the UD-ML optimization strategy, which identifies optimal hyperparameters for extreme event modeling. Collectively, these features enable the UD-ML-Transformer to exhibit exceptional robustness and stability in peak flow prediction.

Qualitative analysis revealed distinct patterns across the six models in peak flow simulation. PSO-optimized models (PSO-Transformer and PSO-RWKV) consistently overestimated peak magnitudes, with PSO-RWKV showing the most pronounced overprediction. In contrast, UD-based models (UD-Transformer and UD-RWKV) moderately underestimated peak discharges, while the integrated UD-ML models achieved the optimal agreement with observed peaks, with minimal systematic bias. Additionally, the UD-ML-Transformer model exhibited superior peak timing accuracy with near-perfect temporal synchronization, whereas PSO-based models suffered from obvious phase shifts. Overall, Transformer-based architectures outperformed RWKV-based architectures in both peak magnitude and timing precision for extreme runoff events. The UD-ML optimization strategy effectively mitigated the underestimation of UD methods and the overestimation of PSO approaches, yielding the most reliable extreme flood simulation among all compared models.

3.3. Hyperparameter Optimization for Transformer-Based Models

Given the better performance of Transformer-based models compared to RWKV-based models, the following section only focused on Transformer-based models. The optimal hyperparameter configurations identified by PSO, UD, and UD-ML for the Transformer architecture on the Rio Pueblo de Taos test set were presented in

Table 3. Optimal hyperparameter configurations for the Transformer-based model with three optimization methods (including PSO, UD, and UD-ML) over the Rio Pueblo de Taos watershed.

Hyperparameters	PSO-Transformer	UD-Transformer	UD-ML-Transformer
Data preprocessing method	Normalization	Normalization	Standardization
Embedding dimension	80	64	80
Number of multi-head attention heads	4	4	2
Encoder–decoder layer count	3	4	2
Fully connected layer neuron count	97	189	197
Batch size	16	256	64
Input sequence length	90	60	30
Learning rate	0.0001	0.001	0.001
Optimizer type	Adam	Adam	Adam
Dropout rate	0.2	0.1	0.1

. Both the UD-ML-Transformer and UD-Transformer employed more fully connected neurons in the hidden layer (197 and 189, respectively) compared to PSO-Transformer (97). This structural difference aligned with the complexity of runoff time series, as the expanded capacity captured long-term dependencies and nonlinear features, ensuring their superior accuracy. These two UD-based models also used larger training batches (64 and 256 vs. 16 for PSO-Transformer), which reduced gradient variance during updates, stabilized training, and enabled convergence to flatter minima, thus improving generalization to unseen data. Additionally, the two UD-based models adopted a higher learning rate (0.001 vs. 0.0001 for PSO-Transformer), accelerating training convergence and aiding in escaping local optima. Furthermore, the UD-ML-Transformer adopted the shortest input sequence (30 days vs. 60 days and 90 days for UD-Transformer and PSO-Transformer, respectively), confirming strong short-term autocorrelation, reducing noise and enhancing the attention mechanism’s focus on relevant time steps.

In summary, each hyperparameter serves a specific functional role during model training. Specifically, the input sequence length determined the range of historical information for runoff prediction, and the dropout was adopted to mitigate overfitting. Despite the black-box characteristics of deep learning framework that excludes explicit hydrological formulas, hyperparameter settings collectively regulated the fitting and generalization abilities of the network.

To evaluate optimization efficiency, the runtime of each model’s complete training cycle (including hyperparameter tuning) was recorded. The PSO-Transformer took 6 h 49 min 22 s to iterate through 100 hyperparameter trial configurations, exhibiting a relatively short runtime but moderate performance with an NSE of 0.867. Under the same hyperparameter search range, UD-Transformer required 12 h 32 min 11 s, yielding improved prediction accuracy with an NSE of 0.903 but at nearly double the runtime. However, to achieve denser sampling in the hyperparameter space (>2000 trials) for higher accuracy, the runtime of UD-Transformer would exceed 251 h, rendering it computationally prohibitive. In contrast, UD-ML-Transformer used a pre-trained meta-learner to simulate 2000 hyperparameter trials in only 24 h 35 min 13 s, while achieving a further improved NSE of 0.906. The meta-learner in UD-ML-Transformer enabled rapid generation of simulated forecasts for a large number of trial configurations. Within the same hyperparameter search space, it achieved denser experimental sampling, reducing the risk of local optima and increasing the likelihood of finding global optimal, thus reconciling optimization efficiency and predictive accuracy.

3.4. Explainability of UD-ML-Transformer Model

The SHAP values were calculated from the Rio Pueblo de Taos test set and visualized as shown in Figure 6. PET exerted the strongest impact on runoff predictions, with the highest mean absolute SHAP value of 0.00044. Specifically, high PET values corresponded to negative SHAP values, implying that increased PET may suppress runoff generation, while low PET values tended to positively promote runoff yield. SWE ranked the second and negative SHAP values under high SWE, indicating its inhibitory effect on immediate runoff generation, which is consistent with cold-region physical mechanisms. In high-altitude regions, snow accumulation stores water in watersheds and delays surface runoff response, which may explain its indirect but critical regulatory role in alpine runoff processes [67]. AT exhibited a distinct bidirectional effect that higher values were linked to positive SHAP values, while low values produced negative SHAP values. this pattern may be attributed to freeze–thaw dynamics in alpine watersheds. Low-temperature conditions accompanied by frozen soil and continuous snow cover may restrict soil infiltration and surface flow formation, thereby suppressing runoff yield. In contrast, rising temperatures promote snowmelt and soil thawing, activating hydrological fluxes and enhancing runoff generation capacity [68]. Additionally, precipitation exerted the weaker influence on runoff predictions, with a low mean absolute SHAP value of 0.00022, and asynchronous runoff-precipitation responses were also observed in Figure 2. This muted precipitation sensitivity may be associated with typical cold-region hydrological characteristics. In high-altitude watersheds, a considerable portion of precipitation occurs as snowfall and is temporarily stored in the snowpack rather than forming immediate runoff. Meanwhile, seasonally frozen soil acts as a low-impermeability layer that restricts infiltration, converting precipitation into evaporation or surface retention instead of effective runoff. Furthermore, sporadic and low-intensity precipitation in alpine watershed may fail to overcome soil storage thresholds, potentially resulting in delayed and indirect streamflow responses to precipitation inputs [69].

4. Model Applicability and Limitation

4.1. Applicability of UD-ML-Transformer Model

To further verify the applicability and generalization capability of the proposed UD-ML-Transformer model beyond its original validation area (the Rio Pueblo de Taos watershed, a high-altitude plateau mountain basin), cross-watershed validation was carried out on the Ford River watershed, a typical lowland river basin (Figure 7). Located in the southern part of the Upper Peninsula of Michigan, USA, the Ford River watershed spans latitudes of 45–46° N and longitudes of 87–88° W, covering a total drainage area of 1165.5 km². The watershed exhibits a gently topographic slope from north to south, with a relatively low main-channel gradient. Characterized by a humid continental climate, the watershed experiences cold and snowy winters, and cool and humid summers. Hydrologically, spring snowmelt produces the dominant annual runoff peak, while summer rainfall events may generate secondary flood peaks, resulting in strong seasonal variability in the annual runoff regime.

To ensure the fairness and comparability of the cross-watershed validation, all experimental settings were kept consistent with those applied in the Rio Pueblo de Taos watershed, including data sources and data time series, preprocessing procedures, hyperparameter search ranges, and evaluation metrics. Model performance in the Ford River watershed was presented in

Table 4. The optimal runoff prediction performance of three Transformer-based models with three optimization methods (including PSO, UD, and UD-ML) in the Ford River watershed.

Models	NSE	MSE	RMSE	MAE
PSO-Transformer	0.783	0.179	0.422	0.187
UD-Transformer	0.854	0.116	0.340	0.154
UD-ML-Transformer	0.890	0.088	0.296	0.141

, with optimal prediction results and linear regression plots shown in Figure 8. The optimal hyperparameter configurations of the three Transformer-based models (PSO-Transformer, UD-Transformer, UD-ML-Transformer) under their respective optimization methods were summarized in

Table 5. Optimal hyperparameter configurations for the Transformer-based models with three optimization methods (including PSO, UD, and UD-ML) in the Ford River watershed.

Hyperparameters	PSO-Transformer	UD-Transformer	UD-ML-Transformer
Data preprocessing method	Normalization	Standardization	Standardization
Embedding dimension	64	96	80
Number of multi-head attention heads	4	8	2
Encoder–decoder layer count	4	3	3
Fully connected layer neuron count	95	89	208
Batch size	32	64	64
Input sequence length	60	60	30
Learning rate	0.0001	0.001	0.001
Optimizer type	Adam	Adam	Adam
Dropout rate	0.3	0.3	0.1

.

The results consistently demonstrated that the UD-ML-Transformer maintained its superiority over the two benchmark models in the Ford River watershed. Consistent with the single-catchment outcomes in Table 1, the UD-ML optimization strategy maintained stable superiority across study sites, though the magnitude of performance improvement varied between the two basins. This finding further demonstrated the catchment characteristics exerted substantial influence on model optimization effects. Specifically, the UD-ML-Transformer achieved the highest NSE of 0.890, outperforming PSO-Transformer (NSE = 0.783) and UD-Transformer (NSE = 0.854) by 13.67% and 4.22%, respectively. It also yielded the lowest MSE of 0.088, MAE of 0.141, and RMSE of 0.296 among these three models. The consistent superiority of the UD-ML-Transformer model across two distinct watershed types (high-altitude mountain basin vs. lowland river basin) highlighted its strong generalization capability.

In terms of hyperparameters, the UD-ML-Transformer employed more fully connected neurons (208 vs. 89 for UD-Transformer and 95 for PSO-Transformer) in the hidden layer, paired with fewer attention heads (2 vs. 8 for UD-Transformer and 4 for PSO-Transformer) and comparable number of encoder–decoder layers (3 vs. 3 for UD-Transformer and 4 for PSO-Transformer). These formed a narrow-and-deep attention structure that enhanced key feature focus and reduced multi-head attention noise. Unlike the normalization adopted by the PSO-Transformer, the UD-ML-Transformer used standardized preprocessing for a more stable gradient environment, while its higher learning rate (0.001 vs. 0.0001 for PSO-Transformer) accelerated training convergence under the Adam optimizer. Additionally, the shortest input sequence length of the UD-ML-Transformer (30 vs. 60 for both UD-Transformer and PSO-Transformer) mitigated redundant information interference, improving training efficiency and feature focus. Moreover, the lower dropout rate (0.1 vs. 0.3 for both UD-Transformer and PSO-Transformer) of the UD-ML-Transformer matched its larger model capacity, which reduced information loss and enabled fully capture of data patterns. In addition to intra-basin comparison of optimization algorithms, the two watersheds exhibited notable differences in optimized hyperparameters, such as input sequence length and dropout rate. These divergence parameter settings reflected varying demands for historical information retention and regularization under different snowmelt-driven hydrological conditions.

In summary, two watersheds with heterogeneous hydrological features were therefore employed to verify the generalization of the UD-ML optimization strategy. Consistent with the results shown in Figure 4, Figure 5 and Figure 8, the UD-ML framework yielded slightly different performance gains across two catchments due to variations in topography and snowmelt hydrological characteristics. Cross-watershed validation confirmed the outstanding applicability and generalization capability of the UD-ML-Transformer model. Its stable predictive superiority under diverse conditions well addressed the transferability limitation of existing deep learning-based hydrological models, supporting its practical extension to broader runoff forecasting scenarios.

4.2. Limitations and Outlooks

In summary, the proposed UD-ML-Transformer model effectively integrated UD and ML hyperparameter optimization within the Transformer architecture, yielding superior predictive accuracy in runoff forecasting. Nevertheless, this study still has several limitations. First, the modeling framework only includes limited hydroclimatic factors. More critical factors such as vegetation coverage, land surface conditions and human activities will be added in future work to improve predictive performance. Second, only three optimization strategies (PSO, UD, UD-ML) and two deep learning architectures (Transformer, RWKV) were compared. More advanced tuning methods like Bayesian optimization will be explored to enhance model efficiency and performance. Notably, the current model comparison was mainly oriented to advanced deep learning architectures, focusing on verifying the superiority of the proposed UD-ML hyperparameter optimization strategy under unified model settings. However, classic and standard baseline models widely used in hydrological forecasting, including LSTM and GRU, as well as conventional process-based hydrological models and hybrid machine learning models, are not incorporated for cross-comparison. This limits the comprehensive evaluation of the generalization and applicability of the proposed optimization framework across diverse models. Furthermore, all experiments in this study were only conducted on the two snow-dominated watersheds in the USA, which restricted the generalizability of the proposed method. Specifically, the runoff generation mechanism of snow-dominated basins is dominated by snowmelt, which differs significantly from rainfall-dominated, arid or semi-arid hydrological systems with strong evaporation and sparse vegetation. In addition, this study did not consider human interference such as reservoir regulation and land use change. Therefore, the optimization performance of the UD-ML framework under complex hydrological conditions and anthropogenic disturbances remains to be further verified. Future work will introduce multi-type catchments covering arid, semi-arid, rainy, and human-influenced basins to comprehensively validate and improve the universal applicability of the proposed method. Finally, inter-model performance statistical significance tests and prediction uncertainty analysis are not covered in this study. Both will be supplemented in future research to improve the reliability and robustness of the method.

5. Conclusions

Runoff Prediction is critical for effective water resource management and flood risk warning and mitigation. This study proposed a novel UD-ML-Transformer model, which integrated the UD and ML methods for hyperparameter optimization into the Transformer architecture. Systematic evaluations were conducted using hydroclimatic data from two distinct watersheds, via four evaluation metrics including NSE, MAE, MSE, and RMSE. The key findings were summarized as follows:

(1)

The proposed UD-ML-Transformer model outperformed five benchmark models in providing reliable and accurate runoff predictions. The hyperparameter optimization strategy integrating UD and ML consistently enhanced the predictive performance of both Transformer and RWKV architectures. Comprehensive performance comparison further verified the stability and superiority of the proposed UD-ML optimization strategy.

(2)

The proposed UD-ML-Transformer model also exhibited robust and reliable performance in peak runoff prediction. Benefiting from the synergistic optimization of optimization strategies, the model can effectively capture the complex nonlinear characteristics of peak runoff, thereby substantially improving its practical applicability for flood simulation and early warning scenarios.

(3)

The proposed UD-ML-Transformer model demonstrated strong regional generalization capability, as verified by cross-watershed validation in the Ford River Watershed. The model maintained stable performance superiority across watersheds with heterogeneous hydrological and topographic characteristics, demonstrating promising extensibility for practical applications in diverse basins.