Learning to Prune Deep Neural Networks via Layer-wise Optimal Brain Surgeon

TL;DR

The paper addresses the problem of compressing deep neural networks without substantial accuracy loss or heavy retraining. It introduces Layer-wise Optimal Brain Surgeon (L-OBS), a per-layer Hessian-based pruning approach that uses a second-order error approximation to select and prune parameters, with a theoretical bound tying layer-wise errors to the overall network error. The main contributions are a tractable, Hessian-based pruning algorithm for fully-connected and convolutional layers, a formal error-propagation bound, and extensive experiments showing high compression with minimal pre-retraining loss and reduced retraining effort compared to baselines like Net-Trim. This work offers a practical, theoretically grounded avenue for slim, accurate models suitable for embedded deployment and real-world applications.

Abstract

How to develop slim and accurate deep neural networks has become crucial for real- world 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. Therefore, there is a guarantee that 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.

Figures (6)

• Figure 1: Illustration of shape of Hessian. For feed-forward neural networks, unit $z_1$ gets its activation via forward propagation: $\mathbf{z}=\mathbf{W}^\top\mathbf{y}$, where $\mathbf{W}\in \mathbb{R}^{4\times 3}$, $\mathbf{y}=[y_1, y_2, y_3, y_4]^\top\in \mathbb{R}^{4\times 1}$, and $\mathbf{z}=[z_1, z_2, z_3]^\top\in \mathbb{R}^{3\times 1}$. Then the Hessian matrix of $z_1$ w.r.t. all parameters is denoted by $\mathbf{H}^{[z_1]}$. As illustrated in the figure, $\mathbf{H}^{[z_1]}$'s elements are zero except for those corresponding to $\mathbf{W}_{*1}$ (the 1st column of $\mathbf{W}$), which is denoted by $\mathbf{H}_{11}$. $\mathbf{H}^{[z_2]}$ and $\mathbf{H}^{[z_3]}$ are similar. More importantly, $\mathbf{H}^{-1}=\hbox{diag}(\mathbf{H}_{11}^{-1},\mathbf{H}_{22}^{-1},\mathbf{H}_{33}^{-1})$, and $\mathbf{H}_{11}=\mathbf{H}_{22}=\mathbf{H}_{33}$. As a result, one only needs to compute $\mathbf{H}_{11}^{-1}$ to obtain $\mathbf{H}^{-1}$ which significantly reduces computational complexity.
• Figure 2: Test accuracy on MNIST using LeNet-300-100 when continually pruning the first layer until pruning ratio is 100%. Comparison on ability to preserve prediction between LWC, ApoZ and our proposed L-OBS.
• Figure 3: Distribution of sensitivity of parameters in LeNet-300-100's first layer. More than 90% of parameters' sensitivity scores are smaller than 0.001.
• Figure 4: Retraining pattern of LWC and L-OBS. L-OBS has a better start point and totally resume original performance after 740 iterations for LeNet-5.
• Figure 5: The corresponding relation between tolerable error threshold and pruning ratio.

Theorems & Definitions (3)

• Lemma 3.1
• Theorem 3.2
• Theorem 5.1