Bayesian Methods for Machine Learning Uncertainty

Beyond Point Estimates: Quantifying Confidence in AI Predictions

Authored by: Loveleen Narang

Date: August 4, 2024

Introduction: Why Uncertainty Matters in ML

Standard machine learning models often provide point estimates – single values representing predictions (e.g., "this image contains a cat," "the stock price will be $105.50"). While useful, these predictions lack a crucial component: a measure of **uncertainty**. How confident is the model in its prediction? Would it predict differently if trained on slightly different data? Knowing the uncertainty associated with a prediction is vital for reliable decision-making, especially in high-stakes applications like medical diagnosis, autonomous driving, and financial modeling. Overly confident incorrect predictions can have severe consequences.

Bayesian methods offer a principled mathematical framework for reasoning about and quantifying uncertainty in machine learning. Instead of learning single point estimates for model parameters, Bayesian approaches aim to learn full probability distributions over parameters, reflecting our beliefs about their plausible values given the observed data. This allows us not only to make predictions but also to understand the confidence associated with those predictions.

Bayesian Inference: The Core Idea

The foundation of Bayesian inference is Bayes' Theorem. Given some observed data \( D \) and model parameters \( \theta \), it describes how to update our prior beliefs about the parameters \( P(\theta) \) into posterior beliefs \( P(\theta | D) \) after observing the data.

Let's break down the terms (Formulas 1-6):

• Posterior Probability \( P(\theta | D) \): Our updated belief about the parameters \( \theta \) after observing data \( D \). This is what we want to compute.
• Likelihood \( P(D | \theta) \): The probability of observing the data \( D \) given a specific set of parameters \( \theta \). This is typically defined by the model structure (e.g., Gaussian likelihood for regression). Often written as \( L(\theta; D) \).
• Prior Probability \( P(\theta) \): Our initial belief about the parameters \( \theta \) *before* observing any data. This allows incorporating prior knowledge or imposing regularization (e.g., assuming weights are close to zero).
• Evidence (Marginal Likelihood) \( P(D) \): The probability of observing the data, averaged over all possible parameter values: \( P(D) = \int P(D | \theta) P(\theta) d\theta \). It acts as a normalization constant, ensuring the posterior integrates to 1. Often computationally intractable.

Since the evidence \( P(D) \) is often hard to compute, we frequently work with the unnormalized posterior: Formula (7):

The goal of Bayesian inference is not just to find a single best \( \theta \) (like Maximum Likelihood Estimation, MLE (Formula 8: \( \theta_{MLE} = \arg\max_\theta P(D|\theta) \)) or Maximum A Posteriori, MAP (Formula 9: \( \theta_{MAP} = \arg\max_\theta P(\theta|D) \))), but to determine the entire posterior distribution \( P(\theta | D) \).

Bayesian Prediction

Instead of predicting with a single point estimate \( \theta_{MAP} \) or \( \theta_{MLE} \), Bayesian prediction involves averaging over the entire posterior distribution of parameters. The posterior predictive distribution for a new data point \( x^* \) is: Formula (10):

This integration accounts for the uncertainty in \( \theta \). The resulting predictive distribution \( P(y^* | x^*, D) \) naturally provides not just a point prediction (e.g., its mean or mode) but also a measure of uncertainty (e.g., its variance or credible intervals). Point estimate prediction: \( \hat{y}^* = f(x^*; \theta_{MAP/MLE}) \) (Formulas 11, 12).

Aleatoric vs. Epistemic Uncertainty

Bayesian methods naturally help distinguish between two fundamental types of uncertainty:

• Aleatoric Uncertainty: Also known as statistical uncertainty, this is the inherent randomness or noise in the data generating process itself. It represents variability that cannot be reduced even with infinite data (e.g., the outcome of a fair coin flip). In models, this is often captured by the variance term in the likelihood function (e.g., \( \sigma^2 \) in \( y = f(x) + \epsilon, \epsilon \sim N(0, \sigma^2) \)). Formula (13): \( \sigma^2 \).
• Epistemic Uncertainty: Also known as model or systematic uncertainty, this arises from our lack of knowledge about the true underlying model or its parameters. It reflects the uncertainty in the model parameters \( \theta \) given the finite training data. This uncertainty *can* be reduced by collecting more data. Bayesian methods explicitly capture epistemic uncertainty through the posterior distribution \( P(\theta | D) \).

The variance of the Bayesian predictive distribution can often be decomposed (approximately) into contributions from both aleatoric and epistemic sources. Formula (14): \( Var(y^*) \approx \underbrace{E_{P(\theta|D)}[\sigma^2(x^*, \theta)]}_{\text{Avg. Aleatoric}} + \underbrace{Var_{P(\theta|D)}[f(x^*, \theta)]}_{\text{Epistemic}} \).

