Structured Bayesian Pruning via Log-Normal Multiplicative Noise

Introduction

Deep neural networks are a flexible family of models which provides state-of-the-art results in many machine learning problems . However, this flexibility often results in overfitting. A common solution for this problem is regularization. One of the most popular ways of regularization is Binary Dropout that prevents co-adaptation of neurons by randomly dropping them during training. An equally effective alternative is Gaussian Dropout that multiplies the outputs of the neurons by Gaussian random noise. In recent years several Bayesian generalizations of these techniques have been developed, e.g. Variational Dropout and Variational Spike-and-Slab Neural Networks . These techniques provide theoretical justification of different kinds of Dropout and also allow for automatic tuning of dropout rates, which is an important practical result.

Besides overfitting, compression and acceleration of neural networks are other important challenges, especially when memory or computational resources are restricted. Further studies of Variational Dropout show that individual dropout rates for each weight allow to shrink the original network architecture and result in a highly sparse model . General sparsity provides a way of neural network compression, while the time of network evaluation may remain the same, as most modern DNN-oriented software can’t work with sparse matrices efficiently. At the same time, it is possible to achieve acceleration by enforcing structured sparsity in convolutional filters or data tensors. In the simplest case it means removing redundant neurons or convolutional filters instead of separate weights; but more complex patterns can also be considered. This way Group-wise Brain Damage employs group-wise sparsity in convolutional filters, Perforated CNNs drop redundant rows from the intermediate dataframe matrices that are used to compute convolutions, and Structured Sparsity Learning provides a way to remove entire convolutional filters or even layers in residual networks. These methods allow to obtain practical acceleration with little to no modifications of the existing software. In this paper, we propose a tool that is able to induce an arbitrary pattern of structured sparsity on neural network parameters or intermediate data tensors. We propose a dropout-like layer with a parametric multiplicative noise and use stochastic variational inference to tune its parameters in a Bayesian way. We introduce a proper analog of sparsity-inducing log-uniform prior distribution that allows us to formulate a correct probabilistic model and avoid the problems that come from using an improper prior. This way we obtain a novel Bayesian method of regularization of neural networks that results in structured sparsity. Our model can be represented as a separate dropout-like layer that allows for a simple and flexible implementation with almost no computational overhead, and can be incorporated into existing neural networks.

Our experiments show that our model leads to high group sparsity level and significant acceleration of convolutional neural networks with negligible accuracy drop. We demonstrate the performance of our method on LeNet and VGG-like architectures using MNIST and CIFAR-10 datasets.

Related Work

Deep neural networks are extremely prone to overfitting, and extensive regularization is crucial. The most popular regularization methods are based on injection of multiplicative noise over layer inputs, parameters or activations . Different kinds of multiplicative noise have been used in practice; the most popular choices are Bernoulli and Gaussian distributions. Another type of regularization of deep neural networks is based on reducing the number of parameters. One approach is to use low-rank approximations, e.g. tensor decompositions , and the other approach is to induce sparsity, e.g. by pruning or $L_{1}$ regularization . Sparsity can also be induced by using the Sparse Bayesian Learning framework with empirical Bayes or with sparsity-inducing priors .

High sparsity is one of the key factors for the compression of DNNs . However, in addition to compression it is beneficial to obtain acceleration. Recent papers propose different approaches to acceleration of DNNs, e.g. Spatial Skipped Convolutions and Spatially Adaptive Computation Time that propose different ways to reduce the number of computed convolutions, Binary Networks that achieve speedup by using only 1 bit to store a single weight of a DNN, Low-Rank Expansions that use low-rank filter approximations, and Structured Sparsity Learning that allows to remove separate neurons or filters. As reported in it is possible to obtain acceleration of DNNs by introducing structured sparsity, e.g. by removing whole neurons, filters or layers. However, non-adaptive regularization techniques require tuning of a huge number of hyperparameters that makes it difficult to apply in practice. In this paper we apply the Bayesian learning framework to obtain structured sparsity and focus on acceleration of neural networks.

Stochastic Variational Inference

