Bayesian Neural Networks

In this tutorial, we demonstrate how one can implement a Bayesian Neural Network using a combination of Turing and Lux, a suite of machine learning tools. We will use Lux to specify the neural network’s layers and Turing to implement the probabilistic inference, with the goal of implementing a classification algorithm.

We will begin with importing the relevant libraries.

using Turing
using FillArrays
using Lux
using Plots
using Tracker
using Functors

using LinearAlgebra
using Random
WARNING: Method definition rrule(typeof(Base.merge), NamedTuple{F1, T} where T<:Tuple, NamedTuple{F2, T} where T<:Tuple) where {F1, F2} in module ChainRules at /home/runner/.julia/packages/ChainRules/hShjJ/src/rulesets/Base/base.jl:313 overwritten in module Lux at /home/runner/.julia/packages/Lux/5xfGO/src/chainrules.jl:16.
ERROR: Method overwriting is not permitted during Module precompilation. Use `__precompile__(false)` to opt-out of precompilation.

Our goal here is to use a Bayesian neural network to classify points in an artificial dataset. The code below generates data points arranged in a box-like pattern and displays a graph of the dataset we will be working with.

# Number of points to generate
N = 80
M = round(Int, N / 4)
rng = Random.default_rng()
Random.seed!(rng, 1234)

# Generate artificial data
x1s = rand(rng, Float32, M) * 4.5f0;
x2s = rand(rng, Float32, M) * 4.5f0;
xt1s = Array([[x1s[i] + 0.5f0; x2s[i] + 0.5f0] for i in 1:M])
x1s = rand(rng, Float32, M) * 4.5f0;
x2s = rand(rng, Float32, M) * 4.5f0;
append!(xt1s, Array([[x1s[i] - 5.0f0; x2s[i] - 5.0f0] for i in 1:M]))

x1s = rand(rng, Float32, M) * 4.5f0;
x2s = rand(rng, Float32, M) * 4.5f0;
xt0s = Array([[x1s[i] + 0.5f0; x2s[i] - 5.0f0] for i in 1:M])
x1s = rand(rng, Float32, M) * 4.5f0;
x2s = rand(rng, Float32, M) * 4.5f0;
append!(xt0s, Array([[x1s[i] - 5.0f0; x2s[i] + 0.5f0] for i in 1:M]))

# Store all the data for later
xs = [xt1s; xt0s]
ts = [ones(2 * M); zeros(2 * M)]

# Plot data points.
function plot_data()
    x1 = map(e -> e[1], xt1s)
    y1 = map(e -> e[2], xt1s)
    x2 = map(e -> e[1], xt0s)
    y2 = map(e -> e[2], xt0s)

    Plots.scatter(x1, y1; color="red", clim=(0, 1))
    return Plots.scatter!(x2, y2; color="blue", clim=(0, 1))
end

plot_data()

Building a Neural Network

The next step is to define a feedforward neural network where we express our parameters as distributions, and not single points as with traditional neural networks. For this we will use Dense to define liner layers and compose them via Chain, both are neural network primitives from Lux. The network nn_initial we created has two hidden layers with tanh activations and one output layer with sigmoid (σ) activation, as shown below.

G Input layer                   Hidden layers                  Output layer cluster_input cluster_hidden1 cluster_hidden2 cluster_output input1 hidden11 input1--hidden11 hidden12 input1--hidden12 hidden13 input1--hidden13 input2 input2--hidden11 input2--hidden12 input2--hidden13 hidden21 hidden11--hidden21 hidden22 hidden11--hidden22 hidden12--hidden21 hidden12--hidden22 hidden13--hidden21 hidden13--hidden22 output1 hidden21--output1 hidden22--output1

The nn_initial is an instance that acts as a function and can take data as inputs and output predictions. We will define distributions on the neural network parameters. 

# Construct a neural network using Lux
nn_initial = Chain(Dense(2 => 3, tanh), Dense(3 => 2, tanh), Dense(2 => 1, σ))

# Initialize the model weights and state
ps, st = Lux.setup(rng, nn_initial)

Lux.parameterlength(nn_initial) # number of paraemters in NN
20

The probabilistic model specification below creates a parameters variable, which has IID normal variables. The parameters vector represents all parameters of our neural net (weights and biases).

# Create a regularization term and a Gaussian prior variance term.
alpha = 0.09
sigma = sqrt(1.0 / alpha)
3.3333333333333335

Construct named tuple from a sampled parameter vector. We could also use ComponentArrays here and simply broadcast to avoid doing this. But let’s do it this way to avoid dependencies.

function vector_to_parameters(ps_new::AbstractVector, ps::NamedTuple)
    @assert length(ps_new) == Lux.parameterlength(ps)
    i = 1
    function get_ps(x)
        z = reshape(view(ps_new, i:(i + length(x) - 1)), size(x))
        i += length(x)
        return z
    end
    return fmap(get_ps, ps)
end
vector_to_parameters (generic function with 1 method)

To interface with external libraries it is often desirable to use the StatefulLuxLayer to automatically handle the neural network states.

const nn = StatefulLuxLayer(nn_initial, st)

# Specify the probabilistic model.
@model function bayes_nn(xs, ts; sigma = sigma, ps = ps, nn = nn)
    # Sample the parameters
    nparameters = Lux.parameterlength(nn_initial)
    parameters ~ MvNormal(zeros(nparameters), Diagonal(abs2.(sigma .* ones(nparameters))))

    # Forward NN to make predictions
    preds = Lux.apply(nn, xs, vector_to_parameters(parameters, ps))

    # Observe each prediction.
    for i in eachindex(ts)
        ts[i] ~ Bernoulli(preds[i])
    end
