Weighted Sum-of-Trees Model for Clustered Data

TL;DR

This work introduces a weighted sum-of-trees model for clustered data, addressing out-of-sample group prediction by first learning group similarity weights from a classifier and then combining predictions from a per-group tree (or forest) ensemble. Predictions for a new observation are a weighted linear combination of the training-group trees, $\hat{y}_t = \sum_{j=1}^J w_j \hat{y}_j(X_t)$, where $w_j$ reflects similarity to group $j$. The approach yields improved predictive performance over traditional DTs, RFs, and LMMs in simulations, and demonstrates competitive, interpretable results on TCGA sarcoma data with clear group-specific insights via VIVI plots. The method offers practical advantages for precision medicine by enabling out-of-sample predictions and providing group-level interpretability without heavy bootstrapping, with future work extending to categorical and survival outcomes.

Abstract

Clustered data, which arise when observations are nested within groups, are incredibly common in clinical, education, and social science research. Traditionally, a linear mixed model, which includes random effects to account for within-group correlation, would be used to model the observed data and make new predictions on unseen data. Some work has been done to extend the mixed model approach beyond linear regression into more complex and non-parametric models, such as decision trees and random forests. However, existing methods are limited to using the global fixed effects for prediction on data from out-of-sample groups, effectively assuming that all clusters share a common outcome model. We propose a lightweight sum-of-trees model in which we learn a decision tree for each sample group. We combine the predictions from these trees using weights so that out-of-sample group predictions are more closely aligned with the most similar groups in the training data. This strategy also allows for inference on the similarity across groups in the outcome prediction model, as the unique tree structures and variable importances for each group can be directly compared. We show our model outperforms traditional decision trees and random forests in a variety of simulation settings. Finally, we showcase our method on real-world data from the sarcoma cohort of The Cancer Genome Atlas, where patient samples are grouped by sarcoma subtype.

Figures (8)

• Figure 1: (A) Consider grouped data such that each training observation falls into some group (e.g., patients clustered by hospital or disease subtype). The test data are similarly structured but come from new groups not seen in the training data. (B) We then construct a classifier that predicts the group that each new observation belongs to, and we extract the predicted probabilities of group membership. (C) Finally, we construct independent trees or forests on each group seen in the training set. The final output is a linear combination of the predictions from all of these trees with the probabilities from (B) as weights.
• Figure 2: Simulation setting 1. Mean squared error (MSE) over a range of noise values $\sigma_\alpha$ for settings where $\mathbf{U} = \mathbf{I}$. The scale of the response variable $y$ is standardized. Each boxplot represents MSE values across 20 simulated datasets. Simulations were performed for the settings $n=10$, $K=40$ (left), $n=K=20$ (center), and $n=40$, $K=10$ (right).
• Figure 3: Simulation setting 2. Mean squared error (MSE) of the various methods tested over a range of random noise values $\sigma_\alpha$ for settings where $\mathbf{U}$ is generated from an inverse-Wishart distribution. The scale of the response variable $y$ is standardized. Each boxplot represents MSE values across 20 simulated datasets. Simulations were performed for the settings $n=10$, $K=40$ (left), $n=K=20$ (center), and $n=40$, $K=10$ (right).
• Figure 4: Simulation setting 3. Mean squared error (MSE) of the various methods tested over a range of values for the number of observations per group, $n$, for settings where data is generated according to one of 3 distinct data generating processes. Throughout all simulations, $K=20$ groups were used. The scale of the response variable $y$ is standardized. Each boxplot represents MSE values across 20 simulated datasets.
• Figure 5: Principal component analysis of the The Cancer Genome Atlas (TCGA) dataset. The points are colored according to that observation's sarcoma subtype. Additionally, the ellipses drawn correspond to the normal data curve fit to each sarcoma subtype with coverage probability of 0.68. Key: DDLPS = dedifferentiated liposarcoma, MFS = myxofibrosarcoma, MPNST = malignant peripheral nerve sheath tumor, SS = synovial sarcoma, STLMS = soft tissue leiomyosarcoma, ULMS = uterine leiomyosarcoma, UPS = undifferentiated pleomorphic sarcoma