Core ML

Bayesian Neural Networks (BNNs)

What if, instead of assigning an exact, rigid number to every parameter, a neural network assigned a flexible probability distribution to every parameter? It could explicitly tell you when it was guessing.

In a Bayesian Neural Network, every single weight is replaced by a probability distribution, allowing the model to mathematically express doubt.

Why Does This Exist?

A standard neural network is an engine of absolute certainty. When it trains, it calculates that a specific weight (the connection between two neurons) should be exactly 1.42. Because the weight is an exact number (a point-estimate), the network's final output will also be an exact, rigid prediction. As we learned in deep-ensembles, this causes models to dangerously hallucinate on unknown data because they have no mathematical mechanism to express doubt.

Bayesian Neural Networks (BNNs) solve this at the fundamental architectural level. A BNN replaces every single weight in the network with a probability distribution.

Instead of saying, "This weight is exactly 1.42," a BNN says, "Based on the training data, I am 95% sure this weight is somewhere between 1.2 and 1.6." By allowing the internal weights to be uncertain, the final output of the network can explicitly express Epistemic Uncertainty (Model Ignorance).

Think of It Like This

Think of a standard neural network like a rigidly built bridge. The steel beams (weights) are exactly 10 feet long. If an unexpected, massively heavy truck (Out-of-Distribution data) drives over it, the bridge cannot adapt. It just snaps.

A Bayesian Neural Network is a bridge built out of shock-absorbing springs. The springs have an average length of 10 feet, but they have the flexibility (variance) to compress or expand when necessary. When the massive truck hits the bridge, the springs flex and absorb the shock, gracefully handling the unknown situation instead of snapping.

How It Actually Works

Transitioning from standard deep learning to Bayesian deep learning changes how we define the parameters of the model.

1. The Architecture

In a standard network, a single weight $w$ is a scalar (e.g., $w = 1.42$ ). In a BNN, a single weight $w$ is defined by a Gaussian distribution $N(\mu, \sigma^2)$ . The network now has to learn two numbers for every connection:

$\mu$ (The mean, or the "best guess" for the weight)
$\sigma$ (The variance, or "how unsure" the network is about that guess)

2. Training (Variational Inference)

Training a BNN is notoriously difficult. You cannot use standard backpropagation easily because you cannot take the derivative of a random probability distribution.

Instead, BNNs use a mathematical trick called the Reparameterization Trick (often via a technique called Variational Inference). The network is penalized not just for making wrong predictions, but for being overly confident when it shouldn't be.

If the network sees 10,000 pictures of dogs, the $\sigma$ (variance) for the "dog-detecting" weights shrinks close to zero. The network becomes confident.
If the network never sees a picture of a car, the $\sigma$ for the relevant weights remains massively wide.

3. Inference (Making Predictions)

When you ask a BNN for a prediction, you don't run it just once. Because every weight is a probability distribution, you run the network 10 or 20 times for a single image. Every time you run it, the weights randomly sample a new number from their distributions.

If the image is a Dog (Known Data), all 20 runs will predict "Dog," because the variance ( $\sigma$ ) on those weights is tiny.
If the image is a Car (Unknown Data), the massive variance in the weights will cause the 20 runs to output wildly different predictions (Dog, Frog, Bird). The high variance of the outputs explicitly alerts you that the data is Out-of-Distribution.

Show Me the Code

Writing a true BNN from scratch is complex, but libraries like bayesian-torch or Pyro abstract the heavy lifting. Here is a conceptual example of a Bayesian Linear Layer.

1import torch2import torch.nn as nn3import torch.nn.functional as F4
5class BayesianLinearLayer(nn.Module):6    def __init__(self, in_features, out_features):7        super().__init__()8        # Instead of one weight matrix, we have a Mean matrix and a Variance matrix9        self.weight_mu = nn.Parameter(torch.Tensor(out_features, in_features))10        self.weight_sigma = nn.Parameter(torch.Tensor(out_features, in_features))11        12        # Initialize them13        nn.init.normal_(self.weight_mu, 0, 0.1)14        nn.init.constant_(self.weight_sigma, -3.0) # Small initial variance15
16    def forward(self, x):17        # 1. Ensure variance is strictly positive18        sigma = torch.exp(self.weight_sigma)19        20        # 2. Sample random noise from a standard normal distribution21        epsilon = torch.randn_like(sigma)22        23        # 3. The Reparameterization Trick! 24        # Create a rigid weight matrix just for this specific forward pass25        sampled_weight = self.weight_mu + (sigma * epsilon)26        27        # 4. Standard linear computation28        return F.linear(x, sampled_weight)29
30# When you call this layer 10 times with the exact same input `x`, 31# it will give you 10 slightly different outputs!

Watch Out For

The Compute and Memory Nightmare

BNNs are the most mathematically pure way to handle uncertainty, but they are almost never used in large-scale production. Why? A standard 7-Billion parameter LLM requires 14GB of RAM. A Bayesian version of that same model would require 28GB of RAM (because every weight needs a $\mu$ and a $\sigma$ ). Furthermore, training them is notoriously unstable and slow. Because of this, the industry usually fakes Bayesian behavior using cheaper approximations like monte-carlo-dropout or deep-ensembles.

The Quick Version

Standard Neural Networks use rigid point-estimates (exact numbers) for weights, causing overconfidence.
Bayesian Neural Networks replace every weight with a probability distribution (a Mean and a Variance).
When encountering unknown data, the high variance in the weights causes the network's predictions to fluctuate wildly.
By running the network multiple times and measuring the fluctuation, you can accurately detect model ignorance (Epistemic Uncertainty).
While mathematically elegant, BNNs double the memory footprint and are incredibly difficult to train, limiting their use in massive modern models.