Optimal way to run multiple chains for a bayesian neural network trained with HMC (tfp.mcmc.HamiltonianMonteCarlo)

Hi,

I'm working on a master's thesis where we want to sample the exact posterior of a bayesian neural network using HMC and the No-U-Turn sampler in regression tasks. The current code implementation does not optimally utilize the fact that several chains can be run simultaneously when calling tfp.mcmc.sample_chain.

I'll walk through the code implementation. First, the weights of the network is stored in a list as

weights = [kernel:0, bias:0, kernel:1, bias:1, ...]

The kernels have shape [num_chains, n, m] and the biases have shape [num_chains, m] to allow for several chains run simultaneously as per the docs of Tensorflow Probability.

The log prior is defined as

def log_prior(weights, lamb=1e-3):
    kernel = weights[::2]
    bias = weights[1::2]
    res = 0
    for w, b in zip(kernel, bias):
        res += tf.reduce_sum(w ** 2, axis=(-1,-2))
        res += tf.reduce_sum(b ** 2, axis=-1)
    return -0.5 * lamb * res

The log likelihood is defined as

def log_likelihood(x, y, weights, activation=tf.nn.relu):
    kernel = weights[::2]
    bias = weights[1::2]
    for w, b in zip(kernel[:-1], bias[:-1]):
        x = activation(tf.matmul(x, w) + b[..., None, :])
    y_pred = tf.matmul(x, kernel[-1]) + bias[-1][..., None, :]
    
    return -0.5 * tf.reduce_sum((y_pred - y) ** 2, axis=(-1,-2))

And the target log probability function is defined by

def get_target_log_prob_fn(x, y):
    def target_log_prob_fn(*weights):
        return log_prior(weights) + log_likelihood(x, y, weights)
    return target_log_prob_fn    

Thus, given training features x of shape [num_points, num_features] and training targets y of shape [num_points, num_outputs], we can extract the target log probability function as

target_log_prob_fn = get_target_log_prob_fn(x, y)

While this code provide adequate results in terms of a proper non-linear regression model, it does not run well when num_chains > 1. In fact, the case num_chains = 1 with num_results = 100 runs significantly faster than num_chains = 10 with num_results = 10, even though the produce the exact same number of results.

As an example of model, we can create it with the following function:

def get_weights(layers, num_chains):
    weights = []
    for n, m in zip(layers[:-1], layers[1:]):
        w = tf.random.normal(shape=(num_chains, n, m))
        b = tf.random.normal(shape=(num_chains, m))
        weights.extend([w, b])
    return weights

And run the chain with adaptive HMC like so:

num_results = 100
num_burnin_steps = 1000
num_chains = 10
layers = [input_sz, 10, 10, output_sz]
weights = get_weights(layers, num_chains)

kernel = tfp.mcmc.HamiltonianMonteCarlo(
    target_log_prob_fn=get_target_log_prob_fn(x_train, y_train),
    num_leapfrog_steps=60,
    step_size=0.01
)

kernel = tfp.mcmc.DualAveragingStepSizeAdaptation(
    inner_kernel=kernel,
    num_adaptation_steps=int(0.8 * num_burnin_steps)
)

chain = sample_chain(
    kernel=kernel,
    num_results=num_results,
    trace_fn=None,
    num_burnin_steps=num_burnin_steps,
    current_state=weights,
)

I'll leave link to a google colab notebook that implements the code and demonstrates the problem: https://colab.research.google.com/drive/1oc5czHsGSsi0XC9T7267EfJUR1OdHtlY?usp=sharing

Is there a better way to structure the model parameters such that it better utilizes the parallelization offered by the tfp.mcmc kernels and tfp.mcmc.sample_chain?

Start by running the linked Google Colab notebook and inspect the calls to tfp.mcmc.sample_chain, tfp.mcmc.HamiltonianMonteCarlo, and DualAveragingStepSizeAdaptation. Compare the one-chain and multi-chain timings, then determine whether the parameter structure can use parallel chains more efficiently; done means a documented approach that improves or explains the observed performance.