Efficient Variational Inference for Sparse Deep Learning with Theoretical Guarantee

Dense neural network (DNN) may face various problems despite its huge successes in AI fields. Larger training sets and more complicated network structures improve accuracy in deep learning, but always incur huge storage and computation burdens. For example, small portable devices may have limited resources such as several megabyte memory, while a dense neural networks like ResNet-50 with 50 convolutional layers would need more than 95 megabytes of memory for storage and numerous floating number computation (Cheng et al.,, 2018). It is therefore necessary to compress deep learning models before deploying them on these hardware limited devices.

In addition, sparse neural networks may recover the potential sparsity structure of the target function, e.g., sparse teacher network in the teacher-student framework (Goldt et al.,, 2019; Tian,, 2018). Another example is from nonparametric regression with sparse target functions, i.e., only a portion of input variables are relevant to the response variable. A sparse network may serve the goal of variable selection (Feng and Simon,, 2017; Liang et al.,, 2018; Ye and Sun,, 2018), and is also known to be robust to adversarial samples against $l_{\infty}$ and $l_{2}$ attacks (Guo et al.,, 2018).

Bayesian neural network (BNN), which dates back to MacKay, (1992); Neal, (1992), comparing with frequentist DNN, possesses the advantages of robust prediction via model averaging and automatic uncertainty quantification (Blundell et al.,, 2015). Conceptually, BNN can easily induce sparse network selection by assigning discrete prior over all possible network structures. However, challenges remain for sparse BNN including inefficient Bayesian computing issue and the lack of theoretical justification. This work aims to resolve these two important bottlenecks simultaneously, by utilizing variational inference approach (Jordan et al.,, 1999; Blei et al.,, 2017). On the computational side, it can reduce the ultra-high dimensional sampling problem of Bayesian computing, to an optimization task which can still be solved by a back-propagation algorithm. On the theoretical side, we provide a proper prior specification, under which the variational posterior distribution converges towards the truth. To the best of our knowledge, this work is the first one that provides a complete package of both theory and computation for sparse Bayesian DNN.

We aim to approximate $f_{0}$ in the generative model (1) by a sparse neural network. Specifically, given a network structure, i.e. the depth $L$ and the width $\boldsymbol{p}$, $f_{0}$ is approximated by DNN models $f_{\theta}$ with sparse parameter vector $\theta\in\Theta=\mathbb{R}^{T}$. From a Bayesian perspective, we impose a spike-and-slab prior (George and McCulloch,, 1993; Ishwaran and Rao,, 2005) on $\theta$ to model sparse DNN.

A spike-and-slab distribution is a mixture of two components: a Dirac spike concentrated at zero and a flat slab distribution. Denote $\delta_{0}$ as the Dirac at 0 and $\gamma=(\gamma_{1},\ldots,\gamma_{T})$ as a binary vector indicating the inclusion of each edge in the network. The prior distribution $\pi(\theta)$ thus follows:

for $i=1,\ldots,T$, where $\lambda$ and $\sigma^{2}_{0}$ are hyperparameters representing the prior inclusion probability and the prior Gaussian variance, respectively. The choice of $\sigma_{0}^{2}$ and $\lambda$ play an important role in sparse Bayesian learning, and in Section 4, we will establish theoretical guarantees for the variational inference procedure under proper deterministic choices of $\sigma_{0}^{2}$ and $\lambda$. Alternatively, hyperparameters may be chosen via an Empirical Bayesian (EB) procedure, but it is beyond the scope of this work. We assume $\mathcal{Q}$ is in the same family of spike-and-slab laws:

for $i=1,\ldots,T$, where $0\leq\phi_{i}\leq 1$.

Comparing to pruning approaches (e.g. Zhu and Gupta,, 2018; Frankle and Carbin,, 2018; Molchanov et al.,, 2017) that don’t pursue sparsity among bias parameter $b_{i}$’s, the Bayesian modeling induces posterior sparsity for both weight and bias parameters.

In the literature, Polson and Rockova, (2018); Chérief-Abdellatif, (2020) imposed sparsity specification as follows $\Theta(L,\boldsymbol{p},s)=\{\theta\mbox{ as in model }(\ref{eqmod2}):||\theta||_{0}\leq s\}$ that not only posts great computational challenges, but also requires tuning for optimal sparsity level $s$. For example, Chérief-Abdellatif, (2020) shows that given $s$, two error terms occur in the variation DNN inference: 1) the variational error $r_{n}(L,\boldsymbol{p},s)$ caused by the variational Bayes approximation to the true posterior distribution and 2) the approximation error $\xi_{n}(L,\boldsymbol{p},s)$ between $f_{0}$ and the best bounded-weight $s$-sparsity DNN approximation of $f_{0}$. Both error terms $r_{n}$ and $\xi_{n}$ depend on $s$ (and their specific forms are given in next section). Generally speaking, as the model capacity (i.e., $s$) increases, $r_{n}$ will increase and $\xi_{n}$ will decrease. Hence the optimal choice $s^{*}$ that strikes the balance between these two is