Given a probabilistic model $p(y\,|\,x,\theta)$ we want to tune parameters $\theta$ of the model using training dataset $\mathcal{D}=\{(x_{i},y_{i})\}_{i=1}^{N}$. The prior knowledge about parameters $\theta$ is defined by prior distribution $p(\theta)$. Using the Bayes rule we obtain the posterior distribution $p(\theta\,|\,\mathcal{D})=p(\mathcal{D}\,|\,\theta)p(\theta)/p(\mathcal{D})$. However, computing posterior distribution using the Bayes rule usually involves computation of intractable integrals, so we need to use approximation techniques.

$L_{D}(\phi)$ is a so-called expected log-likelihood function which is intractable in case of complex probabilistic model $p(y\,|\,x,\theta)$. Following we use the Reparametrization trick to obtain an unbiased differentiable minibatch-based Monte Carlo estimator of the expected log-likelihood. Here $N$ is the total number of objects, $M$ is the minibatch size, and $f(\phi,\varepsilon)$ provides samples from the approximate posterior $q_{\phi}(\theta)$ as a deterministic function of a non-parametric noise $\varepsilon\sim p(\varepsilon)$.

This way we obtain a procedure of approximate Bayesian inference where we solve optimization problem (4) by stochastic gradient ascent w.r.t. variational parameters $\phi$. This procedure can be efficiently applied to Deep Neural Networks and usually the computational overhead is very small, as compared to ordinary DNNs.

If the model $p(y\,|\,x,\theta,w)$ has another set of parameters $w$ that we do not want to be Bayesian about, we can still use the same variational lower bound objective:

This objective corresponds the maximum likelihood estimation $w_{ML}$ of parameters $w$, while finding the approximate posterior distribution $q_{\phi}(\theta)\approx p(\theta\,|\,\mathcal{D},w_{ML})$. In this paper we denote the weights of the neural networks, the biases, etc. as $w$ and find their maximum likelihood estimation as described above. The parameters $\theta$ that undergo the Bayesian treatment are the noisy masks in the proposed dropout-like layer (SBP layer). They are described in the following section.

Group Sparsity with Log-normal Multiplicative Noise

Variational Inference with a sparsity-inducing log-uniform prior over the weights of a neural network is an efficient way to enforce general sparsity on weight matrices . However, it is difficult to apply this approach to explicitly enforce structured sparsity. We introduce a dropout-like layer with a certain kind of multiplicative noise. We also make use of the sparsity-inducing log-uniform prior, but put it over the noise variables rather than weights. By sharing those noise variables we can enforce group-wise sparsity with any form of groups.

In order to train the model, i.e. perform variational inference, we need to choose an approximation family $q_{\phi}$ for the posterior distribution $p(\theta\,|\,\mathcal{D})\approx q_{\phi}(\theta)$.

A common choice of variational distribution $q(\cdot)$ is a fully-factorized Gaussian distribution. However, for this particular model we choose $q(\theta)$ to be a fully-factorized log-normal distribution (10–11). To make this choice, we were guided by the following reasons:

The log-uniform distribution is a specific case of the log-normal distribution when the parameter $\sigma$ goes to infinity and $\mu$ remains fixed. Thus we can guarantee that in the case of no data our variational approximation can be made exact. Hence this variational family has no "prior gap".

We consider a model with multiplicative noise. The scale of this noise corresponds to its shift in the logarithmic space. By establishing the log-uniform prior we set no preferences on different scales of this multiplicative noise. The usual use of a Gaussian as a posterior immediately implies very asymmetric skewed distribution in the logarithmic space. Moreover log-uniform and Gaussian distributions have different supports and that will require establishing two log-uniform distributions for positive and negative noises. In this case Gaussian variational approximation would have quite exotic bi-modal form (one mode in the log-space of positive noises and another one in the log-space of negative noises). On the other hand, the log-normal posterior for the multiplicative noise corresponds to a Gaussian posterior for the additive noise in the logarithmic scale, which is much easier to interpret.

Log-normal noise is always non-negative both during training and testing phase, therefore it does not change the sign of its input. This is in contrast to Gaussian multiplicative noise $\mathcal{N}(\theta_{i}\,|\,1,\alpha)$ that is a standard choice for Gaussian dropout and its modifications . During the training phase Gaussian noise can take negative values, so the input to the following layer can be of arbitrary sign. However, during the testing phase noise $\theta$ is equal to $1$, so the input to the following layer is non-negative with many popular non-linearities (e.g. ReLU, sigmoid, softplus). Although Gaussian dropout works well in practice, it is difficult to justify notoriously different input distributions during training and testing phases.

