Bayesian Neural Network

The defining characteristic of a BNN is that its weights and biases are not single values but probability distributions (e.g., Gaussian). This represents the model's uncertainty about the correct parameter values given the training data.

• Prior Distribution: A starting belief about the weights before seeing data (e.g., a standard normal distribution).
• Posterior Distribution: The updated belief about the weights after observing the training data, calculated using Bayes' theorem. Learning in a BNN is the process of inferring this posterior.
• This framework naturally quantifies epistemic uncertainty—the model's uncertainty due to a lack of knowledge, which decreases as more relevant data is observed.

BNNs provide a principled mathematical framework for separating and measuring different types of uncertainty in predictions, which is critical for safety and reliability.

• Epistemic Uncertainty: Model uncertainty. High for inputs far from the training data. Measured by the variance in predictions from different weight samples from the posterior. Reducible with more data.
• Aleatoric Uncertainty: Data uncertainty from inherent noise or stochasticity. Captured by the model's output distribution (e.g., predicting a mean and variance). Irreducible with more data.
• This allows BNNs to signal low confidence on out-of-distribution inputs, enabling safer deployment in critical applications like medical diagnosis or autonomous systems.

Training a BNN requires computing the posterior distribution over weights, which is analytically intractable for deep networks. Variational Inference (VI) is the standard approximate method.

• Variational Posterior: A simpler, parameterized distribution (e.g., a Gaussian) is defined to approximate the true, complex posterior.
• Evidence Lower Bound (ELBO): The objective function maximized during training. It consists of:
  • A data fidelity term (reconstruction loss).
  • A regularization term: The Kullback-Leibler (KL) Divergence between the variational posterior and the prior, which prevents overfitting by keeping the learned weights close to the prior belief.
• This process balances fitting the data with maintaining calibrated uncertainty.

A landmark result showed that training a standard neural network with Dropout and applying it at test time is equivalent to performing approximate variational inference in a specific BNN.

• Practical Implication: Enables uncertainty estimation with minimal changes to existing neural network architectures and training pipelines.
• Procedure:
  1. Train a network with dropout layers.
  2. At inference, perform T forward passes with dropout active (e.g., T=50).
  3. Treat the mean of the T outputs as the prediction and their variance as the epistemic uncertainty.
• This method, while an approximation, made Bayesian deep learning accessible for many real-world applications.

The Bayesian framework provides inherent protection against overfitting, even with small datasets, and leads to more robust models.

• Built-in Regularization: The KL divergence term in the ELBO acts as a powerful, principled regularizer, penalizing model complexity and preventing the weights from becoming over-specialized to the training noise.
• Ensemble Effect: Predictions are made by integrating over all possible weights (via sampling), which is akin to using an infinite ensemble of networks. This averaging smooths the decision function and improves generalization.
• Calibrated Predictions: BNNs tend to produce better calibrated probabilities, meaning a predicted confidence of 90% corresponds to a 90% accuracy rate, unlike often overconfident standard neural networks.

BNNs are particularly powerful in reinforcement learning (RL) and active learning due to their explicit uncertainty modeling.

• Bayesian Optimization: Uses a BNN as a surrogate model to optimize expensive-to-evaluate functions. It strategically queries points where uncertainty is high (exploration) or predicted performance is high (exploitation).
• Model-Based RL: A BNN can serve as a probabilistic world model. Agents can perform internal simulation (planning) while accounting for model uncertainty, leading to safer and more data-efficient exploration.
• Thompson Sampling: A classic bandit algorithm that directly leverages the Bayesian posterior. The agent samples a weight instance from the posterior, acts optimally according to that single model, then updates the posterior, naturally balancing exploration and exploitation.

Epistemic uncertainty, or model uncertainty, is the reducible uncertainty stemming from a lack of knowledge about the model's parameters or the data distribution. In a Bayesian Neural Network, this is quantified by the variance in the posterior distribution over weights. It is highest in regions of input space with little or no training data and can be reduced by collecting more relevant data. This contrasts with aleatoric uncertainty, which is irreducible noise inherent in the observations.

Variational Inference (VI) is a core technique for approximating the intractable true posterior distribution in Bayesian Neural Networks. Instead of computing the exact posterior, VI introduces a simpler, parameterized distribution (the variational posterior) and optimizes its parameters to minimize its divergence from the true posterior, typically using the Kullback-Leibler (KL) Divergence. This transforms the inference problem into an optimization task, making BNNs computationally feasible for large models and datasets.

Monte Carlo Dropout is a practical and efficient approximation for performing inference in Bayesian Neural Networks. By applying dropout at test time and performing multiple forward passes, the network's stochasticity generates a distribution of predictions. The variance across these samples provides an estimate of epistemic uncertainty. This method establishes a theoretical connection between dropout regularization and approximate variational inference in deep Gaussian processes.

The Evidence Lower Bound (ELBO) is the objective function maximized during variational inference to train a Bayesian Neural Network. It is composed of two terms:

• A reconstruction term (expected log-likelihood) that encourages the model to fit the training data.
• A regularization term (KL divergence) that penalizes the variational posterior for deviating from a prior distribution over weights. Maximizing the ELBO is equivalent to minimizing the KL divergence between the approximate and true posterior.

Thompson Sampling is a classic Bayesian algorithm for solving the exploration-exploitation trade-off in sequential decision problems, such as reinforcement learning and bandits. It works by sampling a model (e.g., a set of neural network weights) from the current posterior belief and then acting optimally according to that sampled model. Bayesian Neural Networks are a natural fit for this paradigm, as their weight distributions provide a direct mechanism for sampling plausible models to guide exploratory behavior.

In Model-Based Reinforcement Learning (MBRL), an agent learns an explicit model of the environment's dynamics (a world model) and uses it for planning. A Bayesian Neural Network is an ideal candidate for this dynamics model because it can quantify its own epistemic uncertainty. This allows the agent to identify which parts of the state-action space are poorly understood, enabling targeted exploration (e.g., via uncertainty-weighted rewards) and more robust, sample-efficient planning while avoiding overconfident predictions in novel situations.