Natural gradient boosting (NGBoost) (Duan et al., 2020) is an advanced model designed for probabilistic forecasting. Figure 1 shows the structure of the NGBoost model. Unlike gradient boosting and XGBoost methods (Friedman, 2001; Chen and Guestrin, 2016) that focus on point predictions, NGBoost emphasizes the entire probability distribution, making it possible to estimate prediction uncertainty. In addition, an iterative method is employed to improve model predictions based on the errors in the previous steps. This approach consists of several key steps:

The natural gradient (Amari, 1998) is computed firstly. In a standard gradient descent, the geometric structure of the parameter space is not considered, resulting in potential instability and inefficiency during optimization. By considering the local curvature, the natural gradient can alleviate these problems. Specifically, the Fisher information matrix measures the local curvature of the parameter space. It shows how changes in pairs of parameters affect the output distribution of the model. For a parameterized distribution model, the Fisher information matrix is defined in Eq. (1): I θ = E ∇ θ ⁡ log ⁡ p y | θ ∇ θ ⁡ log ⁡ p y | θ T (1) where θ represents the parameter of the model; I(θ) represents the Fisher information matrix; E denotes the expectation, which is an average of all possible outcomes; ∇ _θ is the gradient; log p (y | θ) is the log-likelihood of the outcome y given the parameter θ. The natural gradient ∇_natural is then computed by using the Fisher information matrix to correct the standard gradient, as shown in Eq. (2): ∇ natural = I θ − 1 ∇ θ L (2) where L is the loss function, and I(θ)⁻¹ represents the inverse of the Fisher information matrix I(θ). The natural gradient gives a direction to the model optimization so that the prediction errors are minimized more efficiently. This adjustment ensures that parameter updates are in line with the geometry of the parameter space, leading to faster convergence rates and improved robustness.

After calculating the natural gradient, the model residuals can then be determined, which represent the prediction errors in the current iteration. Specifically, y _actual represents the actual measurement, and y pred 1 is the predicted value in the first iteration. Then, the residual in the first iteration is given by r 1 = y actual − y pred 1 . It is worth noting that the residuals in NGBoost include both the mean μ and standard deviation σ of the predicted distribution. Within the framework of NGBoost, “base learner fitting” is a crucial step in the iterative training loop. This step involves selecting and training a base learner to fit the residual r ₁ of the current model. In this study, the decision tree is chosen as the base learner due to its simple structure and high interpretability. In each iteration, the model takes the residual r ₁ as the input and makes predictions through a series of decision nodes. For example, as shown by the red path in Figure 1, the decision tree processes the residuals by examining “Feature 2”. The data point moves to the left when the value of “Feature 2” ≤ 168 and ≤76 and then goes to the right when the value >30. The data point finally reaches the leaf node, which provides the predicted residual r   pred 2 .

Once the predicted residual r pred 2 is obtained from the leaf node, the first prediction y pred 1 can be adjusted accordingly. Specifically, y pred 1 is updated by adding the predicted residual r pred 2 ( y pred 2 = y pred 1 + r   pred 2 ). After n iterations, the model reaches the preset maximum iteration number, and the model training stops. The final prediction accumulates results from all decision trees throughout the whole iterations ( y pred n + 1 = y pred 1 + ∑ n + 1 r pred i ).

Most existing machine learning models are seen as black boxes, with unclear principles and decision processes (Guidotti et al., 2018). To address this issue, this study interprets the model from two main perspectives: global interpretability and local interpretability. For local interpretability, the decision tree-based NGBoost model can visualize the decision-making process to track and explain how an individual prediction is made through the branches of the trees. In addition, the SHAP method (Lundberg and Lee, 2017) is employed to explain the interactions between multiple features and their contributions to the target individual prediction. SHAP utilizes the principles of Shapley values from cooperative game theory to quantify the importance of each feature within the NGBoost model. Within cooperative game theory, Shapley values offer a way to fairly share game rewards, ensuring participants receive compensation in line with their overall contribution to the total gain. When applied in machine learning, SHAP assesses how much each feature contributes to the model prediction, specifically how each feature impacts the final prediction result. Through this process, SHAP uncovers the importance and function of each feature within the internal decision-making mechanism of the model, significantly improving model transparency.

For a given feature x _j , its Shapley value is computed as shown in Eq. (3): ϕ j = ∑ S ⊆ N \ j n ! S ! n − S − 1 ! f S ∪ j − f S (3) where N denotes a set of all features; S is a subset excluding x _j ; n is the total number of features in the set N; |S| is the size (number of elements) of the subset S; ϕ _j represents the Shapley value of feature; N\{j} means the set of all features except x _j ; and f S ∪ j represents the output of a function when the feature subset S and the feature x _j are used together.

