Optimization based Layer-wise Magnitude-based Pruning for DNN Compression

1 Introduction

Deep neural networks (DNNs) have achieved outstanding performance in recent years [Krizhevsky et al., 2012; LeCun et al., 2015]. However, DNNs usually suffer from their enormous number of parameters, which makes them prohibitive to be deployed on platforms with limited memory and processing units, e.g., mobile phones [Kim et al., 2016]. To overcome this obstacle, a variety of approaches have been developed to compress those trained DNNs with overparameterized structures [Denil et al., 2013; Han et al., 2015; Molchanov et al., 2017].

∗This work was supported in part by the National Key Research and Development Program of China (2017YFB1003102), the NSFC (61603367, 61672478), the YESS (2016QNRC001), and the Science and Technology Innovation Committee Foundation of Shenzhen (ZDSYS201703031748284).

Layer-wise magnitude-based pruning (LMP) is an effective DNN compression method and has achieved significant results in many applications [Han et al., 2015; Guo et al., 2016; See et al., 2016]. The idea is to prune connections in each layer separately by removing the connections with absolute weight values lower than a layer-specific threshold. Given a threshold for each layer, LMP can prune connections in parallel, which is especially useful for DNNs with millions or billions connections. With well chosen pruning thresholds, LMP can achieve a significant reduction in the number of parameters, while maintaining a relatively low accuracy loss.

A main difficulty in applying LMP is tuning the pruning thresholds. Such a threshold tuning problem is nontrivial for two reasons. First, the solution space of all threshold combinations can be very large even for a moderatesized DNN. Suppose a DNN has L layers and each layer has W connections, then the possible combinations of all layer-specific thresholds will be of size (W + 1)L. Second, the evaluation of each candidate threshold combination is very time consuming, because it needs to evaluate the performance loss of the pruned model over the training set, e.g., the accuracy loss after retraining. The commonly used approach to tune the thresholds is by hand [Guo et al., 2016; Han et al., 2015], which highly relies on expertise. In order to facilitate non-professional users to use LMP, an automatic threshold adjusting method is required.

It is an appealing idea to tune the thresholds of LMP from an optimization point of view. However, the performance measure of the thresholds is usually a discontinuous and nondifferentiable function, which prevents the direct use of offthe-shelf optimization techniques such as gradient descent methods. An alternative way is to use derivative-free optimization methods [Goldberg, 1989; Brochu et al., 2010; Qian et al., 2015; Yu et al., 2016], which do not require the problem to be either continuous or differentiable.

In this paper, we propose an optimization based approach, namely Optimization based LMP (OLMP), to automatically tune the pruning thresholds for LMP. Concretely, the threshold tuning problem is formulated as a constrained optimization problem, which requires minimizing the size (i.e., the number of connections) of the pruned network, subject to a constraint on the accuracy loss. Then, a powerful derivativefree optimization algorithm is employed to solve this problem. To deal with the costly evaluation of accuracy loss, OLMP evaluates the accuracy loss of the pruned model on a randomly sampled small data set instead of the whole data set. Note that to compress a trained DNN, OLMP is conducted within a new iterative pruning and adjusting pipeline.

To empirically evaluate the performance of OLMP, a recently proposed powerful derivative-free optimization algorithm, negatively correlated search (NCS) [Tang et al., 2016], is adopted. Without incurring any accuracy loss on test data, OLMP can prune 99.66% parameters of LeNet-5 and 99.12% parameters of LeNet-300-100, which are the best results in comparison to a number of state-of-the-art approaches [Han et al., 2015; Guo et al., 2016; Ullrich et al., 2017; Molchanov et al., 2017]. We also apply OLMP to prune an AlexNet-style deep model and 98.78% parameters can be removed without sacrificing accuracy.

The rest of the paper is organized as follows. Section 2 presents the related work. Section 3 introduces the proposed OLMP algorithm. Section 4 presents the empirical studies. Section 5 finally concludes the paper.

2 RelatedWork

Pruning methods [LeCun et al., 1989; Hassibi and Stork, 1992] are widely used for compressing trained DNNs. They usually use some metrics to measure the saliency of parameters and then set thresholds to abandon parameters with low saliency. There are typically two types of pruning methods: neuron pruning [He et al., 2014] and connection pruning [Hassibi and Stork, 1992; LeCun et al., 1989]. This paper focuses on connection pruning.

