Bayesian Hierarchical Modeling in PyMC3 | by Dr. Robert Kübler | Towards Data Science

From time to time, you have the agony of choice when trying to do machine learning on a quite heterogeneous dataset. As an example of what I mean with a heterogeneous dataset, consider a simple height of people dataset.
The really important feature of this data is the height. However, we can also find two different groups here: men and women. The sub-dataset consisting of only men is homogenous on its own, so is the sub-dataset consisting of only women, but putting both together creates a heterogeneous dataset.
Given some heterogeneous data, you can either:
build one big model on your complete dataset, or
build several models on smaller, more homogenous parts of your dataset.
Building one big model, also known as (complete) pooling, is usually the easiest approach: you throw all samples together and forget about the different groups. However, if the big model is too simple it can ignore the different nuances in the data which in turn leads to underfitting. This problem can happen real quick when using highly interpretable models such as linear models. If you can and want to use black-box methods such as gradient boosting, the model can spot and learn the different sub-datasets by itself, but at the cost of reduced interpretability.
Building several smaller models seems like a better idea: you slice your dataset into different more homogenous pieces and build a smaller model for each of them. The resulting model is then called an unpooled model. This way, each model can be good on a small portion of the data, and together they should form a decent model. The obvious drawback with this approach is that you have to fit many models, but an even more severe problem emerges when the small datasets get too small: overfitting.
To illustrate this, just imagine that our height dataset consists of 100 samples: 97 women and 3 men. We want to estimate the mean height in each group now by taking the sample mean of the heights. While the guess for the women should be quite good, the result for the men will be highly unreliable because of the tiny amount of observations. The main problem is the following:
The sub-models are completely independent. Each sub-model that is trained does not pass any information to the other sub-models. In a sense, the total model forgets what it is doing on each sub-dataset.
While training isolated models can make sense in certain situations, sub-models sharing knowledge is better more often than not, in my opinion.
As a real-life example, the waiting time of getting a cheeseburger at The Tavern to the Golden M in Germany takes two minutes on average (I made up this number). If I go to France now, chances are that the average waiting time is not too far off from these two minutes. And probably the same goes for Japan. Sure, some countries are faster, some are slower on average, but waiting for a cheeseburger for an hour should be a rather rare exception, regardless of the country. So, if I am hungry and want a cheeseburger in France, I have an expectation of how long I have to wait, and probably I will be kind of right. I can transfer knowledge learned from one group (Germany) to make predictions for another group (France). The same goes for the heights dataset.
Of course, you can also create situations where it is misleading to remember things that you have learned. The waiting time at the doctor’s office probably does not give you any indication about the duration of the green phase of a traffic light, so you should not try to share the knowledge between these two unrelated variables.
Nevertheless, if we are honest, usually we are not presented with data that comes from completely unrelated and random sources. Usually, it is rather consistent, just with some variations across groups like country, sex, shop, ..., hence
I think sharing knowledge is often a good default.
So, how to give our model a better memory? Or, put another way, how to teach the sub-models to communicate and work together? There are several approaches to this, for example, parameter or weight sharing in the case of neural networks, which I am not going to elaborate on here.
However, as you know from the title of this article, we will go Bayesian instead and learn how to do Bayesian hierarchical modeling as a possible solution to our problem. We will use PyMC3 again, as it is a nice package.
In case you have never heard of Bayesian statistics or PyMC3, and want to learn it, please check out my other introductory articles about this fascinating topic.
towardsdatascience.com
towardsdatascience.com
Let’s go!
First, we will revisit both, the pooled and unpooled approaches in the Bayesian setting because it is
a nice exercise, andthe codebases of the unpooled and the hierarchical (also called partially pooled or multilevel) are quite similar.
a nice exercise, and
the codebases of the unpooled and the hierarchical (also called partially pooled or multilevel) are quite similar.
Before we start, let us create a dataset to play around with. We will create a simple one-dimensional regression problem, i.e. there is a single feature and a single target. There are eight different groups, each with its own slope, and a fixed intercept of zero.
import numpy as npnp.random.seed(0) # to keep it reproduciblemean_slope = 2 # the 8 different slopes have a mean of 2slopes = np.random.normal(mean_slope, size=8)groups = np.array(50*[0, 1, 2, 3, 4, 5, 6] + 5*[7])x = np.random.randn(355)y = slopes[groups] * x + 0.1*np.random.randn(355)
The groups variable contains the group each observation belongs to. There are fifty observations each for the groups 0 to 6. Then, there is the minority group number 7 at the end, with only five observations.
As a small guide, to get the elements belonging to group 2, you can use x[groups==2], y[groups==2] . The slope of this sub-dataset is slopes[2] .
To spice things up, let us add outliers to this minority group.
y[-1] = 30 # the last 5 observations are from minority group 7y[-2] = 30
After importing our favorite Bayesian libraries
import pymc3 as pmimport arviz as az
we can start modeling!
We will start gently by ignoring the groups and treat this dataset as one big block. A simple Bayesian linear regression without intercept in PyMC3 can look like this:
with pm.Model() as pooled_model:    slope = pm.Normal('slope', 0, 20)    noise = pm.Exponential('noise', 0.1)        obs = pm.Normal('obs', slope*x, noise, observed=y)        pooled_trace = pm.sample(return_inferencedata=True)az.plot_posterior(pooled_trace, var_names=['slope'])
The output will be
This is basically useless because we only get a single slope, but this is not surprising. The next best thing to do is to introduce one slope per group.
with pm.Model() as unpooled_model:    slope = pm.Normal('slope', 0, 20, shape=8)    noise = pm.Exponential('noise', 10)        obs = pm.Normal('obs', slope[groups]*x, noise, observed=y)        unpooled_trace = pm.sample(return_inferencedata=True)az.plot_posterior(unpooled_trace, var_names=['slope'])
We get eight different slopes now that were trained independently:
We can also plot only the 94% high-density intervals (HDI), i.e. short credible intervals containing 94% of the posterior mass in a single figure via
az.plot_forest(unpooled_trace, var_names=['slope'], combined=True)
We get
You can see how groups 0 to 6 have small slopes, and group 7 being way out there. But this is totally wrong because all of our slopes should be around a value of two. What happened? Simple: we tricked the model by introducing outliers to the smallest group.
This is the exact problem I was talking about earlier: the sub-model for group 7 has no chance, as it does not know what is going on in the group 0 to 6. It has no clue about the slopes usually being around a value of two.
So, let us fix the problem.
with pm.Model() as hierarchical_model:    mu_slope = pm.Normal('mu_slope', 0, 1) # hyperprior 1    sigma_slope = pm.Exponential('sigma_slope', 13) # hyperprior 2        slope = pm.Normal('slope', mu_slope, sigma_slope, shape=8)    noise = pm.Exponential('noise', 10)        obs = pm.Normal('obs', slope[groups]*x, noise, observed=y)        hierarchical_trace = pm.sample(        return_inferencedata=True,        target_accept=0.995    )az.plot_posterior(hierarchical_trace)
So, what is this now? In the unpooled model, we made up priors for the slopes via
slope = pm.Normal('slope', 0, 20, shape=8)
We told the model that the slopes should be around zero, but with a quite large standard deviation of twenty.
In the hierarchical model, we define so-called hyperpriors to find a better mean and standard deviation for the slope priors. The important part is really only this:
mu_slope = pm.Normal('mu_slope', 0, 1)sigma_slope = pm.Exponential('sigma_slope', 13)    slope = pm.Normal('slope', mu_slope, sigma_slope, shape=8)
We replace the zero with mu_slope and the twenty with sigma_slope , easy as that. Both of these are random variables as well, which we can learn using Bayesian inference. mu_slope and sigma_slope are called hyperpriors, the same way the maximum depth of a decision tree is called a hyperparameter, for example. Both of them are one hierarchy level above slope as they have to be evaluated first before slope can be calculated.
Note: I added target_accept=0.995 to the sample method to improve the sampling, because sampling from this nested posterior is not as easy anymore as for the unpooled model. I can explain this further in another article.
The posteriors look like this:
This is still not perfect, but much closer to the truth than with the unpooled model. That is because the model also tries to figure out the slope mean using the mu_slope hyperprior. The true prior should be two, however, due to the outliers, the model believes that it is about three.
Still, the large slope of fifteen that we have seen in the unpooled model gets pulled towards this mu_slope of around three. This is also true for the other parameters:
In hierarchical modeling, all parameters are being pulled towards a global mean. This effect is known as shrinkage.
All in all, the hierarchical model is the clear winner. As the last step, let us compare the three models visually using the plate notation. This is a quite self-explanatory visual representation of the models that might help the visual learners among you.
PyMC3 has the wonderful function pm.model_to_graphviz(model) to display your model in a nice way.
I like this representation because you can clearly see some kind of evolution of the models. In the pooled model, we start with a single slope. We then increase the number of slopes, indicated by the box with the number 8.
This notation is a shortcut for drawing 8 circles for the slopes. I guess the small box is supposed to be a table, with 8 plates stacked on top of each other that we view from above, hence the name.
The last step of this process is introducing a new layer of variables on top, the hyperpriors.
In this article, we have discussed different approaches to modeling when dealing with heterogeneous datasets. At first, I claimed that there are two ways: pooled and unpooled models. Clearly, I was lying, as we have seen hierarchical modeling as well.
An interesting observation is that you can see hierarchical modeling as a generalization of unpooled modeling. If you set the hyperpriors to some constant random variables, you end up with the unpooled approach again. Otherwise, the hierarchical model does something different by keeping track of global parameters such as mu_slope which serves as the seed for the actual parameters we want to estimate. As a nice side effect, we get estimations for these global parameters as well. And if there is nothing worth sharing among the different groups, the hierarchical model will learn this as well, which is neat. That’s why I believe that you should use the hierarchical approach whenever you can instead of the unpooled approach: usually it does not hurt.
The only thing that can go wrong when using hierarchical modeling instead of unpooled modeling is that there are more parameters to estimate, and estimation also becomes harder because of the nested parameters. You have to use tricks like changing the MCMC sampling parameters or introducing so-called non-centered variables.
A centered (I would call it nested) variable is something like pm.Normal('X', a, b) for two other random variables a and b , this is what we have used above. However, in the special case of a normal distribution, you could also write a + b*pm.Normal('X', 0, 1) , which is equivalent from a statistical point of view, but a huge difference for the MCMC algorithm that computes the posterior distributions.
Still, if you can overcome the computational issues, I think hierarchical modeling is often a path worth trying.
As an additional resource, check out this nice blog post by Thomas Wiecki and Danne Elbers.
I hope that you learned something new, interesting, and useful today. Thanks for reading!
As the last point, if you
want to support me in writing more about machine learning andplan to get a Medium subscription anyway,
want to support me in writing more about machine learning and
plan to get a Medium subscription anyway,
why not do it via this link? This would help me a lot! 😊
To be transparent, the price for you does not change, but about half of the subscription fees go directly to me.
Thanks a lot, if you consider supporting me!
If you have any questions, write me on LinkedIn!