Bayesian Models & Uncertainty Quantification

Bayesian Linear Regression

A simple starting point. Instead of finding single best-fit weights \( \beta \), we place priors on \( \beta \) and the noise variance \( \sigma^2 \) (e.g., Gaussian prior for \( \beta \), Inverse Gamma for \( \sigma^2 \)). Using conjugate priors allows deriving the posterior distributions analytically. The predictive distribution for a new point \( x^* \) is a Student's t-distribution, which naturally has heavier tails than a Gaussian, reflecting parameter uncertainty. Model: \( y = X\beta + \epsilon, \epsilon \sim N(0, \sigma^2 I) \). Priors: \( p(\beta), p(\sigma^2) \). (Formulas 15, 16, 17).

Gaussian Processes (GPs)

A non-parametric Bayesian approach that places a prior directly on the function \( f(x) \) itself, assuming that the function values at any set of points follow a multivariate Gaussian distribution. \( f(x) \sim GP(m(x), k(x, x')) \) (Formula 18), defined by a mean function \( m(x) \) (often zero) (Formula 19) and a covariance (kernel) function \( k(x, x') \) (Formula 20). The kernel encodes prior beliefs about the function's smoothness and other properties. Given training data, the posterior process is also Gaussian, and the predictive distribution for \( y^* \) at \( x^* \) is Gaussian with analytical mean and variance, directly providing uncertainty estimates.

Bayesian Neural Networks (BNNs)

Standard NNs learn point estimates for weights \( W \). BNNs place prior distributions \( P(W) \) over the weights and aim to compute the posterior \( P(W | D) \propto P(D | W) P(W) \) (Formula 21). The predictive distribution \( P(y^* | x^*, D) = \int P(y^* | x^*, W) P(W | D) dW \) requires integrating over this high-dimensional weight posterior.

Computing the exact posterior \( P(W|D) \) for NNs is intractable due to the non-linearity and high dimensionality. Therefore, approximate inference techniques are required.

Approximate Inference Techniques

Since exact Bayesian inference is often intractable for complex models like BNNs, approximation methods are essential.

Markov Chain Monte Carlo (MCMC)

• Idea: Construct a Markov chain whose stationary distribution is the target posterior \( P(\theta | D) \). By running the chain long enough, samples drawn from the chain approximate samples from the posterior.
• Methods: Metropolis-Hastings, Gibbs Sampling, Hamiltonian Monte Carlo (HMC), No-U-Turn Sampler (NUTS). HMC/NUTS are often preferred for high-dimensional problems as they use gradient information to explore the space more efficiently.
• Pros: Asymptotically exact (given infinite time), well-developed theory.
• Cons: Computationally very expensive (requires many sequential model evaluations), diagnosing convergence can be difficult, scaling to massive datasets is challenging.

Variational Inference (VI)

• Idea: Approximate the true (intractable) posterior \( P(\theta | D) \) with a simpler, tractable distribution \( q(\theta; \phi) \) from a chosen family (e.g., fully factorized Gaussian - Mean-Field VI, Formula 22: \( q(\theta) = \prod q_i(\theta_i) \)). The parameters \( \phi \) of \( q \) are optimized to minimize the Kullback-Leibler (KL) divergence between \( q \) and \( p \). Goal: \( \min_\phi D_{KL}(q(\theta; \phi) || P(\theta | D)) \) (Formula 23). Formula (24): \( D_{KL}(q || p) = \int q \log(q/p) \).
• ELBO: Minimizing KL divergence is equivalent to maximizing the Evidence Lower Bound (ELBO): Formula (25): \( \mathcal{L}(\phi) = E_{q(\theta; \phi)}[\log P(D | \theta)] - D_{KL}(q(\theta; \phi) || P(\theta)) \). This turns inference into an optimization problem solvable with gradient-based methods.
• Pros: Much faster than MCMC, scalable to large datasets, leverages standard optimization tools.
• Cons: Provides only an approximation to the posterior (quality depends on the chosen family \( q \)), can underestimate posterior variance, optimization can be challenging.

Monte Carlo Dropout

• Idea: A simpler, heuristic approach. Train a standard neural network with dropout layers. At test time, keep dropout active and perform multiple (\( T \)) forward passes for the same input \( x^* \). The mean of the predictions \( \hat{y}^* \approx \frac{1}{T} \sum_{t=1}^T f(x^*; \hat{W}_t) \) (Formula 26) serves as the prediction, and the variance across these predictions serves as an estimate of uncertainty.
• Connection: Can be shown to be mathematically equivalent to approximate Bayesian inference for a specific type of Gaussian Process or BNN under certain conditions.
• Pros: Easy to implement using standard NN libraries.
• Cons: Provides only an approximation, theoretical grounding is specific, quality of uncertainty estimates can vary.

Laplace Approximation

• Idea: Approximate the posterior distribution \( P(\theta | D) \) with a Gaussian distribution centered at the MAP estimate \( \theta_{MAP} \). The covariance matrix of the Gaussian is derived from the inverse of the Hessian matrix of the negative log-posterior evaluated at the MAP estimate.
• Pros: Relatively simple to compute after finding the MAP estimate.
• Cons: Relies on a Gaussian approximation which may be poor if the true posterior is multi-modal or highly skewed. Requires computing/inverting the Hessian (or approximation).

Comparison of Approximate Inference Techniques

Method	Core Idea	Pros	Cons
MCMC	Sample from posterior via Markov chain	Asymptotically exact, theoretically grounded	Very slow, convergence diagnostics needed, poor scalability
Variational Inference (VI)	Optimize parameters of an approximating distribution (minimize KL)	Much faster than MCMC, scalable, uses optimization tools	Only an approximation, can underestimate variance, choice of family matters
MC Dropout	Average predictions from network with dropout at test time	Very easy to implement, computationally cheap at inference	Heuristic, approximation quality varies, theoretical link specific
Laplace Approximation	Gaussian approximation around MAP estimate using Hessian	Relatively simple post-optimization	Gaussian assumption limiting, Hessian computation/inversion needed

Using and Evaluating Uncertainty Estimates

Quantified uncertainty is valuable for:

• Robust Decision Making: Deferring to humans or fallback systems when model uncertainty is high.
• Active Learning: Querying labels for data points where the model is most uncertain to improve efficiency.
• Out-of-Distribution Detection: Inputs far from the training data often yield high epistemic uncertainty.
• Model Calibration: Checking if predicted probabilities match empirical frequencies. Well-calibrated models have uncertainty estimates that reflect true likelihoods.
• Bayesian Optimization: Using uncertainty to balance exploration and exploitation when optimizing expensive black-box functions.

Evaluating uncertainty itself is challenging. Key aspects include:

• Calibration Plots: Plotting predicted probability vs. actual frequency.
• Entropy/Variance Measures: Higher entropy (Formula 27: \( H(P) = -\sum P(x) \log P(x) \)) or variance indicates higher uncertainty.
• Performance on Downstream Tasks: How well does the uncertainty estimate improve active learning or OOD detection?

Basic formulas often used: Gaussian PDF \( N(x | \mu, \sigma^2) \) (Formula 28), Integral \( \int \) (Formula 29), Expectation \( E[\cdot] \) (Formula 30), Probability \( P(\cdot) \) (Formula 31), Proportionality \( \propto \) (Formula 32), Argmax \( \arg\max \) (Formula 33).

Challenges

• Computational Cost: Exact Bayesian inference is intractable; MCMC is slow; VI optimization can still be complex.
• Prior Specification: Choosing appropriate priors can be difficult and influence the posterior, especially with limited data.
• Scalability: Applying many Bayesian methods (especially MCMC or complex VI) to very large datasets and models remains challenging.
• Approximation Quality: The accuracy of VI or Laplace depends heavily on the appropriateness of the approximating distribution. MC Dropout is more heuristic.
• Evaluation: Standardized and reliable evaluation of the *quality* of uncertainty estimates is still an active research area.

Conclusion

Moving beyond simple point predictions, Bayesian methods provide a robust and theoretically grounded framework for quantifying uncertainty in machine learning. By representing parameters as probability distributions via Bayes' theorem, these techniques capture our state of knowledge and allow us to distinguish between inherent data noise (aleatoric) and model ignorance (epistemic uncertainty). While exact inference is often intractable, approximate methods like Variational Inference, MCMC, and MC Dropout enable practical application to complex models like Bayesian Neural Networks. The resulting uncertainty estimates are invaluable for building more reliable, robust, and trustworthy AI systems, enabling better decision-making, active learning, and outlier detection. Despite computational and evaluation challenges, Bayesian inference remains a cornerstone for principled uncertainty quantification in modern AI.

(Formula count check: Includes Bayes Thm, Post L*P, Likelihood L, Prior P(th), Posterior P(th|D), Evidence P(D), Predictive Dist, Point Est (MAP/MLE), MAP Def, MLE Def, BLR Model, BLR Prior (concept), BLR Noise, GP Prior, GP Mean m, GP Kernel k, BNN Post, VI Goal (min KL), KL Div Def, ELBO Def, MeanField VI, MC Dropout Pred, Pred Var Decomp, Entropy H, Gaussian PDF, Integral, Expectation E, Prob P, Proportional, Argmax. Total > 33).