On the optimization and pruning for Bayesian deep learning

Xiongwen Ke and Yanan Fan

School of mathematics and statistics, UNSW

###### Abstract

The goal of Bayesian deep learning is to provide uncertainty quantification
via the posterior distribution. However, exact inference over the weight space
is computationally intractable due to the ultra-high dimensions of the neural
network. Variational inference (VI) is a promising approach, but naive
application on weight space does not scale well and often underperform on
predictive accuracy. In this paper, we propose a new adaptive variational
Bayesian algorithm to train neural networks on weight space that achieves high
predictive accuracy. By showing that there is an equivalence to Stochastic
Gradient Hamiltonian Monte Carlo(SGHMC) with preconditioning matrix, we then
propose an MCMC within EM algorithm, which incorporates the spike-and-slab
prior to capture the sparsity of the neural network. The EM-MCMC algorithm
allows us to perform optimization and model pruning within one-shot. We
evaluate our methods on CIFAR-10, CIFAR-100 and ImageNet datasets, and
demonstrate that our dense model can reach the state-of-the-art performance
and our sparse model perform very well compared to previously proposed pruning
schemes.

##  1 INTRODUCTION

Bayesian inference (Bishop and Nasrabadi,, 2006) provide an elegant way to
capture uncertainty via the posterior distribution over model parameters.
Unfortunately, posterior inference is intractable in any reasonably sized
neural network models.

Works focussing on scalable inference for Bayesian deep learning over the last
decade can be separated by two streams. One is a deterministic approximation
approach such as variational inference (Graves,, 2011; Blundell et al.,,
2015), dropout (Gal and Ghahramani,, 2016), Laplace approximation (Ritter et
al.,, 2018), or expectation propagation(Hernández-Lobato and Adams,, 2015).
The other stream involve sampling approaches such as MCMC using stochastic
gradient Langevin dynamics (SGLM) (Welling and Teh,, 2011; Chen et al.,,
2014).

Prior to 2019, Bayesian neural networks (BNN) generally struggle with
predictive accuracy and computational efficiency. Recently, a lot of advances
have been made in both directions of research. In deterministic approach,
several authors began to consider using dimension reduction techniques such as
subspace inference (Maddox et al.,, 2019), rank-1 parameterization(Dusenberry
et al.,, 2020), subnetwork inference(Daxberger et al.,, 2021) and node-space
inference(Trinh et al.,, 2022). In the sampling approach, (Zhang et al.,,
2019) propose to use cycles of learning rates with a high-to-low step size
schedule. A large step size in the early stage of the cycle results in
aggressive exploration in the parameter space; as the step size decreases, the
algorithm begins to collect sample around the local mode.

Apart from the progress within these two streams, Wilson and Izmailov, (2020)
show that deep ensembles (Lakshminarayanan et al.,, 2017) can be interpreted
as an approximate approach to posterior predictive distribution. They combine
multiple independently trained SWAG (Gaussian stochastic weight averaging)
approximations (Maddox et al.,, 2019; Izmailov et al.,, 2018) to create a
mixture of Gaussians approximation to the posterior. However, performing
variational inference directly on weight space still produces poor predictive
accuracy and struggles with computational efficiency. Even a simple mean-field
variational inference will involve double the number of parameters of the
neural network (i.e., mean and variance for each weight), which incurs a extra
GPU memory requirement and 2-5 times of the runtime of the baseline neural
network (Osawa et al.,, 2019).

In this paper, we develop an adaptive optimization algorithm for Gaussian
Mean-field variational Bayesian inference that can achieve state-of-the-art
predictive accuracy. We further show that when the learning rate is small and
the update of the posterior variance has been frozen, the algorithm is
equivalent to the SGHMC (Stochastic Gradient Hamiltonian Monte Carlo) with
preconditioning matrix. Therefore, if we exploit the closed form expression of
the gradient for the posterior variances and only keep track of the weight
generated by the algorithm(no need to register the mean and variance parameter
in the code), we can achieve big savings on GPU memory and runtime costs.

Based on the connection to SGHMC, we extend the EM Algorithm for Bayesian
variable selection (Ročková and George,, 2014; Wang et al.,, 2016; Ročková,,
2018) for linear models to neural networks by replacing the Gaussian prior in
BNN with the spike-and-slab group Gaussian prior(Xu and Ghosh,, 2015). Our
method is an MCMC within EM algorithm, which will switch the weight decay
factor between small and large based on the magnitude of each group during
training. Since by construction, there are no exact zeros, we further find a
simple pruning criterion to remove the weights permanently during training. A
sparse model will be trained in one-shot without additional retrain. Our
approach is more computationally efficient than those dynamic pruning
strategies that allows regrow (Zhu and Gupta,, 2017; Dettmers and
Zettlemoyer,, 2019; Lin et al.,, 2020). We will show that this aggressive
approach has no performance loss. Our code is available at github:
https://github.com/z5041294/optimization-and-pruning-for-BNN

##  2 Optimization

###  2.1 Preliminaries on variational Bayesian neural network

Given a dataset $\mathcal{D}=\left\\{x_{i},y_{i}\right\\}_{i=1}^{N}$, a
Bayesian neural network (BNN) is defined in terms of a prior $p(\mathbf{w)}$
on the $p$-dimensional weights, as well as the neural network likelihood
$p(\mathcal{D}|\mathbf{w})$. Variational Bayesian methods approximates the
true posterior $p(\mathbf{w}|\mathcal{D})$ by minimising the KL divergence
between the approximate distribution, $q_{\theta}(\mathbf{w})$, and the true
posterior. This is shown to be equivalent to maximizing the evidence lower
bound (ELBO):

$\mathcal{L}[\theta]=\mathbf{E}_{q_{\theta}}[\log
p(\mathcal{D}\mid\mathbf{w})]-\mathrm{D}_{\mathrm{KL}}(q_{\theta}(\mathbf{w})\|p(\mathbf{w}))$
(1)

where we consider a Bayesian neural net with Gaussian prior $p(\mathbf{w})\sim
N_{p}(0,\boldsymbol{\Sigma}_{0})$ and a Gaussian approximate posterior
$q_{\theta}(\mathbf{w})\sim N_{p}(\boldsymbol{\mu},\boldsymbol{\Sigma})$ where
$\theta=(\boldsymbol{\mu},\boldsymbol{\Sigma})$. To make it scale to large
sized models, we assume both the approximate posterior and prior weights are
independent, such that $p(w_{j})\sim N(0,\delta^{-1})$ and
$q_{\theta_{j}}(w_{j})\sim N(\mu_{j},\sigma_{j}^{2})$. The gradient of ELBO
with respect to $\boldsymbol{\mu}$ and $\boldsymbol{\Sigma}$ is

$\displaystyle\nabla_{\boldsymbol{\mu}}\mathcal{L}$
$\displaystyle=\mathbf{E}_{N_{p}(\boldsymbol{\mu},\boldsymbol{\Sigma})}\left[\nabla_{\mathbf{w}}\log
p(\mathcal{D}\mid\mathbf{w})\right]-\boldsymbol{\Sigma}_{0}^{-1}\boldsymbol{\mu}\approx-\frac{1}{S}\sum_{i=1}^{S}\mathbf{g}_{i}-\delta\mu$
(2) $\displaystyle\nabla_{\boldsymbol{\Sigma}}\mathcal{L}$
$\displaystyle=\frac{1}{2}\mathbf{E}_{N_{p}(\boldsymbol{\mu},\boldsymbol{\Sigma})}\left[\nabla_{\mathbf{w}}^{2}\log
p(\mathcal{D}\mid\mathbf{w})\right]+\frac{1}{2}\boldsymbol{\Sigma}^{-1}-\frac{1}{2}\boldsymbol{\Sigma}_{0}^{-1}$
$\displaystyle\approx-\frac{1}{2S}\sum_{i=1}^{S}\mathbf{g}_{i}^{2}+\frac{1}{2}\mathbf{diag}(\sigma_{j}^{-2})-\frac{1}{2}\delta$

