Probabilistic Linear Regression with Weight Uncertainty
By: Ruben Winastwan
Tags: Bayesian Statistics, Regression, Statistics, Data Science, Machine Learning

Probabilistic Linear Regression with TensorFlow Probability

Thanks to TensorFlow Probability, it is also very easy to build a probabilistic linear regression model. However, you need to install tensorflow_probability library first. You can install it using pip command as follows:

pip install tensorflow_probability

The prerequisite to install this library is that you need to have TensorFlow version 2.3.0. So make sure you upgrade your TensorFlow version before installing TensorFlow Probability.

Building Probabilistic Linear Regression Model for Aleatoric Uncertainty

In this section, we will build a probabilistic linear regression model that takes aleatoric uncertainty into account.

The model is pretty much similar to the deterministic linear regression. However, instead of just using one single dense layer like before, we need to add one more layer as the final layer. This final layer converts the final output value from deterministic into a probability distribution.

In this example, we will create a final layer that converts the output value into probability value that is normally distributed. Below is the implementation of that.

Note that we applied one additional layer in the end with TensorFlow Probability layer. This layer will turn the two outputs of the previous dense layer (one for the mean and one for the standard deviation) into probability value which is normally distributed with a trainable mean (loc) and standard deviation (scale).

We can use RMSprop as the optimizer, but you can use other optimizers if you wish. For the loss function, we need to use the negative log-likelihood.

But why do we use the negative log-likelihood as the loss function?

Negative Log-Likelihood as the Cost Function

In order to fit a distribution to some data, we need to use the likelihood function. With likelihood function, we try to estimate the unknown parameters θ (for example, the mean and the standard deviation of normally distributed data) given the pattern that we’ve seen in the data.

The job for the optimizer in our probabilistic regression model is to find the maximum likelihood estimate for the unknown parameters. In other words, the model is trained to find the most likely parameter value given the pattern from our data.

Maximizing the likelihood estimate is the same as minimizing the negative log-likelihood. In the optimization domain, often the objective is to minimizing the cost instead of maximizing it. That’s why we use the negative log-likelihood as our cost function.

Below is the implementation of the negative log-likelihood as our custom loss function.

Training and Prediction Result of Probabilistic Linear Regression Model for Aleatoric Uncertainty

Now that we have built the model and define the optimizer as well as the loss function, let’s compile and train the model.

Now we can draw samples from the trained model. We can visualize the comparison between the test set and the samples generated from the model with the following code.

Test data v sample generated from probabilistic linear regression model

As you can see from the visualization above, now for any given input value, the model won’t return a deterministic value. Instead, it will return a distribution and it will draw a sample based on that distribution.

If you compare the data points of the test set (blue points) to the data points predicted by the trained model (green points), you’d probably believe that the green points come from the same distribution as the blue points.

Next, we can also visualize the mean and the standard deviation of the distribution generated by the trained model, given the data in the training set. We can do this by applying the following code.

We can see that the probabilistic linear regression model gives us more than the regression line. It also gives the approximation of the standard deviation of the data. It can be seen that roughly 95% of the data points of the test set lie within the two standard deviations.

Building Probabilistic Linear Regression Model for Aleatoric and Epistemic Uncertainties

So far, we’ve built a probabilistic regression model that takes into account the uncertainty that comes from the data or as we call it, aleatoric uncertainty.

However, in reality, we also need to deal with the uncertainty that comes from the regression model itself. Due to the imperfection of the data, there is also uncertainty regarding the weight or slope of the regression parameter. This uncertainty is called epistemic uncertainty.

The probabilistic model that we’ve built so far only considers one deterministic weight. As you’ve seen from the visualization, the model generates only one regression line and often this is not entirely accurate.

In this section, we will improve our probabilistic regression model that takes both aleatoric and epistemic uncertainties into account.We can use the Bayesian perspective to introduce the uncertainty of the regression weights.

First, we need to define our prior belief about what the weight distribution looks like before we see the data. Often, we don’t know what to expect about this, right? To make it simple, let’s assume that the distribution of the weight is normally distributed with a mean equal to 0 and a standard deviation equal to 1.

Since we hard-coded the mean and the standard deviation, this prior belief is not trainable.

Next, we need to define the posterior distribution of the regression weights. The posterior distribution shows how our belief has changed after seeing the pattern in the data. Thus, the parameters in this posterior distribution are trainable. Below is the code implementation to define our posterior distribution.

Now the question is, what is this VariableLayers defined in the posterior function above? The idea behind this variable layers is that we try to approximate the true posterior distribution. Normally, it’s not possible to derive the true posterior distribution, hence we need to approximate it.

After defining the prior and the posterior function, now we can build our probabilistic linear regression model with weight uncertainty. Below is the code implementation for that.

As you may notice, the only difference between this model and the previous probabilistic regression model is just the first layer. Instead of using a normal dense layer, we use the DenseVariational layer. In this layer, we pass our prior and posterior functions as the argument. The second layer is exactly the same as the previous model.

Training and Prediction Result of Probabilistic Linear Regression Model for Aleatoric and Epistemic Uncertainty

Now it’s time for us to compile and train the model.

The optimizer and the cost function are still the same as the previous model. We use RMSprop as the optimizer and the negative log-likelihood as our cost function. Let’s compile and train or model.

Now it’s time for us to visualize the weight or slope uncertainty of our regression model. Below is the code implementation to visualize the result.

In the visualization above, you can see that the linear line (mean) as well as the standard deviation that has been generated by the posterior distribution of the trained model is different in each iteration. All of these lines are a plausible solution to fit the data points in the test set. However, due to epistemic uncertainty, we don’t know which line would be the best one.

Usually, the more the data points we have, the less the uncertainty of the regression line that we will see.