Predicting preterm birth using auto-ML frameworks: a large observational study using electronic inpatient discharge data

Abstract

Background: To develop and compare different AutoML frameworks and machine learning models to predict premature birth.

Methods: The study used a large electronic medical record database to include 715,962 participants who had the principal diagnosis code of childbirth. Three Automatic Machine Learning (AutoML) were used to construct machine learning models including tree-based models, ensembled models, and deep neural networks on the training sample (N = 536,971). The area under the curve (AUC) and training times were used to assess the performance of the prediction models, and feature importance was computed via permutation-shuffling.

Results: The H2O AutoML framework had the highest median AUC of 0.846, followed by AutoGluon (median AUC: 0.840) and Auto-sklearn (median AUC: 0.820), and the median training time was the lowest for H2O AutoML (0.14 min), followed by AutoGluon (0.16 min) and Auto-sklearn (4.33 min). Among different types of machine learning models, the Gradient Boosting Machines (GBM) or Extreme Gradient Boosting (XGBoost), stacked ensemble, and random forrest models had better predictive performance, with median AUC scores being 0.846, 0.846, and 0.842, respectively. Important features related to preterm birth included premature rupture of membrane (PROM), incompetent cervix, occupation, and preeclampsia.

Conclusions: Our study highlights the potential of machine learning models in predicting the risk of preterm birth using readily available electronic medical record data, which have significant implications for improving prenatal care and outcomes.

Keywords: China; administrative data; autoML; machine learning; preterm birth.

Figure 1

Area under the curve (AUC) for different AutoML frameworks and machine learning models. (A) Raincloud plot of the area under the curve (AUC) for three AutoML frameworks (Auto-sklearn, AutoGluon, and H2O AutoML). Each raincloud plot panel consists of three components: a jittered dot plot on the left side, a boxplot in the middle, and a cloud plot of the distribution of AUCs on the right side. (B) Boxplots of AUCs by machine learning models. GBM: Gradient Boosting Machines; GLM: Generalized Linear Models; KNN: K-Nearest Neighbors; LDA: Linear Discriminant Analysis.

Figure 2

Training time in minutes for different AutoML frameworks and machine learning models. (A) Raincloud plot of training time in minutes (training set sample size N = 536,971) for three AutoML frameworks (Auto-sklearn, AutoGluon, and H2O AutoML). Each raincloud plot panel consists of three components: a jittered dot plot on the left side, a boxplot in the middle, and a cloud plot of the distribution of AUCs on the right side. (B) Boxplots of training time in minutes by machine learning models. GBM: Gradient Boosting Machines; GLM: Generalized Linear Models; KNN: K-Nearest Neighbors; LDA: Linear Discriminant Analysis.

Figure 3

Overall feature importance (95% confidence intervals) plots for predicting preterm birth via permutation-shuffling in AutoGluon. PROM: premature rupture of membranes.