Early Detection of Neonatal Infection in NICU Using Machine Learning Models: A Retrospective Observational Cohort Study

2.1. Study Design

This research was conducted as a retrospective observational cohort study using the publicly available MIMIC-III Clinical Database (v1.4). The study aimed to address the research question: Can machine learning models applied within the first 30 and 120 minutes of ICU admission accurately predict neonatal infection?

Figure 1 illustrates the overall workflow of the study. A quantitative analytical approach was used. The study population included 6,539 neonates admitted to the NICU, of whom 4,013 were diagnosed with infection. Patients with multiple ICU admissions or those who died within two hours of admission were excluded.

Physiological and hematological variables (e.g., heart rate, respiratory rate, oxygen saturation, temperature, blood pressure, white blood cell counts) were extracted and aggregated into two time windows (30 and 120 minutes post-admission). Data preprocessing included exclusion criteria, variable aggregation, and imputation of missing values.

The study was observational and non-experimental, based on secondary analysis of de-identified clinical data. Multiple supervised machine learning algorithms were evaluated, and performance was assessed using stratified five-fold cross-validation, with F1-score as the primary metric.

2.2. Data Resource

The data source for this retrospective observational study was the Medical Information Mart for Intensive Care III (MIMIC-III), a relational critical care database developed by the MIT Lab for Computational Physiology. MIMIC-III is a publicly available dataset containing de-identified health data from patients admitted to critical care units at Beth Israel Deaconess Medical Center. It contains detailed health-related data for 46,520 unique critical care patients admitted between 2001 and 2013, covering 58,976 hospital admissions and 38597 ICU stays. The dataset included demographics, clinical diagnoses, procedures, medications, lab results, and other data. MIMIC-III contains over 2 million chart events, 380,000 laboratory measurements, and over 10,000 procedures, making it a comprehensive dataset for critical care research.

2.3. Data Preprocessing

After data collection, preprocessing is crucial for achieving reliable, complete, and high-quality data for prediction tasks [14], [15]. In this research, preprocessing includes collecting relevant data and handling missing values.

This research aimed to develop a machine learning model to predict newborn infections within a short time frame. The first step involved gathering data from patients in the Neonatal Intensive Care Unit (NICU), who were the target population. We excluded patients who had been admitted to the ICU more than once, as well as those who died within two hours of ICU admission, in order to reduce model complexity [12]. Ultimately, we selected 6,539 infants from the NICU cohort, of whom 4,013 were infected.

In the coding section, we extracted data on patients from the Neonatal Intensive Care Unit (NICU) within the MIMIC-III v1.4 database using the following criteria. First, we included only admissions in which the first care unit was recorded as NICU. To ensure adequate observation time, patients were excluded if they died within 120 minutes of ICU admission; those who survived beyond 120 minutes or who did not die during hospitalization were retained. To avoid duplicate information and reduce confounding, we further excluded patients with multiple NICU stays, retaining only those with a single NICU admission. After applying these criteria, demographic information, including gender, was obtained from the patients' table, and additional clinical and administrative details (e.g., admission time, discharge status, ethnicity, religion) were linked through the admissions table.

Various physiological and hematological parameters are essential for early identification and management in predicting neonatal infections. Heart rate is a key indicator; tachycardia may signal systemic infections or sepsis, while bradycardia can indicate severe infection. Respiratory rate is also crucial, with tachypnea reflecting respiratory distress and decreased rates indicating failure [16]. Oxygen saturation (SaO₂) levels reflect respiratory efficiency, particularly during infections, such as pneumonia [17]. Temperature regulation is vital, as fever or hypothermia indicates systemic responses to infection. Elevated temperatures suggest inflammation, and hypothermia often indicates severe sepsis. Blood pressure readings are important for assessing [18] cardiovascular stability, with hypotension being a critical late sign of septic shock [19]. Hematological markers are indispensable; white blood cell counts can indicate infection, while neutrophil counts reflect acute bacterial responses. Lymphocyte counts may decrease, signifying immunosuppression. Thrombocytopenia and abnormal hemoglobin levels may complicate infection management [20]. Immature granulocytes in blood smears indicate a strong bone marrow response to severe infections [16].

