Deep Bayesian Neural Networks¶
What is a Bayesian Neural Network?¶
A standard neural network is trained to find a single set of weights \(\mathbf{w}\) that minimises some loss. This point estimate says nothing about how certain the network is in its predictions, a network can be confidently wrong, especially far from the training data.
A Bayesian Neural Network (BNN) instead treats the weights as random variables and places a prior distribution over them:
After observing data \(\mathcal{D} = \{(\mathbf{x}_i, y_i)\}\), we seek the posterior:
Predictions are then made by integrating over all plausible weight configurations:
This gives us uncertainty-aware predictions: the model can express that it doesn't know the answer, not just what its best guess is.
Example: Fitting a BNN with Bayinx and Equinox¶
Below is a complete example fitting a deep BNN to a 1-D regression problem.
The network is defined with Equinox and embedded inside a Bayinx Model as a stochastic node, so every weight and bias is automatically treated as a latent variable.
import jax
import jax.numpy as jnp
import jax.random as jr
import equinox as eqx
import bayinx as byx
import bayinx.dists as byd
import bayinx.flows as byf
from jaxtyping import Array, PRNGKeyArray, Scalar
class MyNeuralNetwork(eqx.Module):
layers: list
def __init__(self, n_in: int, key: PRNGKeyArray):
k1, k2, k3 = jr.split(key, 3)
self.layers = [
eqx.nn.Linear(n_in, 5, key=k1),
eqx.nn.Linear(5, 5, key=k1),
eqx.nn.Linear(5, 1, key=k2),
]
def __call__(self, x: Array) -> Array:
for layer in self.layers[:-1]:
x = jax.nn.tanh(layer(x))
return self.layers[-1](x).squeeze()
class MyModel(byx.Model):
nn: MyNeuralNetwork = byx.stochastic(
init=MyNeuralNetwork(n_in=1, key=jr.key(0))
)
sigma: Scalar = byx.stochastic(shape=(), lower=0.0)
X: Array = byx.observed(shape=("n_obs", "n_in"))
y: Array = byx.observed(shape="n_obs")
def model(self, target):
# Priors
self.nn << byd.Normal(0.0, 1.0) # constrain weights
self.sigma << byd.Exponential(1.0)
# Likelihood
mu = jax.vmap(self.nn)(self.X)
self.y << byd.Normal(mu, self.sigma)
n_obs, n_in = 1000, 1
X_data = jr.uniform(jr.key(0), (n_obs, n_in), minval=-3.0, maxval=3.0)
y_data = jnp.sin(X_data[:, 0]) + 0.2 * jr.normal(jr.key(1), (n_obs,))
posterior = byx.Posterior(
MyModel,
n_obs=n_obs, n_in=n_in,
X=X_data, y=y_data,
nn=MyNeuralNetwork(n_in=n_in, key=jr.key(99)),
)
posterior.configure([byf.DiagAffine()])
posterior.fit(max_iters=1_000_000, stl=True)
X_test = jnp.array([-jnp.pi / 2, 0, jnp.pi/2])[:, None]
def predict(model: MyModel, ksey: jax.Array) -> Array:
mu = jax.vmap(model.nn)(X_test)
return byd.Normal(mu, model.sigma).sample(mu.shape, key=key)
ppred = posterior.predictive(predict, n_draws=1_000) # shape (1_000, 100)
print("Predictive mean:", ppred.mean(0))
print("Predictive std: ", ppred.std(0))