Bayesian Linear Regression

Linear Regression

Consider a dataset $\mathcal{D} = \{ ( \mathbf x_i, y_i ) \}_{i=1}^n$. Let $\mathbf X \in \reals^{n \times d }$ and $\mathbf y \in \reals^n$. Our goal is to learn $\mathbf w \in \reals^d$ s.t.

In fact, this mean-squared error comes from the fact that we assume

where $\epsilon \sim \mathcal{N}(0, \sigma^2)$ and $\mathbf w^*$ is the oracle weight. With this assumption, we can see that

Thus, we have

where $Z_x$ is the normalizer of the Gaussian distribution. If we assume that $(\mathbf x_i, y_i)$'s are independent and identically distributed, we have

Taking the log and negate the term yields the objective we have just mentioned. The approach based on optimizing $p(\mathcal D | \mathbf w)$ is called maximum likelihood estimation, which yields a point estimate $\mathbf{\hat{w}}$.

Posterior Distribution

However, in many applications, we just do not want only $\mathbf{\hat{w}}$ but rather the distribution of $\mathbf w$ given the data to estimate (epstemic) uncertainty of the prediction. This distribution is called the posterior distribution, and it can be computed via

This is essentially the Bayes' rule. One can see that $p(\mathcal D ) = \int p(\mathcal D| \mathbf w) p(\mathbf w) d\mathbf w$ can be hard to compute; however, this can be done (analytically) with proper assumptions. In particular, we might choose $p(\mathbf w)$ that is structurally comparable with $p(\mathcal D | \mathbf w)$; noting that the technical term is conjugate. For example, if our $p(y|\mathbf x ,\mathbf w)$ is a Gaussian, we might assume that $p(\mathbf w ) = \mathcal{N} (0, \Sigma_w),$ where $\Sigma_w = \sigma_w^2I_d$. With this choice of prior, we have

where $C$ is a constant when marginalizing out $\mathbf w$ . From this, the posterior distribution is in a multivariate Gaussian whose covariance is

When expanding the density of multivariate Gaussian $\mathcal{N}(\mu', \Sigma')$, we know that

Therefore, we can conclude that

Predictive Mean and Variance

Consider a new sample $\mathbf x^*$. We would like to know its prediction according to the posterior (i.e. averaging across all possible $\mathbf w )$ and its variance which is the uncertainty of the prediction. In particular, we know that

where $(\dagger)$ is based on the assumptions that (1) the new prediction and the data $\mathcal D$ given the model's parameters and (2) our model's parameters are independent of the new sample $\mathbf x^*$ given the data $\mathcal D$. Because both terms in the integral are Gaussian, this distribution of the prediction is also a Gaussian. Let assume $p(y_*|\mathbf x_*, \mathcal D) = \mathcal N (\mu_*, \sigma_*^2)$. Writing the two term together, we get

Expanding the term inside the exponent yields

Thus, we have

We can see that the exponent term inside the integrate is the form of Gaussian; in particular, this is

with an assumption that $L^{-1}$ exists. Because of the Gaussian form, we thus have

Because $L^{-1}$ is symmetric, i.e. $L^{-1} = (L^{-1})^T$, we can simplify $\xi ^T L \xi$ , yielding

Combining all the results together we have

where we have

These results can be further simplified by using the Sherman–Morrison formula (1). Denote $a=1/\sigma^2$. We have

Therefore,

Similary, we can follow the same steps for $\mu_{y_*}$,

These are the predictive variance and predictive mean for logistic regression with Gaussian prior or ridge regression.

We can look closer at $\sigma_{y_*}^2$. Here the first term $\sigma^2$ is a constrant we assume; more precisely it tells us about aleatoric uncertainty, which is the uncertainy due to noise in measurement. On the other hand, the second term $\mathbf x_*^T \Sigma' \mathbf x_*$ is what we are interested in if we make prediction. It captures epistemic uncertainty, which indicates the level of knowledge one does not have in the problem or the model. Therefore, if one is interested in the uncertainty of her/his model $f$'s prediction, one can determine the uncertainy by

Example

Now, it is time to put things together. We take a dataset and train a linear regression model on four different subsets. We assume that all train samples are in the range $[-1, 1]$, while test samples are $[-1.5, 1.5]$.

Fig. 1: Ridge regression trained with data with different sizes; the more training data the more certain prediction it is, especially in the extrapolation regime.

From Fig. 1, we see the effect of extrapolation in the range that our training data does not cover. Without the posterior distribution, we would not know how much uncertainty we had if we relied the point-estimate of the solution (i.e. solutions of ML or MAP). Nevertheless, as we have more training samples, our model does not only do well in the interpolation regime (where training data is covered) but also the extrapolation regime.

Conclusion and References

Perhaps, in the next step, we shall look at classification tasks and how to estimate uncertainty of the prediction in such situations.

This post is my recap of Yarin Gal's tutorial at SMILES 2019. I also consulted mathematicalmok's Youtube channel.

The figure is generated from this Google Colab.