Learning to Prune Deep Neural Networks via Layer-wise Optimal Brain Surgeon
How to develop slim and accurate deep neural networks has become crucial for realworld applications, especially for those employed in embedded systems. Though previous work along this research line has shown some promising results, most existing methods either fail to significantly compress a well-trained deep network or require a heavy retraining process for the pruned deep network to re-boost its prediction performance. In this paper, we propose a new layer-wise pruning method for deep neural networks. In our proposed method, parameters of each individual layer are pruned independently based on second order derivatives of a layer-wise error function with respect to the corresponding parameters. We prove that the final prediction performance drop after pruning is bounded by a linear combination of the reconstructed errors caused at each layer. By controlling layer-wise errors properly, one only needs to perform a light retraining process on the pruned network to resume its original prediction performance. We conduct extensive experiments on benchmark datasets to demonstrate the effectiveness of our pruning method compared with several state-of-the-art baseline methods. Codes of our work are released at: https://github.com/csyhhu/L-OBS.
Summary:
This paper adapts Optimal Brain Surgeon (OBS) method to a local version, and modified the objective function to be the target activation per each layer. Similar to OBS, it uses an approximation to compute Hessian inverse by running through the dataset once. Compare to prior methods, it finishes compression with much less retraining iterations. A theoretical bound on the total error based on local reconstruction error is provided.
Pros:
- The paper explores a local version of OBS and shows effectiveness of proposed method in terms of less time cost for retraining the pruned network.
- Leveraging local second order information is a much clever way than solving/retraining a network through brute force.
- The proposed method achieves good compression ratio without losing performance, with much less retraining iterations.
Comments:
- I recommend the authors to include the basics of Optimal Brain Surgeon (OBS) as a background section. This will make the flow of the paper much better. Many of the techniques used later (e.g. Lagrange multiplier of the MSE error function, iterative algorithm for computing the Hessian inverse) should be part of the background.
- The formulation of layer-wise objective is difficult to understand. Z^l is also a function of Theta^l, and since Z^l = Y^(l-1) W^l, then E^l is actually not even a function of Theta^l, since both terms cancel out in Z^l and hat(Z)^l. This is another reason why the - paper should have included OBS in the background, since here it uses Z^l as a local target activation value, and no longer treated as a function of Theta^l. I had read section 3.2 over many times to realize this point.
- Line 129: “For a well-trained network, gradient w.r.t. Parameters is small enough to be ignored as well”. This is not true. For the original OBS, it is true since the error is the total loss function of the network. For a local MSE function, the gradient of E^l wrt. Theta^l is literally zero (not just small enough).
- I recommend the authors move the proof Theorem 3.2 to the Appendix. This bound is good to know (that the total error is bounded by each local errors) but not very useful and hurts the flow of the paper.
- I recommend the authors edit section 3.4 into algorithmic format. It is hard to read an algorithm in a paragraph. Same for section “prune convolutional layer”.
- Line 191: Again, I don’t understand why “well-trained” is part of the clause here. For any function, if you perturb the input a little bit, the output will be quite close. In the original OBS, it is the final output vs. the target, so “well-trained” is necessary. But here, the target is itself.
- The paper gives some asymptotic bound on time complexities, but I am still curious of the actual wall clock timing compared to related work (DNS, Net-Trim), especially the time cost to compute the Hessian inverse.
- The paper claims that Net-Trim has memory limitation. But maybe it is worth trying a mini-batch version of Net-Trim?
- The paper puts ResNet-101 as future work. It would have been better if the authors can give preliminary results to give insights on how expensive it is to compute Hessian inverse on a deeper network like this. Based on the time complexity, I don’t see why running ResNet-101 with L-OBS is an issue.
- The model achieves good results with much less number of retraining iterations. However, in the AlexNet experiment, it is still losing to DNS in terms of compression ratio. Although, the paper claims to have faster compression time, one major use case for network compression/sparsification deployment on embedded systems, in which the compression ratio is more of a concern than time cost to retrain the network. For completeness, I also recommend the authors to implement DNS and report numbers on VGG-16.
Overall, I feel the paper has some good contribution, but the paper clarity needs to be improved by a lot (Background section, local error function formulation). Some experiments need to be clarified (e.g. wall clock time comparison), and completed (DNS numbers on VGG). Based on these points, I recommend marginal accept, conditioned on a major revision.
After rebuttal note: I am glad that DNS numbers on VGG and ResNet is provided. After seeing the rebuttal I decided to upvote my rating to 7. I hope the authors will reorganize the paper as promised in the rebuttal.