The final loss function is presented in equation (12) and is essentially the original variational lower bound (4).

where $\mu$ and $\sigma$ are the variatianal parameters, and $W$ denotes all other trainable parameters of the neural network, e.g. the weight matrices, the biases, batch normalization parameters, etc.

Note that we can optimize the variational lower bound w.r.t. the parameters $\mu$ and $\sigma$ of the log-normal noise $\theta$. We do not fix the mean of the noise thus making our variational approximation more tight.

2 Problems of Variational Inference with Improper Log-Uniform Prior

A common way to tackle this problem is to consider the density of the log-uniform distribution to be equal to $\frac{C}{\theta}$ and to treat $C$ as some finite constant. This trick works well for the case of a Gaussian posterior distribution . The KL divergence between a Gaussian posterior and a log-uniform prior has an infinite gap, but can be calculated up to this infinite constant in a meaningful way . However, for the case of the log-normal posterior the KL divergence is infinite for any finite values of variational parameters, and is equal to zero for a fixed finite $\mu$ and infinite $\sigma$. As the data-term (3) is bounded for any value of variational parameters, the only global optimum of the variational lower bound is achieved when $\mu$ is finite and fixed, and $\sigma$ goes to infinity. In this case the posterior distribution collapses into the prior distribution and the model fails to extract any information about the data. This effect is wholly caused by the fact that the log-uniform prior is an improper (non-normalizable) distribution, which makes the whole probabilistic model flawed.

3 Variational Inference with Truncated Approximation Family

Due to the improper prior the optimization problem becomes ill-posed. But do we really need to use an improper prior distribution? The most common number format that is used to represent the parameters of a neural network is the floating-point format. The floating-point format is only able to represent numbers from a limited range. For example, a single-point precision variable can only represent numbers from the range $-3.4\times 10^{38}$ to $+3.4\times 10^{38}$, and the smallest possible positive number is equal to $1.2\times 10^{-38}$. All of probability mass of the improper log-uniform prior is concentrated beyond the single-point precision (and essentially any practical floating point precision), not to mention that the actual relevant range of values of neural network parameters is much smaller. It means that in practice this prior is not a good choice for software implementation of neural networks.

We propose to use a truncated log-uniform distribution (14) as a proper analog of the log-uniform distribution. Here $I_{[a,b]}(x)$ denotes the indicator function for the interval $x\in[a,b]$. The posterior distribution should be defined on the same support as the prior distribution, so we also need to use a truncated log-normal distribution (14).

Our final model then can be formulated as follows.

Note that all the nice facts about the log-normal posterior distribution from the Section 4.1 are also true for the truncated log-normal posterior. However, now we have a proper probabilistic model and the Stochastic Variational Inference can be preformed correctly. Unlike (13), now the KL divergence term (16–17) can be calculated correctly for all valid values of variational parameters (see Appendix A for details).

where $\alpha_{i}=\frac{a-\mu_{i}}{\sigma_{i}}$, $\beta_{i}=\frac{b-\mu_{i}}{\sigma_{i}}$, $\phi(\cdot)$ and $\Phi(\cdot)$ are the density and the CDF of the standard normal distribution.

The reparameterization trick also can still be performed (18) using the inverse CDF of the truncated normal distribution (see Appendix B).

The final loss and the set of parameters is the same as described in Section 4.1, and the training procedure remains the same.

4 Sparsity

The SNR can be computed analytically, the derivation can be found in the appendix. It has a simple interpretation. If the SNR is low, the corresponding neuron becomes very noisy and its output no longer contains any useful information. If the SNR is high, it means that the neuron output contains little noise and is important for prediction. Therefore we can remove all neurons or filters with a low SNR and set their output to constant zero.

5 Implementation details

We perform a minibatch-based stochastic variational inference for training. The training procedure looks as follows. On each training step we take a minibatch of $M$ objects and feed it into the neural network. Consider a single SBP layer with input $X^{M\times I}$ and output $Y^{M\times I}$. We independently sample a separate noise vector $\theta^{m}\sim q(\theta)$ for each object $x_{m}$ and obtain a noise matrix $\theta^{M\times I}$. The output matrix $Y^{M\times I}$ is then obtained by component-wise multiplication of the input matrix and the noise matrix: $y_{mi}=x_{mi}\cdot\theta^{m}_{i}$.

