Gradient Boosted Mixed Models: Flexible Estimation of Mean and Variance Components for Clustered Data

Abstract

We introduce a novel way to combine gradient boosting with mixed effects models, whereby the mean and variance components are learned jointly as functions of covariates via likelihood-based gradients. Gradient Boosted Mixed Models (GBMixed) estimates a nonparametric fixed effects function characterizing the overall mean of the response, while also allowing the random effects covariance matrix along with the residual variance to depend on covariates in a flexible manner. We demonstrate how GBMixed facilitates covariate-dependent random effect predictions, and subsequently point predictions and prediction intervals for individual treatment effects, that can adapt between population-level and cluster-level information. Experiments and applications to two real-world datasets show that GBMixed can accurately recover complex nonlinear fixed effect functions and covariate-dependent covariances in a linear mixed model, while also improving point and probabilistic predictive performance compared with several existing approaches such as parametric linear mixed models, Natural Gradient Boosting, and Gaussian Process Boosting. In simulations where the variance components are designed to vary as a function of covariates, GBMixed reduces the mean squared error in recovering the random effects variance function by a factor of eight relative to linear mixed models and Gaussian Process Boosting.