where $\mathbf{g}_{i}=-\nabla_{\mathbf{w}}\log
p(\mathcal{D}\mid\mathbf{w_{i}})$ and
$\mathbf{w}_{i}\sim\prod_{j=1}^{p}N(\mu_{j},\sigma_{j}^{2})$ are Monte Carlo
samples. In addition, we also have the second order derivative of $\mu$

$\displaystyle\nabla_{\boldsymbol{\mu}}^{2}\mathcal{L}$
$\displaystyle=-\mathbf{E}_{N_{p}(\boldsymbol{\mu},\boldsymbol{\Sigma})}\left[\nabla^{2}_{\mathbf{w}}\log
p(\mathcal{D}\mid\mathbf{w})\right]+\boldsymbol{\Sigma}_{0}^{-1}\approx\frac{1}{S}\sum_{i=1}^{S}\mathbf{g}_{i}^{2}+\delta$

Using Gaussian back-propagation and reparameterization trick, an alternative
MC approximation (Khan et al.,, 2018) can be used, such that
$-\mathbf{E}_{N_{p}(\boldsymbol{\mu},\boldsymbol{\Sigma})}\left[\nabla^{2}_{\mathbf{w}}\log
p(\mathcal{D}\mid\mathbf{w})\right]\approx\frac{1}{S}\sum_{i=1}^{S}\mathbf{g}_{i}\odot\frac{\mathbf{\epsilon}_{i}}{\mathbf{\sigma}}$,
where $\mathbf{\epsilon}_{i}\sim N(0,\mathbf{1}_{p})$. In practice, to reduce
the runtime complexity, we often use $S=1$ as long as the batch size in one
iteration is not too small.

###  2.2 Bayesian versions of adaptive algorithm

In Gaussian mean-field variational Bayesian inference, the natural-gradient
descent(Khan et al.,, 2018; Zhang et al.,, 2018; Osawa et al.,, 2019)
optimizes the ELBO bound by updating

$\displaystyle\boldsymbol{\mu}_{t+1}$
$\displaystyle=\underset{\boldsymbol{\mu}\in\mathbf{R}^{p}}{\operatorname{argmin}}\left\\{\langle\nabla_{\boldsymbol{\mu}}\mathcal{L},\boldsymbol{\mu}\rangle+\frac{1}{2l_{t}}(\boldsymbol{\mu}-\boldsymbol{\mu}_{t})^{T}\mathbf{diag}(\mathbf{\sigma}^{-2})^{\alpha}(\boldsymbol{\mu}-\boldsymbol{\mu}_{t})^{T}\right\\}$
$\displaystyle=\boldsymbol{\mu}_{t}-l_{t}(\boldsymbol{\sigma}_{t}^{2})^{\alpha}\odot\nabla_{\boldsymbol{\mu}_{t}}\mathcal{L}$

where $l_{t}$ is a learning rate, $\alpha\in\left\\{\frac{1}{2},1\right\\}$
and $\frac{1}{\boldsymbol{\sigma}_{t}^{2}}$ is updated with momentum
$\frac{1}{\boldsymbol{\sigma}_{t}^{2}}=\frac{1-\lambda\gamma}{\boldsymbol{\sigma}_{t-1}^{2}}+\gamma([\boldsymbol{g}_{t}\odot\boldsymbol{g}_{t}]+\frac{\delta}{N})$
($0<\gamma<1$ and $\lambda$ is another learning rate). When
$\alpha=\frac{1}{2}$, this is similar to the Adam (Kingma and Ba,, 2014)
algorithm and when $\alpha=1$, the algorithm is very close to second-order
optimization with the Hessian matrix given by $\boldsymbol{\Sigma}^{-1}$ when
$\nabla_{\boldsymbol{\Sigma}}\mathcal{L}=0$ in equation (2).

There are two concerns for this version of Bayesian adaptive algorithm: First,
similar to Adam, the variance of adaptive learning rate is problematically
large in the early stages of the optimization (Liu et al.,, 2019). A warm-up
stage is recommended for traditional Adam. As an example, consider a ReLu
neural network with a single hidden layer and binary cross entropy loss, then
if a normal initialization of the weight with mean zero has been used, then
the variance of the adaptive learning will diverge at the beginning of the
training period (See detailed discussion in the Appendix A). Second, when the
number of Monte Carlo samples for the gradient is $S=1$ and the learning rate
is small, the injected Gaussian noise may dominate the gradient. We see that

$\displaystyle\mathbf{w}_{t}-\mathbf{w}_{t-1}$
$\displaystyle=\boldsymbol{\mu}_{t}-\boldsymbol{\mu}_{t-1}+(\boldsymbol{\sigma}_{t}\odot\boldsymbol{\epsilon}-\boldsymbol{\sigma}_{t-1}\odot\boldsymbol{\epsilon}^{\prime})$
$\displaystyle=-l_{t}(\boldsymbol{\sigma}_{t}^{2})^{\alpha}\odot\boldsymbol{m}_{t}+\boldsymbol{\epsilon}\odot\sqrt{\boldsymbol{\sigma}_{t}^{2}+\boldsymbol{\sigma}_{t-1}^{2}}$

where $\boldsymbol{m}_{t}$ is the momentum of gradient. For
$\alpha=\frac{1}{2}$, when the learning rate is small, the gradient may be
erased by the injected noise. For $\alpha=1$, this could also happen when
$\boldsymbol{\sigma}_{t}^{2}\ll 1$. Empirically, we observe that both
$\mathbf{w}\ll 1$ and $\boldsymbol{\sigma}^{2}\ll 1$ in Bayesian deep
learning. So a warm up strategy that starts with a very small learning rate
may not work for variational BNN. On the contrary, when the learning rate is
large, the size of the injected noise will be relatively small, the model
could overfit. In summary, there is a delicate balance between the size of the
gradient and the size of the injected noise. This argument will become more
lucid when we show the connection between our algorithm and SGHMC.

###  2.3 Constrained variational Adam

To remedy the issue highlighted above, we reparameterize
$\boldsymbol{\sigma}_{i}$ as a product of global and local parameters such
that $\boldsymbol{\sigma}_{i}=\alpha\boldsymbol{\tau}_{i}$ where
$\boldsymbol{\tau}_{i}\in(0,1)$ and $\alpha>0$. We treat $\alpha$ as a hyper-
parameter, which can be used in an annealing scheme. Now, the variance of the
adaptive learning rate is upper bounded by $\frac{\alpha^{2}}{12}$. The
updated posterior mean is modified to

$\displaystyle\boldsymbol{\mu}_{t+1}$
$\displaystyle=\underset{\boldsymbol{\mu}\in\mathbf{R}^{p}}{\operatorname{argmin}}\left\\{\langle\nabla_{\boldsymbol{\mu}}\mathcal{L},\boldsymbol{\mu}\rangle+\frac{1}{2l_{t}}(\boldsymbol{\mu}-\boldsymbol{\mu}_{t})^{T}\mathbf{diag}(\boldsymbol{\tau}^{-1})(\boldsymbol{\mu}-\boldsymbol{\mu}_{t})^{T}\right\\}$
$\displaystyle=\boldsymbol{\mu}_{t}-l_{t}\boldsymbol{\tau}_{t}\odot\nabla_{\boldsymbol{\mu}_{t}}\mathcal{L}$