Not all measurements were collected simultaneously for each patient, and some measurements were taken multiple times. We aggregated the data for each patient into two time windows: 30 minutes and 120 minutes. For the 30-minute dataset, we calculated the minimum, maximum, and mean values for each variable within the first 30 minutes after patients were admitted to the ICU. We then applied the same aggregation process to the 120-minute dataset.

2.4. Missing Value Imputation

In a hospital environment, particularly in ICUs, many variables are measured. However, these measurements are not always collected and available simultaneously for each patient, leading to the common issue of missing data, especially during the early hours of admission. Fig. 2 illustrates the percentage of missing data in the raw dataset, highlighting that handling missing values is a critical challenge. Proper management of missing data is essential to ensure the validity and generalizability of study findings [21], [22].

There are three types of missing data: Missing Completely at Random (MCAR), Missing at Random (MAR), and Missing Not at Random (MNAR). Understanding these data types is essential for determining whether imputation methods, sensitivity analyses, or model adjustments are needed to avoid drawing misleading conclusions from research findings [23].

We explored the nature of missingness using a multi-step approach. First, we employed Missingno visualizations, a matrix plot to reveal patterns of missing data across observations and variables, and a heatmap to identify correlations in missingness [24]. Next, we performed Little’s MCAR test by label encoding categorical variables and applying mean imputation to create a complete covariance matrix. We compared this matrix with that of the fully observed subset and used the resulting statistics to determine whether the data were missing completely at random (MCAR). A non-significant p-value (above 0.05) supported the MCAR hypothesis, while a significant result indicated that the data might follow missing at random (MAR) or missing not at random (MNAR) mechanisms [25]. To differentiate between MAR and MNAR, we conducted chi-square tests for each variable with missing data, creating a binary indicator to evaluate associations with other observed variables; significant results suggested MAR [26]. Finally, we performed two-sample t-tests to determine if observed values differed between the “missing” and “non-missing” groups for each variable. A significant difference (p < 0.05) suggested that missingness could be MNAR, implying dependence on unobserved information rather than solely on observed data.

We examined two strategies for addressing missing data, row removal based on a pre-defined missing threshold (θ) and three imputation methods (Simple Imputer, KNN Imputer, and Iterative Imputer) [27], [28]. We systematically varied θ from 0.95 to 0.6 in increments of 0.05, discarding rows that surpassed each threshold’s proportion of missing values to ensure that highly incomplete rows did not compromise imputation results. Subsequently, we artificially masked 10% of the numeric entries to create a known “ground truth” for performance evaluation, comparing original values to the imputed ones.

We assessed imputation quality using mean squared error (MSE) and mean absolute error (MAE), calculated over all features where masking occurred. Hyperparameters, such as k-values for KNN and iteration counts for the iterative imputer, were tuned via grid search, providing a robust benchmark of each method’s effectiveness. By systematically quantifying the trade-offs between excluding rows and employing various imputation techniques, this approach provides a clearer understanding of optimal strategies for handling incomplete datasets.

2.5. Model Selection

We explored a diverse range of classification algorithms to capture a broad spectrum of decision boundaries, incorporating models from multiple learning paradigms, such as Linear and Statistical Models (Logistic Regression, Naïve Bayes), Tree-Based Methods (Decision Tree, Random Forest, Gradient Boosting, AdaBoost, XGBoost, LightGBM, Extra Trees, CatBoost, HistGradientBoosting, and Bagging Classifier), Instance-Based Learning (k-Nearest Neighbors), Neural Networks (Multilayer Perceptron), and Support Vector Machine [29].

This comprehensive set of models ensured that both linear and non-linear decision boundaries were considered, allowing for a more robust evaluation of classification performance.

