Probabilistic Programming with Pyro and a Kitchen Scale
By: Boris Shabash
Tags: Statistics, Pyro, Data Science, Pytorch

One of the techniques which has become of interest lately is the use of probabilistic programming. If you’ve never heard of probabilistic programming before, you might conclude it is just programming with probabilities. That would be correct, but only part of the story.

At its core, probabilistic programming is about finding underlying distributions which affect processes of interest. You define a process and how it works, and then allow your Probabilistic Programming Language (PPL) of choice to infer what are the likely values for factors which influence your process, and how likely they are.

This is explained best with an example, and for this example, I’ll be using Pyro as the PPL of choice, as well as a simple case study largely inspired from the official Pyro examples.

Building Our Process — A Magical Kitchen Scale

Imagine you have a kitchen scale, you put an object on it, and the scale tells you its weight. But imagine the scale is quite old, and is not very accurate. Every time you put the same object on it, it gives you a slightly different measurement. However, you observed that this scale, in some magical way, errs in a way that forms a normal distribution around the true weight of the item, with a standard deviation of 0.1kg

Let’s describe this process (we’ll be using PyTorch, which is what Pyro uses as its underlying tensor calculation engine).

First, we will import our required libraries

Let’s define our kitchen scale, or more specifically, the process it employs to give us observations

Now let’s ask ourselves an easy question;

“If we put a 0.5kg item on this scale, what are some possible values we will get back?”

tensor(0.6541)

tensor(0.4707)

tensor(0.2821)

We got back a collection of 3 values. Each time we ask our scale for an answer, we’ll get a slightly different value. Sometimes the values will be very close to 0.5 (this will happen most of the time), and sometimes they will be very far. This was easy enough to implement though, nothing new in terms of programming there.

Let’s try a slightly more difficult question;

“What is the probability of observing a value above 0.63kg, if I place an item which weighs 0.5kg on my scale?”

We can answer that programmatically as well. We’ll just generate many values, count all the values we generated, count all the values above 0.63, and find their ratio:

rough estimate: 0.097

reasonable estimate: 0.101

good estimate: 0.09691

true estimate: 0.0968004845856103

If we run the process enough times, we’ll get an estimate that is good for all intents and purposes. This took a bit more work, and some more time, but overall this workflow wasn’t too hard either.

Let’s try an even more difficult question:

“I’ve observed a value of 0.63kg, what is the likeliest true weight of the item?”

This might take you by surprise, but after a bit of thinking, you might recall your magical inaccurate-yet-well-behaved-scale reports a normal distribution around the true weight of the item, so with no further information, the most likely weight is in fact 0.63kg.

However, this was largely easy because we know the distribution is normal around the true weight of the scale. What if it wasn’t normal? what if it wasn’t even something pre-built into Python?

Now let’s try a hard question that will motivate our use of probabilistic programming:

“Let’s imagine we got the following set of observations, from our scale: 0.74kg, 0.98kg, 0.66kg, 0.75kg, 0.84kg, and 0.74kg, what is the likeliest true weight of the object?”

The true weight I used to generate those values is 0.8kg, but let’s see how we might find that out

Building Our Process in Pyro

The question just in the last section is exactly the type of question probabilistic programming is designed to answer. The general form is “given some observations, and a good understanding of the process, can I infer something about the hidden values or initial conditions in my process?”

Now we would like to know, what is the likeliest weight of our item. We’ll now employ Pyro to answer this question.

The mean of the observations is: 0.7850000262260437.

First, let’s define our process again, but with a few changes. We’ll also rename the function from process to model since that is usually how Pyro code is written:

This might seem like a lot, so let’s unpack how that function works.

We are trying to find a distribution over possible values of the weight, given the observations we have.

We first define a prior distribution on those values, our best guess for what the probability distribution might be, at (1).

Then, we sample from that distribution at (2). However, this is now a named distribution (we named it "weight1" ). This is our way of telling Pyro to optimize the distribution from which the weight came from. Behind the scenes, Pyro will modify this distribution to be more in-line with our observations.

). This is our way of telling Pyro to optimize the distribution from which the weight came from. Behind the scenes, Pyro will modify this distribution to be more in-line with our observations. At (3), we define the way the scale works. Recall that to produce measurements, the scale returns values from a normal distribution centered at the weight, and with a standard deviation of 0.1.