For global interpretability, SHAP summary plots are used to help understand the overall trends and patterns in model predictions by assessing the average contribution of features across all data points, revealing which features have the greatest impact on the model prediction. In addition, the decision tree-based NGBoost model can also show the global importance analysis of each feature on the target prediction, including the mean and the standard deviation of the prediction, aiming to provide a global analysis method based on the model itself for assessing prediction accuracy and uncertainty. The importance of each feature is evaluated using Eq. (4): F I x j = ∑ i ∈ nodes I i × Error parent i − Error child i (4) where FI(x _j ) denotes the importance of feature x _j ; I(i) is an indicator function, which is equal to one if the node i splits and 0 otherwise; Error parent i and Error child i represent the error of the parent node before splitting and the error of the child node after splitting, respectively.

The proposed algorithm for shield movement performance has five main steps: data acquisition, data preprocessing, model construction, model evaluation, and model interpretability, as shown in Figure 2. For the tunnel boring excavation, data are acquired from the installed sensors on the shield machine equipment. Typical parameters include the total thrust, torque, advance rate, etc. These data are sequentially stored in a secured database, either locally or in the cloud. In the data preprocessing step, the downtime data and outliers are first removed. The reserved dataset is then divided into a training set and a test set at a ratio of 8:2. In the model construction step, the NGBoost model is constructed using a decision tree as a base learner. Decision trees are chosen because they are highly interpretable (Quinlan, 1986). The model undergoes an iterative training loop that consists of natural gradient calculation for optimization, fitting of base learners, and adaptive model updating. Specifically, the NGBoost model uses the natural gradient descent to adjust the model parameters. The base decision tree fits the residuals from the previous iteration and serves as the current base learner. A natural gradient is then computed from the loss function MSE (mean squared error) and model parameters, guiding an adaptive parameter update. This process is repeated until a predefined maximum iteration number is met. Next, the trained NGBoost model can provide probabilistic prediction results due to its inherent capability to estimate the distribution of the target variable. Optimal hyperparameters and the superior model are determined through comparative analysis. Finally, based on the prediction results, the importance and influence of each feature on the target variable are analyzed. By visualizing the decision-making process, the transparency of the model can be improved.

A shield tunneling project in Hangzhou is used to illustrate the proposed method. The dataset has real-time tunneling information from 100 rings, all gathered by the monitoring system. Specifically, this tunneling dataset is collected from seven subsystems: the propulsion system, cylinder stroke system, mud and water silo pressure system, air cushion silo pressure system, cutter system, grouting system, and shield attitude system, totaling 50 tunneling parameters. Among these parameters, the total thrust and torque directly reflect the shield movement performance. The values of total thrust and torque at the next moment are chosen as the output for the prediction model (Ntoutsi et al., 2020). Advance rate is the distance that the shield machine tunnels within a unit of time. Penetration is an important indicator for assessing the propulsive capacity of a shield machine. A higher penetration means the shield can traverse hard geological conditions more easily. Ring number is usually used to label segments of the tunnel, and it is an indicator of construction progress. The aforementioned three parameters (advance rate, penetration, and ring number) are closely related to the shield movement performance. Therefore, they are selected as input parameters for the predictive model. In addition, as this task is a time-series prediction, the future states are influenced by the past states. Therefore, the torque and total thrust at the current moment are also selected as the input parameters for the model. The monitoring system automatically logs data every 10 s, resulting in approximately 350,000 data samples.

Data preprocessing is crucial to ensure data quality and model accuracy (Liu et al., 2019; Zhang et al., 2022b). The data preprocessing in this study includes removing the downtime data and the outliers.

In practical engineering projects, there is a significant amount of downtime data for shield machines due to cutterhead replacement, electrical circuit inspections, lubrication system maintenance, and so on. Figure 3 shows the change in advance rate during the shield tunneling process. A complete shield tunneling operational cycle, as illustrated in the upper subfigure, consists of a rising phase, a stable phase, and a declining phase. When the advance rate = 0, it indicates that the shield machine has stopped tunneling. In this case, the downtime data cannot provide practical value for predicting the shield movement performance, therefore, they are removed, and 66,224 data samples are then reserved.