To be fully Bayesian, one would also sample and average over different dropout masks $\theta$ during testing, i.e. perform Bayesian ensembling. Although this procedure can be used to slightly improve the final accuracy, it is usually avoided. Bayesian ensembling essentially requires sampling of different copies of neural networks, which makes the evaluation $K$ times slower for averaging over $K$ samples. Instead, during the testing phase in most dropout-based techniques the noise variable $\theta$ is replaced with its expected value. In this paper we follow the same approach and replace all non-pruned $\theta_{i}$ with their expectations (20) during testing. The derivation of the expectation of the truncated log-normal distribution is presented in Appendix C.

We tried to use Bayesian ensembling with this model, and experienced almost no gain of accuracy. It means that the variance of the learned approximate posterior distribution is low and does not provide a rich ensemble.

Throughout the paper we introduced the SBP dropout layer for the case when input objects are represented as one-dimensional vectors $x$. When defined like that, it would induce general sparsity on the input vector $x$. It works as intended for fully-connected layers, as a single input feature corresponds to a single output neuron of a preceding fully-connected layer and a single output neuron of the following layer. However, it is possible to apply the SBP layer in a more generic setting. Firstly, if the input object is represented as a multidimensional tensor $X$ with shape $I_{1}\times I_{2}\times\dots\times I_{d}$, the noise vector $\theta$ of length $I=I_{1}\times I_{2}\times\dots\times I_{d}$ can be reshaped into a tensor with the same shape. Then the output tensor $Y$ can be obtained as a component-wise product of the input tensor $X$ and the noise tensor $\theta$. Secondly, the SBP layer can induce any form of structured sparsity on this input tensor $X$. To do it, one would simply need to use a single random variable $\theta_{i}$ for the group of input features that should be removed simultaneously. For example, consider an input tensor $X^{H\times W\times C}$ that comes from a convolutional layer, $H$ and $W$ being the size of the image, and $C$ being the number of channels. Then, in order to remove redundant filters from the preceding layer (and at the same time redundant channels from the following layer), one need to share the random variables $\theta$ in the following way:

Similarly, the SBP layer can be applied in a DropConnect fashion. One would just need to multiply the weight tensor $W$ by a noise tensor $\theta$ of similar shape. The training procedure remains the same. It is still possible to enforce any structured sparsity pattern for the weight tensor $W$ by sharing the random variables as described above.

Experiments

