Uncertainty Estimation in BNNs

Standard deep learning models provide point estimates for predictions, offering no inherent measure of their confidence. This can be problematic in applications where understanding the reliability of a prediction is important. Bayesian Neural Networks (BNNs) address this by learning a distribution over the model parameters (weights and biases, $w$) rather than single point estimates. This approach allows BNNs to quantify uncertainty. The posterior distribution, $p(w | \mathcal{D})$, where $\mathcal{D}$ represents the training data, is the foundation for quantifying uncertainty in BNN predictions.

Understanding and quantifying uncertainty allows us to build more trustworthy and informative models. In the context of BNNs, uncertainty isn't a single concept but typically decomposes into two primary types:

Aleatoric Uncertainty

Aleatoric uncertainty, sometimes called data uncertainty, captures the inherent noise or randomness in the data generating process itself. It represents uncertainty that cannot be reduced even if we had infinite data, because it stems from irreducible variability in the observations.

Consider a regression task where multiple different output values $y$ are possible for the exact same input features $x$. This variability is aleatoric. In BNNs, aleatoric uncertainty is typically modeled via the likelihood function, $p(y | x, w)$. For example, using a Gaussian likelihood for regression:

$$p(y | x, f_w(x)) = \mathcal{N}(y | f_w(x), \sigma^2)$$

Here, $f_w(x)$ is the output of the neural network with parameters $w$, and $\sigma^2$ represents the observation noise variance.

• Homoscedastic Uncertainty: Assumes the noise variance $\sigma^2$ is constant for all inputs. It can be treated as a parameter to be learned or fixed beforehand.
• Heteroscedastic Uncertainty: Allows the noise variance to depend on the input $x$. The BNN can be designed to predict both the mean output $f_w(x)$ and the variance $\sigma^2(x)$. This is useful when the noise level varies across the input space.

Epistemic Uncertainty

Epistemic uncertainty, also known as model uncertainty or knowledge uncertainty, reflects our ignorance about the true model parameters. It captures the uncertainty arising from having limited data to constrain the posterior distribution $p(w | \mathcal{D})$. As we gather more relevant data, epistemic uncertainty should decrease because the posterior distribution becomes more concentrated around the optimal parameter values.

In BNNs, epistemic uncertainty arises directly from the fact that we have a distribution over weights $p(w | \mathcal{D})$ instead of a single set of weights. Different plausible weight configurations (sampled from the posterior) will produce different predictions for the same input $x$. The variation in these predictions reflects the model's uncertainty about the function it has learned.

Quantifying Predictive Uncertainty

To obtain a prediction for a new input $x^*$, we are interested in the predictive distribution $p(y^* | x^*, \mathcal{D})$. This involves marginalizing out the model parameters:

$$p(y^* | x^*, \mathcal{D}) = \int p(y^* | x^*, w) p(w | \mathcal{D}) dw$$

This integral averages the predictions $p(y^* | x^*, w)$ over all possible parameter values weighted by their posterior probability $p(w | \mathcal{D})$. The variance of this predictive distribution captures the total uncertainty, encompassing both aleatoric and epistemic sources.

Since the posterior $p(w | \mathcal{D})$ is generally intractable, we rely on the approximation methods discussed previously:

1. MCMC Methods: Techniques like Stochastic Gradient HMC provide samples $w^{(1)}, w^{(2)}, ..., w^{(S)}$ from the posterior $p(w | \mathcal{D})$. We can approximate the predictive distribution by generating predictions for each sample and forming an empirical distribution:

   $$p(y^* | x^*, \mathcal{D}) \approx \frac{1}{S} \sum_{s=1}^S p(y^* | x^*, w^{(s)})$$

   For regression with a Gaussian likelihood, this means obtaining $S$ predicted means $f_{w^{(s)}}(x^*)$ and potentially variances $\sigma^2_{w^{(s)}}(x^*)$. The mean of the predictive distribution can be estimated by $\frac{1}{S} \sum_s f_{w^{(s)}}(x^*)$, and the variance (total uncertainty) can be estimated from the variance of the samples $\{f_{w^{(s)}}(x^*)\}_{s=1}^S$ plus the average predicted aleatoric variance.

2. Variational Inference (VI): VI methods like Bayes by Backprop learn an approximate posterior $q(w; \phi)$. The predictive distribution is approximated as:

   $$p(y^* | x^*, \mathcal{D}) \approx \int p(y^* | x^*, w) q(w; \phi) dw$$

   This integral is often approximated using Monte Carlo sampling: draw samples $w^{(s)} \sim q(w; \phi)$ and average the predictions $p(y^* | x^*, w^{(s)})$ similar to the MCMC approach.

3. Monte Carlo Dropout: A computationally cheaper, widely used technique involves training a standard neural network with dropout applied before weight layers. At test time, dropout is kept active, and multiple forward passes ($S$ times) are performed for the same input $x^*$. Each forward pass yields a different output due to the stochastic nature of dropout. This set of outputs $\{f_{w^{(s)}}(x^*)\}_{s=1}^S$ (where $w^{(s)}$ represents the network configuration in the $s$-th pass) can be treated as samples from an approximate predictive distribution. This method approximates Bayesian inference in deep Gaussian processes.

Interpreting and Using Uncertainty

The spread (e.g., variance or interquartile range) of the approximated predictive distribution gives us a measure of the model's confidence. High variance suggests low confidence.

• Distinguishing Uncertainty Types: While the total predictive variance combines both effects, techniques exist to disentangle aleatoric and epistemic uncertainty. For instance, in regression using Monte Carlo samples $\{f_{w^{(s)}}(x^*), \sigma^2_{w^{(s)}}(x^*)\}_{s=1}^S$:
  • Epistemic Uncertainty: Can be estimated by the variance of the sampled means: $Var_w(f_w(x^*))$.
  • Aleatoric Uncertainty: Can be estimated by the average of the predicted variances: $\mathbb{E}_w[\sigma^2_w(x^*)]$.
  • Total Uncertainty: Approximated by $Var_w(f_w(x^*)) + \mathbb{E}_w[\sigma^2_w(x^*)]$.

Knowing the source of uncertainty is valuable. High epistemic uncertainty suggests the model is unsure due to lack of data in that region of the input space, indicating where collecting more data might be beneficial. High aleatoric uncertainty suggests inherent limits to predictability based on the given features and noise level.

Example of a BNN's predictive mean and total uncertainty (shaded region, e.g., Mean +/- 2 Std Dev) for a 1D regression problem. Uncertainty is lower near training data points and higher in regions far from the data, reflecting increased epistemic uncertainty.

Applications of Uncertainty Estimation

Reliable uncertainty estimates from BNNs are beneficial in numerous areas:

• Safety-Critical Systems: Identifying predictions where the model is uncertain is essential in domains like autonomous driving or medical diagnosis, allowing for human intervention or alternative actions.
• Active Learning: Selecting new data points where the model is most uncertain (high epistemic uncertainty) can lead to more efficient data acquisition.
• Out-of-Distribution Detection: BNNs often exhibit higher uncertainty for inputs that are significantly different from the training data distribution.
• Reinforcement Learning: Uncertainty estimates can guide exploration strategies, encouraging agents to explore states or actions with higher uncertainty.
• Model Diagnostics: Analyzing the spatial distribution of epistemic uncertainty can reveal limitations in the training data or model architecture.

By providing a principled way to represent and quantify what the model doesn't know, BNNs offer a significant advantage over standard deep learning models, enabling more trustworthy AI systems.