The raw dataset may include outliers caused by factors such as sensor errors or human operational mistakes, which could negatively affect the training of the model. To solve this problem, this paper firstly divides the 66,224 samples into a training set and a test set at a ratio of 8:2. Boxplot analysis has been proved as an effective method for identifying and addressing outliers in shield tunneling, as shown in Hou et al. (2022), Ma et al. (2024), and Chen et al. (2024). Given its proven efficacy, this technique is utilized to discern and remove outliers from the training set. Specifically, the boxplot analysis is used to identify and eliminate the outliers in the training set to improve the data quality. Figure 4 illustrates the outlier removal for the total thrust and torque. In the boxplot, the central horizontal line rep-resents the median, while the bottom and top of the box denote the first and third quartiles, respectively. The upper and lower edges of the boxplot represent the maxi-mum and minimum values of the data. The scattered points represent the outliers. There are 1782 outliers for total thrust and 1,508 outliers for torque in the training set. To improve data quality, records including these outliers are removed (Qin et al., 2023).

Based on the preprocessed dataset, a shield movement performance prediction model is developed. The hyperparameters of the NGBoost model are determined through random search, and the last 20% of the training set is used as a validation set for model evaluation. The optimal hyperparameters are summarized in Table 1. The “n_estimators” parameter is set to 300 iterations. The normal distribution is selected as the target distribution for the NGBoost model. This is because when the number of independent random variables is large enough, their sum tends to approximate a normal distribution. In regression problems, the target variable can be considered as the sum of multiple factors (Devore, 2011). To determine the learning rate and the base learner combination, a systematic comparison of model performance is then conducted. Mean absolute error (MAE), R-squared (R ²), and Root mean square error (RMSE) shown in Eqs 5–7 are used as the evaluation metrics. R 2 = 1 − ∑ y i − y ^ i 2 ∑ y i − y ¯ i 2 (5) MAE = 1 n ∑ i = 1 n y i − y ^ i (6) RMSE = 1 n ∑ i = 1 n y i − y ^ i 2 (7) where n represents the number of measurements, y _i is the actual value for the ith measurement. y ^ i is the prediction for the ith measurement; y ¯ i is the mean of the measurements.

The selection of base learner and learning rate has significant impact on the model performance. Figure 5 shows the model performance for the total thrust and torque. The decision tree is a simple and effective model for making decisions based on previous data, and random forest is an ensemble method that aggregates multiple decision trees. The two models both show steady performances across different learning rates. Extreme gradient boosting (XGBoost) improves predictions by adding weak learners iteratively, but it does not perform well in this case. Specifically, for the total thrust shown in Figure 5A, the decision tree with a learning rate of 0.001 achieves the highest R ² value of 0.948. In contrast, XGBoost with a learning rate of 0.1 results in the lowest R ² of 0.840. For the torque shown in Figure 5B, the decision tree with a learning rate of 0.001 leads to the best R ² value of 0.933. In contrast, XGBoost with a learning rate of 0.05 has the lowest R ² of 0.852. Therefore, the base learner is selected as the decision tree, with a learning rate of 0.001.

Probabilistic prediction is helpful for risk management and decision-making as it provides decision makers with a prediction interval besides a single prediction value. Figure 6 displays the prediction results of the NGBoost model. The 95% confidence interval (CI) is plotted, which means that it includes the actual dataset parameter in about 95 out of 100 cases when sampled repeatedly. Overall, the measurements predominantly reside within the 95% confidence interval. This indicates that the model has a high level of predictive accuracy. Specifically, between the 3000th and 7500th sample points, the CI is considerably wider, which indicates a larger prediction uncertainty in this region. According to the geological survey report, this region is transitioning from a uniform layer to an upper-soft and lower-hard composite stratum, leading to heightened prediction challenges. In this complex geological situation, it is particularly important to control the operational parameters of the shield machine with greater precision and caution. In practice, the shield control system can automatically adjust the advance rate and rotation speed of the shield machine based on the confidence interval to adapt to the current geological conditions. When the model provides significant uncertainty in complex strata, such as the area between the 3000th and 7500th sample points, the system may suggest reducing the advance rate and increasing torque to ensure that the cutting head can effectively break through hard strata without causing mechanical overload.

In addition, the width of the confidence interval is also suddenly increased at the area marked by the red box, where no anomalies are identified by the geological survey. The possible reason is that the geological survey report is not detailed enough to reflect the local variations in strata. The scatter plot shows the comparison between the model predictions and the actual monitored values. Most of the points are densely clustered around the diagonal line. This indicates that the model can predict the total thrust force and torque accurately. However, Figure 6 also displays several noticeable outliers, where the model overestimates the monitored values. Specifically, these outliers mainly appear during the descending phase of the shield tunneling operational cycle. The appearance of outliers during the descending phase can be attributed to various factors, such as geological shifts, equipment malfunctions, or regular inspections of the shield machine. These factors introduce significant randomness to abrupt changes, which makes the prediction more challenging.