Since $\boldsymbol{\tau}_{i}\in(0,1)$ are the parameters we need to learn, the
objective function for updating the posterior variance change to:

$\displaystyle\boldsymbol{\tau}_{t+1}$
$\displaystyle=\underset{\boldsymbol{\tau}\in(0,1)^{p}}{\operatorname{argmin}}\left\\{\langle\nabla_{\boldsymbol{\tau}}\mathcal{L},\boldsymbol{\tau}\rangle+\frac{1}{2\eta}\|\boldsymbol{\tau}-\boldsymbol{\tau}_{t}\|^{2}\right\\}$
$\displaystyle=\underset{\boldsymbol{\tau}\in(0,1)^{p}}{\operatorname{argmin}}\|(\boldsymbol{\tau}_{t}-\eta\nabla_{\boldsymbol{\tau}}\mathcal{L})-\tau\|^{2}$
$\displaystyle=\pi_{(0,1)^{P}}(\boldsymbol{\tau}_{t}-\eta\nabla_{\boldsymbol{\tau}}\mathcal{L})$

where $\eta$ is a learning rate, $\pi_{(0,1)^{P}}$ denotes the Euclidean
projection onto $(0,1)^{P}$. Now we replace the Euclidean distance of
$\frac{1}{2}\|\boldsymbol{\tau}_{t+1}-\boldsymbol{\tau}_{t}\|^{2}$ to Bregman
distance $D_{G}(\boldsymbol{\tau}_{t+1},\boldsymbol{\tau}_{t})$, where
$D_{G}(\cdot,\cdot)$ is generated by a strictly convex twice-differentiable
function $G(\cdot)$

$D_{G}(\boldsymbol{\tau},\boldsymbol{\tau}^{\prime}):=G(\boldsymbol{\tau})-G(\boldsymbol{\tau}^{\prime})-\langle\nabla
G(\boldsymbol{\tau}^{\prime}),\boldsymbol{\tau}-\boldsymbol{\tau}^{\prime}\rangle.$

If $G=\frac{1}{2}\|\cdot\|^{2}$, we recover the squared Euclidean distance.
This modification leads to a more general version of gradient descent
algorithm which is known as mirror descent.

The idea is that we want to find a distance function which can better reflect
the geometry of $(0,1)^{p}$. Now our objective function can be written as a
more general form:

$\boldsymbol{\tau}_{t+1}=\underset{\boldsymbol{\tau}\in(0,1)^{p}}{\operatorname{argmin}}\left\\{\langle\nabla_{\boldsymbol{\tau}}\mathcal{L},\boldsymbol{\tau}\rangle+\frac{1}{\eta}D_{G}(\boldsymbol{\tau},\boldsymbol{\tau}_{t})\right\\}$
(3)

The choice of $G(\cdot)$ will define the geometry and the distance metric of
the primal space and dual space. Values in $(0,1)$ may be interpreted as a
probability. We define the distance
$D_{G}(\boldsymbol{\tau},\boldsymbol{\tau}_{t})$ as negative binary cross
entropy loss:

$G(\tau)=\tau\log\tau+(1-\tau)\log(1-\tau)$

Then parameter in the dual space is defined as

$\rho=\nabla_{\tau}G(\tau)=g(\tau)=\log(\tau)-\log(1-\tau)$ (4)

with the inverse function from dual space to the primal space, which is the
logistic sigmoid function

$\tau=\nabla_{\rho}H(\rho)=h(\rho)=g^{-1}(\rho)=\frac{1}{1+\exp(-\rho)}$ (5)

The mirror descent allows us to perform gradient descent in the dual space,
which is unconstrained in our case, as shown in equation (4), and finally move
back to the primal space by equation (5). Algorithm 1 provides the psudo-code
for the constrained variational Adam. (Appendix B provides a quick
introduction of mirror descent and the comparison with the reparameterization
trick).

Algorithm 1 Constrained variational Adam

Input: $\beta=(\beta_{1},\beta_{2})$,$\delta\text{(weight decay)}$

$N\text{(training size)}$, $B\text{(batch size)}$

Initialization: $\boldsymbol{m}$, $\boldsymbol{\rho}$, $\boldsymbol{\mu}$,
$\boldsymbol{\tau}$

for $t=1\>\cdots\>$ do

$l\leftarrow s_{1}(t)$ $\triangleright$ learning rate schedule

$\eta\leftarrow s_{2}(t)$ $\triangleright$ learning rate schedule

$\alpha\leftarrow s_{3}(t)$$\triangleright$ annealing schedule

$\textbf{Sample}\>\boldsymbol{\epsilon}\sim N(0,\mathbf{1}_{p})$

$\mathbf{w}\leftarrow\boldsymbol{\mu}+\alpha\boldsymbol{\tau}\odot\boldsymbol{\epsilon}$

$\boldsymbol{g}\leftarrow-\frac{1}{B}\sum_{i\in\mathcal{B}}\nabla_{\mathbf{w}}\log(\mathcal{D}_{i}|\mathbf{w})+\frac{\delta}{N}\boldsymbol{\mu}$

$\boldsymbol{m}\leftarrow\beta_{1}\boldsymbol{m}+\beta_{2}\boldsymbol{g}$

$\boldsymbol{\rho}\leftarrow\boldsymbol{\rho}+\left(\frac{\eta}{\tau}-\eta\alpha^{2}\delta\boldsymbol{\tau}\right)-\eta\alpha\boldsymbol{\epsilon}\odot\boldsymbol{g}$

$\boldsymbol{\tau}\leftarrow 1/(1+e^{-\boldsymbol{\rho}})$

$\boldsymbol{\mu}\leftarrow\boldsymbol{\mu}-l\boldsymbol{\tau}\odot\boldsymbol{m}$

end for

###  2.4 Connection to SGHMC

If we set the global parameter of posterior standard deviation
$\alpha=kl_{t}^{3/4}$ for some positive $k$ and $\beta_{1}=1-hl_{t}^{1/2}$ for
some positive $h$, then based on the following two assumptions:

  * •

the learning rate is small and $l_{t}\approx l_{t-1}$

  * •

the update of $\tau$ is frozen or it has already converged

then the constrained variational Adam algorithm is equivalent to SGHMC(Chen et
al.,, 2014)

$\displaystyle d\mathbf{w}_{t}$
$\displaystyle=D_{\boldsymbol{\tau}}\boldsymbol{r}_{t}dt$ (6) $\displaystyle
d\boldsymbol{r}_{t}$
$\displaystyle=-\beta_{2}\boldsymbol{g}_{t}dt-h\boldsymbol{r}_{t}dt+\sqrt{2+2\beta_{1}^{2}}kd\boldsymbol{B}_{t}$

where $D_{\boldsymbol{\tau}}=\mathbf{Diag}(\boldsymbol{\tau})$ is a
preconditioning matrix and $d\boldsymbol{B}_{t}$ is Brownian motion. In
Appendix C, we show this relationship by setting $\sqrt{l_{t}}=\Delta t$. The
above dynamic produces the distribution proportional to
$\exp\left(-\frac{\beta_{2}L(\mathbf{w})}{k^{2}}\right)$ where
$L(\mathbf{w})=-\frac{1}{N}\sum_{i=1}^{N}\log
p(y_{i}|x_{i},\mathbf{w})-\frac{1}{N}\log p(\mathbf{w})$. By setting
$\beta_{2}=1$ and $k=\frac{1}{\sqrt{N}}$, we recover the posterior of the
weight $p(\mathbf{w}|\mathcal{D})$.

