Enabling Retrain-free Deep Neural Network Pruning Using Surrogate Lagrangian Relaxation

1 Introduction

Deep neural network (DNN)-based statistical models are increasingly taxing of computational and storage resources, with costs proportional to the model size (i.e. the number of parameters in a model). This concern is especially pressing for embedded or IoT devices [Krizhevsky et al., 2012; Simonyan and Zisserman, 2014]. By reducing model size, one reduces both storage costs and computation costs when evaluating a model. Various techniques exist for reducing model size while maintaining performance of the model, e.g. weight pruning, sparsity regularization, quantization, and clustering. These techniques are collectively known as model compression [Dai et al., 2017; Yang et al., 2016; Molchanov et al., 2017; Guo et al., 2016; Tung et al., 2017; Luo et al., 2017; Ding et al., 2017; Park et al., 2017; Frankle and Carbin, 2018; He et al., 2018; Liu et al., 2019].

∗These authors contributed equally.

These works leverage the observation that training a compact model from scratch is more difficult and less performant than retraining a pruned model [Frankle and Carbin, 2018; Liu et al., 2019]. Therefore, a typical three-stage pipeline, i.e., training (large model), pruning and retraining (also called "fine-tuning") is required. The pruning process is to set the redundant weights to zero and keep the important weights to best preserve the accuracy. The retraining process is necessary since the model accuracy dramatically drops after hardpruning. However, this three-stage weight pruning technique significantly increases the overall training cost. For example, although the state-of-the-art weight pruning methods achieve very high compression rate while maintaining the prediction accuracy on many DNN architectures, the retraining process takes a longer time, e.g., 80 epochs for ResNet-18 on ImageNet, which is 70% of the original training epochs using Alternate Direction Method of Multipliers (ADMM) [Zhang et al., 2018; Ren et al., 2019].

Given the pros and cons of current weight pruning-method, this paper aims to answer the following questions: Is there an optimization method that can achieve a high model accuracy even at the hard-pruning stage and can reduce retraining trails significantly? Given a limited budget of retraining epochs, is there an optimization method that can quickly recover the model accuracy (much faster than the state-of-the-art methods)?

The major obstacle when answering these questions is the discrete nature of the model compression problems caused by "cardinality" constraints, which ensure that a certain proportion of weights is pruned.

In this paper, to address this difficulty, we develop a weight-pruning optimization approach based on recent Surrogate Lagrangian relaxation (SLR) [Bragin et al., 2015], which overcomes all major convergence difficulties of standard Lagrangian relaxation. Within the SLR approach, Lagrangian multipliers approach their optimal values much faster as compared to those within other methods (e.g., ADMM). We summarize our contributions/findings as:

• Tailoring the SLR-based approach to overcome difficulties caused by the discrete nature of the weight-pruning problem while ensuring fast convergence.
• Further acceleration of the SLR convergence by using quadratic penalties. The method possesses nice convergence properties inherited from fast accelerated reduction of constraint violation due to quadratic penalties and from the guaranteed convergence thereby leading to unparalleled performance as compared to other methods. Therefore, model parameters obtained by SLR are much closer to their optimal values as compared to those obtained by other state-of-the-art methods.
• Convergence proof of the SLR method for weight pruning problems. Existing coordination-based weight pruning methods do not converge when solving non-convex problems. Other coordination methods (e.g., ADMM) are not designed to handle discrete variables and other types of non-convexities.
• The achievement of high model accuracy even at the hard-pruning stage through our SLR-based weightpruning optimization approach; given a limited budget of retraining epochs, the new method quickly recovers the model accuracy.

2 RelatedWorks onWeight Pruning

Since lots of researchers have investigated that some portion of weights in neural networks are redundant, weight pruning is proposed to remove these less important coefficient values and it achieves model compression with similar performance compared to uncompressed one. Structured and non-structured (irregular) are two mainstream weight pruning methods.

In this work, we focus on irregular pruning which can achieve much higher accuracy than structured pruning [Wen et al., 2016] due to its flexibility in selecting weights.