LMP extends magnitude-based pruning (MP) [LeCun et al., 1989] with respect to the multi-layer nature of DNNs. Given a set of connections and a positive threshold, MP prunes connections with absolute weight values smaller than the threshold. Considering the non-linear connected nature of DNNs, LMP applies MP with different thresholds to different layers. Popular LMP methods always tune the thresholds by experience, e.g., iterative pruning and retraining (ITR) [Han et al., 2015], dynamic surgery (DS) [Guo et al., 2016] and layer-wise optimal brain surgeon [Dong et al., 2017].

There also exist some connection pruning methods without the need of tuning layer-specific thresholds. SWS [Ullrich et al., 2017] extends the idea of soft weight-sharing [Nowlan and Hinton, 1992] to DNN compression. It fits a mixture of multiple Gaussian models over the weights with one Gaussian model fixed to zero mean, and then prunes the weights which are under the zero-mean distribution. Sparse VD [Molchanov et al., 2017] uses a technique called variational dropout, which provides Bayesian interpretation to Gaussian dropout. It assigns an individual dropout rate to each weight and explores all possible values of dropout rates to find a sparse model.

Pruning followed by adjusting is a popular approach to getting better performance than applying pruning alone. For example, ITR and DS both use an additional model adjusting phase after pruning. DS first prunes the network drastically based on a variant of LMP, and then retrains the pruned model and randomly recovers the incorrectly removed connections by an operation called splicing. ITR iteratively applies LMP and then retrains the pruned model until the model cannot retrieve the original accuracy. Similar ideas of ITR can also be found in [Castellano et al., 1997].

The proposed OLMP adopts a heuristic optimization algorithm NCS [Tang et al., 2016] to optimize the pruning thresholds. Heuristic optimization algorithms have been used to evolve the network structure [Yao and Liu, 1997], and to prune the connections [Reed, 1993]. However, most of them are effective only for shallow networks, and are computational impossible for DNNs with millions of parameters. Recently, Google Research [Real et al., 2017] successfully applied genetic algorithms to design the structures of DNNs by using clusters of GPU servers.

3 The Proposed Approach

In this section, we introduce OLMP, the whole framework of DNN compression with OLMP, and a realistic implementation of OLMP, respectively.

3.1 OLMP

Suppose a trained DNN model with L + 1 layers can be represented as W = {W l i,j | W l i,j ̸= 0, 1 ≤ l ≤ L, 1 ≤ i ≤ nl , 1 ≤ j ≤ nl+1 }, where W l i,j denotes the connection weight between the ith neuron in layer l and the jth neuron in layer l+1, and nl denotes the number of neurons in layer l. Note that W l i,j = 0 indicates that the corresponding connection does not exist. Let f(W) denote the evaluation function which measures the accuracy of model W, and let δ denote the userdefined tolerance on accuracy loss. For W, the task of finding the best pruned model with respect to δ can be formulated as

Note that W ′ is a subset of W (namely a pruned model), |W ′ | denotes the size of W ′ , and W ∗ denotes the best pruned model with the smallest size while satisfying the constraint.

In LMP, the generation of the pruned model W ′ is decided by the pruning thresholds at each layer. That is, the pruning in layer l of W is described by MP(W, l, εl) = {W l i,j | |W l i,j | ≥ εl, 1 ≤ i ≤ nl , 1 ≤ j ≤ nl+1 }, where |W l i,j | denotes the absolute value of W l i,j and εl denotes the pruning threshold at layer l. For many LMP methods, εl is always determined by a heuristic function g which contains hyper-parameters to be tuned. Here, we only consider one hyper-parameter cl to be tuned for each layer l, denoted as εl = g(cl). Therefore, the pruned model by LMP can be represented as follows:

where c = (c1, c2, . . . , cL). By combining Eqs. (1) and (2), the optimization based LMP (OLMP) can be formulated as:

The formulated problem Eq. (3) is usually difficult, which can be discontinuous and non-differentiable. Thus, we need to employ a powerful derivative-free optimization algorithm to solve it. After getting the output hyper-parameter vector c∗ , we can get a corresponding pruned model LMP(W, c∗ ). Note that OLMP can be used in DNN compression together with other techniques, like accuracy adjustment.

3.2 DNNCompression with OLMP