Based on this connection, it is natural to consider using the cyclical
learning rate (Zhang et al.,, 2019) with high-to-low learning rate in each
cycle. A high learning rate at the beginning can quickly find a local mode and
a small learning rate at the end ensures an accurate simulation around the
mode.

From Gaussian variational inference perspective, $\alpha$ is the annealing
parameter, which controls the temperature of the approximated posterior.
Unlike the original Bayesian backpropagation (Blundell et al.,, 2015), we
actually introduce an extra inductive bias into the approximate posterior by
only training the local variance parameter $\tau$ and defining the global
variance parameter $\alpha$. We believe that, in high dimension, the magnitude
of posterior variance for individual weight should be very small.

When the size of the model is large, to improve the predictive performance, we
use the cold posterior (Heek and Kalchbrenner,, 2019; Zhang et al.,, 2019;
Wenzel et al.,, 2020) by setting $k=\frac{1}{N}$ rather than
$\frac{1}{\sqrt{N}}$. By defining
$\boldsymbol{v}_{t}=\boldsymbol{r}_{t-1}\Delta t$, (Chen et al.,, 2014), we
derive the numerical scheme shown in algorithm 2.
Remark: The implementation of constrained variational Adam(CV-Adam) algorithm
is always based on Algorithm 2. In addition, we found that it is not necessary
to freeze the update of $\tau$ as it converges well in practice.

Algorithm 2 SGHMC

Input $\beta=(\beta_{1},\beta_{2})$, $\delta\text{(weight decay)}$

$N\text{(training size)}$, $B\text{(batch size)}$

Initialization $\boldsymbol{\rho}$,$\mathbf{w}$,$\boldsymbol{v}$

for $t=1\>\cdots\>$ do

$\Delta t\leftarrow s_{1}(t)$ $\triangleright$ learning rate schedule

$\eta\leftarrow s_{2}(t)$ $\triangleright$ learning rate schedule

$k\leftarrow s_{3}(t)$ $\triangleright$ annealing schedule

$\alpha\leftarrow k*\Delta t^{1.5}$

$\boldsymbol{g}\leftarrow-\frac{1}{B}\sum_{i\in\mathcal{B}}\nabla_{\mathbf{w}}\log(\mathcal{D}_{i}|\mathbf{w})+\frac{\delta}{N}\mathbf{w}$

$\boldsymbol{\rho}\leftarrow\boldsymbol{\rho}+\left(\frac{\eta}{\boldsymbol{\tau}}-\eta\alpha^{2}\delta\boldsymbol{\tau}\right)-\eta\alpha\boldsymbol{\epsilon}\odot\boldsymbol{g}$

$\boldsymbol{\epsilon}\leftarrow N(0,\mathbf{1}_{p})$

$\boldsymbol{\tau}\leftarrow 1/(1+e^{-\boldsymbol{\rho}})$

$\boldsymbol{v}\leftarrow{\beta_{1}\boldsymbol{v}-\beta_{2}\boldsymbol{g}(\Delta
t)^{2}+\sqrt{2+2\beta_{1}^{2}}k(\Delta t)^{1.5}\boldsymbol{\epsilon}}$

$\mathbf{w}\leftarrow\mathbf{w}-\boldsymbol{\tau}\odot\boldsymbol{v}$

end for

##  3 Pruning

In this section, we describe how we construct an efficient pruning scheme to
obtain a sparse network for Bayesian deep learning.

###  3.1 Spike-and-Slab prior

We apply the group spike and slab prior to the weight parameters in the neural
network. For simplicity, we ignore the subscript for the layers and consider
the weights between any two layers. For convolutional layer, let
$w_{ij}=(w_{ij1},...,w_{ijK^{2}})$ denote the group of $K^{2}$ parameters from
the $i$th input channel to $j$th output channel, and $K\times K$ is the size
of the kernel. Then the prior for $w_{ijk}$ conditioned on a binary inclusion
parameter $\gamma_{ij}$ for the entire group of weights, follows an
independent normal distribution such that:

$\pi\left(w_{ijk}\mid\gamma_{ij}\right)=\begin{cases}\mathrm{N}\left(0,\delta_{1}^{-1}\right)&\text{
if }\gamma_{ij}=1\\\ \mathrm{N}\left(0,\delta_{0}^{-1}\right)&\text{
otherwise, }\end{cases}$ (7)

where $\pi(\gamma_{ij})$ are i.i.d binary distribution such that

$\pi\left(\gamma_{ij}\right)=\begin{cases}1-p_{ij}&\text{ if }\gamma_{ij}=1\\\
p_{ij}&\text{ if }\gamma_{ij}=0.\end{cases}$

The basic assumption is that a priori, within the same kernel all the weight
parameters have the same distribution. For a fully connected layer, we can
also make groupings based on whether they share the same input or output unit,
for instance, if a group is made based on the input unit, then this is similar
to variable/feature selection. Alternatively, we can assign the spike-and-
slab-Gaussian prior to all the weights individually, this is related to
unstructured pruning.

###  3.2 The EM algorithm

We now derive a conditional EM algorithm that returns a MAP estimator of
$\gamma_{ij}$. With the likelihood from the neural network model and prior
distributions specified above, we write the full posterior distribution as
follows:

$\pi\left(\boldsymbol{\gamma},\mathbf{w}\mid\mathbf{y}\right)\propto\underset{ijk}{\prod}p\left(\mathbf{y}\mid
w_{ij}\right)\pi(w_{ijk}\mid\boldsymbol{\gamma}_{ij})\pi(\boldsymbol{\gamma}_{ij}).$

E-Step:The objective function $Q$ at the $t$th iteration in an EM algorithm is
defined as the integrated logarithm of the full posterior with respect to
$w_{ij}$ $\begin{aligned}
Q\left(\gamma_{ij}\mid\gamma_{ij}^{(t-1)}\right)&=\mathbf{E}_{\pi(\mathbf{w}\mid\gamma_{ij}^{(t-1)},\gamma_{-(ij)},\mathbf{y})}\log\pi(\boldsymbol{\gamma},\mathbf{w}\mid\mathbf{y})\\\
&=\mathbf{E}_{\pi(\mathbf{w}\mid\gamma_{ij}^{(t-1)},\gamma_{-(ij)},\mathbf{y})}\log\pi(y\mid
w)\\\
&-\frac{1}{2}E_{\pi(\mathbf{w}\mid\gamma_{ij}^{(t-1)},\gamma_{-(ij)},\mathbf{y})}\sum_{k}[(\delta_{1}-\delta_{0})\gamma_{ij}+\delta_{0}]w_{ijk}^{2}\\\
&+\frac{K^{2}}{2}\log[\delta_{0}+(\delta_{1}-\delta_{0})\gamma_{ij}]+\log\pi(\gamma_{ij})+C\end{aligned}$
M-Step: Set $\gamma_{ij}=0$, if
$Q\left(\gamma_{ij}=0\mid\gamma_{ij}^{(t-1)}\right)\geq
Q\left(\gamma_{ij}=1\mid\gamma_{ij}^{(t-1)}\right)$,that is,

$\displaystyle\frac{E_{\pi(\mathbf{w}\mid\gamma_{ij}^{(t-1)},\gamma_{-(ij)},\mathbf{y})}[\sum_{k}w_{ijk}^{2}]}{K^{2}}\leq\frac{1}{\delta_{0}-\delta_{1}}\left[\log\frac{\delta_{0}}{\delta_{1}}+\frac{2}{K^{2}}\log\frac{p_{ij}}{1-p_{ij}}\right]=\lambda_{1}$

and 1 otherwise. Since the term
$\log\left(\frac{p_{ij}}{1-p_{ij}}\right)\in\mathbf{R}$ and $0<p_{ij}<1$ is
the hyper-parameter,instead of turning $p_{ij}$, we can turn the threshold
$\lambda_{1}$ on the right hand side directly.

