Modeling Uncertainty in Neural Networks

This article is a personal summary of Y. Gal's MLSS 2019 Moscow, Bayesian Deep Learning (Slide Deck 2).

Background

Consider a dataset $\mathcal D = \{ (\mathbf x_i, y_i)\}_{i=1}^n$ where $\mathbf x \in \mathcal X \subseteq \reals^d$ and $y \in \mathcal Y$. Consider the underlying data generation process $f : \mathcal X \rightarrow \mathcal Y$ has the following form

where $\phi: \mathcal X \rightarrow \reals^h$ is a feature extractor that is assumed to be given. Consider regression problems, i.e. $\mathcal Y \subseteq \Reals$, we assume that the observed target is corrupted by noise $\epsilon$ . That is

where $\epsilon \sim \mathcal N(0, \sigma^2)$. In other words, $p(y|\mathbf x, \mathbf w) = \mathcal {N} (\mathbf w^T\phi(\mathbf x), \sigma^2)$. Let define $\mathbf X = (\mathbf x_1, \dots, \mathbf x_n)^T$ and $\mathbf y = (y_1, \dots, y_n)^T$. Using Bayes' rule, we have

If we choose $p(\mathbf w) = \mathcal N(0, \sigma_w^2I_h)$. This setup leads to ridge regression, and the predictive distribution has a Gaussian form.

Let's consider a $K$-category classification problem, i.e. $\mathcal Y = \{0, 1\}^K$ and assume the underlying data generation process is

where $W \in \reals^{K \times h}$ and $\text{softmax}: \reals^K \rightarrow [0, 1]^K$ whose exact form is

Define $\mathbf Y = (\mathbf y_1, \dots, \mathbf y_n)^T$. Therefore, we can write the likelihood as

However, in this classification setting and the assumption of the prior remains Gaussian, the posterior cannot be found because the evidence becomes intractable:

In other words, the category distribution (the output of the softmax function) is not conjugate with the Gaussian prior.

Approximating the Posterior : Variational Inference

Instead of computing $p(W|\mathcal D)$ directly, we approxiate it using a variational distribution $q_\theta(W)$. Typically, we use a simple distribution like Gaussian for $q_\theta$; in this case, $\theta = \{ \mu_\text{VI}, \Sigma_\text{VI}\}$. We can then find $\theta$ by minimizing the reverse Kullback-Leibler divergence

Expanding the KL, we get

Thus, we have

where the first term corresponds to how we predict the data, and the second term is how well our approximated distribution aligns with the prior.

Stochastic Approximate Inference

Let's consider the individual likelihood in classification settings; the log likelihood is

Then, we have

which is not integratable.

Detour: Monte Carlo Estimation and Re-parameterization Trick

Consider a density model $p(x)$ and a function of interest $f: \mathcal X \rightarrow \reals$. We assume that $x \sim p(x)$ can be done easily and assume that $\mathbb E[ f(X)]$ is difficult to compute (no analytical solution). Instead, one can sample $x_i \sim p(x)$ for $i = [1, n]$ and compute

It can be shown that $\mathbb{\hat{E}} [ f(X)]$ is an unbiased estimator, i.e.

Let's assume $w \sim \mathcal N(\mu, \sigma^2)$ and consider the following function

which has a closed-form solution (Link to derivation):

We can see that $\frac{\partial \mathcal L }{\partial \mu } = 1 + 2 \mu$ and $\frac{\partial \mathcal L }{\partial \sigma } = 2 \sigma$.

However, if we take $\hat w \sim \mathcal N(\mu, \sigma^2)$ and compute the derivatives, we get $\frac{ \partial }{\partial \mu } (\hat w + \hat w^2) = \frac{ \partial }{\partial \sigma } (\hat w + \hat w^2) = 0$. This is because the dependency of $\mu$ and $\sigma$ is via $w$.

Instead, we can re-parameterize $\hat w$ using

where $\epsilon \sim \mathcal N(0, 1)$. Thus, we have

We can see that

which are unbiased estimators of what we have derived previously. This trick is also known as pathwise derivative estimator, infinitesimal perturbation analysis, and stochastic backpropagation.

Approximating the Posterior with Stochastic Variational Inference

With the re-parameterization trick, we can then construct an objective function to learn $q_\theta(W)$ with $\theta = \{ \mu_\text{VI}, \Sigma_\text{VI} \}$. We assume $\epsilon \sim \mathcal {N}(0, I_{Kh})$ and $q_\theta = \mathcal N(\mu_\text{VI}, \Sigma_\text{VI})$ where