To apply OLMP in DNN compression, we use a new iterative pruning and adjusting pipeline, as shown in the right part of Figure 1. Suppose the whole data set D contains |D| batches of data, and is split into |D|/K small sets, where each set contains K batches. In each iteration of the pipeline, one split data set is selected for pruning and adjusting. OLMP is first used to prune the reference model based on one batch extracted from the selected split data set, and then the pruned model is adjusted on the rest batches of the split data set to recover its accuracy. Note that the adjusting phase utilizes retraining with splicing to recover the incorrect pruned connections just the same as DS does [Guo et al., 2016]. The adjusting will last one epoch on the K - 1 batches of the split data set and the refined model will then be treated as a new reference model for pruning in the next loop. This procedure will be repeated until a limit number of loops (denoted as pruning loops) is reached. Finally, the pruned model will be retrained with splicing on the whole data until it converges.

One thing should be noted here is how to combine splicing into the pipeline. In DS, for a layer l, there are two parameters (al, bl) where al is the pruning threshold and bl controls the splicing operation. Through an investigation of the released

θl and σl are the mean of absolute values and standard deviation of weights in layer l, and cl is a layer-specific hyperparameter. We use the same value for the pruning threshold g(cl) in OLMP and the pruning threshold al in splicing, i.e., g(cl) = al. Thus, the hyper-parameter tuning of splicing is unified into the optimization of Eq. (3).

The iterative pipeline used in the framework is a combination of ITR and DS with the layer-wise hyper-parameters tuned by OLMP. Instead of using the whole data set in adjusting of each iteration (e.g., ITR), the proposed pipeline only uses a fraction of data for speeding up the compression. By adopting different optimization algorithms for OLMP or selecting different values for hyper-parameters (i.e., δ, K and pruning loops), this framework of DNN compression can have different implementations.

3.3 OLMP Implementation

In this subsection, we are to provide an effective implementation of the proposed DNN compression framework by using a powerful derivative-free optimization algorithm for OLMP and giving some guidelines on hyper-parameter tuning.

Here, an instantiated OLMP method based on NCS [Tang et al., 2016] is established. NCS is a population-based heuristic optimization algorithm which achieves the state-of-theart performance on many multimodal optimization problems, and does not require derivative information. It uses negative correlations to increase diversities among solutions and to encourage them to search different areas of the solution space.

To use NCS to solve the problem Eq. (3), a fitness function is needed to evaluate the qualities of solutions. A simple and direct choice is to maximize the following function

infeasible solutions (i.e., the constraint is violated) is assigned by the worst value 0. Note that maximizing |W\W ′ | |W | is equivalent to minimizing |W ′ | in Eq. (3). However, such a fitness function is usually difficult to be solved. Taking pruning LeNet-300-100 as an example, Figure 2 shows the landscape of solutions, when iterative pruning and retraining is applied. We can observe that the optimal solution is surrounded by some "plateau" structures on the landscape, which cannot provide useful searching information. Thus, we propose to minimize the fitness function Eq. (4), which prefers infeasible solutions with less constraint violation degree instead of directly assigning them by the same worst value.

Figure 3 shows the negative fitness landscape of Eq. (4), which becomes less flat and will be easier to be optimized.

The procedure of OLMP with NCS is shown in the left of Figure 1. A reference model W, a user defined δ and a sampled data set are packed into a fitness evaluation module. In each iteration, for newly generated alternatives of c, the evaluation module prunes the reference model by LMP(W, c) and returns the fitness values; NCS then searches for new solutions based on these solutions and their fitness values.

As NCS is adopted, its hyper-parameters are involved into the compression process. Although NCS contains adaptive settings for its parameters, we find some manual selections for initial parameter values can speed up the search process. The parameters of NCS that we tune are σ, popN and Tmax where σ decides the stride in searching solutions, popN decides the number of search directions and Tmax decides the epochs of NCS. Therefore, the overall hyper-parameters consist of those used by the proposed framework (i.e., δ,K and pruning loops) and those introduced by the employed optimization algorithms (e.g., σ, popN and Tmax for NCS).

In practice, we can first fix δ and σ to a small value (e.g., 0.1) and tune other parameters. The sequence of tuning was K → pruning loops → Tmax → popN. For each parameter, we can use grid search to get a value at which the performance stops increasing. Finally, we can tune δ and σ based on the sensitivity analysis, as provided in Section 4.3.

4 Experiments

