Training Sparse Neural Networks

11/21/2016 ∙ by Suraj Srinivas, et al. ∙ indian institute of science 0

Deep neural networks with lots of parameters are typically used for large-scale computer vision tasks such as image classification. This is a result of using dense matrix multiplications and convolutions. However, sparse computations are known to be much more efficient. In this work, we train and build neural networks which implicitly use sparse computations. We introduce additional gate variables to perform parameter selection and show that this is equivalent to using a spike-and-slab prior. We experimentally validate our method on both small and large networks and achieve state-of-the-art compression results for sparse neural network models.

READ FULL TEXT VIEW PDF
POST COMMENT

Comments

There are no comments yet.

Authors

page 1

page 2

page 3

page 4

This week in AI

Get the week's most popular data science and artificial intelligence research sent straight to your inbox every Saturday.

Introduction

For large-scale tasks such as image classification, large networks with many millions of parameters are often used [Krizhevsky, Sutskever, and Hinton2012], [Simonyan and Zisserman2015], [Szegedy et al.2015]. However, these networks typically use dense computations. Would it be advantageous to use sparse computations instead? Apart from having fewer number of parameters to store ( to )111For a matrix of size with non-zero elements, sparse computations also decrease feedforward evaluation time ( to )222

For matrix-vector multiplies with a dense vector of size

. Further, having a lower parameter count may help in avoiding overfitting.

Regularizers are often used to discourage overfitting. These usually restrict the magnitude () of weights. However, to restrict the computational complexity of neural networks, we need a regularizer which restricts the total number of parameters of a network. A common strategy to obtain sparse parameters is to apply sparsity-inducing regularizers such as the penalty on the parameter vector. However, this is often insufficient in inducing sparsity in case of large non-convex problems like deep neural network training as shown in [Collins and Kohli2014]. The contribution of this paper is to be able to induce sparsity in a tractable way for such models.

The overall contributions of the paper are as follows.

  • We propose a novel regularizer that restricts the total number of parameters in the network. (Section 2)

  • We perform experimental analysis to understand the behaviour of our method. (Section 4)

  • We apply our method on LeNet-5, AlexNet and VGG-16 network architectures to achieve state-of-the-art results on network compression. (Section 4)

Problem Formulation

To understand the motivation behind our method, let us first define our notion of computational complexity of a neural network.

Let be a set of vectors. This represents an -layer dense neural network architecture where is a vector of parameter indices for the layer, i.e; . Here, each layer contains elements. Zero indicates absence of a parameter and one indicates presence. Thus, for a dense neural network, is a vector of all ones, i.e.; . For a sparse parameter vector, would consist of mostly zeros. Let us call as the index set of a neural network.

For these vectors, our notion of complexity is simply the total number of parameters in the network.

Definition 1.

The complexity of a -layer neural network with index set is given by .

We now aim to solve the following optimization problem.

(1)

where denotes the weights of the neural network, and the index set.

denotes the loss function, which depends on the underlying task to be solved. Here, we learn both the weights as well as the index set of the neural network. Using the formalism of the index set, we are able to penalize the total number of network parameters. While easy to state, we note that this problem is difficult to solve, primarily because

contains elements .

Figure 1: Our strategy for sparsifying weight matrices. First, we sample / threshold the gate variables. We then multiply the resulting binary matrix with , to yield a sparse matrix .

Gate Variables

How do we incorporate the index set formalism in neural networks? Assume that the index set ( in Fig. 1) is multiplied pointwise with the weight matrix. This results in a weight matrix that is effectively sparse, if the index set has lots of zeros rather than ones. In other words, we end up learning two sets of variables to ensure that one of them - weights - becomes sparse. How do we learn such binary parameters in the first place ?

To facilitate this, we interpret index set variables (

) as draws from a bernoulli random variable. As a result, we end up learning the real-valued bernoulli parameters (

in Fig. 1), or gate variables rather than index set variables themselves. Here the sampled binary gate matrix corresponds exactly to the index set, or the matrix described above. To clarify our notation, and stand for the real-valued gate variables, while the superscript indicates binary sampled variables.