We perform an evaluation on different supervised classification tasks and with different architectures of neural networks including deep VGG-like architectures with batch normalization layers. For each architecture, we report the number of retained neurons and filters, and obtained acceleration. Our experiments show that Structured Bayesian Pruning leads to a high level of structured sparsity in convolutional filters and neurons of DNNs without significant accuracy drop. We also demonstrate that optimization w.r.t. the full set of variational parameters $(\mu,\sigma)$ leads to improving model quality and allows us to perform sparsification in a more efficient way, as compared to tuning of only one free parameter that corresponds to the noise variance. As a nice bonus, we show that Structured Bayesian Pruning network does not overfit on randomly labeled data, that is a common weakness of non-bayesian dropout networks. The source code is available in Theano and Lasagne, and also in TensorFlow (https://github.com/necludov/group-sparsity-sbp).

The truncation parameters $a$ and $b$ are the hyperparameters of our model. As our layer is meant for regularization of the model, we would like our layer not to amplify the input signal and restrict the noise $\theta$ to an interval $$. This choice corresponds to the right truncation threshold$b$set to . We find empirically that the left truncation parameter$a$does not influence the final result much. We use values$a=-20$and$b=0$ in all experiments.

2 More Flexible Variational Approximation

Usually during automatic training of dropout rates the mean of the noise distribution remains fixed. In the case of our model it is possible to train both mean and variance of the multiplicative noise. By using a more flexible distribution we obtain a tighter variational lower bound and a higher sparsity level. In order to demonstrate this effect, we performed an experiment on MNIST dataset with a fully connected neural network that contains two hidden layers with 1000 neurons each. The results are presented in Fig. 2.

3 LeNet5 and Fully-Connected Net on MNIST

We compare our method with other sparsity inducing methods on the MNIST dataset using a fully connected architecture LeNet-500-300 and a convolutional architecture LeNet-5-Caffe. These networks were trained with Adam without any data augmentation. The LeNet-500-300 network was trained from scratch, and the LeNet-5-CaffeA modified version of LeNet5 from . Caffe Model specification: https://goo.gl/4yI3dL network was pretrained with weight decay. An illustration of trained SNR for the image features for the LeNet-500-300Fully Connected Neural Net with 2 hidden layers that contains 500 and 300 neurons respectively. network is shown in Fig. 2. The final accuracy, group-wise sparsity levels and speedup for these architectures for different methods are shown in Table 1.

4 VGG-like on CIFAR-10

To prove that SBP scales to deep architectures, we apply it to a VGG-like network that was adapted for the CIFAR-10 dataset. The network consists of 13 convolutional and two fully-connected layers, trained with pre-activation batch normalization and Binary Dropout. At the start of the training procedure, we use pre-trained weights for initialization. Results with different scaling of the number of units are presented in Table 2. We present results for two architectures with different scaling coefficient $k\in\{1.0,1.5\}$ . For smaller values of scaling coefficient $k\in\{0.25,0.5\}$ we obtain less sparse architecture since these networks have small learning capacities. Besides the results for the standard StructuredBP procedure, we also provide the results for SBP with KL scaling (StructuredBPa). Scaling the KL term of the variational lower bound proportional to the computational complexity of the layer leads to a higher sparsity level for the first layers, providing more acceleration. Despite the higher error values, we obtain the higher value of true variational lower bound during KL scaling, hence, we find its another local maximum.

5 Random Labels

A recent work shows that Deep Neural Networks have so much capacity that they can easily memorize the data even with random labeling . Binary dropout as well as other standard regularization techniques do not prevent the networks from overfitting in this scenario. However, recently it was shown that Bayesian regularization may help . Following these works, we conducted similar experiments. We used a Lenet5 network on the MNIST dataset and a VGG-like network on CIFAR-10. Although Binary Dropout does not prevent these networks from overfitting, SBP decides to remove all neurons of the neural network and provides a constant prediction. In other words, in this case SBP chooses the simplest model that achieves the same testing error rate. This is another confirmation that Bayesian regularization is more powerful than other popular regularization techniques.

Conclusion

We propose Structured Bayesian Pruning, or SBP, a dropout-like layer that induces multiplicative random noise over the output of the preceding layer. We put a sparsity-inducing prior over the noise variables and tune the noise distribution using stochastic variational inference. SBP layer can induce an arbitrary structured sparsity pattern over its input and provides adaptive regularization. We apply SBP to cut down the number of neurons and filters in convolutional neural networks and report significant practical acceleration with no modification of the existing software implementation of these architectures.

We would like to thank Christos Louizos and Max Welling for valuable discussions. Kirill Neklyudov and Arsenii Ashukha were supported by HSE International lab of Deep Learning and Bayesian Methods which is funded by the Russian Academic Excellence Project ’5-100’. Dmitry Molchanov was supported by the Ministry of Education and Science of the Russian Federation (grant 14.756.31.0001). Dmitry Vetrov was supported by the Russian Science Foundation grant 17-11-01027.

References

Appendix A KL divergence for the truncated log-normal distribution

Now we can calculate the KL divergence between a truncated log-normal distribution $q(\theta^{i})$ and a log-uniform distribution $p(\theta^{i})$ with a bounded support $\theta^{i}\in[e^{a},e^{b}]$.

Entropy for the truncated normal distribution is

Appendix B Sampling from the truncated log-normal distribution

Appendix C Mean of the truncated log-normal ditribution

Let’s derive the expected value of $\theta$. In order to this, let’s first find the PDF of the truncated log-normal distribution.

The obtained PDF is very similar to the log-normal distribution PDF. Hence,

Appendix D Signal-to-noise ratio of the truncated log-normal distribution

We already have the expression for $p(a)$

For the rest two summands we introduce a new variable $t$.

Next, we use $t$ in a variable substitution.

Finally, we obtain the signal-to-noise ratio

Appendix E Stable Computation of Statistics

we can rewrite equations (19), (20) in the following form