In this section, we empirically investigate the performance of OLMP. Firstly, we compare OLMP with some state-of-theart connection pruning methods on LeNet-5 and LeNet-300100 [Han et al., 2015], which are commonly used by all comparative methods. Secondly, OLMP is applied to compress a real deep network, i.e., an AlexNet [Krizhevsky et al., 2012] style model called AlexNet-Caltech. Thirdly, we perform a sensitivity analysis of hyper-parameters. Finally, we study how OLMP performs without the proposed iterative pipeline.

For the data sets and models, LeNet-5 and LeNet-300-100 are trained on MNIST, and AlexNet-Caltech is trained on Caltech-256 [Griffin et al., 2006]. Note that the output dimension of AlexNet-Caltech is 257 instead of 1000 in AlexNet, since Caltech-256 contains 257 classes (one class for unrecognized images); and this is the only difference in structures compared with original AlexNet.

In the context of all experiments, pruning ratio (PR) indicates |W| |W ′ | , i.e., the ratio between the number of connections of the reference model and the pruned model. All of the experiments are based on Caffe [Jia et al., 2014] and released projects of DS and NCS, and run on a workstation with one Titan X pascal and dual Intel E5-2683 v3@2.0 GHz CPUs.

4.1 Comparison on LeNet-Style Models

In this subsection, we are to show that OLMP can achieve the best performance on DNN compression. We use the reference modes of LeNet-300-100 and LeNet-5 that were used in testing DS [Guo et al., 2016]. Besides, SGD with the same experimental settings (e.g., base learning rate and batch size) in training these reference models is used to retrain the pruned models in the adjusting phase of OLMP. The final pruned models need to be retrained with splicing for additional 15, 000 iterations to converge. In this experiment, we set the hyper-parameters (K, pruning loops, popN, Tmax) to (1000, 15, 10, 160). For LeNet-300-100, (δ, σ) is set to (8%, 5); for LeNet-5, δ is set to 5%, and σ is set to 5 in the first 10 pruning loops and 0.5 since then.

The results are shown in Table 1. We can observe that OLMP achieves the best PRs on LeNet-300-100 and LeNet5 without degrading test accuracy. Note that ITR [Han et al., 2015] and DS [Guo et al., 2016] are two LMP methods which manually set pruning thresholds; and SWS [Ullrich et al., 2017] and Sparse VD [Molchanov et al., 2017] are two non-LMP connection pruning methods.

4.2 Application to AlexNet

In this subsection, we apply OLMP to prune a real deep network, AlexNet-Caltech with 58.33 million parameters. The reference model is trained from scratch for 10, 000 iterations using SGD which mainly follows the experimental settings in [Zeiler and Fergus, 2014] except that the batch size is 256 here; and the same settings of SGD are also used during retraining. The reference model achieves a validation accuracy of 40.18% which is similar to [Zeiler and Fergus, 2014] (i.e., 38.8 ± 1.4%). The final pruned model needs to be retrained with splicing for additional 5, 000 iterations to converge. (K, pruning loops, popN, Tmax, δ) is set to (250, 40 , 8, 200, 8%), and σ is set to 5 in the first 28 pruning loops and 0.5 since then.

The results are shown in Table 2. We can see that OLMP can achieve the PR value of 82 without accuracy loss. We also give the the reduction in each layer, as shown in Table 3. We can observe that the percentages of remaining parameters in the first two layers and the last layer are much larger than that in other layers, which suggests that these three layers are less redundant than other layers.

For the computational cost, it takes 108 minutes to train the reference model, and 605 minutes to compress (433 minutes for OLMP and 172 minutes for adjusting), which is 5.6 (i.e., 605/108) times than training the reference model. Note that NCS can be easily implemented in parallel, and the cost of NCS can be reduced by a factor of popN, i.e., the cost of 433 minutes can be reduced to 54 (i.e., 433/8) minutes if enough GPUs are available.

4.3 Analysis of Sensitivities to δ and σ

In this subsection, we investigate the impact of (δ, σ) on the performance of OLMP. LeNet-300-100 with the same settings as in subsection 4.1 is used as an example. We use δ ∈ {1%, 6%, 11%, . . . , 51%} and σ ∈ {0.05, 5.05, 10.05, . . . , 50.05}, and prune the model with each possible setting of (δ, σ). There are totally 121 settings of (δ, σ). For each setting, we independently run OLMP 30 times and record the number of pruned models which can recover accuracy (i.e., no accuracy loss compared with the reference model). Meanwhile, for those models which can recover accuracy, we also record their average PR. The results are shown in Table 4.