When we draw from a bernoulli distribution, we have two choices - we can either perform a

unbiased draw (the usual sampling process), or we can perform a so-called maximum-likelihood (ML) draw. The ML draw involves simply thresholding the values of at . To ensure determinism, we use the ML draw or thresholding in this work.

Promoting Sparsity

Given our formalism of gate variables, how do we ensure that the learnt bernoulli parameters are low - or in our case - mostly less than ? One plausible option is to use the or the regularizer on the gate variables. However, this does not ensure that there will exist values greater than . To accommodate this, we require a bi-modal regularizer, i.e; a regularizer which ensures that some values are large, but most values are small.

To this end, we use a regularizer given by . This was introduced by [Murray and Ng2010] to learn binary values for parameters. However, what is important for us is that this regularizer has the bi-modal property mentioned earlier, as shown in Fig. 1(a)

(a) Bi-modal regularizer
(b) Spike and slab prior
Figure 2: (a) The bi-modal regularizer used in our work. Note that this encourages values to be close to and , in contrast to the regularizer. (b) An example of a spike-and-slab prior similar to the one used in this work (except for a constant).

Our overall regularizer is simply a combination of this bi-modal regularizer as well the traditional or regularizer for the individual gate variables. Our objective function is now stated as follows.

(2)

where denotes the gate parameter in the layer. Note that for , the second term in Eqn. 2 vanishes and the third term becomes , thus reducing to Eqn.1.

An Alternate Interpretation

Now that we have arrived at the objective function in Eqn.2, it is natural to ask the question - how do we know that it solves the original objective in Eqn.1 ? We shall now derive Eqn.2 from this perspective.

Assuming the formulation of gate variables, we can re-write the objective in Eqn.1 as follows.

(3)

where is the sampled version of gate variables . Note that Eqn.3 is a stochastic objective function, arising from the fact that is a random variable. We can convert this to a real-valued objective by taking expectations. Note that expectation of the loss function is difficult to compute. As a result, we approximate it with a Monte-Carlo average.

where . While this formulation is sufficient to solve the original problem, we impose another condition on this objective. We would like to minimize the number of Monte-Carlo evaluations in the loss term. This amounts to reducing for a fixed

, or reducing the variance of the loss term. This is done by reducing the variance of

, the only random variable in the equation. To account for this, we add another penalty term corresponding to . Imposing this additional penalty and then using gives us back Eqn.2.

Relation to Spike-and-Slab priors

We observe that our problem formulation closely resembles spike-and-slab type priors used in Bayesian statistics for variable selection

[Mitchell and Beauchamp1988]. Broadly speaking, these priors are mixtures of two distributions - one with very low variance (spike), and another with comparatively large variance (slab). By placing a large mass on the spike, we can expect to obtain parameter vectors with large sparsity.

Let us consider for a moment using the following prior for weight matrices of neural networks.

(4)

Here, denotes the dirac delta distribution, and denotes the normalizing constant, and is the mixture coefficient. Also note that like [Mitchell and Beauchamp1988], we assume that for some . This is visualized in Fig. 1(b). Note that this is a multiplicative mixture of distributions, rather than additive. By taking negative logarithm of this term and ignoring constant terms, we obtain

(5)

Note that the first term in this expression corresponds exactly to the number of non-zero parameters, i.e; the term of Eqn. 1. The second term corresponds to the usual regularizer on the weights of the network (rather than gates). As a result, we conclude that Eqn. 4 is a spike-and-slab prior which we implicitly end up using in this method.

Estimating gradients for gate variables

How do we estimate gradients for gate variables, given that they are binary stochastic variables, rather than real-valued and smooth? In other words, how do we backpropagate through the bernoulli sampling step? Bengio

et al. bengio2013estimating investigated this problem and empirically verified the efficacy of different possible solutions. They conclude that the simplest way of computing gradients - the straight-through estimator works best overall. Our experiments also agree with this observation.

The straight-through

estimator simply involves back-propagating through a stochastic neuron as if it were an identity function. If the sampling step is given by

, then the gradient is used.

