## Bayesian Neural Networks Background with code snippets

BNNs represent the integration of a hierarchical Bayesian framework together with a neural network structure composed of recursive applications of linear weighted functions followed by nonlinear transformations. With prior distributions on weights, we are able to approximate their posterior distributions and perform posterior prediction via a [variational inference](https://en.wikipedia.org/wiki/Variational_Bayesian_methods) framework.
Consider a model that returns a probability distribution over an output $y$ given an input $x$ and parameter $\theta$, that is $p(y|x,\theta)$. The goal is to learn these parameters from the observed data $D = \{D_i\}_{i=1}^N = \{x_i,y_i\}_{i=1}^N$. Following a Bayesian approach, we put a prior distribution $p(\theta)$ over $\theta$ and aim to obtain the posterior distribution

$$p(\theta|D)\propto p(\theta) p(D|\theta) = p(\theta) \prod_{i=1}^N p(y_i|\theta, x_i). $$

In most cases, the posterior distribution is intractable and an approximation is required. Consider using a variational family distribution $q_v(\theta)$ with parameters $v$ to approximate the posterior by minimizing the KL divergence, i.e.,

\begin{equation}\label{eq:minKL}
\min_v D_{KL}{q_v(\theta)}{p(\theta|D)}
\end{equation}

where,
$$
D_{KL}{q_v(\theta)}{p(\theta|D)} = \mathbb{E}_{q} \bigg[ {\log\frac{q_v(\theta)}{p(\theta|D)}} \bigg] \\
$$

$$
= \mathbb{E}_{q} \bigg[ {\log\frac{q_v(\theta) p(D)}{p(D|\theta)p(\theta)}} \bigg] 
$$

$$
= - \left\{\mathbb{E}_{q}{ [\log p(D|\theta) ]}- D_{KL}{q_v(\theta)}{p(\theta)}\right\} + \log p(D)
$$

where the first term represents the log-likelihood of the model predictions, while the second part serves as a regularizer KL divergence between the approximate posterior $q_v(\theta)$ and the prior distribution $p(\theta)$. In the case where $\log p(D)$ is constant given prior distribution $p(\theta)$, minimizing the KL divergence is equivalent to maximizing the objective function or the Evidence Lower BOund (ELBO) given by
\begin{equation}
L(v) = \sum_{i=1}^N \mathbb{E}_{q}{\log p(y_i|\theta,x_i)} - \beta D_{KL}{q_v}{p}, \label{eq:Lv}
\end{equation}

```python

#############################################################################
##
## CNN_BNN_Model.py
##
#############################################################################


# %% Loss
# Convert the target to a one-hot tensor of shape (batch_size, 10) and
# with a on-value of 1 for each one-hot vector of length 10.
labels_distribution = tfd.Categorical(logits=logits,name='label_dist')


# Compute the -ELBO as the loss, averaged over the batch size.
neg_log_likelihood = -tf.reduce_mean(input_tensor=labels_distribution.log_prob(label))

kl = sum(model.losses) / N

KLscale=1.0
Loss_ = neg_log_likelihood + KLscale * kl


```



```python
#############################################################################
##
## CIFAR10_BNN.py
##
#############################################################################


kl_regularizer = FLAGS.kl_regularizer

# Compute the -ELBO as the loss. The kl term is annealed from 0 to 1 over
# the epochs specified by the kl_annealing flag.
log_likelihood = labels_distribution.log_prob(labels)
neg_log_likelihood = -tf.reduce_mean(input_tensor=log_likelihood)
kl = sum(model.losses) / len(x_train) * tf.minimum(1.0, kl_regularizer)

loss = neg_log_likelihood + kl

```

- [losses](https://www.tensorflow.org/probability/api_docs/python/tfp/layers/Convolution2DFlipout)

  + Upon being built, this layer adds losses (accessible via the losses property) representing the divergences of kernel and/or bias surrogate posteriors and their respective priors. When doing minibatch stochastic optimization, make sure to scale this loss such that it is applied just once per epoch (e.g. if kl is the sum of losses for each element of the batch, you should pass kl / num_examples_per_epoch to your optimizer).


The $\beta$ is a hyperparameter for tuning the degree of regularization during training of the neural network. Typically the $\beta$ hyperparameter is introduced to address the  difficulty of the {KL-term vanishing} which is often observed while training Variational Auto Encoders (VAEs). This hyperparameter was introduced by Bowman et.al [1] as *KL-annealing* where the $\beta$ parameter can be varied from 0 to 1 over the course of training.

Recently an analysis on the scheduling strategies of $\beta$ was presented by Liu et.al [2]. Once the BNN model is trained the inference procedure is to compute the predictive probability distribution $p(y^{*}|x^{*})$ given by

\begin{equation}
p(y^{*}|x^{*}) = \int p_{w}(y^{*}|x^{*})p(w) dw \, ,
\label{Eq:InferPost}
\end{equation}

where $x^{*}$ is the unseen (test) sample and $y^{*}$ is the corresponding predicted class.
Finding a closed form solution to the above integral for non conjugate pair of distribution is not possible hence the integral can be approximated as an expectation by sampling from $q_v(w|D)$ as

\begin{equation}
\mathbb{E}_{q}{p(y^{*}|x^{*})} = \int q_v(\theta | D) p_{w}(y|x) dw \approx \frac{1}{S} \Sigma_{i=1}^{S} p_{w_i}(y^{*}|x^{*}) 
\end{equation}

where $S$ are the number of samples or Monte-Carlo iterations.


# Training Procedure

- The procedure is summarize to optimize the loss function described in minimization of KL divergence. The backpropagation algorithm is at the heart of training a neural network and  relies of the calculation of gradients of the loss function which are subsequently used to update the weights layer by layer. There are several approaches available for computing the gradients such as  score-function gradient estimators shown in [3,4] reparameterization  gradient  estimators  described in [5,6], or a combination of the two as described in Naesseth et.al [7]. These methods provide an unbiased stochastic gradient which is used to calculate the optimum for the loss function.

- In training the model reparametrization procedure for computing the gradients is used.  
We use the Gaussian distribution for the variational posterior and initilize the weights, $\theta$, by sampling from the unit Gaussian with mean $\mu$ and standard-deviation $\sigma$. 

- The hyperparameters for the variational posterior are therefore $v = (\mu, \sigma)$. The reparametization trick is to parameterize weights as a function, $t$, where $\theta = t(v, \epsilon) = \mu + {\rm{log}}(1 + exp(\sigma)) \circ \epsilon $, where $\circ$ represents an element-wise multiplication and the parameter free noise is defined by $\epsilon \sim \mathit{N}(0, I)$. Eq. for Loss $L$ defined earlier can be rewritten in terms of $L(\theta,v)$ and the gradients, $\Delta_{\mu} L(\theta,v)$ and $\Delta_{\sigma} L(\theta,v)$, are computed to update the hyperparamters $ \mu  \leftarrow  \mu + \eta \Delta_{\mu} L(\theta,v) $ and  $ \sigma  \leftarrow  \sigma + \eta \Delta_{\sigma} L(\theta,v) $.


# Distributed Training Using Horovod

+ Setting the distributed optimizer,with scaling the learning rare with horovod size.


```python
##############
#
# Demo Code
#
###############
# Horovod: adjust learning rate based on number of GPUs.
opt = tf.train.AdamOptimizer(FLAGS.learning_rate * hvd.size())

# Horovod: add Horovod Distributed Optimizer.
opt = hvd.DistributedOptimizer(opt)

```

+ Hooks to setup the monitored training session.


```python
##############
#
# Demo Code
#
###############


hooks = [
    # Horovod: BroadcastGlobalVariablesHook broadcasts initial variable states
    # from rank 0 to all other processes. This is necessary to ensure consistent
    # initialization of all workers when training is started with random weights
    # or restored from a checkpoint.
    hvd.BroadcastGlobalVariablesHook(0),

    # Horovod: adjust number of steps based on number of GPUs.
    tf.train.StopAtStepHook(last_step= Num_iter // hvd.size()),

    tf.train.LoggingTensorHook(tensors={'step': global_step, 'loss': Loss_},every_n_iter=500),
]
```

+ Config for horovod, more details can be seen in horovod tutorials.


```python

##############
#
# Demo Code
#
###############
# # Horovod: pin GPU to be used to process local rank (one GPU per process)
config.gpu_options.allow_growth = True
config.gpu_options.visible_device_list = str(hvd.local_rank())


checkpoint_dir = os.path.join(dirmake,'checkpoints') if hvd.rank() == 0 else None

```

+ Monitored training session to take cares of initialization and hooks to create a session to run various operations in the tensorflow graph etc. The while loop contains the training code.

```python
##############
#
# Demo Code
#
###############

# The MonitoredTrainingSession takes care of session initialization,
# restoring from a checkpoint, saving to a checkpoint, and closing when done
# or an error occurs.
with tf.train.MonitoredTrainingSession(checkpoint_dir=checkpoint_dir,hooks=hooks,config=config) as mon_sess:
    #Train Time start 
    start_Train_time = time.time()
    iter_num = 0

    while not mon_sess.should_stop():
        # Training code
        

```


+ **EXAMPLE BNN Network Summary**

```
Model: CIFAR-10 Bayesian Model "sequential"
_________________________________________________________________
Layer (type)                 Output Shape              Param #   
=================================================================
Conv_I (Conv2DFlipout)       (None, 32, 32, 32)        1760      
_________________________________________________________________
Conv_II (Conv2DFlipout)      (None, 32, 32, 32)        18464     
_________________________________________________________________
Max_I_ (MaxPooling2D)        (None, 16, 16, 32)        0         
_________________________________________________________________
Conv_III (Conv2DFlipout)     (None, 16, 16, 64)        36928     
_________________________________________________________________
Conv_IV (Conv2DFlipout)      (None, 16, 16, 64)        73792     
_________________________________________________________________
Max_II (MaxPooling2D)        (None, 8, 8, 64)          0         
_________________________________________________________________
Conv_V (Conv2DFlipout)       (None, 8, 8, 128)         147584    
_________________________________________________________________
Conv_VI (Conv2DFlipout)      (None, 8, 8, 128)         295040    
_________________________________________________________________
Max_III (MaxPooling2D)       (None, 4, 4, 128)         0         
_________________________________________________________________
flatten (Flatten)            (None, 2048)              0         
_________________________________________________________________
Dense_I (DenseFlipout)       (None, 128)               524416    
_________________________________________________________________
Dense_II (DenseFlipout)      (None, 10)                2570      
=================================================================
Total params: 1,100,554
Trainable params: 1,100,554
Non-trainable params: 0
_________________________________________________________________
```


# References

[1] Samuel R Bowman, Luke Vilnis, Oriol Vinyals, Andrew M Dai, Rafal Jozefowicz, and SamyBengio. Generating sentences from a continuous space.arXiv preprint arXiv:1511.06349,2015.

[2]  Xiaodong Liu, Jianfeng Gao, Asli Celikyilmaz, Lawrence Carin, et al.  Cyclical annealingschedule: A simple approach to mitigating kl vanishing.arXiv preprint arXiv:1903.10145,2019.  

[3] John Paisley, David Blei, and Michael Jordan. Variational bayesian inference with stochas-tic search.arXiv preprint arXiv:1206.6430, 2012.

[4] Rajesh  Ranganath,  Sean  Gerrish,  and  David  M  Blei.   Black  box  variational  inference.arXiv preprint arXiv:1401.0118, 2013.

[5] Charles  Blundell,  Julien  Cornebise,  Koray  Kavukcuoglu,  and  Daan  Wierstra.   Weightuncertainty in neural networks.arXiv preprint arXiv:1505.05424, 2015
 
[6] Diederik P Kingma and Max Welling. Auto-Encoding Variational Bayes. Technical report
  
[7] CA  Naesseth,  FJR  Ruiz,  SW  Linderman,  and  DM  Blei.   Reparameterization  gradientsthrough  acceptance-rejection  sampling  algorithms.   InProceedings  of  the  20th  Interna-tional Conference on Artificial Intelligence and Statistics, AISTATS 2017, 2017.

**Author: Himanshu Sharma, Argonne National Laboratory**


```python

```

d61cc87709e16fdfafad465b41ea47d38efd61a3
ipynb
Jupyter Notebook
src/VariationalInferenceNotebook.ipynb
Himscipy/bnn_hvd
c564c076904585600987eda16e96804060f529f2
MIT
2020-05-18T22:07:48.000Z
2022-01-15T19:31:35.000Z
src/VariationalInferenceNotebook.ipynb
Himscipy/bnn_hvd
c564c076904585600987eda16e96804060f529f2
MIT
src/VariationalInferenceNotebook.ipynb
Himscipy/bnn_hvd
c564c076904585600987eda16e96804060f529f2
MIT
2020-05-18T22:07:53.000Z
2020-08-04T21:21:24.000Z
Qwen/Qwen-72B
1. YES
2. YES