From the distribution of bold values (i.e., the largest PR in each row) in Table 4, the PRs keep to a small value (i.e., 6) and the pruned model can always restore accuracy, when σ is very low (i.e., σ = 0.05). When σ is not very low (i.e., σ ≥ 5.05), we can see that a larger δ value will possibly lead to a larger PR value. This is expected since a larger δ value implies a larger feasible solution space and thus a better optimal objective function value, i.e., a larger PR value. However, a large δ may also make the pruned model difficult to recover accuracy, which means that users need to rerun OLMP many times for finding a pruned model without accuracy loss. Thus, a large δ is prohibitive for DNNs with expensive computational cost. When δ is not very large (i.e., δ ≤ 36%), the relationship between the performance of OLMP and σ is relatively stable for different δ values. That is, as σ increases, the PR tends to decrease while the number of pruned models without accuracy loss keeps within a certain range.

The above observations give us some practical guidance on tuning δ and σ. First, a suitable σ can be selected with a small fixed δ value, since the pattern between σ and the performance of OLMP is stable for small δ values. Then, we can carefully select δ for a large PR. Thus, combining with the analysis in subsection 3.3, we have shown how to tune the hyper-parameters of OLMP.

4.4 OLMP without Iterative Pruning

In this subsection, we test OLMP without iterative pruning and adjusting, to show the process of optimization in a single pruning. We also compare it with the baseline algorithm MP [LeCun et al., 1989] to investigate whether optimization really brings some benefit.

Different from previous experiments, all the models used here are trained on a training set from scratch, pruned on a validation set and tested on an unseen test set. Since MNIST and Caltech-256 only have two separate data sets, a manual division is applied. For MNIST, a validation set containing 6, 000 samples is selected from the training set (60, 000 in total) uniformly at random; the remaining samples form the new training set; the test set is untouched (10, 000 in total). For Caltech-256, the training set is untouched (15, 420 in total) while the validation set (15, 187 in total) is uniformly randomly divided into a new validation set (7, 530 in total) and a new test set (7, 657 in total). All the settings of training the reference models are the same as those used in subsections 4.1 and 4.2.

During pruning, δ is set to 1%. MP tests 100 thresholds which prune 0%, 1%, . . . , 99% parameters of the reference model, respectively, and then selects the pruned model with the largest PR among those without accuracy loss. For OLMP without iterative pruning and adjusting, (popN, σ, Tmax) is set to (4, 5, 400). Figure 4 shows the optimization process of OLMP which involves both the model accuracy and the

-fitness value (i.e., the value of eval in Eq. (4)) of the best solution (i.e., the solution with the smallest fitness value) found in each iteration of NCS. Table 5 summarizes the results.

From the right plot in Figure 4, we can observe that OLMP can gradually improve the solution quality. To be specific, in the early stage, the found solutions always fit the accuracy constraint and the fitness values continue to increase. In the later stage, the trend of the solution quality is still becoming better, while the solutions can suddenly become worse (i.e., the curve decreases) sometimes. The reason might be that, at the early stage of pruning, the model has a lot of redundant connections and thus better solutions are easily generated; while as the optimization goes forward, the model is less redundant and it becomes easier to prune some important connections that will dramatically decrease the fitness value.

From Table 5, we can see that OLMP without iterative pruning and adjusting always achieves better PRs than MP, which implies that introducing optimization can lead to better performance. For the accuracy loss, neither OLMP nor MP has significant larger loss in testing than in validation, which shows that the pruned models on the validation data set can generalize well on the unseen test set.

5 Conclusion

State-of-the-art LMP based DNN compression methods require to manually select the pruning thresholds on each layer, and thus are hard to be applied in practice, particularly for end-users with limited expertise. This paper proposes OLMP, an optimization based LMP method which can choose thresholds automatically. OLMP first formulates the task of threshold selection as a constrained optimization problem, i.e., maximizing the pruning ratio under a constraint on the accuracy loss, and then uses a derivative-free optimization algorithm to solve the problem. To compress DNNs exhaustively, OLMP is applied within a new iterative pruning and adjusting pipeline. Empirical results show the superiority of OLMP over several state-of-the-art connection pruning methods.