For neural networks,
$\pi(\mathbf{w}|\boldsymbol{y},\boldsymbol{\gamma})\propto\pi(\boldsymbol{y}|\mathbf{w})\pi(\mathbf{w}|\boldsymbol{\gamma})$
has no closed form, an approximation can be obtained by using the samples
produced by our constrained variational Adam algorithms:
$E_{\pi(\mathbf{w}|\gamma_{ij}^{(t-1)},\gamma_{-(ij)},y)}[\sum_{k}w_{ijk}^{2}]\approx\sum_{k}w_{ijk}^{2}$.

Our method can be viewed as a hybird MCMC or variational with EM algorithm,
which will adaptively change the weight decay factor for each group based on
the magnitude. A large weight decay factor will be assigned to the group of
weights with small magnitude via $\frac{\delta_{0}}{N}$. This will further
push them toward zeros. On the contrary, the group with strong signal will
assign a small weight decay factor, via $\frac{\delta_{1}}{N}$.

###  3.3 Dynamic forward pruning (DFP)

Since our EM-MCMC algorithm produces no exact zeros, a hard threshold is
required to prune the weight and maintain sparsity. Recently there are large
number of dynamic pruning methods that maintain sparse weights through a
prune-regrowth cycle during training (Zhu and Gupta,, 2017; Bellec et al.,,
2017; Dettmers and Zettlemoyer,, 2019; Mostafa and Wang,, 2019; Lin et al.,,
2020; Kusupati et al.,, 2020). Some of these methods can be inserted into our
framework directly. For example, we can incorporate dynamic pruning with
feedback (DPF) (Lin et al.,, 2020), which evaluates a stochastic gradient for
the pruned model
$\widetilde{\mathbf{w}}_{t}=\mathbf{m}_{t}\odot\mathbf{w}_{t}$ (where
$\mathbf{m}_{t}$ is a mask matrix) and apply it to the (simultaneously
maintained) dense model

$\mathbf{w}_{t+1}:=\mathbf{w}_{t}-\gamma_{t}\mathbf{g}\left(\mathbf{m}_{t}\odot\mathbf{w}_{t}\right)=\mathbf{w}_{t}-\gamma_{t}\mathbf{g}\left(\widetilde{\mathbf{w}}_{t}\right)$
(8)

and $\mathbf{g}\left(\mathbf{m}_{t}\odot\mathbf{w}_{t}\right)$ is the gradient
evaluated at the sparse model. Although these methods are one-shot, some
computational budget for weight regrow is still required.

For linear models, it has been argued that forward or backward selection
strategy may be trapped into a bad local mode (Hastie et al.,, 2009), thus a
variable selection algorithm allowing parameter regrow is preferred.

However, Gur-Ari et al., (2018) argued that only a tiny subspace of the
parameters have non-zero gradients. Empirically, when we run algorithm 2, we
find that the gradients will vanish everywhere except in a small number of
input and output channels. In addition, during optimization, many of
parameters in the input and out channels will concentrate to zero. These
concentration is quite stable even under the two following perturbations from
the algorithm:
1\. A cyclical learning rate, which can jump from small to large. See Figure
1-(f).
2\. At each iteration, adding an injected noise
$\sqrt{2+2\beta_{1}^{2}}k(\Delta t)^{1.5}\boldsymbol{\epsilon}$.

(a) Layer 1 Input channel 3

(b) Layer 6 Input channel 20

(c) Layer 6 Output channel 14

(d) Layer 12 Output channel 50

(e) Layer 16 Output channel 100

(f) Cyclical learning rate

Figure 1: Sample paths of all weights from select channels of ResNet18 in
CIFAR10 during training with cyclical learning rate $l_{t}$ and
$k=\frac{1}{N}$ from algorithm 2. (a)-(e) are grouped by either input and
output channels, (f) is the cyclical learning rate used.

We observe the concentration phenomenon by visualizing the dynamic
distributions of the weights in Figure 1, which shows the sample paths of all
weights from select channels during optimization. Note that, at this stage, we
just ran our algorithm 2 with cyclical learning rate. We can see that the
contractions occur regardless of whether the grouping is based on input or
output channels. In addition, we rarely observe regrowth of the weight once
the whole channel concentrates to zero, see Figure 2. This phenomenon is very
similar to the posterior contraction in high dimensional sparse Bayesian
linear regression (Castillo et al.,, 2015), where only a sub-model has
substantial probability mass. Finally, Figure 1$(e)$ shows the behavior of
weights in those channels which are not concentrated around zero.

Motivated by the above observations, we propose a structured pruning rule
based on a simple concentration metric, where, within each layer, two for
loops will be used to scan through all the input channels and output channels.

If
$\displaystyle\,\,\max_{j,k}|W_{ijk}|-\min_{j,k}|W_{ijk}|<\lambda_{2}\,\,\text{set}\,\,W_{i}=0;$
(9) If
$\displaystyle\,\,\max_{i,k}|W_{ijk}|-\min_{i,k}|W_{ijk}|<\lambda_{2}\,\,\text{set}\,\,W_{j}=0;$

where $\lambda_{2}>0$ is a hard threshold which we need to tune. Algorithm 3
provides the pseudo code. The SoftMask returned by the EM algorithm will
control the shrinkage of the weight. It is used to separate the noise and
signal as in sparse linear regression (Castillo et al.,, 2015). The group of
weights with very small magnitude will be identified as noise and a strong
shrinkage (a large weight decay) will be applied to it in the next iteration
and vice versa. The HardMask returned by the pruning criterion rule will
remove the weight. Although we recommend using DFP, where regrowth is not
allowed, we also provide the user with DPF as another option in Algorithm 3.
Figure 2 shows no significant differences between the DFP and DPF, when
comparing their test accuracy and sparsity ratios.
Remark: The pruning criterion from equation (9) assumes that if the weights
within the input channels or output channels concentrate together, it will
concentrate to zero. However, the strong concentration phenomenon we observed
in CNN as shown in Figure 1 are observed in feedforward neural networks, we
find that even when the problem is known to be sparse, this criterion can lead
to a lack of convergence. Instead, we propose the following $L_{2}$ magnitude
pruning criterion

$\text{If}\,\,\sum_{i=1}^{k_{o}}w_{ij}^{2}/k_{o}\leq\lambda_{1},\text{set}\,\,w_{j}=0$
(10)

where $i$ is output units, $j$ is input units and $k_{o}$ is the number of
output units.

Figure 2: Resnet18 pruned using algorithm 3 with the cyclical learning rate
under DFP (blue) and DPF (orange), under the same hyper-parameter setting. The
left panel shows testing accuracy. The middle panel shows the percentage of
the weights, which use large weight decay factor as suggested by EM algorithm
(soft sparsity ratio). The right panel shows the sparsity ratio of the model
during training. Algorithm 3 EM-MCMC

Input $\beta=(\beta_{1},\beta_{2})$, $(\delta_{0},\delta_{1})\text{(weight
decay)}$

$N\text{(training size)}$, $B\text{(batch size)}$

$\lambda_{1}\text{(soft threshold)}$, $\lambda_{2}\text{(hard threshold)}$

$\text{pruning with feedback}=\text{False}$

Initialization
${\boldsymbol{\rho},\boldsymbol{w},\boldsymbol{v},\text{SoftMask}\leftarrow
1,\text{HardMask}\leftarrow 1}$

for $t=1\>\cdots\>$ do

$\delta\leftarrow\delta_{1}$

$\delta[\text{SoftMask}==0]\leftarrow\delta_{0}$