• $\mu \in \reals^{Kh}$ and
• $\Sigma_\text{VI} = \text{diag}(\sigma_{\text{VI}_1}^2, \dots, \sigma_{\text{VI}_{Kh}}^2) \in \reals^{Kh\times Kh}$.

In other words, we have $\hat W = \mu_\text{VI} + \Sigma_{\text{VI}}^{1/2} \epsilon$. Therefore, our learning objective is

whose $\mathbb {E}_{p(\epsilon)}[ \mathcal{\hat{L}}_\text{ELBO} ] = \mathcal{{L}}_\text{ELBO}$ and $\mathbb {E}_{p(\epsilon)}[ \nabla_\theta \mathcal{\hat{L}}_\text{ELBO} ] = \nabla_\theta \mathcal{{L}}_\text{ELBO}$ .

Uncertainty in Classification

For a multinomial distribution with $K$ classes whose mass is $p_{c_k}$ for $k \in [1, K]$, the uncertainty of the distribution is indicated by the entropy $H$

One observation is that the entropy is highest when $p_c = p_{c'} \forall c, c' \in [1, K]$. In other words, this is the case when we have a uniform distribution over $K$ classes, indicating absolute ambiguity in the prediction. Noting that, here we use the natural logarithm; hence, $H$ is measured the unit of nats.

Because the output of neural networks for classification is a parameterized multinomial distribution, we can use the entropy as a measure of uncertainty, this is known as predictive entropy

where $p(y_*=c_k|\mathbf x_*, \mathcal D)$ can be approximated using Monte Carlo. More precisely, let $\hat W_t \sim q_\theta(W)$ for $t=[1, T]$, we have

Because the ambiguity comes from two sources: 1) noisy measurements and 2) model parameters, the predictive entropy captures both aleatoric and epistemic uncertainties. To compute the epistemic uncertainty, we can look at the mutual information between the predicted label and the weights

where the first term is the predictive entropy and the second term is the average uncertainty of the prediction w.r.t the posterior. In other words, the second term captures the epistemic uncertainty, and it can be computed by

where $\hat W_t \sim q_\theta(W)$. Moreover, because entropy is non-negative, we have the following condition

Stochastic Approximate Inference in Neural Networks

From the previous example, we only consider the posterior for the last layer; however, we can use the variational approach to relax this constraint by considering the following model

where $W^1 \in \reals^{h \times d}, W^2 \in \reals^{K \times h}$, $b^1 \in \reals^h, b^2 \in \reals^K$ . We then assume that $w_{ij}^1, w_{ij}^2 \sim \mathcal N(0, 1)$. Define $\mathbf W = \{ W^{1}, W^{2}\}$. We thus have

where (1) we assume that $W^1$ and $W^2$ are independent; this is known as mean-field variational inference. However, we see that finding $q_\theta(\mathbf W)$ requires a doubled number of parameters (each $w$'s distribution is mean and variance); hence, the approach is not suitable for big models.

Fig. 1: Uncertainty from Bayesian Logistic Regression and (1 hidden layer) Neural Network; PyMC3's ADVI is used for the inference.

Dropout as Approximation Inference

Dropout is one of the regularization technique used in deep learning. In standard settings, during training, we randomly set activations in the network with probability $\rho$, i.e. the dropout activity is $\text{Ber}(\rho)$. Then, during test time, we turn off the dropout activity and multiply each neuron with $\frac{1}{1 -\rho}$ to compensate the dropout activity. These two steps are called stochastic and deterministic forward passes respectively.

Let's take a closer look into this using the previous model. Define $\epsilon^1_{ij} \sim \text{Ber}(1-\rho^1), \epsilon^2_{ij} \sim \text{Ber}(1-\rho^2)$ and the parameters are $M^1 \in \reals^{h \times d}, M^2 \in \reals^{K \times h}, b^1 \in \reals^d, b^2 \in \reals^K$, We have

We can see that we can write $\hat W^{1} = M^1 \epsilon^1$ and $\hat W^{2} = M^2 \epsilon^2$. Gal (2016) shows that the KL term can be approximated

where $C$ is a constant. Define $\hat \omega_\epsilon = \{\hat W^1, \hat W^2\}$. Therefore, the loss function is

Here, it should be noted that $\rho^1$ and $\rho^2$ also have to chosen properly; for example, one might consider using cross-validation.

Conclusion

Prediction is only part of the whole story. In the real world, we also need to know when our predictive models are uncertain; this is critical in high stake applications, such as automous vehicles and heathcare. Recent development in variational inference and sampling methods allow us to approximate the posterior, hence enabling extracting uncertainty from the model.

Figures are maded from Google Colab.