Gradient boosted trees with individual explanations: An alternative to logistic regression for viability prediction in the first trimester of pregnancy

Abstract

Background: Clinical models to predict first trimester viability are traditionally based on multivariable logistic regression (LR) which is not directly interpretable for non-statistical experts like physicians. Furthermore, LR requires complete datasets and pre-established variables specifications. In this study, we leveraged the internal non-linearity, feature selection and missing values handling mechanisms of machine learning algorithms, along with a post-hoc interpretability strategy, as potential advantages over LR for clinical modeling.

Methods: The dataset included 1154 patients with 2377 individual scans and was obtained from a prospective observational cohort study conducted at a hospital in London, UK, from March 2014 to May 2019. The data were split into a training (70%) and a test set (30%). Parsimonious and complete multivariable models were developed from two algorithms to predict first trimester viability at 11-14 weeks gestational age (GA): LR and light gradient boosted machine (LGBM). Missing values were handled by multiple imputation where appropriate. The SHapley Additive exPlanations (SHAP) framework was applied to derive individual explanations of the models.

Results: The parsimonious LGBM model had similar discriminative and calibration performance as the parsimonious LR (AUC 0.885 vs 0.860; calibration slope: 1.19 vs 1.18). The complete models did not outperform the parsimonious models. LGBM was robust to the presence of missing values and did not require multiple imputation unlike LR. Decision path plots and feature importance analysis revealed different algorithm behaviors despite similar predictive performance. The main driving variable from the LR model was the pre-specified interaction between fetal heart presence and mean sac diameter. The crown-rump length variable and a proxy variable reflecting the difference in GA between expected and observed GA were the two most important variables of LGBM. Finally, while variable interactions must be specified upfront with LR, several interactions were ranked by the SHAP framework among the most important features learned automatically by the LGBM algorithm.

Conclusions: Gradient boosted algorithms performed similarly to carefully crafted LR models in terms of discrimination and calibration for first trimester viability prediction. By handling multi-collinearity, missing values, feature selection and variable interactions internally, the gradient boosted trees algorithm, combined with SHAP, offers a serious alternative to traditional LR models.

Keywords: First trimester viability; Gradient boosted tree; Logistic regression; Machine learning; Post-hoc interpretability; Shapley value.

Fig. 1

Study flowchart n represents the number of unique patients and n scans the number of scans. All repeated scans of a given patient were strictly allocated to either the training or the test set. GA: gestational age; LMP: last menstruation period; LFTU: lost to follow-up, TOP: termination of pregnancy.

Fig. 2

Longitudinal assessment of the performance metrics for the parsimonious models. AUC, calibration slope and calibration in the large are displayed depending on the GA by LMP at the date of scan using a time window of 30 days around the GA. Both LR and LGBM display similar profiles in terms of discrimination and calibration performance. Note that the proportion of pregnancies remaining at risk of miscarriage naturally decreases with time, which partly explains the increase of AUC for higher GA. AUC: area under the curve; GA: gestational age; LGBM: light gradient boosted machine; LMP: last menstruation period; LR: logistic regression.

Fig. 3

Raw features importance measured with averaged absolute SHAP values per variable. Variables that contribute significantly to the model's predictions for many patients have a high importance depicted as a large averaged absolute SHAP value. Main effects are colored in blue, interaction effects (first order variable interaction) are colored in orange. For LR models, the correlated SHAP approach (first column) takes variables collinearity into account when computing the SHAP values. The resulting feature importance is more balanced among correlated variables than the independent SHAP approach (second column) which directly reflects the LR coefficients. LGBM are reported with main effect only (third column) and with first order interactions (fourth and fifth columns, only the top 20 features). Despite similar performances, the algorithms have a different internal use of the same set of features. For example, the interaction MSD*FH is the main driving force of the LR model but appears as the 7th most important variable (second interaction term) in LGBM. CRL: crown-rump length; GA: gestational age; LGBM: light gradient boosted machine; LMP: last menstruation period; LR: logistic regression; MSD: mean sac diameter; PUQE: Pregnancy-Unique Quantification of Emesis.

Fig. 4

Longitudinal features importance measured as averaged absolute SHAP values per variable. In the LR model, under the independent SHAP values computation, the interaction term MSD*FH is much more important in the beginning than in the second half of the examined period, although it remains from far the main driving variable throughout all gestational ages. Under the correlated approach, this interaction term remains the most important variable, but the credit is now shared among other correlated variables. In the LGBM model, the CRL variable stays the main variable while the importance of MSD decreases with GA. On the other hand, the difference in estimated GA becomes more important in the second half of the examined period. CRL: crown-rump length; FH: fetal heart; GA: gestational age; LGBM: light gradient boosted machine; LMP: last menstruation period; LR: logistic regression; MSD: mean sac diameter; PUQE: Pregnancy-Unique Quantification of Emesis.

Fig. 5

Decision path plots. Example of individual prediction explained with the SHAP values for 3 instances of the test set (dashed line = correlated SHAP, plain line = independent SHAP). The SHAP values attributed to each variable fixed at their current value (as indicated below the x-axis) are gradually summed from left to right to explain the departure of the current prediction from the study population prevalence (gray dashed line), acting as a baseline viability probability. The bigger the magnitude of the SHAP value, the larger the variable contribution to the final prediction. The variables are ordered from lower to higher importance at the prediction level. For example, in the top graph, small values of CRL (2.6 mm) and MSD (4.8 mm) in combination with an older maternal age (41 years) and a large difference in estimated GA (17.2 mm) produce a low chance of viability from both models. CRL: crown-rump length; GA: gestational age; LGBM: light gradient boosted machine; LMP: last menstruation period; LR: logistic regression; MSD: mean sac diameter; Nan: missing; PUQE: Pregnancy-Unique Quantification of Emesis.

Fig. 6

Top three of interaction effects learned from LGBM and measured with the SHAP framework. The x-axis variable represents the first variable of the interaction. The color bar on the right indicates the value of the second variable present in the interaction. The SHAP value of the interaction is reported on the y-axis. (A). MSD*CRL, intrauterine pregnancies without visible embryo (CRL equals to zero) have a higher predicted risk of miscarriage (depicted as negative SHAP values) when MSD increases; (B). MSD*FH, pregnancies with high MSD without FH present have a higher predicted risk of miscarriage than with FH (see the negative SHAP values of yellow dots for MSD greater to 15 mm); and (C). Maternal Age * difference in GA: a large discrepancy in GA is modeled as a protective variable in young women while it becomes a risk factor in older women. CRL: crown-rump length; FH: fetal heartbeat, GA: gestational age; MSD: mean sac diameter.