Another issue of consideration is that of ensuring that always lies in so that it is a valid bernoulli parameter. Bengio et al.

bengio2013estimating use a sigmoid activation function to achieve this. Our experiments showed that clipping functions worked better. This can be thought of as a ‘linearized’ sigmoid. The clipping function is given by the following expression.

The overall sampling function is hence given by , and the straight-through estimator is used to estimate gradients overall.

Comparison with LASSO

LASSO is commonly used method to attain sparsity and perform variable selection. The main difference between the above method and LASSO is that LASSO is primarily a shrinkage operator, i.e.; it shrinks all parameters until lots of them are close to zero. This is not true for the case of spike-and-slab priors, which can have high sparsity and encourage large values at the same time. This is due to the richer parameterization of these priors.

(a) Effect of Gate initialization
(b) Effect of varying
(c) Effect of varying
Figure 3: (a) We vary the initialization of the gate variables and observe it’s effect on sparsity. The dotted blue lines denote the variance of sparsity in case of the sampling-based method. (b) seems to have a stabilizing effect on sparsity whereas (c) increasing seems to increase sparsity.

Practical issues

In this section we shall discuss some practical issues pertaining to our method. Our method ironically uses twice the number of parameters as a typical neural network, as we have two sets of variables - weights and gates. As a result, model size doubles while training. However, we multiply the two to result in sparse matrices which considerably reduces model size at test time. Essentially we do not have to store both sets of parameters while testing,only a element-wise product of the two is required. Even though the model size doubles at train time, we note that speed of training / feedforward evaluation is not affected due to the fact that only element-wise operations are used.

Our method can be applied to both convolutional tensors as well as fully connected matrices. However while performing compression, we note that convolutional layers are less susceptible to compression that fully connected layers due to the small number of parameters they possess.

Related Work

There have been many recent works which perform compression of neural networks. Weight-pruning techniques were popularized by LeCun et al.lecun1989optimal and Hassibi et al.hassibi1993second, who introduced Optimal Brain Damage and Optimal Brain Surgery respectively. Recently, Srinivas and Babu BMVC2015_31 proposed a neuron pruning technique, which relied on neuronal similarity. In contrast, we perform weight pruning based on learning, rather than hand-crafted rules.

Previous attempts have also been made to sparsify neural networks. Han et al.han2015learning create sparse networks by alternating between weight pruning and network training. A similar strategy is followed by Collins and Kohli DBLP:journals/corr/CollinsK14. On the other hand, our method performs both weight pruning and network training simultaneously. Further, our method has considerably less number of hyper-parameters to determine () compared to the other methods, which have thresholds to be set for each of the layers in a neural network.

Many methods have been proposed to train models that are deep, yet have a lower parameterisation than conventional networks. Denil et al.denil2013predicting demonstrated that most of the parameters of a model can be predicted given only a few parameters. At training time, they learn only a few parameters and predict the rest. Yang et al.yang2014deep propose an Adaptive Fastfood transform, which is an efficient re-parametrization of fully-connected layer weights. This results in a reduction of complexity for weight storage and computation. Novikov et al.novikov2015tensorizing use tensor decompositions to obtain a factorization of tensors with small number of parameters. Cheng et al.cheng2015exploration make use of circulant matrices to re-paramaterize fully connected layers. Some recent works have also focussed on using approximations of weight matrices to perform compression. Gong et al.gong2014compressing use a clustering-based product quantization approach to build an indexing scheme that reduces the space occupied by the matrix on disk. Note that to take full advantage of these methods, one needs to have fast implementations of the specific parameterization used. One the other hand, we use a sparse parameterization, fast implementations of which are available on almost every platform.

Our work is very similar to that of Architecture Learning [Srinivas and Babu2015b], which uses a similar framework to minimize the total number of neurons in a neural network. On the other hand, we minimize the total number of weights.

