Applied Bayesian Inference with PyMC3 and Bambi Part 3 | by Ani Madurkar | Towards Data Science

This is the final part of my Applied Bayesian Inference series. In part 1, I introduced the conditional world and probabilistic programming:
towardsdatascience.com
In part 2, I showed how you can scale your models up to “real” data by modeling Nike and Adidas shoes. This piece showed how you can run more effective A/B tests and group comparisons:
towardsdatascience.com
In this part, part 3, I will show why Bayesian modeling is so incredible by gently introducing Linear Regression in PyMC3 and then taking it further into Hierarchical Models, Generalized Linear Models, and Out-of-Sample Prediction.
This story will take a look at Spotify data, namely the Top 200 songs from 2020–2021. You can find the dataset here: https://www.kaggle.com/sashankpillai/spotify-top-200-charts-20202021. We’ll attempt to model the posterior of Streams for each song based on features related of each song (Danceability, Artist Followers, etc.)
Dtypes are a great place to start:
This will definitely need some cleaning up. Furthermore, because they’re objects, they’re likely to have weird text to clean up to make numeric. Between that, the biggest change I saw needed was the Genre column. Each song seems to fit ≥1 genre which is represented as a list. For our use case, we’ll want to One Hot Encode these column. This can easily be done with sklearn’s MultiLabelBinarizer.
That should be much better! Now we want to separate out the columns by type to 1. fix the dtype and 2. filtering down for a place to start. From there we can plot the Streams distribution to see what we’re targeting.
With Regressions, it’s common practice to center or standardize your values and it’s no different here. In fact, standardizing your values can be especially significant with MCMC models as the sampler will have a hard time converging otherwise. Standardizing can have another advantage for us which is being able to use weak informative priors since the Intercept will always be zero and the slope around -1 and 1. It also allows us to talk in terms of Z-scores, identifying how many standard deviations values are from the mean.
0
We’ll only work with these numeric fields for now; a later implementation of this can easily scale up to include the categorical columns we OneHotEncoded earlier (genre).
This is a common distribution shape we see in plenty of scenarios. It’s the consequence of having a lot of people at “lower” values and exponentially less as you go farther out in the tail to the right. It’s interesting we find it so strong even for the Top 200 songs, but this will be an important piece to remember as we go to model this posterior.
To start, we can take the plot the standardized log of Streams:
Still not the best, but much better. We clearly have some outliers that have seen some massive virality in the last year.
The best way to learn is to do, and when things are complex it’s best to start as simple as possible and then slowly scale up. The simplest example here is Linear Regression, so let’s see we wanted to model Streams as a function of Artist Followers. Maybe we have some premonition that the artist’s followers has some impact on Streams.
It’s good practice to first get an understanding of relationships; how does Artist’s Followers correlate to Streams for the Top 200?
Barely a correlation which seems to make sense as these are already highly popular songs — there could even be a reverse correlation occurring here (as a new song comes up to Top 200, that artist’s followers significantly increases). These kinds of insights are incredibly important to understand not only before modeling, but also when it comes time to draw conclusions from results. Everyone should understand by now that correlation != causation, but which direction may the correlation flow? And why? These can give amazing insight into how to fix problems in your model or what kind of data you may want to curate to improve your model.
What’s interesting here is that we are creating a prior for beta (the weight of x, the Artist’s Followers) and the bias and then creating a linear function (alpha + beta*x) to model the likelihood. Remember, we’re attempting to model the posterior of the likelihood as a function of priors, but when we do it this way we’re able quantify a distribution of values based on the uncertainty in the observed values. We want to wrap mu in a deterministic function because we don’t want it to be different each time; it should always follow the consequence of alpha + beta*x, but alpha and beta have stochastic priors that allow the model to scan a range of options for them.
Furthermore, we can create some deterministic metrics to measure our output pretty simply as well. Since we’re dealing with Regression, we can also calculate the r2 value in the model build itself. This creates a really easy way to identify how well the fitted line fits the data.
Now we can sample the data. NUTS (No-U-Turn Sampler) is a fantastic choice for most use cases. Although it can be slower per step than Metropolis, it is able to traverse diagonal spaces with less difficulties. Furthermore, it usually needs far fewer steps to get a reasonable approximation to the posterior.
We don’t see an noticeable divergences in the traces and the standardization/centering has significantly helped our model find appropriate distributions for each song. Let’s see our r2 and plot the findings.
So, our model performs great but we get quite a low correlation which should be obvious when seeing this scatterplot. There really isn’t much correlation to begin with , and this seems to pick up on it. There are a number of things we can do to improve on this model though. First, let’s use a StudentT distribution for our likelihood and see if we can make it more robust to outliers while also predicting out-of-sample.
We get a dramatic improvement in the line fitting to our data — ~.16! Technically, since we standardized the data, the slope and the Pearson correlation coefficient have the same value but this is a really interesting way to bake it into your modeling. Still not the best predictor of Streams, but we can see how the line is less likely to be picked up by potential outliers in our dataset. Although it would do us little good, we can also easily update our data and predict the posterior using this model on the test dataset.
Mean Squared Error: 0.6710557105827462
It is amazing how easy it can be to not only find a test point (mean value of each posterior distribution) but also a distribution of values to estimate the error! We can now take these skills to take this simple model and keep building on top of it.
Enter BAyesian Model-Building Interface — Bambi.
As before, if we’re going to incorporate more variables we should first visualize relationships and ensure minimum correlation between features and between features/predicted variable. We can do this effectively with my favorite visualization, the scatterplot matrix.
It seems ‘Energy’ and ‘Loudness’ are more correlated than we’d prefer. This can cause a lot of issues when attempting to model a relationship between these variables so we can remove one.
Note the last row is the Log_Streams value, so we don’t need to evaluate its correlations to Streams as we’re not modeling that (yet).
Using Bambi, we can easily write up our formula and cook up the model like so:
I’ll skip the traceplots and forestplots to analyze the data here, but make sure to use those to ensure your model is converging on each feature properly! Traceplots are fantastic ways to immediately visualize the outputs of a MCMC model.
Mean Squared Error: 0.5720862510645325
Note, we’re sampling 2 chains here on the posterior predictive to model the test predictions so this is why we’re taking two column-wise means — one for the mean of the two chains and two for the mean of the posterior distribution of each prediction (the test value).
The MSE is significantly better as we add more features of a particular song (Danceability, Valence, Energy, and more). From assessing the means in the fitted model’s summary, it is interesting that Followers appears to have the strongest effect to increase Log_Streams out of the other values. Features such as Danceability and Speechiness appear to have a negative effect on Log_Streams, but it is important to take these features with a grain of salt as I’m not entirely sure how Spotify records this information (let alone how someone may have manipulated it to get it in Kaggle).
An important question that you may have realized at this point is what about the variation of songs by artists? We are currently modeling each song individually (unpooled), but we also could model each artist’s songs (pooled). The issue with each of these is that you have to give something up for it. With unpooled, you lose the ability to see the nuance across artists; with pooled, you lose the granularity of each song’s distribution. So what can we do? Hierarchical partial pooling to the rescue!
This looks much more complicated, but it’s actually beautifully simple. The whole concept of Bayesian inference that we have learned so far is that you are able to bake in prior knowledge and model a posterior distribution which is suited to be updated as new knowledge is formed. So far we have only modeled the immediate priors (the mu and sigma priors on a Normal likelihood, for example), but what if we have a hierarchical nature to our data? Wouldn’t it be great if we could model the distribution of song metadata per artist? Short answer is that we can, and easily. This is a hierarchical model and it involves you adding hyperpriors to your priors; in this case, we’re adding grouping the song metadata features by Artist and adding hyperpriors to our priors along with our likelihood.
In the code block, we can signify we want a hierarchical model by adding the “(1|...)” to each feature and then simply adding Priors to values appropriately. The syntax can be broken down as such:
Common predictors are your main effects. For example, Danceability.
Group predictors can be identified by three different ways: 1. (Danceability|Artist) indicates we want to allow each artist to have his/her own slope and intercept for Danceability, 2. (1|Artist) indicates we wanted different intercepts for Artist but not slopes, and 3. (0+Danceability|Artist) indicates we only wanted slopes specific to each artist without including an artist specific intercept.
Obviously the less knowledge we may have, the wider priors we want to implement just the same as before. In any case, how did the model perform?
Unfortunately, Bambi does not have current functionality (version 0.6.3) to do out-of-sample predictions when you’re involving categorical grouping hierarchically yet, but our model does seem to learn the distribution just as well. The benefit of doing partial pooling is that we can not only see the posterior distribution in this regard, but also measure the impact each artist’s song attribute distribution had. Here’s an example of just a few for Popularity:
Although hierarchical modeling is extremely fascinating, it’s quite visible that we are not able to fully model the observed distribution as effectively. The reality in that is due to it’s non-gaussian shape. We’re attempting to model it with linear models that assume a level of normality which in reality, we just don’t see. And furthermore, our Streams don’t go below 0 ever so our Log_Streams will never go below 0 (the lowest Stream count possible is 1 and discrete). How can we model this?
Generalized Linear Models are incredibly easy to implement in Bambi, and they are the solution to our issue. Our Streams are not normally distributed by nature, and so we should think about how to effectively model their real nature. A good start is the guide I outlined in Part 1 — how do we believe this data is collected? Streams are a counting measure; they reflect the number of discrete times a song has been played on Spotify. Luckily for us, there is a popular GLM type that models exactly this: Poisson. Here’s how we’d implement it with Bambi and the results (Note how we’re modeling Streams directly now).
Important disclaimer: at time of writing I used the Streams field that was already standardized which isn’t correct to model as a Poisson distribution since this makes discrete count data into continuous data that lies below 0 (shoutout to Susia for catching this). Until I’m able to fix this, please make sure to know that the intended field to be modeled is the raw count of Streams.
Mean Squared Error: 0.2989611604295006
What a dramatic improvement! But why did the MSE get so much better? Actually, I’d wager that we could get it even lower but how? The answer to this question is the essence of Bayesian Machine Learning.
The reason the model performed dramatically better is because we carefully thought about the data and what it represents. We thought about the nuances of it and how it is collected. Throwing endless modeling techniques at a problem could improve your metrics, but the model doesn’t truly sing until you do the hard work of strategic design. Bayesian ML gives us the ability to do this really flexibly and also gets stronger with robust and precise priors, which is another thing we did/improve on. Allowing your priors a general space to learn in can create a really informative model that is able to converge quickly and learn effectively.
To wrap this up with a nice bow, let’s do model comparison to see how each performed a bit more clearly.
By default AriZ uses leave one out cross-validation. Another option is the widely applicable information criterion (WAIC). Since the results are in the log scale, the better out-of-sample predictive fit is given by the model with the highest value, which is the Generalized Linear Model as we saw earlier.
This wraps up my 3 part series on Applied Bayesian Inference and we have been able to do quite a lot in that time. Honestly, I started this series knowing very little about the field and knew I wanted to learn it deeply so I figured the best way would be to build with all the knowledge I’m finding as I read around. And what better addition to then document my journey for others to learn to?
I hope this series not only helped you learn this craft and dive into it yourself, but also opened your eyes to the beautiful world of Machine Learning. Learning Bayesian methods for ML gives you a greater appreciation for how a strong foundation in statistics can take you so far here. From modeling the uncertainty in your world, making predictions in spite of it, and updating your knowledge in light of new evidence.
The end of this series marks the official start for me with designing my brand in this space. Applied content that is meant to help you not only get started, but also understand concepts just enough to enrich your corner of the world. Follow along for more content related to the art & science of Machine Learning.
References
[1] Osvaldo Martin, Bayesian Analysis with Python
[2] PyMC3, GLM in PyMC3: Out-Of-Sample Predictions
[3] PyMC3, (Generalized) Linear and Hierarchical Linear Models in PyMC3
[4] PyMC3, GLM: Poisson Regression
[5] PyMC3, Hierarchical Partial Pooling
[6] PyMC3, A Primer on Bayesian Methods for Multilevel Modeling
[7] Bambi, Wald and Gamma Regression (Australian insurance claims 2004–2005)
[8] Bambi, Multi-level Regression
[9] Bambi, Logistic Regression and Model Comparison with Bambi and ArviZ
[9] Bambi, Hierarchical Logistic regression with Binomial family