end
UndefVarError: UndefVarError(:StatefulLuxLayer)
UndefVarError: `StatefulLuxLayer` not defined
Stacktrace:
 [1] top-level scope
   @ ~/work/docs/docs/tutorials/03-bayesian-neural-network/index.qmd:226

Inference can now be performed by calling sample. We use the NUTS Hamiltonian Monte Carlo sampler here.

setprogress!(false)
# Perform inference.
N = 2_000
ch = sample(bayes_nn(reduce(hcat, xs), ts), NUTS(; adtype=AutoTracker()), N);
UndefVarError: 
UndefVarError: `bayes_nn` not defined
Stacktrace:
 [1] top-level scope
   @ ~/work/docs/docs/tutorials/03-bayesian-neural-network/index.qmd:259

Now we extract the parameter samples from the sampled chain as θ (this is of size 5000 x 20 where 5000 is the number of iterations and 20 is the number of parameters). We’ll use these primarily to determine how good our model’s classifier is.

# Extract all weight and bias parameters.
θ = MCMCChains.group(ch, :parameters).value;
UndefVarError: 
UndefVarError: `ch` not defined
Stacktrace:
 [1] top-level scope
   @ ~/work/docs/docs/tutorials/03-bayesian-neural-network/index.qmd:273

Prediction Visualization

We can use MAP estimation to classify our population by using the set of weights that provided the highest log posterior.

# A helper to run the nn through data `x` using parameters `θ`
nn_forward(x, θ) = nn(x, vector_to_parameters(θ, ps))

# Plot the data we have.
fig = plot_data()

# Find the index that provided the highest log posterior in the chain.
_, i = findmax(ch[:lp])

# Extract the max row value from i.
i = i.I[1]

# Plot the posterior distribution with a contour plot
x1_range = collect(range(-6; stop=6, length=25))
x2_range = collect(range(-6; stop=6, length=25))
Z = [nn_forward([x1, x2], θ[i, :])[1] for x1 in x1_range, x2 in x2_range]
contour!(x1_range, x2_range, Z; linewidth=3, colormap=:seaborn_bright)
fig
UndefVarError: UndefVarError(:ch)
UndefVarError: `ch` not defined
Stacktrace:
 [1] top-level scope
   @ ~/work/docs/docs/tutorials/03-bayesian-neural-network/index.qmd:295

The contour plot above shows that the MAP method is not too bad at classifying our data.

Now we can visualize our predictions.

\[ p(\tilde{x} | X, \alpha) = \int_{\theta} p(\tilde{x} | \theta) p(\theta | X, \alpha) \approx \sum_{\theta \sim p(\theta | X, \alpha)}f_{\theta}(\tilde{x}) \]

The nn_predict function takes the average predicted value from a network parameterized by weights drawn from the MCMC chain.

# Return the average predicted value across
# multiple weights.
function nn_predict(x, θ, num)
    num = min(num, size(θ, 1))  # make sure num does not exceed the number of samples
    return mean([first(nn_forward(x, view(θ, i, :))) for i in 1:10:num])
end
nn_predict (generic function with 1 method)

Next, we use the nn_predict function to predict the value at a sample of points where the x1 and x2 coordinates range between -6 and 6. As we can see below, we still have a satisfactory fit to our data, and more importantly, we can also see where the neural network is uncertain about its predictions much easier—those regions between cluster boundaries.

# Plot the average prediction.
fig = plot_data()

n_end = 1500
x1_range = collect(range(-6; stop=6, length=25))
x2_range = collect(range(-6; stop=6, length=25))
Z = [nn_predict([x1, x2], θ, n_end)[1] for x1 in x1_range, x2 in x2_range]
contour!(x1_range, x2_range, Z; linewidth=3, colormap=:seaborn_bright)
fig
UndefVarError: UndefVarError(:θ)
UndefVarError: `θ` not defined
Stacktrace:
 [1] (::var"#20#21")(::Tuple{Float64, Float64})
   @ Main.Notebook ./none:0
 [2] iterate
   @ ./generator.jl:47 [inlined]
 [3] collect(itr::Base.Generator{Base.Iterators.ProductIterator{Tuple{Vector{Float64}, Vector{Float64}}}, var"#20#21"})
   @ Base ./array.jl:834
 [4] top-level scope
   @ ~/work/docs/docs/tutorials/03-bayesian-neural-network/index.qmd:348

Suppose we are interested in how the predictive power of our Bayesian neural network evolved between samples. In that case, the following graph displays an animation of the contour plot generated from the network weights in samples 1 to 1,000.

# Number of iterations to plot.
n_end = 500

anim = @gif for i in 1:n_end
    plot_data()
    Z = [nn_forward([x1, x2], θ[i, :])[1] for x1 in x1_range, x2 in x2_range]
    contour!(x1_range, x2_range, Z; title="Iteration $i", clim=(0, 1))
end every 5
UndefVarError: UndefVarError(:θ)
UndefVarError: `θ` not defined
Stacktrace:
 [1] (::var"#22#23"{Int64})(::Tuple{Float64, Float64})
   @ Main.Notebook ./none:0
 [2] iterate
   @ ./generator.jl:47 [inlined]
 [3] collect(itr::Base.Generator{Base.Iterators.ProductIterator{Tuple{Vector{Float64}, Vector{Float64}}}, var"#22#23"{Int64}})
   @ Base ./array.jl:834
 [4] macro expansion
   @ ~/work/docs/docs/tutorials/03-bayesian-neural-network/index.qmd:367 [inlined]
 [5] top-level scope
   @ ~/.julia/packages/Plots/du2dt/src/animation.jl:251

This has been an introduction to the applications of Turing and Lux in defining Bayesian neural networks.

Back to top