Layers Initial
Params Final
Params Sparsity (%)
conv1 0.5K 0.04K 91
conv2 25K 1.78K 92.8
fc1 400K 15.4K 96.1
fc2 5K 0.6K 86.8
Total 431K 17.9K 95.84
Table 1: Compression results for LeNet-5 architecture.
Method Params (P+I)* Accuracy(%) Compression Rate(%)
Baseline 431K 99.20 1x
SVD(rank-10)[Denton et al.2014] 43.6K 98.47 10x
AL [Srinivas and Babu2015b] 40.9K 99.04 10.5x
AF-1024 [Yang et al.2014] 38.8K 99.29 11x
Han et al.[Han et al.2015] 36K (72k)* 99.23 12x
Method-1 18K (36k)* 99.19 24x
Method-2 22K (44k)* 99.33 19x
Table 2: Comparison of compression performance on LeNet-5 architecture. *Count of number of parameters and the indices to store them. This is the effective storage requirement when using sparse computations.

Experiments

In this section we perform experiments to evaluate the effectiveness of our method. First, we perform some experiments designed to understand typical behaviour of the method. These experiments are done primarily on LeNet-5 [LeCun et al.1998]

. Second, we use our method to perform network compression on two networks - LeNet-5 and AlexNet. These networks are trained on MNIST and ILSVRC-2012 dataset respectively. Our implementation is based on Lasagne, a Theano-based library.

Analysis of Proposed method

We shall now describe experiments to analyze the behaviour of our method. First, we shall analyze the effect of hyper-parameters. Second, we study the effect of varying model sizes on the resulting sparsity.

For all analysis experiments, we consider the LeNet-5 network. LeNet-5 consists of two convolutional layers with 20 and 50 filters, and two fully connected layers with 500 and 10 (output layer) neurons. For analysis, we only study the effects sparsifying the third fully connected layer.

Layers Initial
Params Final
Params Sparsity (%)
conv(5 layers) 2.3M 2.3M -
fc6 38M 1.3M 96.5
fc7 17M 1M 94
fc8 4M 1.2M 70
Total 60.9M 5.9M 90
Table 3: Layer-wise compression performance on AlexNet
Layers Initial
Params Final
Params Sparsity(%)
conv1_1 to conv4_3
(10 layers) 6.7M 6.7M -
conv5_1 2M 2M -
conv5_2 2M 235K 88.2
conv5_3 2M 235K 88.2
fc6 103M 102K 99.9
fc7 17M 167K 99.01
fc8 4M 409K 89.7
Total 138M 9.85M 92.85
Table 4: Layer-wise compression performance on VGG-16

Effect of hyper-parameters

In Section 2.1 we described that we used maximum likelihood sampling (i.e.; thresholding) instead of unbiased sampling from a bernoulli. In these experiments, we shall study the relative effects of hyper-parameters on both methods. In the sampling case, sparsity is difficult to measure as different samples may lead to slightly different sparsities. As a result, we measure expected sparsity as well the it’s variance.

Our methods primarily have the following hyper-parameters: and the initialization for each gate value. As a result, if we have a network with layers, we have hyper-parameters to determine.

First, we analyze the effects of and . We use different combinations of initializations for both and look at it’s effects on accuracy and sparsity. As shown in Table 5, both the thresholding as well as the sparsity-based methods are similarly sensitive to the regularization constants.

In Section 2.3, we saw that roughly controls the variance of the bernoulli variables while penalizes the mean. In Table 5, we see that the mean sparsity for the pair is high, while that for is considerably lower. Also, we note that the variance of is smaller than that of , confirming our hypothesis that controls variance.

Overall, we find that both networks are almost equally sparse, and that they yield very similar accuracies. However, the thresholding-based method is deterministic, which is why we primarily use this method.

To further analyze effects of and , we plot sparsity values attained by our method by fixing one parameter and varying another. In Figure 2(b) we see that , or the variance-controlling hyper-parameter, mainly stabilizes the training by reducing the sparsity levels. In Figure 2(c) we see that increasing increases the sparsity level as expected.