The performance of the NGBoost model is compared to five other machine learning models. The hyperparameter settings for the five models are determined through random search, as detailed in Table 2. Three popular decision-tree-based models are used for comparison, specifically, random forest consists of a collection of decision trees. Each tree is trained on a random subset of the data and gives its own prediction. The random forest model then averages these results to produce a final outcome. Extreme gradient boosting (XGBoost) is a decision tree-based ensemble model that uses a gradient boosting framework. It iteratively adds new trees and corrects errors made by previously trained trees. Light gradient boosting machine (LightGBM) also follows the gradient boosting framework. Instead of simply growing trees level by level, it prioritizes splitting the most optimal leaf nodes. In addition, to further demonstrate the effectiveness of the established NGBoost model, two commonly used machine learning models for the shield tunneling problem, namely extreme learning machine (ELM) and K-nearest neighbors (KNN), are also included for comparison. ELM is a feedforward neural network with a single layer of hidden nodes, where the weights are randomly assigned and do not require iterative tuning. KNN considers the “k” nearest data points to estimate the value of a data point based on the average of neighbors. Table 3 summarizes the performance of all the models for predicting the total thrust. The LightGBM and NGBoost models show superior accuracy, with an R ² value of 0.9477. The ELM model has the worst performance, with an R ² of 0.3633. In terms of MAE, NGBoost has the lowest MAE of 361.60 kN, which is less than 1% of the total thrust during the stable phase, indicating that the average deviation of its predictions from the measurements is small. Table 4 shows the model performance for predicting the torque. The NGBoost model performs the best with an R ² value of 0.9329 and an MAE of 432.22 kN m. For both torque and total thrust, the NGBoost model performs best in term of the R ², MAE, and RMSE. The predictive accuracy of the LightGBM model is close to that of the NGBoost. However, the NGBoost model can provide probabilistic predictions for the shield movement performance. Based on these probabilistic results, decision-makers can better understand the uncertainty of predictions and make more reasonable adjustments accordingly.

Sections 3.4, 3.5 show the accuracy of the decision tree-based NGBoost model. In practical engineering applications, an interpretable model is preferable because it improves transparency and acceptance of the model. The paper visualizes the decision pathways of NGBoost and utilizes the SHAP technique to explore the influence of different features on the prediction of shield movement performance.

Each decision tree path can translate into clear rules to improve model interpretability and transparency of its decisions. This study combines this highly interpretable feature with the NGBoost model to explain the decision path changes during each iteration of NGBoost. Figure 7 shows the decision path of the 1st and 300th iterations in the NGBoost model for predicting the total thrust. Node 0 is the root node of the decision tree; “Samples” denotes the number of samples under the path; and “Correction value” denotes an adjustment to the model prediction. The decision path directs samples from the root node to the leaf nodes based on their feature values and node rules. Samples go left if they meet the rule; otherwise, they go right. For example, during the 300th iteration as shown in Figure 7B, given the input conditions of penetration = 16 mm/r, advance rate = 25 mm/min, and total thrust = 5,000 kN, the path proceeds left to Node 1 as the penetration 16 mm/r ≤ 17 mm/r meets the condition at Node 0. At Node 1, since the advance rate 25 mm/min ≤28 mm/min is met, the path continues left to Node 2. However, failing to meet the condition of advance rate ≤12 mm/min at Node two (i.e., 25 ≥ 12), it finally reaches Node 4 with a correction value of −1.7 kN. Note that the correction from the 300th iteration is just one example. The actual prediction is based on the cumulative sum of the corrections from 300 iterations.

In the first iteration of the NGBoost model, only the feature “total thrust” is used as the criterion for decision-making. This indicates that in this task of predicting the total thrust, the total thrust of the current moment is the most crucial feature for forecasting the total thrust of the next moment. The large correction values in the first iteration indicate that the model initially tends to make a wide range of parameter adjustments to move quickly towards the target value. In contrast, in the 300th iteration of NGBoost, more features are used for parameter tuning. The corrections are small and are mainly used to fine-tune the model. These results indicate that in the 300th iteration, the model has gradually stabilized and attempts to reduce the prediction error using as many features as possible.

The importance of input features for total thrust and torque at the next moment is shown in Figure 8. Overall, “total thrust” and “torque” show the highest levels of importance when predicting their respective future values, and the value of importance is significantly higher than all other input features. The torque and total thrust significantly affect the standard deviation of their respective predicted outcomes (see Figures 8A, B). This is because the torque mainly overcomes the friction between the cutter and the soil, and the total thrust pushes the shield machine forward. Excessive thrust with insufficient torque may cause the cutting tools to slip or jam, while insufficient thrust can lead to slow progress or machine stalling. Therefore, an increase in one parameter leads to a higher demand for the other, resulting in increased uncertainty in the predictive results. The ring number in shield tunneling significantly affects the standard deviation of the predicted values for total thrust and torque. The possible reason might be the variation in geological layers traversed by shield machine at different ring numbers, which also leads to increased predictive uncertainty.

