Mixed Effects Random Forests in Python | by Sourav Dey | Towards Data Science

This blog post introduces an open source Python package for implementing mixed effects random forests (MERFs). The motivation for writing this package came from the models we have been building at Manifold. Much of the data we come across is clustered, e.g. longitudinal data from individuals, data clustered by demographics, etc. In many of these circumstances, we’ve found that using MERFs provide substantial improvements compared to vanilla random forests.
MERFs are great if your model has non-negligible random effects, e.g. there are large idiosyncrasies by cluster. You can pip install our package off of PyPi by typing:
pip install merf
The source code is available here. Contribute to it! The package is based on the excellent published work of Prof. Larocque from the Department of Decision Sciences at HEC and Prof. Ahlem from the Department of Marketing at l’UQAM.
Lots of data in the wild has a clustered structure. The most common example we see is longitudinal clustering, where there are multiple measurements per individual of a phenomena you wish to model. For example, say we want to model math test scores as a function of sleep factors but we have multiple measurements per student. In this case, the specific student is a cluster. Another common example is clustering due to a categorical variable. Continuing the example above, the specific math teacher a student has is a cluster. Clustering can also be hierarchical. For example, in the the example above there is a student cluster contained within a teacher cluster contained within a school cluster. When making this model, we want to learn the common effect of sleep factors on the math test scores — but want to account for the idiosyncrasies by student, teacher, and school.   There are four sensible model building strategies for clustered data:
Make a separate models per cluster. This strategy is unsatisfying for two reasons: (a) sharding by cluster makes it so that any given model does not have much training data and (b) the model can’t learn across the population.Make one global model for all clusters. With this strategy, the model can learn across the population, but won’t learn anything idiosyncratic about particular clusters. For example, some students may have higher mathematical ability, irrespective of sleep factors. It would be nice to model that effect explicitly, rather than lumping it into the sleep factors.Make one global model, but put in cluster id as a feature. This is a sensible strategy, but has scaling issues as the cluster size grows. Specifically, in a random forest, a single categorical variable can only have a small cardinality before the splitting decision becomes unwieldy. In R that is about n_categories = 30. In H20, that is n_categories = 64. The random forest implementation in Python Scikit-learn doesn’t even support non-binary categorical variables. One hot encoding cluster is an option, but this has known bad performance as explained in this excellent blog post (and as we’ll show below). Long story short, directly using high cardinality categorical variables as features in a model sucks.Mixed effect model. This is the right way to attack clustered data. In a mixed effect model, each cluster gets a random effect that is learned but drawn from a prior that is itself learned from the data. As explained below, this is not a new idea, statisticians have been doing this forever. But mixed effects random forests are novel — they combine the best of linear mixed effects models with the power of non-parametric modeling, where you don’t need to understand the “physics” of the problem.
Make a separate models per cluster. This strategy is unsatisfying for two reasons: (a) sharding by cluster makes it so that any given model does not have much training data and (b) the model can’t learn across the population.
Make one global model for all clusters. With this strategy, the model can learn across the population, but won’t learn anything idiosyncratic about particular clusters. For example, some students may have higher mathematical ability, irrespective of sleep factors. It would be nice to model that effect explicitly, rather than lumping it into the sleep factors.
Make one global model, but put in cluster id as a feature. This is a sensible strategy, but has scaling issues as the cluster size grows. Specifically, in a random forest, a single categorical variable can only have a small cardinality before the splitting decision becomes unwieldy. In R that is about n_categories = 30. In H20, that is n_categories = 64. The random forest implementation in Python Scikit-learn doesn’t even support non-binary categorical variables. One hot encoding cluster is an option, but this has known bad performance as explained in this excellent blog post (and as we’ll show below). Long story short, directly using high cardinality categorical variables as features in a model sucks.
Mixed effect model. This is the right way to attack clustered data. In a mixed effect model, each cluster gets a random effect that is learned but drawn from a prior that is itself learned from the data. As explained below, this is not a new idea, statisticians have been doing this forever. But mixed effects random forests are novel — they combine the best of linear mixed effects models with the power of non-parametric modeling, where you don’t need to understand the “physics” of the problem.
Linear mixed effects (LME) modeling is a classic technique. Let’s look at this in some detail because it motivates the MERF model. The LME model assumes a generative model of the form:
In the equation above:
y is the target variable.
X is the fixed effect features. X is assumed to be p dimensional, e.g. there are p features.
Z is the random effect features. Z is assumed to be q dimensional, e.g. there are q features.
e is independent, identically distributed (iid) noise. It is distributed as N(0, sigma_e2)
a is the fixed effect coefficients. They are the same for all clusters.
i is the cluster index. We assume that there are k clusters in the training.
bi is the random effect coefficients. They are different per cluster i but are assumed to be drawn from the same distribution. In the classic LME problem, this distribution is assumed to be N(0, Sigma_b) where Sigma_b is learned from the data.
The LME is a special case of the more general hierarchical Bayesian model. These models assume that the fixed effect coefficients are unknown constants but that the random effect coefficients are drawn from some unknown distribution. The random effect coefficients and prior are learned together using iterative algorithms. This article is not about hierarchical Bayesian models though. If you want to learn more about them this is a great resource.   Though hierarchical Bayesian modeling is a mature field, they require you to specify a functional form to the regression, i.e. you need to know the physics of the problem. Unfortunately, in many situations, particularly in complex systems, it’s hard to specify a functional form. This is where the random forest shines. Sinces it cuts up feature space, it can act as a universal function approximator. Our work at Manifold led us on a search to combine random forests with the power of mixed effects.
We found the answer in the excellent work of Prof. Larocque’s group at the Department of Decision Sciences at HEC Montreal. In a series of papers, they have illustrated a methodology to combine random forests with linear random effects. More importantly, Prof. Larocque has been very generous with his feedback to make the MERF Python package a reality.
Similar to LME, the MERF model assumes a generative model of the form:
Note all the terms are the same as the LME, except the linear fixed effect a*X is replaced with a general non-linear function f(.). In the paper, this non-linear function is learned using a random forest. More generally, f(.) can be any non-linear regression model like gradient boosting trees or a deep neural network.  An expectation-maximization (EM) technique is used to fit the MERF. There are four parameters to fit:
f(.)
bi for all known clusters.
Sigma_b
sigma_e
We won’t go into the math here — the paper does it much better — but the basic intuition behind the EM algorithm is alternative optimization, e.g. optimize one parameter while keeping the others fixed. You keep doing this iteratively until convergence. Specifically, to fit MERF the steps are:
Fix all the bi and compute y* as y — bi*Z. Fit a random forest, f(X) to y* globally across all samples.Fix f(), Sigma_b, sigma_e. Optimize to find bi*. There is a closed form solution assuming a linear random effect and Gaussian prior.Fix f(), bi. Optimize to find Sigma_b and sigma_e. There is a closed form solution assuming a linear random effect and Gaussian prior.
Fix all the bi and compute y* as y — bi*Z. Fit a random forest, f(X) to y* globally across all samples.
Fix f(), Sigma_b, sigma_e. Optimize to find bi*. There is a closed form solution assuming a linear random effect and Gaussian prior.
Fix f(), bi. Optimize to find Sigma_b and sigma_e. There is a closed form solution assuming a linear random effect and Gaussian prior.
The generalized log likelihood (GLL), which can be computed at every iteration, is a measure of training loss. As the GLL goes down the fit improves. Once it flattens out we can usually stop iterating — the fit has converged. There is not much danger of overfitting with MERF — no more than the classic random forest which is fairly robust to overfitting.  Once fit the MERF can be used to predict. For data in “known” clusters, that the MERF saw in training, the prediction includes the random effect correction:
For data in “new” clusters, that the MERF did not see in training, the prediction only includes the fixed effect:
Our contribution was to implement a MERF in Python and open source it. We implemented so that it adheres, as much as possible, to the scikit-learn model interface. We’ll illustrate with an example below. This example can be downloaded as a runnable Jupyter notebook here.   We have three matrices containing our n_samples =500 of training data:
X. Matrix containing three fixed effect features. Dimension = 500 x 3.
y. Vector containing the single target variable. Dimension = 500 x 1.
clusters. Vector containing the cluster_id for each sample. Dimension = 500 x1. We have k = 100 unique clusters in the training data.
In this example, there is not an explicit Z matrix. We create one to model a random mean for each cluster. It is a matrix of all 1s with dimension = 500 x 1.  Given the X, y, clusters, and Z matrices we can fix a MERF model to this data by instantiations a MERF model and running the fit method. The MERF model currently has two hyperparameters:
n_estimators. The number of trees in to train for the random forest.
max_iterations. The maximum number of iterations to run for the EM model.
We eventually plan on implementing early stopping for the EM algorithm as well. This will add an additional two hyper parameters, i.e. whether or not to early stop and what the patience on that should be.   Here’s what the output of fit looks like. Note that fitting usually takes some time to run — especially if you run for many iterations. This is unavoidable because each EM iteration requires a random forest fit. Even with parallelization of the random forest fit a single MERF fit can take many minutes.
Once the model is trained we can access the final random forest and the final matrix of trained b_is through two public variables:
mrf.trained_rfmrf.trained_b
We can look at the distribution of the b’s. It looks approximately Gaussian with variance Sigma_b as expected:
The model also host the history of the EM iterations. We can look at the generalized log likelihood (GLL) at every step, the estimates of Sigma_b, sigma_e, and all the bi’s.
Once fit, the model can be used to predict on new samples given X, Z, and clusters. The predict code handles whether or not to apply the random effect correction based on if the cluster_id of the new sample was seen in training or not.
We tested the performance of MERF on synthetic data. The experimental setup mimics that of the published MERF paper. We generate data from this function:
In this equation:
m is a fixed constant of our choosing.
the bi are drawn from normal distribution N(0, sigma_b). They are constant per cluster i. sigma_b is a fixed constant of our choosing.
the error, e, is drawn IID from normal distribution N(0, sigma_e). sigma_e is a fixed constant of our choosing.
For our experiments, we choose various values for m, sigma_e, and sigma_b and generated 500 samples from 100 different unbalanced clusters. We summarize the values in the table below, but the specific values are not that important. What matters are two derived parameters, the percentage total effects variance (PTEV) and the percentage random effects variance (PREV). They are defined in detail in the paper, but intuitively they are:
PTEV = the amount of total variance that is not due to random noise. The higher it is, the more predictable the output y is.
PREV = the amount of PTEV that is due to random effect. The higher is it, the larger random effect there is.
We did an experiment comparing MERF to two competitors:
a random forest that does not take random effects into account, i.e. fitting a random forest on only X.
a random forest that includes one hot encoded cluster id as features, i.e. we create k new one hot encoded features, X_ohe, and join that with X and train a random forest on the combined [X, X_ohe] feature matrix.
We varied PTEV and PREV for the experiments and assessed performance on two different test sets:
test set consists of known clusters that the model trained on
test set consisted of new (unknown) clusters that the model had no knowledge of at training time
We use the average mean squared error (MSE) over 10 runs as our error metric. We define the MERF gain as the percentage gain over the compared algorithm:
The table and chart below summarize the results. MERF has marked performance gains over both of the compared models — particularly when the PREV is high (over 50%!). In addition, there is negligible loss in performance when predicting new clusters. MERF is a strictly better model when there are significant random effects in the data. It’s a winner.   We have included the data generator as a utility module in merf package. The source code for the experiment is in a Jupyter notebook here. Check it out and run your own experiments!
Manifold is working on a thermal modelling problem for one of our clients. The problem is to predict, everyday, when a building is going to reach its target temperature based on a number of factors, e.g. outdoor temperature, outdoor humidity, indoor temperature the previous day. We have historical data that is clustered longitudinally by building. The clusters are highly imbalanced. For some buildings we have many years of data; for others we only have a few weeks. This problem is begging for a mixed effects model. We want to build one model that learns across the population, but, at the same time, accounts for the idiosyncrasies of each building.
Similar to the experiment above, we compared MERF against:
a random forest that does not take random effects into account, i.e. fitting a random forest on only the fixed effect features.
a random forest that includes one hot encoded building id as features, i.e. we create k new one hot encoded features, X_ohe, and join that with the fixed effect features and train a random forest on the combined [X, X_ohe] feature matrix.
Because of outliers, we changed the error metric to the median absolute error (MAE) rather than the MSE. We compare the performance using 5-fold cross validation on a small N~1000 data set. The results of the experiment are summarized below.
For this data set, there is a small MERF gain, but it’s not that significant. Looking at the MERF convergence plot we get some clues.
The GLL is flat even though we run for 100 iterations and the other variables are converging. It’s not clear what that means. In addition, the estimated error variance, sigma_e, is very significant. It seems like the data has a very low PTEV (because of outliers?). The one hot encoding (OHE) model is doing ok because the number of buildings is small — k = 15. As the number of buildings scale the OHE model performance will suffer more.
All of these factors, coupled with the fact that there is not that much data, leaves this experiment a bit inconclusive. MERF is not performing poorer than it’s counterparts — so it doesn’t hurt — and in fact has a slight gain on average. This experiment needs more analysis — when we do we’ll publish them on this blog.
Mixed effects models are powerful — and lots of data has a structure that is amenable to using them. Add them to your repertoire of models. Mixed effects random forests (MERFs) are a piece of the puzzle. They combine the best of random forests with the best of mixed effects models. We hope you find them useful.   More importantly, we’re looking for help extending the MERF package! We would love more contributors. We have a bunch of open issues to work on:
Add early stopping.
Add more learners, e.g. gradient boosting trees, xgboost, etc.
Add generalization to classification using GMERFs. See paper here.
Conduct novel research into using a nonlinear mixed effect, rather than a purely linear one.
Send us a pull request!
Manifold is an AI studio with offices in Silicon Valley and Boston. Leading companies work with us on new AI solutions when the stakes are high. We accelerate time to market, create new IP, and foster in-house capabilities. Manifold’s team is experienced in conceiving of and deploying production machine learning systems and data platforms at places like Google, Facebook, Qualcomm, MIT, and successful venture-backed startups. Our domain expertise spans electronics (consumer and industrial), wireless, healthcare, and marketing. We have an extended network of experts that we draw on across industry and engineering sub-specialties.
Wanna work on stuff like this? Reach out at sdey@manifold.com! We’re hiring data scientists.