Training Sparse Neural Networks

Introduction

For large-scale tasks such as image classification, large networks with many millions of parameters are often used (?), (?), (?). 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 ($\mathcal{O}(mn)$ to $\mathcal{O}(k)$)For a matrix of size $m\times n$ with $k$ non-zero elements, sparse computations also decrease feedforward evaluation time ($\mathcal{O}(mnp)$ to $\mathcal{O}(kp)$)For matrix-vector multiplies with a dense vector of size $p$. Further, having a lower parameter count may help in avoiding overfitting.

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 $\Phi=\{g^{s}_{1},g^{s}_{2},...,g^{s}_{m}\}$ be a set of $m$ vectors. This represents an $m$-layer dense neural network architecture where $g^{s}_{i}$ is a vector of parameter indices for the $i^{th}$ layer, i.e; $g^{s}_{i}=\{0,1\}^{n_{i}}$. Here, each layer $g^{s}_{i}$ contains $n_{i}$ elements. Zero indicates absence of a parameter and one indicates presence. Thus, for a dense neural network, $g_{i}$ is a vector of all ones, i.e.; $g^{s}_{i}=\{1\}^{n_{i}}$. For a sparse parameter vector, $g^{s}_{i}$ would consist of mostly zeros. Let us call $\Phi$ 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.

The complexity of a $m$-layer neural network with index set $\Phi$ is given by $\|\Phi\|=\sum\limits_{i=1}^{m}n_{i}$.

We now aim to solve the following optimization problem.

How do we incorporate the index set formalism in neural networks? Assume that the index set ($G^{s}$ 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 ?

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 $G$ at $0.5$. To ensure determinism, we use the ML draw or thresholding in this work.

Promoting Sparsity

To this end, we use a regularizer given by $w\times(1-w)$. This was introduced by (?) 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. 2(a)

where $g_{i,j}$ denotes the $j^{th}$ gate parameter in the $i^{th}$ layer. Note that for $g_{i,j}\in\{0,1\}$, the second term in Eqn. 2 vanishes and the third term becomes $\lambda\|\Phi\|$, 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.

where $g^{s}$ is the sampled version of gate variables $g$. Note that Eqn.3 is a stochastic objective function, arising from the fact that $g^{s}$ 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.

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 (?). 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.

Here, $\delta(\cdot)$ denotes the dirac delta distribution, and $Z$ denotes the normalizing constant, and $\alpha$ is the mixture coefficient. Also note that like (?), we assume that $w_{i}\in[-k,k]$ for some $k>0$. This is visualized in Fig. 2(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

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. (?) 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 $g^{s}\sim bernoulli(g)$, then the gradient $\frac{dg^{s}}{dg}=1$ is used.

Another issue of consideration is that of ensuring that $g$ always lies in $$ so that it is a valid bernoulli parameter. Bengio et al. (?) 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 $g^{s}\sim bernoulli(clip(g))$, 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.

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.(?) and Hassibi et al.(?), who introduced Optimal Brain Damage and Optimal Brain Surgery respectively. Recently, Srinivas and Babu (?) 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.(?) create sparse networks by alternating between weight pruning and network training. A similar strategy is followed by Collins and Kohli (?). 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 ($\lambda_{1},\lambda_{2}$) compared to the other methods, which have $n$ thresholds to be set for each of the $n$ 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.(?) 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.(?) 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.(?) use tensor decompositions to obtain a factorization of tensors with small number of parameters. Cheng et al.(?) 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.(?) 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 (?), 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.

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 (?). 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.

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 $5\times 5$ 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.

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 3(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.

We test compression performance on three different network architectures - LeNet-5, AlexNet (?) and VGG-16.

For LeNet-5, we simply sparsify each layer. As shown in Table Related Work, 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 $(\lambda_{1},\lambda_{2})=(0.001,0.05)$, while for Proposed Method-2, we used $(\lambda_{1},\lambda_{2})=(0.01,0.1)$. 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 (?). 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-$k\%$ weights in the $W$ 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 ($95\%$) and convolutional layers with very little sparsity. This means that $95\%$ of $g^{s}$ parameters are zero, and rest are one. For $g^{s}=1$, the underlying gate values were $g=1$ and for $g^{s}=0$, we used $g=0.49$. 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 $88\%$ sparsity.

We run fine-tuning on AlexNet for 30k iterations ($\sim 18$ hours), and VGG for 40k ($\sim 24$ hours) iterations before stopping training based on the combination of compression ratio and validation accuracy. This is in contrast with (?), 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.

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.