Using tree-based ensemble methods to produce a population-based mortality risk score in Ontario, Canada

Abstract

Introduction Risk adjustment is critical in observational epidemiology to control for confounding of the exposure-outcome relationship. Accurate prediction of outcomes, such as mortality, can improve risk adjustment. In the present study, we compared logistic regression with a range of tree-based ensemble methods to predict 1-year mortality in the general population of Ontario, Canada. Methods Ontario adults (age 18 years and older) who were alive as of January 1, 2022 were included. Using a window of up to 3 years, various measures of health and healthcare utilization were captured from administrative databases. To predict 1-year mortality, we applied logistic regression, random forests, extremely randomized trees, adaptive boosting, gradient boosting, extreme gradient boosting, Newton boosting, and CatBoost. All models also included age and sex. Performance was evaluated using the area under the ROC curve (AUROC), the area under the precision-recall curve (PR-AUC), the Brier score, and a quantile-based version of the Integrated Calibration Index (ICI), reported in the 30% test set. Feature importance was assessed using CatBoost’s internal model structure, supplemented with permutation feature importance, explainable boosted machines, and marginal effects. Results A total of 12,080,801 Ontarians were included and 121,951 (1.0%) died within 1 year. Logistic regression showed excellent discrimination (AUROC 0.926; PR-AUC 0.256) and acceptable calibration (ICI 0.0022). The best model was CatBoost, which had the best discrimination (AUROC 0.933, PR-AUC 0.280) and calibration (ICI 0.0003). In sensitivity analyses of the CatBoost model, including more detailed definitions of cancer (to include its subtype) and chronic kidney disease (defined using serum creatinine instead of diagnostic codes) produced modest improvements in PR-AUC (0.290), along with substantially improved calibration amongst the highest-risk (70–100%) individuals. The most influential model-building feature was age. Residence in long-term care and receipt of palliative care was associated with the largest marginal effects. Conclusion The machine learning model CatBoost yielded the most accurate predictive model for 1-year mortality using individual comorbidities and additional measures of healthcare utilization for the general population. These findings demonstrate that machine learning methods can enhance risk adjustment efforts in observational studies, leading to more accurate confounder control and better support for health policy and epidemiologic research.