$\Delta t\leftarrow s_{1}(t)$; $\eta\leftarrow s_{2}(t)$; $k\leftarrow T(t)$;

$\alpha\leftarrow k*\Delta t^{1.5}$

$\boldsymbol{g}\leftarrow-\frac{1}{B}\sum_{i\in\mathcal{B}}\nabla_{\mathbf{w}}\log(\mathcal{D}_{i}|\mathbf{w})+\frac{\delta}{N}\mathbf{w}$

$\boldsymbol{\rho}\leftarrow\boldsymbol{\rho}+\left(\frac{\eta}{\boldsymbol{\tau}}-\eta\alpha^{2}\delta\boldsymbol{\tau}\right)-\eta\alpha\boldsymbol{\epsilon}\odot\boldsymbol{g}$

$\boldsymbol{\tau}\leftarrow 1/(1+e^{-\boldsymbol{\rho}})$

$\boldsymbol{\epsilon}\leftarrow N(0,\mathbf{1}_{p})$

${\boldsymbol{v}\leftarrow\beta_{1}\boldsymbol{v}-\beta_{2}g(\Delta
t)^{2}+\sqrt{2+2\beta_{1}^{2}}k(\Delta t)^{1.5}\epsilon}$

$\mathbf{w}\leftarrow\mathbf{w}-\tau\odot\boldsymbol{v}$

if $t>\text{warm up}$ then

$\text{SoftMask}\leftarrow\textbf{EM}(\mathbf{w},\lambda_{1})$

$\text{HardMask}\leftarrow\textbf{Pruning Criterion}(\mathbf{w},\lambda_{2})$

$\mathbf{w}[\text{HardMask}]\leftarrow 0$

if pruning with feedback is False then

$\text{Freeze}\,\,\mathbf{w}[\text{HardMask}]$

end if

end if

end for

##  4 Experiments

###  4.1 Experimental setup for image classfication task

In this section, we conduct experiments to show the effectiveness of our
proposed algorithms in terms of the test accuracy and sparsity ratio by
comparing against both dense and sparse methods in the literature for image
classfication. Apart from the baseline, which is SGD with momentum, we include
two SOAT Bayesian methods: Rank-1 BNN (Dusenberry et al.,, 2020), a
variational inference approach with batch ensemble (Wen et al.,, 2020); and an
MCMC approach cSGLD (Zhang et al.,, 2019). We also include three popular
dynamic sparse training methods: SM (Dettmers and Zettlemoyer,, 2019);
DSR(Mostafa and Wang,, 2019) and DPF (Lin et al.,, 2020). We repeat each
experiment three times and average them to obtain stable results.

The Rank-1 BNN has some advantage in this experimental comparison as their
training approach is more computationally intensive. Since we don’t have
enough TPU to implement their method, we extracted their experimental result
from the original paper as indicated by * in the table.

For the other methods, we use 200 epochs to train models in CIFAR10 and
CIFAR100 datasets with single RTX3090 node and 200 epochs to train model in
ImageNet dataset with single A100 node. The cyclical learning rate with 4
cycles as shown in Figure 1-(f) is used for cSGLD (Zhang et al.,, 2019) and
CV-Adam. We collect three samples at the end of each cycle, which gives us 12
samples in total for CIFAR10 for CIFAR100 datasets, while for ImageNet, the
first 50 epochs is used as warm up, therefore, we only collect 9 samples. The
posterior predictive distribution is calculated by averaging all the samples
as an ensemble.

Since DPF equation (8) has no official code release, but easy to insert into
Algorithm 3, we implement it by ourselves. For the baseline model, we follow
the hyper-parameter settings from Xie et al., (2022). Their setting allows us
to train a strong baseline for comparison. Details of hyper-parameter settings
is provided in Appendix D.

###  4.2 Results for image classfication task

Table 1: Comparison of test accuracy with different target sparsity ratios in CIFAR10. | Dense Method |  | Sparse Method
---|---|---|---
Model | Baseline | Rank-1 BNN | cSGLD | CV-Adam |  | SM | DSR | DFP | DPF | Sparsity ratio
ResNet18 | 95% | - | 95.7% | 95.5% |  | 92.8% | 93.1% | 94.7% | 94.8% | 70%
ResNet18 | 95% | - | 95.7% | 95.5% |  | 91.1% | 91.2% | 94.5% | 94.5% | 80%
ResNet18 | 95% | - | 95.7% | 95.5% |  | 89.7% | 90.0% | 93.9% | 93.8% | 90%
WRN-28-10 | 96.3% | 96.5%* | 96.2% | 96.5% |  | 95.4% | 95.8% | 95.5% | 95.8% | 90%
WRN-28-10 | 96.3% | 96.5%* | 96.2% | 96.5% |  | 95.3% | 95.4% | 95.3% | 95.4% | 95%
WRN-28-10 | 96.3% | 96.5%* | 96.2% | 96.5% |  | 94.5% | 94.5% | 94.7% | 94.5% | 99%
Table 2: Comparison of test accuracy for different target sparsity ratios for in CIFAR100 | Dense method |  | Sparse method
---|---|---|---
Model | Baseline | Rank-1 BNN | cSGLM | CV-Adam |  | SM | DSR | DFP | DPF | Sparsity ratio
ResNet34 | 78.5% | - | 79.7% | 79.4% |  | 76.4% | 76.6% | 77.6% | 77.4% | 50%
ResNet34 | 78.5% | - | 79.7% | 79.4% |  | 73.2% | 73.8% | 75.4% | 75.1% | 70%
ResNet34 | 78.5% | - | 79.7% | 79.4% |  | 72.1% | 72.3% | 75.1% | 75.3% | 90%
WRN-28-10 | 81.7% | 82.4%* | 82.7% | 83.4% |  | 78.0% | 78.1% | 78.5% | 78.2% | 70%
WRN-28-10 | 81.7% | 82.4%* | 82.7% | 83.4% |  | 77.0% | 76.5% | 77.7% | 77.9% | 80%
WRN-28-10 | 81.7% | 82.4%* | 82.7% | 83.4% |  | 76.5% | 76.4% | 77.5% | 77.4% | 90%
Table 3: Comparison of test accuracy for different target sparsity ratios for ImageNet | Dense method |  | Sparse method
---|---|---|---
Model | Baseline | Rank-1 BNN | cSGLM | CV-Adam |  | SM | DSR | DFP | DPF | Sparsity ratio
ResRet50 | 76.5% | 77.3%* | 76.7% | 76.8% |  | 73.8% | 73.3% | 74.5% | 74.6% | 80%
ResNet50 | 76.5% | 77.3%* | 76.7% | 76.8% |  | 72.3% | 72.0% | 73.5% | 73.2% | 90%
Table 4: Comparison of sparse methods for three simulated examples | Example 1 |  | Example 2 |  | Example 3
---|---|---|---|---|---
Method | MSE | FDR | FNDR | $\hat{S}$ |  | MSE | FDR | FNDR | $\hat{S}$ |  | Accuracy | FDR | FNDR | $\hat{S}$
SPLBNN | 1.4 | 0 | 0 | 4 |  | 2.43 | 0 | 0 | 5 |  | 91.2% | 0 | 0 | 5
SVBNN | 1.6 | 38.3% | 0 | 6.3 |  | 3.12 | 51% | 0 | 10.2 |  | 86% | 39.9% | 0 | 8.2
DFP | 1.32 | 0 | 0 | 4 |  | 2.46 | 0 | 0 | 5 |  | 94.70% | 0 | 0 | 5
DPF | 1.29 | 0 | 0 | 4 |  | 2.49 | 0 | 0 | 5 |  | 94.50% | 0 | 0 | 5