Early work [Han et al., 2015] proposed an iterative and static magnitude-based weight pruning to explore the redundancy. Later, the compression rate has been further improved by integrating the ADMM [Boyd et al., 2011], an optimization algorithm that breaks optimization problems into smaller subproblems, where dynamic penalty is applied on all targeted weights each of which is then solved iteratively and more easily [Zhang et al., 2018; Niu et al., 2020]. [Louizos et al., 2018] proposed a framework for L0 norm regularization for neural networks, which aim to prune the network during training by choosing weights setting them to exactly zero. Lottery ticket hypothesis was proposed by [Frankle and Carbin, 2018], which observes that a subnetwork of randomly-initialized network can replace the original network with the same performance.

3 Weight Pruning using SLR

Consider a DNN with N layers indexed by n ∈ 1, ...,N, where the weights at layer n are denoted byWn.

The objective is to minimize the loss function subject to constraints on the cardinality of weights within each layer n (the number of nonzero weights should be less than or equal to the predefined number ln): min Wn f (Wn ) +

∑Nn=1 gn(Wn), where the first term represents the nonlinear smooth loss function and the other represents the nondifferentiable "cardinality" penalty term [Zhang et al., 2018] with gn(·) being the indicator function:

In its entirety, the problem cannot be solved either analytically or by using the stochastic gradient descent. To enable the decomposition into smaller manageable subproblems, duplicate variables are introduced and the problem is equivalently rewritten as: min Wn f (Wn ) +

∑Nn=1 gn(Zn), subject to Wn = Zn, n = 1, . . . ,N. To solve the problem, constraints are first relaxed by introducing Lagrangian multipliers and their violations are penalized by using quadratic penalties. The resulting Augmented Lagrangian function [Boyd et al., 2011; Zhang et al., 2018] of the above optimization problem is this given by

where Λn is a matrix of Lagrangian multipliers (dual variables) corresponding to constraints Wn = Zn, and has the same dimension as Wn. The positive scalar ρ is the penalty coefficient, tr(·) denotes the trace, and ‖ · ‖2F denotes the Frobenius norm.

In the following, motivated by decomposibility enabled by SLR [Bragin et al., 2015], which overcame all major difficulties of standard Lagrangian Relaxation, with much alleviated zigzagging and guaranteed convergence, the relaxed problem will be decomposed into two manageable subproblems, and the subproblems will then be efficiently coordinated by Lagrangian multipliers.

Step 1: Solve "Loss Function" Subproblem for Wn by using Stochastic Gradient Decent. At iteration k, for given values of multipliers Λkn , the first "loss function" subproblem is to minimize the Lagrangian function, while keeping Zn at previously obtained values Zk-1 n as

Since the regularizer is a differentiable quadratic norm, and the loss function is differentiable, the subproblem can be solved by stochastic gradient descent (SGD) [Bottou, 2010]. To ensure that multipliers updating directions are "proper," the following "surrogate" optimality condition needs to be satisfied following [Bragin et al., 2015, p. 179, eq. (12)]:

If (3) is satisfied, multipliers are updated following [Bragin et al., 2015, p. 179, eq. (15)] as:

where stepsizes are updated as [Bragin et al., 2015, p. 180, eq. (20)] as

Step 2: Solve "Cardinality" Subproblem for Zn through Pruning by using Projections onto Discrete Subspace. The second "cardinality" subproblem is solved with respect to Zn while fixing other variables at values Wkn as

Since gn(·) is the indicator function, the globally optimal solution of this problem can be explicitly derived as [Boyd et al., 2011]:

where ΠSn (·) denotes the Euclidean projection onto the set Sn = {Wn | card(Wn) ≤ ln}, n = 1, . . . , N.

To ensure that multipliers updating directions are "proper," the following "surrogate" optimality condition needs to be satisfied:

Once (8) is satisfied, 1 multipliers are updated as:

where stepsizes are updated as

where stepsize-setting parameters [Bragin et al., 2015, p. 188, eq. (67)] are:

The algorithm of the new method is presented below:

1 If condition (8) is not satisfied, the subproblems (2) and (6) are solved again by using the latest available values for Wn and Zn.

Algorithm 1 Surrogate Lagrangian Relaxation

