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 of this retrospective observational study is 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 has detailed health information for 46,520 unique critical care patients who were admitted to the hospital between 2001 and 2013. This includes 58,976 hospital admissions and 38,597 ICU stays. The dataset includes demographics, clinical diagnoses, procedures, medications, lab results, etc. 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 perfect data for prediction tasks [14, 15]. In this research, preprocessing includes gathering relevant data and dealing with missing values.

The goal of this research was to develop a machine learning model that can quickly predict infections in newborns. The initial step involved assembling patients in the NICU, specifically the neonatal patients. Patients who had been in the ICU more than once and patients who died within two hours of being admitted to the ICU were excluded to make the model less complicated [12]. We finally selected 6539 infants from all NICU patients, of whom 4013 infants were infected.

In the coding part, we extracted neonatal intensive care unit (NICU) patients from 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 reveal 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 are collected simultaneously for each patient, and some may be collected multiple times. We gathered the data for one patient into two time windows: 30 minutes and 120 minutes. In the 30-minute dataset, we aggregated the data to minimum, maximum, and mean for each variable in the first 30 minutes after patients were admitted to the ICU. Then, a similar method was applied for the 120 minutes dataset.

2.4. Missing Value Imputation

In hospitals, especially in ICUs, different clinical parameters are measured. However, these tests are not always performed simultaneously and available at the same time for every patient; therefore, data is often missing, especially in the early hours of admission. Figure 2 shows the missing percentage of the raw dataset, which indicates that handling missing values is the most important issue. Appropriate handling of missing data is crucial to ensure the validity and generalizability of study findings [21, 22].

There are three missing types: missing completely at random, missing at random, and missing not at random. It is important to know the types of missing data to determine whether imputation methods, sensitivity analyses, or model adjustments are necessary to avoid misleading conclusions in research findings [23].

We explored the nature of missingness through a multi-step approach. First, we employed missingno visualizations, which is a plot to reveal patterns of missing data in observations and variables, and a heatmap to indicate correlations in missingness [24]. Next, we performed Little’s MCAR test by label encoding categorical variables and applying mean imputation to derive a complete covariance matrix. We compared this covariance matrix to that of the fully observed subset and used the resulting statistics to ascertain 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 row removal based on a pre-defined missing threshold (θ) and three imputation methods (Simple Imputer, KNN Imputer, and Iterative Imputer) for addressing missing data [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 used mean squared error (MSE) and mean absolute error (MAE) to assess the quality of the imputation. They were calculated over all features where masking happened. Grid search was used to adjust hyperparameters like k-values for KNN and iteration counts for the Iterative Imputer. This gave a strong benchmark for how well each method worked. This method gives a better idea of the best ways to deal with incomplete datasets by systematically measuring the trade-offs between leaving out rows and using different imputation methods.

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.

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

All models’ predictive performance was assessed using the following 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, ensuring stability in performance evaluation.

The best-performing model was determined based on its mean F1-score because 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 re-trained on the entire imputed dataset (X_imputed, y) to produce a final, optimized predictive model for validation and deployment.

2.6. Hyperparameter Optimization

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

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

For the evaluation of each hyperparameter combination, we utilized k-fold cross-validation again. Similarly, the training data were split into k folds (e.g., k=5), and each fold in turn was treated as a temporary validation set while the remaining folds were used for model training. Five folds offered the best trade-off between statistical reliability and computation cost. Performance metrics, such as accuracy, precision, recall, or F1-score, were computed on the validation fold, and the process was repeated for each combination of parameters. 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 Gradient Boosting, 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 (and optionally standard deviation) of the performance metrics across folds was recorded for each parameter setting. Computational parallelization (using multiple CPU cores or threads, and the CPU is AMD Ryzen 9 3900X) was leveraged where it was 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, run time) or a tiebreaker (e.g., validation loss) was used to distinguish between comparably performing solutions.

2.7. Feature Selection

The built-in feature importance mechanism of the chosen model was used to give each variable a number that showed how much it helped make predictions more accurate. We used the full feature matrix to train baseline tree-based ensembles like Random Forest, Extra Trees, Gradient Boosting, HistGradient Boosting, LightGBM, XGBoost, and CatBoost. The native importance metric for each model is the mean decrease in impurity for bagging trees and the total gain for boosting algorithms. We min-max scaled these scores to 0–1 and then averaged them to make a consensus ranking that showed which variables were considered important across models. After that, we made a ranked list of the variables, with the most important ones at the top. We then used the SHAP (SHapley Additive exPlanations) framework to figure out how much each feature added to each prediction. This gave us both summary plots and insights at the class level. This interpretability analysis helped us figure out which variables had the most significant effect on the model's predictions.