Therefore, one needs to develop a selection criteria for $\widehat{s}$ such that $\widehat{s}\approx s^{*}$. In contrast, our modeling directly works on the whole sparsity regime without pre-specifying $s$, and is shown later to be capable of automatically attaining the same rate of convergence as if the optimal $s^{*}$ were known.

In this section, we will establish the contraction rate of the variational sparse DNN procedure, without knowing $s^{*}$. For simplicity, we only consider equal-width neural network (similar as Polson and Rockova, (2018)).

The following assumptions are imposed:

Condition 4.2 is very mild, and includes ReLU, sigmoid and tanh. Note that Condition 4.3 gives a wide range choice of $\lambda$, even including the choice of $\lambda$ independent of $s^{*}$ (See Theorem 4.1 below).

We first state a lemma on an upper bound for the negative ELBO. Denote the log-likelihood ratio between $p_{0}$ and $p_{\theta}$ as $l_{n}(P_{0},P_{\theta})=\log(p_{0}(D)/p_{\theta}(D))=\sum^{n}_{i=1}\log(p_{0}(D_{i})/p_{\theta}(D_{i}))$. Given some constant $B>0$, we define

Recall that $r_{n}(L,N,s)$ and $\xi_{n}(L,N,s)$ denote the variational error and the approximation error.

Noting that $\mbox{KL}(q(\theta)||\pi(\theta|\lambda))+\int_{\Theta}l_{n}(P_{0},P_{\theta})q(\theta)(d\theta)$ is the negative ELBO up to a constant, we therefore show the optimal loss function of the proposed variational inference is bounded.

The next lemma investigates the convergence of the variational distribution under the Hellinger distance, which is defined as

In addition, let $s_{n}=s^{*}\log^{2\delta-1}(n)$ for any $\delta>1$. An assumption on $s^{*}$ is required to strike the balance between $r_{n}^{*}$ and $\xi^{*}$:

The above two lemmas together imply the following guarantee for VB posterior:

The $\varepsilon_{n}^{*2}$ denotes the estimation error from the statistical estimator for $P_{0}$. The variational Bayes convergence rate consists of estimating error, i.e., $\varepsilon_{n}^{*2}$, variational error, i.e., $r_{n}^{*}$, and approximation error, i.e., $\xi_{n}^{*}$. Given that the former two errors have only logarithmic difference, our convergence rate actually strikes the balance among all three error terms. The derived convergence rate has an explicit expression in terms of the network structure based on the forms of $\varepsilon_{n}^{*}$, $r_{n}^{*}$ and $\xi_{n}^{*}$, in contrast with general convergence results in Pati et al., (2018); Zhang and Gao, (2019); Yang et al., (2020).

To conduct optimization of (4) via stochastic gradient optimization, we need to find certain reparameterization for any distribution in $\mathcal{Q}$. One solution is to use the inverse CDF sampling technique. Specifically, if $\theta\sim q\in\mathcal{Q}$, its marginal $\theta_{i}$’s are independent mixture of (7). Let $F_{(\mu_{i},\sigma_{i},\phi_{i})}$ be the CDF of $\theta_{i}$, then $\theta_{i}\stackrel{{\scriptstyle d}}{{=}}F^{-1}_{(\mu_{i},\sigma_{i},\phi_{i})}(u_{i})$ holds where $u_{i}\sim\mathcal{U}(0,1)$. This inverse CDF reparameterization, although valid, can not be conveniently implemented within the state-of-art python packages like PyTorch. Rather, a more popular way in VB is to utilize the Gumbel-softmax approximation.

We rewrite the loss function $\Omega$ as

Since it is impossible to reparameterize the discrete variable $\gamma$ by a continuous system, we apply continuous relaxation, i.e., to approximate $\gamma$ by a continuous distribution. In particular, the Gumbel-softmax approximation (Maddison et al.,, 2017; Jang et al.,, 2017) is used here, and $\gamma_{i}\sim\mbox{Bern}(\phi_{i})$ is approximated by $\widetilde{\gamma}_{i}\sim\mbox{Gumbel-softmax}(\phi_{i},\tau)$, where

$\tau$ is called the temperature, and as it approaches 0, $\tilde{\gamma}_{i}$ converges to $\gamma_{i}$ in distribution (refer to Figure 4 in the supplementary material). In addition, one can show that $P(\widetilde{\gamma}_{i}>0.5)=\phi_{i}$, which implies $\gamma_{i}\stackrel{{\scriptstyle d}}{{=}}1(\widetilde{\gamma}_{i}>0.5)$. Thus, $\widetilde{\gamma}_{i}$ is viewed as a soft version of $\gamma_{i}$, and will be used in the backward pass to enable the calculation for gradients, while the hard version $\gamma_{i}$ will be used in the forward pass to obtain a sparse network structure. In practice, $\tau$ is usually chosen no smaller than 0.5 for numerical stability. Besides, the normal variable $\mathcal{N}(\mu_{i},\sigma^{2}_{i})$ is reparameterized by $\mu_{i}+\sigma_{i}\epsilon_{i}$ for $\epsilon_{i}\sim\mathcal{N}(0,1)$.