• 1: InitializeW0n ,Z0n ,Λ0n and s0
• 2: while Stopping criteria are not satisfied do
• 3: 1 solve subproblem (2),
• 4: if surrogate optimality condition (3) is satisfied then
• 5: keepWkn , Zkn , and update multipliers Λkn per (4) ,
• 6: else
• 7: keep Wkn , Zkn , do not update multipliers Λkn ,
• 8: end if
• 9: 2 solve subproblem (6),
• 10: if surrogate optimality condition (8) is satisfied then
• 11: keepWkn , Zkn , and update multipliers Λkn per (9) ,
• 12: else
• 13: keep Wkn , Zkn , do not update multipliers Λkn ,
• 14: end if
• 15: end while

The theoretical results are summarized below:

Theorem. Sufficient Condition for Convergence of the Method: Assuming for any integer number κ there exists k > κ such that surrogate optimality conditions (3) and (8) are satisfied, then under the stepsizing conditions (5) and (10)-(11), the Lagrangian multipliers converge to their optimal values Λ∗n that maximize the following dual function:

Proof. The proof will be based on that of [Bragin et al., 2015]. The major difference in the original SLR method [Bragin et al., 2015] and the SLR method of this paper is the presence of quadratic terms within the Lagrangian function (1).

It should be noted that the weight pruning problem can be equivalently rewritten in a generic form as:

where X collectively denotes the decision variables {Wn, Zn} and

The feasible set of (13) is equivalent to that of the original model compression problem. Feasibility requires that Wn = Zn, which makes the term ρ2 ‖Wn - Zn‖2F within (13) disappear. Therefore, the Lagrangian function that corresponds to (13) is the Augmented Lagrangian function (1) to the original model compression problem. Furthermore, the surrogate optimality conditions (3) and (8) are the surrogate optimality conditions that correspond to the Lagrangian function F(X) +ΛG(X) that corresponds to the problem (13). Therefore, since within the original SLR [Bragin et al., 2015, Prop. 2.7, p. 188] convergence was proved under conditions on stepsizes (5) and (10)-(11) and the satisfaction of surrogate optimality conditions, both of which are satisfied here, multipliers converge to their optimal values for the model compression under consideration as well.

This is a major result in machine learning since this is the first application of the SLR method to guarantee theoretical convergence of the model compression problem while handling discrete variables as well as the quadratic terms. In fact, owing to quadratic terms, the method inherits nice convergence properties similar to those of Augmented Lagrangian Relaxation (ALR) (fast reduction of constraint violations) and to those of SLR (fast and guaranteed convergence without much of zigzagging of multipliers and the need of the socalled "optimal dual value" [Bragin et al., 2015]) thereby leading to unparalleled performance as compared to other methods.

The SLR method enjoys the benefits of efficient subproblem solution coordination with guaranteed convergence made possible by stepsizes (5) and (10)-(11) approaching zero (without this requirement, multipliers (9) would not exhibit convergence), and by the satisfaction of surrogate surrogate optimality conditions (3) and (8) ensuring that multipliers are updated along "good" directions. In Section 4, it will be verified empirically that there always exists iteration κ after which the "surrogate" optimality conditions are satisfied thereby ensuring that the multipliers approach their optimal values during the entire iterative process.

The SLR method also benefits from the independent and systematic adjustment of two hyper-parameters: penalty coefficient ρ and the stepsize sk . In contrast, other coordination methods are not designed to handle discrete variables and other types of non-convexities. For example, ADMM do not converge when solving non-convex problems [Boyd et al., 2011, p. 73] because stepsizes ρ within the method does not converge to zero. Lowering stepsizes to zero within ADMM would also lead to decrease of the penalty coefficient, thereby leading to slower convergence.

4 Evaluation

In this section, we discuss our experimental results for image classification and object detection tasks.

4.1 Experimental Setup

All of the baseline models we use and our code in image classification tasks are implemented with PyTorch 1.6.0 and Python 3.6. For our experiments on COCO 2014 dataset, we used Torch v1.6.0, pycocotools v2.0 packages. For our experiments on TuSimple lane detection benchmark dataset 2 , we used Python 3.7 with Torch v1.6.0, and SpConv v1.2 package. We conducted our experiments on Ubuntu 18.04 and using Nvidia Quadro RTX 6000 GPU with 24 GB GPU memory. We used 4 GPU nodes to train our models on the ImageNet dataset.

We start by pruning the pretrained models through SLR training. Afterwards, we perform hard-pruning on the model, completing the compression phase. We report the overall compression rate (or the percentage of remaining weights) and prediction accuracy.

