Predicting Neonatal Mortality Using Demographic and Health Survey Data: Comparing Logistic Regression and Machine Learning under Severe Class Imbalance
Erick Cheruiyot Kirui1* and Stephen Muthii Wanjohi2
Abstract
Neonatal mortality continues to be a significant public health problem, especially in low and middle-income countries, where improvements in neonatal mortality have been slower than progress in overall child survival. Logistic regression has been applied in the past to determine the factors associated with neonatal mortality from the Demographic and Health Survey (DHS) data, but machine learning techniques have been gaining popularity as alternative methods for predictive modelling. Evidence comparing these approaches based on nationally representative household survey data, however, is limited, especially when considering the case of severe class imbalance. This study evaluated alternative strategies to address class imbalance and examined the effect of outcome-endogenous predictors to compare the predictive performance of logistic regression, random forest and gradient boosting classifiers with respect to the prediction of neonatal mortality. Secondary analysis was performed with data from 18,978 live births with 422 neonatal deaths (2.22%) and 27 predictors related to the mother, household, and child. The data were split with stratified sampling into training set (75%) and testing set (25%). Three strategies of handling imbalanced data were used in this study: baseline (no adjustment), class weighting, and Synthetic Minority Oversampling Technique (SMOTE). The area under the receiver operating characteristic curve (ROC-AUC), the area under the precision-recall curve (PR-AUC), the sensitivity, the specificity, the precision, the F1 score, the calibration curve and the permutation feature importance were used to assess model performance. Class-weighted gradient boosting had the best overall predictive performance (ROC-AUC = 0.921, PR-AUC = 0.311, sensitivity = 83.0%) and logistic regression had similar discrimination ability but was easier to interpret (ROCAUC = 0.920). Explicitly dealing with class imbalance significantly outperformed unadjusted models for sensitivity across all of the algorithms. The number of living children under five years, number of births in the last five years and gestation length were the most important variables identified by the feature importance analysis. When the mechanically outcome-related predictor, "number of living children under 5 years," was removed, the ROC-AUC dropped from 0.921 to 0.700, suggesting that a significant amount of the apparent predictive performance was due to information leakage and not to the predictor itself. These results show that the solution for class imbalance problem is more important than using different classification methods, and emphasize the importance of selecting predictors for temporal and definitional relationship with the outcome. The study gives methodological directions for creating accurate predictions to neonatal mortality and other rare health phenomena based on survey data nationally representative.
