Group Probability-Weighted Tree Sums for Interpretable Modeling of Heterogeneous Data

TL;DR

This paper introduces G-FIGS, a novel instance-weighted tree-sum method that improves predictive accuracy and interpretability for heterogeneous data, especially in healthcare, by leveraging group membership probabilities.

Contribution

The paper presents G-FIGS, a new method combining group probability estimation with decision tree sums to enhance modeling of diverse data while maintaining interpretability.

Findings

G-FIGS outperforms CART and FIGS in clinical datasets.

G-FIGS increases specificity by up to 10% at fixed sensitivity.

Final models remain interpretable with fewer than 16 rules.

Abstract

Machine learning in high-stakes domains, such as healthcare, faces two critical challenges: (1) generalizing to diverse data distributions given limited training data while (2) maintaining interpretability. To address these challenges, we propose an instance-weighted tree-sum method that effectively pools data across diverse groups to output a concise, rule-based model. Given distinct groups of instances in a dataset (e.g., medical patients grouped by age or treatment site), our method first estimates group membership probabilities for each instance. Then, it uses these estimates as instance weights in FIGS (Tan et al. 2022), to grow a set of decision trees whose values sum to the final prediction. We call this new method Group Probability-Weighted Tree Sums (G-FIGS). G-FIGS achieves state-of-the-art prediction performance on important clinical datasets; e.g., holding the level of sensitivity fixed at 92%, G-FIGS increases specificity for identifying cervical spine injury by up to 10% over CART and up to 3% over FIGS alone, with larger gains at higher sensitivity levels. By keeping the total number of rules below 16 in FIGS, the final models remain interpretable, and we find that their rules match medical domain expertise. All code, data, and models are released on Github.