To ensure a reliable assessment of model generalizability, we adopted a Stratified K-Fold cross-validation approach with five folds (n_splits = 5). This technique maintains consistency in class distribution across folds, mitigating potential class imbalance and ensuring fair evaluation across all candidate models [30].

Each model's predictive performance was assessed using four key metrics, computed through cross-validation:

• Accuracy: The proportion of correctly classified instances.

• Precision (weighted): The class-weighted mean of precision across all categories.

• Recall (weighted): The class-weighted means of recall, ensuring balanced sensitivity.

• F1-Score (weighted): The harmonic means of precision and recall, providing a robust measure of model effectiveness.

The mean scores for each metric were computed across all five folds to ensure stable performance evaluation.

The best-performing model was determined based on its mean F1-score, as this metric provides a balanced evaluation of both precision and recall, making it particularly suitable for imbalanced datasets [31]. Once the top-performing model was identified, it was retrained on the entire imputed dataset (X_imputed, y) to produce a final, optimized predictive model ready for validation and deployment.

2.6. Hyperparameter Optimization

We rigorously evaluated a diverse set of algorithms by employing this multi-step framework, comprising imputation, cross-validation, and metric-based model selection. This approach ensured that the final chosen model exhibited strong predictive performance across multiple evaluation metrics, enhancing its reliability for real-world applications [32].

A predefined grid of hyperparameter values is established, covering a range of model-specific configurations, such as learning rate, regularization strength, and complexity parameters. The grid systematically enumerates all possible combinations to be tested, ensuring that various configurations are explored comprehensively.

For the evaluation of each hyperparameter combination, a k-fold cross-validation procedure was employed. Specifically, the training data were split into k folds (e.g., k=5), and each fold was treated as a temporary validation set, while the remaining folds were used for model training. Five-fold offered the best trade-off between statistical reliability and computation cost. Performance metrics, such as accuracy, precision, recall, and F1-score, were computed on the validation fold, and the process was repeated for each parameter combination. This approach mitigated the risk of overfitting to a single train–validation partition and provided more robust estimates of out-of-sample performance [33].

All hyperparameter combinations in the predefined grid were iterated over, with each model built and evaluated using the cross-validation scheme described above. For GradientBoosting, the grid spanned n_estimators (50, 100, 150), learning_rate (0.01, 0.05, 0.1), max_depth (3, 5, 7), and subsample (0.6, 0.8, 1) to balance ensemble size, step size, tree complexity, and sampling noise. For CatBoost, we varied iterations (100, 200, 300), learning_rate (0.01, 0.05, 0.1), depth (4, 6, 8), l2_leaf_reg (1, 3, 5), and border count (32, 64, 128), allowing fine-grained control over boosting rounds, tree depth, regularization strength, and stochasticity. The mean (optionally standard deviation) of the performance metrics across folds was recorded for each parameter setting. Computational parallelization (using multiple CPU cores or threads; CPU: AMD Ryzen 9 3900X) was leveraged where feasible to expedite the search process. The configuration achieving the highest mean performance metric (e.g., the highest mean F1-score) was identified as optimal. In certain cases, a secondary performance measure (e.g., model complexity or runtime) or a tiebreaker (e.g., validation loss) was used to distinguish between comparably performing solutions.

2.7. Feature Selection

The selected model’s built-in feature importance mechanism was first used to assign each variable a numerical score reflecting its contribution to predictive accuracy. The full feature matrix was used to train baseline tree-based ensembles, including Random Forest, Extra Trees, Gradient Boosting, HistGradient Boosting, LightGBM, XGBoost, and CatBoost. The native importance metric of each model, mean decrease in impurity for bagging trees and total gain for boosting algorithms, was then extracted. These scores were min-max scaled to 0–1 and averaged to create a consensus ranking that highlighted variables deemed influential across models. We then generated a ranked list of variables in descending order of importance. Subsequently, we employed the SHAP (SHapley Additive exPlanations) framework to quantify each feature’s contribution to individual predictions, producing both summary plots and class-level insights. This interpretability analysis helped clarify which variables most strongly influenced the model’s predictive outcomes.