2 https://github.com/TuSimple/tusimple-benchmark

4.2 Evaluation on Image Classification Tasks

Models and Datasets. We used ResNet-18, ResNet-50, ResNet-56, ResNet-110 [He et al., 2016] and VGG-16 [Simonyan and Zisserman, 2014] on CIFAR-10. On ImageNet ILSVRC 2012 benchmark, we use ResNet-18, ResNet-50 [He et al., 2016] and MobileNetV2 [Sandler et al., 2018]. We use the pretrained ResNet models on ImageNet from Torchvision's "models" subpackage. The accuracy of the pretrained baseline models we used are listed in Table 1.

Training Settings. In all experiments we used ρ = 0.1. In CIFAR-10 experiments, we used a learning rate of 0.01, batch size of 128 and ADAM optimizer during training. On ImageNet, we used a learning rate of 10-4 , batch size of 256 and SGD optimizer. For a fair comparison of SLR and ADMM methods, we used the same number of training epochs and sparsity configuration for both methods in the experiments.

Table 1 shows our comparison of SLR and ADMM on CIFAR-10 and ImageNet benchmark. Here, SLR parameters are set as M = 300, r = 0.1 and s0 = 10-2 . After SLR and ADMM training, final hardpruning is performed and the hardpruning accuracy is reported without any additional retraining, given a limited budget of training epochs. According to our results, SLR outperforms ADMM method in terms of accuracy under the same compression rate. With higher compression rates, CIFAR-10 results show higher gap in accuracy between SLR and ADMM.

Figure 1 shows the hardpruning accuracy during SLR vs. ADMM on CIFAR-10 and ImageNet, corresponding to Table 1. During training, hardpruning accuracy is checked periodically. If the hardpruning accuracy meets the accuracy criteria, the training is stopped. As seen in the figures, SLR quickly converges and reaches the desired accuracy, almost 3× faster than ADMM on CIFAR-10. Moreover, in Figure 1c, ADMM is still below the desired accuracy even after 300 epochs of training on VGG-16, while SLR completes training in 80 epochs. Similarly, as can be seen in Figures 1f and 1g, ADMM cannot output the desired accuracy after 60 and 50 epochs of training on ImageNet, while SLR reaches the threshold quickly.

Table 2 shows our comparison of SLR with other recent model compression works on CIFAR-10 benchmark. We report the percentage of parameters pruned after SLR training

• (d) ResNet-56 on CIFAR-10.
• (e) ResNet-110 on CIFAR-10.
• (f) ResNet-18 on ImageNet.
• (g) ResNet-50 on ImageNet.

10 15 20 25 30 35 40 45 50

Table 2: SLR performance comparison with VGG-16, ResNet-18, ResNet-50 and ResNet-56 on CIFAR-10.

and the final accuracy. We start training the networks with a learning rate of 0.1 decreasing the learning rate by a factor of 10 at epochs 80 and 100. On ResNet-18, we compare our result at only 20k iterations. For VGG-16 and ResNet-50, we observe that SLR can achieve up to 60% pruning with less than 1% accuracy drop.

4.3 Evaluation on Object Detection Tasks

Models and Datasets. We used YOLOv3 and YOLOv3tiny models [Redmon and Farhadi, 2018] on COCO 2014 benchmark. We used and followed the publicly available Ultralytics repository 3 for YOLOv3 and its pretrained models. For lane detection experiment, we used the pretrained model from Ultra-Fast-Lane-Detection [Qin et al., 2020] on TuSimple lane detection benchmark dataset.

Training Settings. In all experiments we used ρ = 0.1. We set SLR parameters as M = 300, r = 0.1 and s0 = 10-2 . We follow the same training settings provided by the repositories we use. Finally, we use the same number of training epochs and sparsity configuration for ADMM and SLR.

Testing Settings. On YOLOv3 models, we calculate the COCO mAP with IoU = 0.50 with image size of 640 for testing. On lane detection experiments, evaluation metric is "accuracy", which is calculated as ∑clip Cclip ∑clip Sclip , where Cclip is the number of lane points predicted correctly and Sclip is the total number of ground truth in each clip.

Our comparison of SLR and ADMM methods on COCO dataset is shown in Table 3. We have compared the two methods under 3 different compression rates for YOLOv3-tiny and tested YOLOv3-SPP pretrained model with a compression rate of 1.98×. We can see that the model pruned with SLR method has higher accuracy after hard-pruning in all cases.