Finally, at (4), we sample from the distribution we defined at (3). Notice these are also named samples, but they are conditioned on the observations. This is where we tell Pyro how our value for "weight1" relates to the observations.

Approximating Our Weight Using Sampling Techniques

Now we are going to use our model to find a suitable estimate for the weight of the item that gave us our measurements. We’re going to use a family of algorithms called Markov Chain Monte Carlo (MCMC), and a particular instance called Hamiltonian Monte Carlo (HMC). For an in-depth explanation of MCMC and the original Metropolis sampling algorithm (on which HMC is built), I encourage you to look at this talk by Zach Anglin

In essence, this algorithm will build a probability distribution of our named distribution “weight1” based on the observations we defined that distribution affects (the variables we named "obs_{i}" ).

Sample: 100%|██████████| 20100/20100 [02:03, 163.05it/s, step size=1.55e+00, acc. prob=0.902]

Let’s unpack what happened here as well:

At (1) we clear all named samples from our Pyro storage. Pyro uses a special, built-in, dictionary-like object to keep track of its estimates for items we request. It’s important to clear that so if we run this code several times, previous runs won’t skew our results.

At (2) and (3) we define an MCMC algorithm which uses an HMC implementation to estimate the distribution of the named variable "weight1" . The MCMC algorithm runs for a certain number of iterations, in this case 20,100 iterations. However, the first few estimates it draws are very sensitive to initial random conditions, so we tend to discard these estimates (we will discard the first 100). All remaining 20,000 samples will be used to build our posterior estimate of "weight1"

. The MCMC algorithm runs for a certain number of iterations, in this case 20,100 iterations. However, the first few estimates it draws are very sensitive to initial random conditions, so we tend to discard these estimates (we will discard the first 100). All remaining 20,000 samples will be used to build our posterior estimate of Finally at (4) we run our MCMC algorithm. Notice we have to send our observations into the model function (as we defined it before to be accepting a series of observations). This is where we send all variables the model function accepts into it.

Now, let’s draw all the samples we got when estimating "weight1" , convert those samples to a NumPy array (they come back as a PyTorch tensor), and plot them as a histogram.

array([0.8106951 , 0.77260906, 0.79576313, ..., 0.86141765, 0.7857265 , 0.7469458 ], dtype=float32)

The Sample Distribution for Weight Values Generated by the HMC Algorithm with a Normal Prior

We can also plot the summary and see what is the mean value of this distribution to find the most likely value, as well as a 90% credibility interval.

mean std median 5.0% 95.0% n_eff r_hat

weight1 0.78 0.04 0.79 0.72 0.85 29855.87 1.00 Number of divergences: 0

So we see the mean value is about 0.78kg, and we’re 90% confident the value for the weight of the item is between 0.72kg — 0.85kg (this is also room for us to judge if we’re happy with that 90% range).

You might now say “Well, you had a pretty good prior, your mean was very close to the true value of 0.8kg. What if you had used a less informative prior?”

Let’s explore that scenario as well. This time we’ll use a uniform distribution between 0.0 (the weight can’t be negative), and 2kg (seems like a good upper threshold).

Sample: 100%|██████████| 20100/20100 [03:20, 100.21it/s, step size=8.57e-01, acc. prob=0.978]

mean std median 5.0% 95.0% n_eff r_hat

weight2 0.78 0.04 0.78 0.72 0.85 1059.35 1.00



Number of divergences: 0

Having this less informative prior did not make much of a difference. The algorithm is 90% confident the weight lies between 0.72kg and 0.85kg, and the mean of those values is still centered at 0.78kg (very close 0.8kg which was used).

We can grab the samples for this new approximation again, convert them to a NumPy array again, and plot the results as a distribution. You will see the distributions overlap very strongly, even though two very different priors were used.