To further investigate the relationship between feature subset size and model performance, we adopted an incremental coverage scheme ranging from 10% to 80% (in 5% increments). Specifically, if there were F total features, a coverage level of c% dictated retaining the top [c%×F] features according to the importance ranking. For example, 10% coverage preserved only the top 10% of features, whereas 80% coverage retained 80% of the most important ones. At each coverage level, we trained a new CatBoost model (using the same hyperparameters) and conducted a 5-fold stratified cross-validation. The mean and standard deviation of accuracy for each subset provided insight into both performance and stability across different feature inclusion thresholds.

3.1. Row Removal Threshold Result

Figure 1 illustrates how varying the missingness threshold (0.60≤θ≤0.95) affects the number of retained and discarded rows in the 30-minute and 120-minute datasets, respectively. As the threshold becomes more permissive (increasing from 0.60 to 0.95), the number of retained and discarded rows changes as follows:

• 30-Minute Data: At the lowest threshold (0.60), only about 1,500 rows are retained, while approximately 5,000 rows are discarded, indicating that many rows exceed 60% missingness. As the threshold increases, the strict row-removal criteria are relaxed, allowing more rows to remain in the dataset. By a threshold of 0.95, nearly 6,500 rows are retained and fewer than 500 are excluded (Fig. 3).

• 120-Minute Data: A similar trend is observed, although even at θ = 0.60, a larger portion of rows (over 4,000) is retained compared to the 30-minute dataset. This suggests that the 120-minute dataset contains fewer high-missingness rows or generally more complete measurements. By θ = 0.95, over 6,300 rows are retained, and only a small fraction are removed (Fig. 4).

Overall, these results demonstrate that raising the missingness threshold substantially increases the number of retained samples, especially for the 30-minute data (which generally exhibit higher missing rates). Investigators must balance the desire to retain more data with the risk that overly incomplete records may introduce noise or degrade imputation quality.

3.2. Best Imputation Method

To identify the optimal approach to handling missing data in both the 30-minute and 120-minute datasets, we systematically varied the row-removal threshold, defined as the maximum allowable proportion of missing features per row, and compared multiple imputation strategies. Performance was assessed using mean squared error (MSE) and mean absolute error (MAE) relative to a “ground truth” created by artificially masking 10% of the values. For the 120-minute dataset, a threshold of 95% missingness discarded only eight rows while yielding the lowest MSE (968,311) and MAE (91.65) under Iterative Imputer (max_iter=12), demonstrating that even highly incomplete rows could be retained without severely compromising imputation accuracy. By contrast, the 30-minute data benefited from a more stringent threshold of 75%, removing 3,720 rows, to achieve the lowest MSE (962,903) and MAE (91.42) with the same imputation method. Despite differing optimal thresholds, Iterative Imputer consistently outperformed simpler approaches in both time windows, underscoring the utility of modeling inter-variable relationships to achieve robust imputation results.

3.3. Best Machine Learning Model

After applying the various classification algorithms to the 120-minute dataset, as shown in Fig. (5), Gradient Boosting emerged as the top performer, achieving the highest average F1-score (0.7983) across the cross-validation folds. In terms of overall classification metrics, Accuracy, Precision, Recall, and F1, Gradient Boosting consistently outperformed established ensemble methods, such as Random Forest and CatBoost, as well as linear/logistic models (e.g., Logistic Regression) and instance-based methods (KNN). This result indicates that, over a longer 120-minute window of recorded physiological and laboratory parameters, a boosted ensemble approach can successfully capture complex nonlinear patterns that distinguish infected from non-infected neonates.