To further investigate the relationship between feature subset size and model performance, we adopted an incremental coverage scheme ranging from 10% to 80% (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, a coverage of 10% 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 for the 30-minute dataset and 120-minute dataset, respectively. As the threshold became more permissive (moving from 0.60 to 0.95), the number of retained and discarded rows changed as follows:

• 30-Minute Data: At the lowest threshold (0.60), only about 1,500 rows were retained, and around 5,000 rows were discarded, indicating that many rows exceeded 60% missingness. As the threshold increased, stricter row-removal criteria were relaxed, so more rows remained in the dataset. By a threshold of 0.95, nearly 6,500 rows were retained and fewer than 500 were excluded (Fig. 3).

• 120-Minute Data: A similar trend appeared, although even at θ=0.60, a larger portion of rows (over 4,000) was retained compared to the 30-minute dataset. This suggests that the 120-minute data had fewer high-missingness rows or more complete measurements overall. At θ=0.95, over 6,300 rows remained, and only a small fraction were removed (Fig. 4).

These results demonstrated that raising the missingness threshold substantially increased the number of retained samples, especially for the 30-minute data (which generally exhibit higher missing rates). Investigators need to find a balance between wanting to keep more data and the risk that records that are too incomplete could add noise or lower the quality of imputation.

3.2. Best Imputation Method

We systematically changed the row-removal threshold, which is the highest percentage of missing features that can be in a row, to find the best way to deal with missing data in both the 30-minute and 120-minute datasets. We then compared different imputation methods. 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. Even though the optimal thresholds differed, Iterative Imputer consistently performed better than simpler methods in both time windows. This shows how useful it is to model the relationships between variables to obtain strong 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, CatBoost was found to be the best algorithm in the 30-minute data by 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, as shown in Fig. (6). Although certain algorithms (e.g., Logistic Regression) also demonstrated robust Accuracy, CatBoost’s specialized handling of categorical features and iterative boosting led to more balanced gains in Precision and Recall, rendering it particularly effective in this shorter time window, where missing data and limited observation periods could constrain model inputs.

Notably, the F1-scores for Gradient Boosting (in the 120-minute data) and CatBoost (in the 30-minute data) were sufficiently close that either method could be considered a strong candidate for modeling neonatal infection risk. The next step in the experiment was to use both CatBoost and Gradient Boosting on both time windows while optimizing the hyperparameters more carefully, ensuring that the final model selection takes into account the slight differences in how each method works in different time contexts. This two-model strategy was used to figure out which boosting framework makes the most accurate and clinically useful predictions.

We conducted a 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) under a configuration involving 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 attained a CV accuracy near 0.7903, reflecting competitive but slightly lower performance relative to 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. It was found that the optimal Gradient Boosting configuration differed from the 30-minute scenario; it favored a shallower tree (max_depth=3) with more estimators (150) and a subsampling rate of 0.8, emphasizing consistent gains in predictive performance across folds. Both models yielded robust accuracy on the fully trained dataset (ranging between 0.82 and 0.86 when evaluated via confusion matrices). However, the slight differences between the final accuracy and CV ranking indicate the importance of cross-validation for choosing hyperparameters. 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

We performed an incremental coverage analysis where only the top-ranked features (by model-specific importance) were evaluated to assess how classification performance is influenced by the number of features retained 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. It 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 a near-competitive accuracy being observed 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 (losing important signals) and too many (risking noise or redundancy). By identifying this “spot,” we can streamline model complexity while preserving high predictive performance. Table 2 presents a statistical summary of all variables across the two datasets, together with their respective missing rates.

OVERVIEW OF FINDINGS

In the 30-minute dataset, CatBoost emerged as the top-performing model with the highest F1-score (~0.7634), surpassing Gradient Boosting (~0.7628) and other tree-based or linear methods. The marginal gap between CatBoost and Gradient Boosting suggests that, when data are relatively sparse (i.e., at only 30 minutes into an ICU stay), subtle differences in how ensemble algorithms handle missingness, outlier measurements, and imputation steps can substantially shape final performance metrics. CatBoost’s specialized handling of categorical data, alongside robust boosting iterations, may have conferred an advantage in the “ultra-early” context, where variable availability is inconsistent, and the physiologic signals can be less stable.

When the time window was extended to 120 minutes, Gradient Boosting models exhibited the highest average F1-score (~0.7983) and displayed excellent calibration. This improvement likely stems from the greater data richness at 120 minutes; additional laboratory results, updated vital signs, and time for more hemodynamic fluctuations can allow gradient-boosted ensembles to better capture complex, nonlinear relationships. Nonetheless, CatBoost, Random Forest, and certain neural network models also demonstrated strong discriminatory ability (F1-scores typically above 0.76), implying that multiple algorithmic families can achieve clinical usefulness once enough data points become available. The key difference was that Gradient Boosting reached slightly higher and more consistent cross-validation scores across accuracy, precision, and recall metrics.

INTERPRETATIONS AND CLINICAL IMPLICATIONS

Our findings highlight the significance of time-dependent data accessibility in infection-risk modeling. The 30-minute window allows to intervene very early, but it also means that the data is limited, with fewer blood gas analyses or lab results that are not completed. So, advanced methods work well with missing data or can use categorical inputs well.

The 120-minute dataset, on the other hand, provides a better overall picture of the patient's health. Taking more vital signs, conducting more lab tests, and monitoring the heart and brain for more extended periods of time can all help to make risk estimates more accurate. Gradient Boosting works better in the 120-minute window. This shows that ensemble-based methods can leverage detailed information to model more complex relationships between variables.

These results indicate a two-stage or adaptive methodology: (1) an ultra-early model appropriate for the initial 30 to 60 minutes of admission, utilizing robust management of missing data (CatBoost or a comparable boosted technique), and (2) a refined model (e.g., Gradient Boosting) re-trained with updated data at approximately 120 minutes, providing a more precise prediction for ensuing critical-care decisions.

COMPARISON WITH EXISTING LITERATURE

Several prior investigations have compared machine learning models (e.g., Random Forest, XGBoost, and neural networks) against established clinical scoring systems for sepsis or infection-related mortality. Consistent with those reports, we found that ensemble-based algorithms offer better calibration and higher F1 scores than classical logistic regression or basic scoring systems (e.g., APACHE II or simplified risk scores). Notably, the incremental gains are especially relevant in the “intermediate” time frame (1–2 hours), when important physiologic trends start to manifest yet remain absent at baseline.