The Sample Distributions for Weight Values Generated by the HMC Algorithm with a Normal Prior, as well as the HMC Algorithm with a Uniform Prior

Approximating Our Weight Using Stochastic Variational Inference

If you run the code above you will notice that drawing all the samples using an MCMC method takes quite a bit of time. Pyro offers an alternative method of evaluating our hidden parameters that is faster, and can provide equally good results for some problems: Stochastic Variational Inference (SVI).

SVI methods treat this inference problem as an optimization one. Under the SVI method, we choose another distribution which is defined by some parameters (for example, the normal distribution is defined by its mean and standard deviation), and we optimize those parameters until our chosen distribution provides results which are “close enough” to our observations.

This alternative method of evaluating a posterior requires an additional function, which is called a guide. This function defines our parameters, as well as our proposed distribution to match the observations.

Let’s define our model again:

And now, let’s define our guide function:

First, the guide and the model functions must accept the same inputs, even if the guide doesn’t use any of them at all.

At (2) and (3) we say take a normal distribution, and adjust its parameters (the mean and standard deviation) so that it is more likely to produce the observed observation. In (2) we have to provide initial guesses for those parameters (notice the pyro.param statement, which is different than the pyro.sample in that it registers a parameter, not a distribution).

statement, which is different than the in that it registers a parameter, not a distribution). At (4) we draw a weight from the distribution we defined in (3). We don’t actually have to store that value into a python variable, we only need the pyro.sample statement which is named exactly the same as it is in the model function. This is only the case for unobserved samples (those that don't have the obs= additional argument)

Now this becomes an optimization problem, we simply need to optimize "mu" and "sigma" such that the normal distribution defined by them is closer to producing our observed values.

iter: 9900, loss: -2.031

After performing the optimization, we can ask Pyro for the found values of "mu" and "sigma"

"mu" is 0.7845805883407593

"sigma" is 0.040342673659324646

We can also ask for the same statistics as before about “weight3” , such as the 90% credibility interval, but we have to do a bit more work for this.

We need to use our found values to generate new samples (MCMC works by generating samples, so we don’t need that extra step after the algorithm runs then):

weight_mean weight_5% weight_95%

0 0.784121 0.717712 0.850685

Run time for SVI algorithm and result extraction: 19.51 seconds

We can see we are getting very similar results to those we observed when employing an MCMC model, but at a fraction of the run time.

SVI approaches are designed to perform faster. However, since they require a pre-defined distribution as the guide, those can be difficult to choose for some distributions (e.g. multimodal). For distributions that can be approximated by pre-defined distributions, they are ideal since they offer a substantial performance boost.

We can also compare the resulting distribution we obtained by sampling from the SVI algorithm against our previous samples:

The Sample Distributions for Weight Values Generated by both HMC Algorithms, as well as the One Produced by the SVI Algorith,

Employing an Autoguide

One final functionality I would like to introduce here is the use of autoguides in Pyro. Pyro offers the ability to employ pre-built auto-guide functions, when constructing your own guide can be a laborious task.

Optimizing...

iter: 9900, loss: -1.88

Parameters:

"mu" is [0.77290195]

"sigma" is [0.04055462]

Predicting...

weight_mean weight_5% weight_95%

0 0.773199 0.706324 0.8403

Run time for SVI algorithm and result extraction: 23.17 seconds

We can also plot these results and compare the distributions of the two SVI approaches:

The Sample Distributions for Weight Values Generated by an SVI Algorithm with a Custom Guide, as well as with an Autogude

This post showcases the basic functionality Pyro offers in performing inference about hidden parameter values in a model.

I encourage everyone to explore this talk by Zach Anglin for a discussion on probabilistic programming applications (the examples are explored using PyMC3, but should be easier to translate to Pyro now), as well as this talk by Chi Nhan Nguyen which showcases how to employ Pyro to introduce uncertainty for overconfident neural networks (the talk uses an older version of Pyro).

In future posts, I will explore how to employ Pyro for more realistic applications, but when starting with any new library, an easy example which is easy to grasp can be incredibly instrumental in deciphering how the library works.

For a notebook version of this post, visit here