3 https://github.com/ultralytics/yolov3

At a glance at YOLOv3-tiny results, we observe that the advantage of SLR is higher with an increased compression rate.

In compression rate 3.33× on YOLOv3-tiny, given a limit of 25 epochs, we can observe that the gap between ADMM and SLR is much higher, which is the due to the faster convergence of SLR as shown in Figure 2b. Similarly, Figure 2a shows the mAP progress of YOLOv3 during SLR and ADMM training for 50 epochs, pruned with 2× compression. SLR reaches the mAP threshold only at epoch 15.

Table 4 reports our result for the Lane Detection task on TuSimple lane detection benchmark after 40 epochs of training and 5 epochs of masked-retraining. We conducted experiments under 8 different compression rates. Similarly, Figure 3 illustrates the accuracy gap between ADMM and SLR methods after hard pruning as compression rate increases.

From Figure 3, our observation is that for a small compression rate such as 1.82×, SLR has little advantage over ADMM in terms of hardpruning accuracy. However, as the compression rate increases, SLR starts to perform better than ADMM. For example, SLR survives 77.67× compression with slight accuracy degration and results in 89.72% accuracy after hardpruning while ADMM accuracy drops to 46.09%. This demonstrates that our SLR-based training method has a greater advantage over ADMM especially in higher compression rates, as it achieves compression with less accuracy loss and reduces the time required to retrain after hard-pruning.

Finally, in Figure 4, we show the difference between weights of one layer before and after pruning with SLR and ADMM. In Figure 4a, we show the initial (non-pruned) weights and then show the sparsity of weights under the same compression rate (77×) with SLR and ADMM. Initially, the layer has low sparsity. After training with SLR and ADMM, we can see an increased number of zeroed-weights. SLR moves towards the desired sparsity level faster than ADMM. In Figure 4b, we compare the sparsity of weights under the same accuracy (89.0%). It can be observed that SLR significantly reduced the number of non-zero weights and ADMM has more non-zero weights remaining compared with SLR.

4.4 Ablation Studies

We conducted several experiments to observe SLR behavior with respect to SLR parameters ρ, s0, r and M on ResNet18 model (93.33% accuracy) and CIFAR-10. We pruned the model through SLR training for 50 epochs with a compression rate of 8.71× and observed the hardpruning accuracy every 10 epochs. Figure 5 shows the accuracy of the model through SLR training based on the different values of s0, M and r. Based on the hardpruning accuracy throughout training, it can be seen that, even though the parameters do not have a great impact on the end result, choice of s0 can impact the convergence of the model. From Figure 5a, we can state that s0 = 10-2 provides higher starting accuracy and converges quickly. Figure 5b and Figure 5c demonstrate the impact of M and r on the hardpruning accuracy respectively. Figure 6 demonstrates that there exists iteration κ (as required in the Theorem) so that the surrogate optimality condition, the high-level convergence criterion of the SLR method, is satisfied during training with s0 = 10-2 , ρ = 0.1 thereby signifying that "good" multiplier-updating directions are always found. For example, after the conditions are violated at epoch 9, there exits κ = 10 so that at iteration 11, after κ = 10, the surrogate conditions are satisfied again.

Figure 5: Hardpruning accuracy of ResNet-18 on CIFAR-10 during SLR training with respect to different values of s0, M and r.

5 Conclusions

In this paper, we presented the DNN weight-pruning problem as a non-convex optimization problem by adopting the cardinality function to induce weight sparsity. By using the SLR method, the relaxed weight-pruning problem is decomposed into subproblems, which are then efficiently coordinated by updating Lagrangian multipliers, resulting in fast convergence. We conducted weight-pruning experiments on image classification and object detection tasks to compare our SLR method against ADMM. We observed that SLR has a significant advantage over ADMM under high compression rates and achieves higher accuracy during weight pruning. SLR reduces the accuracy loss caused by the hard-pruning and so shortens the retraining process. With the effective optimization capabilities through coordination with clear advantages shown from several examples, the SLR method has a strong potential for more general DNN-training applications.

Acknowledgments

This research was supported by the National Science Foundation under Grants S&AS-1849246 and ECCS-1810108.