Performance for CIFAR10 and CIFFAR100 dataset: We report the testing accuracy
and sparsity ratio for sparse method in Table 1 and Table 2. For the dense
method, comparing with two SOTA Bayesian methods and baseline, the CV-Adam of
Algorithm 1 is very competitive in both CIFAR10 and CIFAR100 datasets.

When pruning is applied, there is almost no performance loss with both DFP and
DPF approaches compared with the dense methods in ResNet18, and a much higher
sparsity ratio can be achieved by WRN-28-10 in the CIFAR10 dataset. For
CIFAR100 dataset, we found that it is harder to target the high sparsity ratio
while at the same time keep the performance loss negligible. Overall, our
pruning algorithms can achieve high sparsity ratio while sacrificing a small
amount in accuracy.

Finally, there is no evidence that using a dynamic forward pruning with (no
regrowth suffers performance loss compared to DPF.

Performance for ImageNet: As shown in Table 3, CV-Adam produced a modest
performance gain compared with the baseline, but performed worse than Rank-1
BNN. For cSGLM(Zhang et al.,, 2019), we failed to reproduce the $77\%$
predictive accuracy as claim in the original paper. We believe the performance
for both our CV-Adam and cSGLM can be further improved by tuning the hyper-
parameters carefully. For sparse method, our pruning algorithms out performed
the other two methods SM and DSR. Again, there is still no evidence that using
DFP will incur performance loss.

###  4.3 Simulated Examples: variable selection for nonlinear regression

Sun et al., (2021) applied the Spike-and-Slab Gaussian prior to prune the
neural network (SPLBNN). They use a Laplace approximation-based approach to
approximate the marginal posterior inclusion probability, which requires the
user to train the dense model to find the local mode first followed by pruning
and retraining the sparse model, finally they use the Bayesian evidence to
elicit sparse DNNs in multiple runs with different initialization. Another
related approach is variational BNN with Spike-and-Slab prior (Bai et al.,,
2020) (SVBNN). Both approaches are not scalable to large models, but they work
well in high dimensional sparse nonlinear regression, so we carry out a
comparison with these two approaches in this setting. Here, we follow the same
data generate process and examples as described in Sun et al., (2021):

  * •

Simulate $e,z_{1},...,z_{p}$ independently from the truncated standard normal
distribution on the interval $[-10,10]$.

  * •

Set $x_{i}=\frac{e+z_{i}}{\sqrt{2}}$ for $i=1,...,p$

Then, all the predictor $x_{i}$ fall into a compact set and are mutually
correlated with a correlation coefficient of about 0.5. Based on this setting,
we generate three toy examples from Sun et al., (2021) where each example
consists of 10000 training samples and 1000 testing samples. To fit the data
and make comparison, we follow the network structure described in Sun et al.,
(2021).
Example 1:

$\displaystyle y=$
$\displaystyle\tanh\left(2\tanh\left(2x_{1}-x_{2}\right)\right)+2\tanh\left(\tanh\left(x_{3}-2x_{4}\right)\right.$
$\displaystyle\left.-\tanh\left(2x_{5}\right)\right)++0x_{6}+\cdots+0x_{1000}+\varepsilon,$

Network Structure: 1000-5-3-1 with ReLu activation

Example 2:

$y=\frac{5x_{2}}{1+x_{1}^{2}}+5\sin\left(x_{3}x_{4}\right)+2x_{5}+0x_{6}+\cdots+0x_{2000}+\varepsilon$

Network Structure: 2000-6-4-3-1 with ReLu activation

Example 3:

$y=\begin{cases}1&e^{x_{1}}+x_{2}^{2}+5\sin\left(x_{3}x_{4}\right)+0x_{5}+\cdots+0x_{1000}>3\\\
0&\text{ otherwise. }\end{cases}$

Network Structure: 2000-6-4-3-1 with ReLu activation

where $\epsilon\sim N(0,1)$. We only compare our algorithm with the method
from Sun et al., (2021) and (Bai et al.,, 2020). The other two sparse methods
we used in the image task didn’t work in these three examples. We use equation
(10) as the pruning criterion in Algorithm 3. The performance of variable
selection is measured by the false discover rate (FDR) and false non-discover
rate (FNDR). The predictive mean square error is used for examples 1 and 2 and
the prediction accuracy is used for example 3. $\hat{S}$ is the number of the
selected variables returned by the competing methods. We repeat the
experiments 10 times and report the averaged results in Table 4. Details of
hyper-parameter setting is provided in the Appendix D.

Performance analysis: In terms of variable selection, SPLBNN, DFP and DPF
perfectly selected all the relevant variables and removed all irrelevant
variables in all three examples. SVBNN always selected more variables but was
able to keep the FNDR at zero. The predictive performance for examples 1 and 3
are good for SPLBNN, DFP and DPF. But for example 2, the mean square error are
high (compared to the oracle value of 1), A better network structure may
improve this performance. One of the surprising results is that DFP can still
identify all relevant predictors and was never trapped in a bad local mode.
Comparing with the SPLBNN from Sun et al., (2021), which requires the train-
prune-fine tune process multiple times, our algorithm is very aggressive and
efficient and it outperforms the SPLBNN in example 1 and 3.

##  5 Discussion of future work

In this paper, we showed the connection between our variational BNN algorithm
with stochastic gradient Hamiltonian Monte Carlo. The cyclical learning rate
has been used with the hope that we can jump to different modes by suddenly
changing the learning rate from small to large. Using the posterior predictive
distributions in an ensemble approach can improve the predictive accuracy. One
question arises after observing the concentration phenomenon from CNN is do we
really need to train the full model after each cycle? Our experimental
findings suggest that only a small number of the subnetworks needs to be
trained after one cycle. Maybe the mode jump is due to the change of a small
proportion of the weights from subnetworks. If this is true, it implies that
training can be accelerated by developing the dynamic forward freeze algorithm
to gradually freeze the parameters at the end of each cycle. In addition, we
can use the subnetwork ensemble approach where only small proportion of the
weights are different among different models and majority of the weights are
shared cross all the models. This approach can save storage requirements and
speed up computation.

## References

  * Bai et al., (2020) Bai, J., Song, Q., and Cheng, G. (2020).  Efficient variational inference for sparse deep learning with theoretical guarantee.  Advances in Neural Information Processing Systems, 33:466–476.
  * Bellec et al., (2017) Bellec, G., Kappel, D., Maass, W., and Legenstein, R. (2017).  Deep rewiring: Training very sparse deep networks.  arXiv preprint arXiv:1711.05136.
  * Bishop and Nasrabadi, (2006) Bishop, C. M. and Nasrabadi, N. M. (2006).  Pattern recognition and machine learning, volume 4.  Springer.
  * Blundell et al., (2015) Blundell, C., Cornebise, J., Kavukcuoglu, K., and Wierstra, D. (2015).  Weight uncertainty in neural network.  In International conference on machine learning, pages 1613–1622. PMLR.
  * Castillo et al., (2015) Castillo, I., Schmidt-Hieber, J., and Van der Vaart, A. (2015).  Bayesian linear regression with sparse priors.  The Annals of Statistics, 43(5):1986–2018.
  * Chen et al., (2014) Chen, T., Fox, E., and Guestrin, C. (2014).  Stochastic gradient hamiltonian monte carlo.  In International conference on machine learning, pages 1683–1691. PMLR.
  * Daxberger et al., (2021) Daxberger, E., Nalisnick, E., Allingham, J. U., Antorán, J., and Hernández-Lobato, J. M. (2021).  Bayesian deep learning via subnetwork inference.  In International Conference on Machine Learning, pages 2510–2521. PMLR.
  * Dettmers and Zettlemoyer, (2019) Dettmers, T. and Zettlemoyer, L. (2019).  Sparse networks from scratch: Faster training without losing performance.  arXiv preprint arXiv:1907.04840.
  * Dusenberry et al., (2020) Dusenberry, M., Jerfel, G., Wen, Y., Ma, Y., Snoek, J., Heller, K., Lakshminarayanan, B., and Tran, D. (2020).  Efficient and scalable bayesian neural nets with rank-1 factors.  In International conference on machine learning, pages 2782–2792. PMLR.
  * Gal and Ghahramani, (2016) Gal, Y. and Ghahramani, Z. (2016).  Dropout as a bayesian approximation: Representing model uncertainty in deep learning.  In international conference on machine learning, pages 1050–1059. PMLR.
  * Graves, (2011) Graves, A. (2011).  Practical variational inference for neural networks.  Advances in neural information processing systems, 24.
  * Gur-Ari et al., (2018) Gur-Ari, G., Roberts, D. A., and Dyer, E. (2018).  Gradient descent happens in a tiny subspace.  arXiv preprint arXiv:1812.04754.
  * Hastie et al., (2009) Hastie, T., Tibshirani, R., Friedman, J. H., and Friedman, J. H. (2009).  The elements of statistical learning: data mining, inference, and prediction, volume 2.  Springer.
  * Heek and Kalchbrenner, (2019) Heek, J. and Kalchbrenner, N. (2019).  Bayesian inference for large scale image classification.  arXiv preprint arXiv:1908.03491.
  * Hernández-Lobato and Adams, (2015) Hernández-Lobato, J. M. and Adams, R. (2015).  Probabilistic backpropagation for scalable learning of Bayesian neural networks.  In International conference on machine learning, pages 1861–1869. PMLR.
  * Izmailov et al., (2018) Izmailov, P., Podoprikhin, D., Garipov, T., Vetrov, D., and Wilson, A. G. (2018).  Averaging weights leads to wider optima and better generalization.  arXiv preprint arXiv:1803.05407.
  * Khan et al., (2018) Khan, M., Nielsen, D., Tangkaratt, V., Lin, W., Gal, Y., and Srivastava, A. (2018).  Fast and scalable bayesian deep learning by weight-perturbation in adam.  In International Conference on Machine Learning, pages 2611–2620. PMLR.
  * Kingma and Ba, (2014) Kingma, D. P. and Ba, J. (2014).  Adam: A method for stochastic optimization.  arXiv preprint arXiv:1412.6980.
  * Kusupati et al., (2020) Kusupati, A., Ramanujan, V., Somani, R., Wortsman, M., Jain, P., Kakade, S., and Farhadi, A. (2020).  Soft threshold weight reparameterization for learnable sparsity.  In International Conference on Machine Learning, pages 5544–5555. PMLR.
  * Lakshminarayanan et al., (2017) Lakshminarayanan, B., Pritzel, A., and Blundell, C. (2017).  Simple and scalable predictive uncertainty estimation using deep ensembles.  Advances in neural information processing systems, 30.
  * Lin et al., (2020) Lin, T., Stich, S. U., Barba, L., Dmitriev, D., and Jaggi, M. (2020).  Dynamic model pruning with feedback.  arXiv preprint arXiv:2006.07253.
  * Liu et al., (2019) Liu, L., Jiang, H., He, P., Chen, W., Liu, X., Gao, J., and Han, J. (2019).  On the variance of the adaptive learning rate and beyond.  arXiv preprint arXiv:1908.03265.
  * Maddox et al., (2019) Maddox, W. J., Izmailov, P., Garipov, T., Vetrov, D. P., and Wilson, A. G. (2019).  A simple baseline for bayesian uncertainty in deep learning.  Advances in Neural Information Processing Systems, 32.
  * Mostafa and Wang, (2019) Mostafa, H. and Wang, X. (2019).  Parameter efficient training of deep convolutional neural networks by dynamic sparse reparameterization.  In International Conference on Machine Learning, pages 4646–4655. PMLR.
  * Osawa et al., (2019) Osawa, K., Swaroop, S., Khan, M. E. E., Jain, A., Eschenhagen, R., Turner, R. E., and Yokota, R. (2019).  Practical deep learning with bayesian principles.  Advances in neural information processing systems, 32.
  * Ritter et al., (2018) Ritter, H., Botev, A., and Barber, D. (2018).  A scalable laplace approximation for neural networks.  In 6th International Conference on Learning Representations, ICLR 2018-Conference Track Proceedings, volume 6. International Conference on Representation Learning.
  * Ročková, (2018) Ročková, V. (2018).  Particle em for variable selection.  Journal of the American Statistical Association, 113(524):1684–1697.
  * Ročková and George, (2014) Ročková, V. and George, E. I. (2014).  EMVS: The EM approach to bayesian variable selection.  Journal of the American Statistical Association, 109(506):828–846.
  * Sun et al., (2021) Sun, Y., Song, Q., and Liang, F. (2021).  Consistent sparse deep learning: Theory and computation.  Journal of the American Statistical Association, pages 1–15.
  * Trinh et al., (2022) Trinh, T. Q., Heinonen, M., Acerbi, L., and Kaski, S. (2022).  Tackling covariate shift with node-based bayesian neural networks.  In International Conference on Machine Learning, pages 21751–21775. PMLR.
  * Wang et al., (2016) Wang, J., Liang, F., and Ji, Y. (2016).  An ensemble EM algorithm for bayesian variable selection.  arXiv preprint arXiv:1603.04360.
  * Welling and Teh, (2011) Welling, M. and Teh, Y. W. (2011).  Bayesian learning via stochastic gradient langevin dynamics.  In Proceedings of the 28th international conference on machine learning (ICML-11), pages 681–688. Citeseer.
  * Wen et al., (2020) Wen, Y., Tran, D., and Ba, J. (2020).  Batchensemble: an alternative approach to efficient ensemble and lifelong learning.  arXiv preprint arXiv:2002.06715.
  * Wenzel et al., (2020) Wenzel, F., Roth, K., Veeling, B. S., Świątkowski, J., Tran, L., Mandt, S., Snoek, J., Salimans, T., Jenatton, R., and Nowozin, S. (2020).  How good is the bayes posterior in deep neural networks really?  arXiv preprint arXiv:2002.02405.
  * Wilson and Izmailov, (2020) Wilson, A. G. and Izmailov, P. (2020).  Bayesian deep learning and a probabilistic perspective of generalization.  Advances in neural information processing systems, 33:4697–4708.
  * Xie et al., (2022) Xie, Z., Wang, X., Zhang, H., Sato, I., and Sugiyama, M. (2022).  Adaptive inertia: Disentangling the effects of adaptive learning rate and momentum.  In International Conference on Machine Learning, pages 24430–24459. PMLR.
  * Xu and Ghosh, (2015) Xu, X. and Ghosh, M. (2015).  Bayesian variable selection and estimation for group lasso.  Bayesian Analysis, 10(4):909–936.
  * Zhang et al., (2018) Zhang, G., Sun, S., Duvenaud, D., and Grosse, R. (2018).  Noisy natural gradient as variational inference.  In International Conference on Machine Learning, pages 5852–5861. PMLR.
  * Zhang et al., (2019) Zhang, R., Li, C., Zhang, J., Chen, C., and Wilson, A. G. (2019).  Cyclical stochastic gradient mcmc for bayesian deep learning.  arXiv preprint arXiv:1902.03932.
  * Zhu and Gupta, (2017) Zhu, M. and Gupta, S. (2017).  To prune, or not to prune: exploring the efficacy of pruning for model compression.  arXiv preprint arXiv:1710.01878.