In contrast, (Fig. 6) shows that CatBoost was the best-performing algorithm on the 30-minute data based on F1-score (0.7634), narrowly surpassing Gradient Boosting (0.7628) and other tree-based methods (e.g., Random Forest, LightGBM) in terms of balanced predictive performance. Although some algorithms (e.g., Logistic Regression) also demonstrated robust accuracy, CatBoost’s specialized handling of categorical features and iterative boosting resulted in more balanced gains in precision and recall, making it particularly effective in this shorter time window, where missing data and limited observation periods can constrain model inputs.

Notably, the F1-scores for Gradient Boosting (on the 120-minute data) and CatBoost (on the 30-minute data) were sufficiently close that either method could be considered a strong candidate for modeling neonatal infection risk. Consequently, the next stage of experimentation applied both CatBoost and Gradient Boosting to both time windows with more targeted hyperparameter optimization, ensuring that the final model selection reflects the nuances of each approach’s performance in different temporal contexts. This two-model strategy aims to determine which boosting framework ultimately provides the most consistent and clinically relevant predictions.

We conducted an extensive grid search to optimize CatBoost and Gradient Boosting models for both the 30-minute and 120-minute datasets. Each experiment varied key parameters, such as tree depth, number of iterations (or estimators), learning rate, and subsampling rates, with cross-validation (CV) accuracy serving as the primary criterion for ranking configurations. For the 30-minute dataset, CatBoost achieved the highest CV accuracy (approximately 0.7946) with a moderate tree depth (depth=4) and 300 iterations, at a learning rate of 0.05. By contrast, the best Gradient Boosting model on the same dataset achieved a CV accuracy of 0.7903, reflecting competitive but slightly lower performance than CatBoost. Inspection of confusion matrices and classification reports confirmed that these top-ranked parameter sets offered balanced improvements in both Recall and Precision.

In the 120-minute setting, the situation was reversed; Gradient Boosting attained a marginally higher CV accuracy (about 0.8024) (Fig. 6), surpassing CatBoost’s best of approximately 0.8006. Notably, the optimal Gradient Boosting configuration differed from the 30-minute scenario; it favored a shallower tree (max_depth=3), more estimators (150), and a subsampling rate of 0.8, resulting in consistent gains in predictive performance across folds. Although both models yielded robust accuracy on the fully trained dataset (ranging from 0.82 to 0.86 when evaluated using confusion matrices), the slight discrepancy between the final accuracy and the CV ranking underscores the importance of cross-validation in guiding hyperparameter selection. In particular, the alignment of strong CV scores with high out-of-sample accuracy highlights the capacity of boosted ensemble methods to capture subtle nonlinearities and interactions within newborn infection data, even under demanding early-time-window conditions.

3.4. Incremental Coverage Analysis

To evaluate how classification performance is influenced by the number of features retained, we performed an incremental coverage analysis, retaining only the top-ranked features (by model-specific importance) in successive subsets. For the 30-minute dataset (using CatBoost), coverage levels ranged from 10% to 80% of the total feature set, with each subset evaluated via 10-fold cross-validation. As shown in the results table, the highest mean accuracy (0.7896) was achieved at 50% coverage, suggesting that roughly half of the most important features were sufficient to obtain robust predictive performance. Notably, while lower coverage (e.g., 10% or 25%) produced a slightly lower accuracy (0.7754 and 0.7815, respectively), including too many features at 75–80% coverage again lowered accuracy, likely due to the introduction of less informative or redundant variables (Fig. 7).

A similar procedure was applied to the 120-minute dataset with a Gradient Boosting model, which was also tested at 10% to 80% coverage. In this case, the peak accuracy (0.8058) occurred at 80% coverage, indicating that a relatively larger proportion of features contributed a meaningful signal in the extended 120-minute window. Despite near-competitive accuracy at 70% coverage (0.8045), the model benefited from retaining a broader set of features, perhaps reflecting the richer data collected over a longer timeframe. Overall, these findings underscore that an optimal balance exists between too few features, which may omit important signals, and too many, which may introduce noise or redundancy. By identifying this “optimal point,” we can streamline model complexity while preserving high predictive performance. Table 2 presents a statistical summary of all variables across the two datasets, along with their respective missing rates.