Sparsity (%)
[T] Avg.Sparsity (%)
[S] Variance (%)
[S]
0 0 54.5 53.1 16.1
1 1 98.3 93.7 3.3
1 0 62.1 57.3 5.4
0 1 99.0 92.7 4.1
Table 5: Effect of parameters on sparsity. [T] denotes the threshold-based method, while [S] denotes that sampling-based method.
Method Params (Params + Indices) Top-1 Accuracy(%) Compression Rate
Baseline 60.9M 57.2 -
Neuron Pruning [Srinivas and Babu2015a] 39.6M 55.60 1.5x
SVD-quarter-F [Yang et al.2014] 25.6M 56.18 2.3x
Adaptive FastFood 32 [Yang et al.2014] 22.5M 57.39 2.7x
Adaptive FastFood 16 [Yang et al.2014] 16.4M 57.1 3.7x
ACDC [Moczulski et al.2015] 11.9M 56.73 5x
Collins & Kohli [Collins and Kohli2014] 8.5M (17M) 55.60 7x
Han et al.[Han et al.2015] 6.7M (13.4M) 57.2 9x
Proposed Method 5.9M (11.8M) 56.96 10.3x
Table 6: Comparison of compression performance on AlexNet architecture
Method Params (Params + Indices) Top-1 Accuracy(%) Compression Rate
Baseline 138M 68.97 -
Han et al.[Han et al.2015] 10.3M (20.6M) 68.66 13x
Proposed Method 9.8M (19.6M) 69.04 14x
Table 7: Comparison of compression performance on VGG-16 architecture

We now study the effects of using different initializations for the gate parameters. We initialize all gate parameters of a layer with the same constant value. We also tried stochastic initialization for these gate parameters (Eg. from a Gaussian distribution), but we found no particular advantage in doing so. As shown in Figure

2(a), both methods seem robust to varying initializations, with the thresholding method consistently giving higher sparsities. This robustness to initialization is advantageous to our method, as we no longer need to worry about finding good initial values for them.

Compression Performance

We test compression performance on three different network architectures - LeNet-5, AlexNet [Krizhevsky, Sutskever, and Hinton2012] and VGG-16.

For LeNet-5, we simply sparsify each layer. As shown in Table 1, we are able to remove about 96% of LeNet’s parameters and only suffer a negligible loss in accuracy. Table 2 shows that we obtain state-of-the-art results on LeNet-5 compression. For Proposed Method - 1, we used , while for Proposed Method-2, we used . These choices were made using a validation set.

Note that our method converts a dense matrix to a sparse matrix, so the total number of parameters that need to be stored on disk includes the indices of the parameters. As a result, we report the parameter count along with indices. This is similar to what has been done in [Collins and Kohli2014]. However, for ASIC implementations, one need not store indices as they can be built into the circuit structure.

For AlexNet and VGG-16, instead of training from scratch, we fine-tune the network from pre-trained weights. For such pre-trained weights, we found it be useful to pre-initialize the gate variables so that we do not lose accuracy while starting to fine-tune. Specifically, we ensure that the gate variables corresponding to the top- weights in the matrix are one, while the rest are zeros. We use this pre-initialization instead of the constant initialization described previously.

To help pruning performance, we pre-initialize fully connected gates with very large sparsity () and convolutional layers with very little sparsity. This means that of parameters are zero, and rest are one. For , the underlying gate values were and for , we used . This is to ensure good accuracy by preserving important weights while having large sparsity ratios. The resulting network ended up with a negligible amount of sparsity for convolutional layers and high sparsity for fully connected layers. For VGG-16, we pre-initialize the final two convolutional layers as well with sparsity.

We run fine-tuning on AlexNet for 30k iterations ( hours), and VGG for 40k ( hours) iterations before stopping training based on the combination of compression ratio and validation accuracy. This is in contrast with [Han et al.2015], who take about 173 hours to fine-tune AlexNet. The original AlexNet took 75 hours to train. All wall clock numbers are reported by training on a NVIDIA Titan X GPU. As shown in Table 6 and Table 7, we obtain favourable results when compared to the other network compression / sparsification methods.

Conclusion

We have introduced a novel method to learn neural networks with sparse connections. This can be interpreted as learning weights and performing pruning simultaneously. By introducing a learning-based approach to pruning weights, we are able to obtain the optimal level of sparsity. This enables us to achieve state-of-the-art results on compression of deep neural networks.

References