A SHAP summary plot shown in Figure 9 is used to illustrate how the different features of shield tunneling affect model predictions. The SHAP values on the horizontal axis show how much each feature contributes to the prediction, and the features are listed on the vertical axis in order of their importance, from most to least. The color indicates the magnitude of the feature values, with blue colors for lower values and green colors for higher values. Overall, the two subfigures show distinct orders of impact for different prediction targets. For the total thrust at the next moment, the order of impact is total thrust > penetration > ring number > advance rate > torque. For the torque at the next moment, the order of impact is torque > penetration > ring number > advance rate > total thrust. Specifically, Figure 9A explains that the total thrust at the next moment is mostly influenced by its previous value, highlighting the significance of past conditions on future predictions. Other factors like penetration rate, advance rate, and torque have a more minor effect, and they usually reduce the total thrust at the next moment. Although the impact of the ring number is similarly small, higher ring numbers tend to reduce the total thrust for the next ring more significantly. Figure 9B discusses the torque required at the next moment and shows that, except for a significant impact of torque, all other features have minor impacts and no apparent trends.

Shield tunneling is a highly intricate process involving multiple variables and complex interactions. The contribution of a single feature to the tunneling performance is not isolated but depends on the other features. Partial dependence plots for double features are employed to examine the influence of these feature on the predictive accuracy, which illustrate how two features interact to affect the predictions of the NGBoost model. The SHAP value quantifies the contribution of each feature to the predictions made by the NGBoost model. The larger the value, the greater the contribution. The horizonal axis represents the value of the most important feature, and the vertical axis represents the SHAP values contributed by the main feature. The color intensity of each point represents the value of the second important feature, with darker colors indicating larger values. Figure 10A shows the interplay between the advance rate and the total thrust and how they affect the torque at the next moment. When the advance rate ranges from 0 mm/min to 1.3 mm/min, it makes a negative contribution to the predicted value of torque at the next moment. The negative contribution increases as the advance rate is gradually decreased. This phenomenon indicates that the shield machine faces less resistance at a lower advance rate, leading to a lower torque. When the range of the advance rate is between 1.3 mm/min and 3.2 mm/min, a total thrust of 70,000 kN to 75,000 kN causes no increase in the torque. Similarly, Figure 10B shows the interplay between the penetration and the torque and how they affect the torque at the next moment. When the value of penetration is less than 1 mm/r, it results in a negative contribution. However, when it exceeds 1 mm/r, the contribution becomes positive, peaking at around a SHAP value of 3. Figure 10C shows that as the total thrust increases, its negative contribution to the torque at the next moment decreases to about 0. This means that the total thrust can be adjusted to reduce the torque, but not to increase the torque. Figure 10D demonstrates that the SHAP values remain unchanged with the changes in the ring number. The ring number has little effect on the subsequent torque. Based on these interpretable results, the proposed algorithm can adjust various tunneling parameters more precisely, thereby significantly enhancing shield movement performance. The partial dependence plots and SHAP value analysis provide a scientific basis for this predictive adjustment. For example, these analyses reveal that the relationship between increasing the advance rate, torque and the total thrust at the next moment is not simply linear but exhibits complex non-linear dynamics, with significant boosts in some ranges and reductions in others. This not only helps maintain the shield machine in optimal working condition but also effectively minimizes downtime caused by inappropriate adjustments, thereby improving the overall construction efficiency and safety of the project.

Figure 11 shows the partial dependence plots for the total thrust at the next moment. Overall, the trends observed in Figure 11 are similar to those for the next-moment torque shown in Figure 10. Specifically, Figure 11C demonstrates that as the penetration increases to about 1.7 mm/r, its negative impact on the next-moment thrust decreases to 0. As the penetration increases further, its impact suddenly changes, dropping to about −0.7 mm/r and continuing to decline to −3 mm/r. This pattern of change is different from that for the next-moment torque shown in Figure 10C. In Figure 10C, as the penetration increases, the negative impact on the torque changes from negative to positive, eventually increasing to about +2 mm/r. This difference may be due to a change in the penetration mechanism from cutting to squeezing upon reaching a certain penetration. After the transition, the shield tunneling needs to overcome greater resistance, thus negatively impacting the total thrust.