The complete variational inference procedure with Gumbel-softmax approximation is stated below.

The Gumbel-softmax approximation introduces an additional error that may jeopardize the validity of Theorem 4.1. Our exploratory studies (refer to Section B.3 in supplementary material) demonstrates little differences between the results of using inverse-CDF reparameterization and using Gumbel-softmax approximation in some simple model. Therefore, we conjecture that Gumbel-softmax approximation doesn’t hurt the VB convergence, and thus will be implemented in our numerical studies.

We evaluate the empirical performance of the proposed variational inference through simulation study and MNIST data application. For the simulation study, we consider a teacher-student framework and a nonlinear regression function, by which we justify the consistency of the proposed method and validate the proposed choice of hyperparameters. As a byproduct, the performance of uncertainty quantification and the effectiveness of variable selection will be examined as well.

For all the numerical studies, we let $\sigma^{2}_{0}=2$, the choice of $\lambda$ follows Theorem 4.1 (denoted by $\lambda_{opt}$): $\log(\lambda^{-1}_{opt})=\log(T)+0.1[(L+1)\log N+\log\sqrt{n}p]$. The remaining details of implementation (such as initialization, choices of $K$, $m$ and learning rate) are provided in the supplementary material. We will use VB posterior mean estimator $\widehat{f}_{H}=\sum^{H}_{h=1}f_{\theta_{h}}/H$ to assess the prediction accuracy, where $\theta_{h}\sim\widehat{q}(\theta)$ are samples drawn from the VB posterior and $H=30$. The posterior network sparsity is measured by $\widehat{s}=\sum^{T}_{i=1}\phi_{i}/T$. Input nodes who have connection with $\phi_{i}>0.5$ to the second layer is selected as relevant input variables, and we report the corresponding false positive rate (FPR) and false negative rate (FNR) to evaluate the variable selection performance of our method.

Our method will be compared with the dense variational BNN (VBNN) (Blundell et al.,, 2015) with independent centered normal prior and independent normal variational distribution, the AGP pruner (Zhu and Gupta,, 2018), the Lottery Ticket Hypothesis (LOT) (Frankle and Carbin,, 2018), the variational dropout (VD) (Molchanov et al.,, 2017) and the Horseshoe BNN (HS-BNN) (Ghosh et al.,, 2018). In particular, VBNN can be regarded as a baseline method without any sparsification or compression. All reported simulation results are based on 30 replications (except that we use 60 replications for interval estimation coverages). Note that the sparsity level in methods AGP and LOT are user-specified. Hence, in simulation studies, we try a grid search for AGP and LOT, and only report the ones that yield highest testing accuracy. Furthermore, note that FPR and FNR are not calculated for HS-BNN since it only sparsifies the hidden layers nodewisely.

We proposed a variational inference method for deep neural networks under spike-and-slab priors with theoretical guarantees. Future direction could be investigating the theory behind choosing hyperparamters via the EB estimation instead of deterministic choices.

Furthermore, extending the current results to more complicated networks (convolutional neural network, residual network, etc.) is not trivial. Conceptually, it requires the design of structured sparsity (e.g., group sparsity in Neklyudov et al., (2017)) to fulfill the goal of faster prediction. Theoretically, it requires deeper understanding of the expressive ability (i.e. approximation error) and capacity (i.e., packing or covering number) of the network model space. For illustration purpose, we include an example of Fashion-MNIST task using convolutional neural network in the supplementary material, and it demonstrates the usage of our method on more complex networks in practice.

We believe the ethical aspects are not applicable to this work. For future societal consequences, deep learning has a wide range of applications such as computer version and natural language processing. Our work provides a solution to overcome the drawbacks of modern deep neural network, and also improves the understanding of deep learning.

The proposed method could improve the existing applications. Specifically, sparse learning helps apply deep neural networks to hardware limited devices, like cell phones or pads, which will broaden the horizon of deep learning application. In addition, as a Bayesian method, not only a result, but also the knowledge of confidence or certainty in that result are provided, which could benefit people in various aspects. For example, in the application of cancer diagnostic, by providing the certainty associated with each possible outcome, Bayesian learning would assist the medical professionals to make a better judgement about whether the tumor is a cancer or a benign one. Such kind of ability to quantify uncertainty would contribute to the modern deep learning.

We would like to thank Wei Deng for the helpful discussion and thank the reviewers for their thoughtful comments. Qifan Song’s research is partially supported by National Science Foundation grant DMS-1811812. Guang Cheng was a member of the Institute for Advanced Study in writing this paper, and he would like to thank the institute for its hospitality.