Probabilistic Linear Regression with Weight Uncertainty | by Ruben Winastwan | Towards Data Science

Linear regression is probably the first statistical approach that you’ll ever encounter when you’re learning data science and machine learning. So, I’d take my chance to guess that this is not the first time you’re dealing with linear regression. Thus in this article, I want to talk about probabilistic linear regression instead of the typical/deterministic linear regression.
But before that, let’s briefly discuss the concept of deterministic linear regression to get us up to the speed to the main talking point of this article.
Linear regression is a fundamental statistical approach to model the linear relationship between one or multiple input variables (or independent variables) with one or multiple output variables (or dependent variables).
In the above equation, a is called the intercept, and b is called the slope. x is our independent variable, and y is our dependent variable, which the value we try to predict.
The value for a and b need to be optimized with a gradient descent algorithm. Then, we obtain a regression line that shows the best fit between independent and dependent variables. With the regression line, we can predict the value of ywith any given input of x . Those are the steps on how the typical or deterministic linear regression algorithm is normally built.
However, this deterministic linear regression algorithm doesn’t really tell the complete story of both the data and the model. Why is that?
In reality, there are two types of uncertainty that arise when we do linear regression analysis:
Aleatoric uncertainty, which is the uncertainty that arises from the data.
Epistemic uncertainty, which is the uncertainty that arises from the regression model.
I’m going to elaborate more on these uncertainties as we go through the article. To take these uncertainties into account, probabilistic linear regression should be used instead of deterministic linear regression.
In this article, we will talk about probabilistic linear regression and how it differs from the deterministic linear regression. We will first see how deterministic linear regression is built in TensorFlow, and then we will move on to build a probabilistic linear regression model with TensorFlow probability.
First, let’s start with loading the dataset that we will use in this article.
The dataset that will be used in this article is car’s MPG dataset. As usual, we can load the data with Pandas.
Below is the statistical summary of the data.
Next, we can look at the correlation between variables in the dataset with the code below.
Now if we look at the correlation, the car’s Miles per Gallon (MPG) and the car’s weight has a strong negative correlation.
In this article, I’m going to do a simple linear regression analysis for visualization purpose. The independent variable will be the car’s weight and the dependent variable will be the car’s MPG.
Now, let’s split the data into training data and test data with Scikit-learn. After splitting the data, we can now scale both of the dependent and independent variables. This is to make sure that both variables are going to be on the same scale and this will also improve the convergence rate of our linear regression model.
Now if we visualize the training data, we get the following visualization:
Awesome! Next, let’s move on to build our deterministic linear regression model with TensorFlow.
It is very easy to build a simple linear regression model with TensorFlow. All we need to do to build the model is one single dense layer without any activation function. For the cost function, the mean squared error is normally used. In this example, I will use RMSprop as the optimizer and the model will be trained in 100 epochs. We can build and train the model in just a few lines of code as below.
After we trained the model, let’s take a look at the loss history of the model to check the loss convergence.
It seems like the loss has converged. Now if we use the trained model to predict the test set, we can see the following regression line.
And that’s it. We’re done!
As I mentioned earlier, it is very easy to build a simple linear regression model with TensorFlow. With the regression line, we can now approximate the car’s MPG at any given input of the car’s weight. As an example, let’s say that the car’s weight after feature scaling is 0.64. We can get the corresponding value of the car’s MPG by passing this value to the trained model as follows.
Now as you can see, the model predicts that the car’s MPG would be 0.21. Simply put, for any given car’s weight, we get a single deterministic value of the car’s MPG
However, that output value doesn’t really tell the whole story. There are two things that we should pay attention here. First, we only have a limited amount of data points. Second, as we can see from the linear regression plot, most of the data points don’t really lie on the regression line.
Although we get the output value of 0.21, we know that the actual car’s MPG is not precisely 0.21. It could be slightly below that, could be slightly above that. In other words, there is an uncertainty that needs to be taken into account. And this uncertainty is called aleatoric uncertainty.
Deterministic linear regression fails to capture this aleatoric uncertainty of the data. To capture this aleatoric uncertainty, the probabilistic linear regression can be applied instead.
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.
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?
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.
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.
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.
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.
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.
And that’s it! Now you’ve seen how the probabilistic linear regression differs from the deterministic linear regression. With probabilistic linear regression, two types of uncertainty that arise from both the data (aleatoric) and the regression model (epistemic) can be taken into account.
Taking these uncertainties into account is very important if we want to build a deep learning model where the inaccurate predictions lead to very serious negative consequences, for example in the field of autonomous driving and medical diagnosis.
Normally, as we have more data points, the epistemic uncertainty of the model will be reduced.