MobilePrune: Neural Network Compression via ℓ₀ Sparse Group Lasso on the Mobile System

Many pruning algorithms for deep learning models achieve slim neural networks by introducing sparse-inducing norms to the models. ${\ell }_1$ regularization [31,32,33] and ${\ell }_0$ regularization [8] were applied to induce sparsity on each individual weight. However, such individual sparsity has arbitrary structures, which cannot be utilized by software and hardware. Wei et al. [24] applied group sparsity to prune filters or channels, which can reduce the matrix size used in GEMM in cuDNN. Because the pruned models are compatible with cuDNN, they achieved large speedups. There are methods [34,35] aiming to find sparse models at both individual and group levels, which is similar to our goal. However, they all used ${\ell }_1$ norm to induce individual sparsity in addition to group sparsity. We have performed a comprehensive comparison and demonstrated that our MobilePrune method is the best in inducing sparsity at both individual and group levels for pruning deep learning models in Section 5.

Recent studies [12,13] demonstrate that ${\ell }_0$ norm is the best sparsity-inducing penalty comparing to lasso ${\ell }_1$ [14], elastic net [15], SCAD [16], and MCP [17] models. Remarkably, even current ${\ell }_0$ optimization techniques cannot achieve the global optimal solution, these ${\ell }_0$-based methods with sub-optimal solutions still significantly outperform ${\ell }_1$ and elastic net models, which can be solved to global optima [13,18]. Therefore, the machine learning community has shown great interest in using ${\ell }_0$ norm to compress the large-scale DNN models [8,19,20,21,22,23,24]. However, some ${\ell }_0$-based DNN pruning approaches [8,25] rely on different relaxation strategies to conquer the non-convex and non-differentiable challenges, which does not fully exploit the strength of ${\ell }_0$ regularization. The other methods [19,20,21,22,23] rely on the alternating direction method of multipliers (ADMM) whose global convergence has not proved so far when applied to ${\ell }_0$ norm optimization [27,28,29].

In this paper, we aim to design an algorithm to make the pruned DNN models compatible with cuDNN library [5] and hardware accelerator architecture that uses the sparse systolic tensor array [11,36]. cuDNN is the GPU-accelerated library used in deep learning models [5]. As shown in Figure 1a, convolution used in the convolutional neural network is lowered to matrix multiplication. Therefore, the size of the filter matrix can be reduced when inducing group sparsity column-wise as shown in Figure 1(b.3),(b.4), which can reduce the memory of the DNN models to achieve practical performance improvement. The systolic tensor array is an efficient hardware accelerator for structured sparse matrix as shown in Figure 1(c.4). Specifically, each column in Figure 1(c.4) is sparse. To achieve a pruned DNN model that is compatible with cuDNN and the systolic tensor array, sparsity needs to be induced on both the group level and within-group level. We will show how we achieve this in the following sections.

We aim to prune a generic (deep) neural network, which includes fully connected feed-forward networks (FCN) and convolutional neural networks (CNN). Assuming the generic neural network has N neurons in FCN and M channels in CNN. Let $W^i$ denote the outgoing weights of the ith neuron in FCN and $T^j$ represent the 3D tensor of all filters in the jth channel, which can come from different layers. The training objective for the neural network is given as follows: where $W=\{W^1,\dots ,W^N\}$, $T=\{T^1,\dots ,T^M\}$, $\mathcal{D}={\{x_i,y_i\}}_{i=1}^P$ is a training dataset with P instances, $\mathcal{L}$ is an arbitrary loss function parameteized by W and T, ${\Omega }_{\lambda }^{\eta }\left(W\right)$ and ${\Gamma }_{\beta ,\gamma }^{\alpha }\left(T\right)$ represent the ${\ell }_0$ sparse group lasso penalties for neurons and channels, respectively. Specifically, ${\Omega }_{\lambda }^{\eta }\left(W\right)$ is defined as where $\eta \ge 0$ and $\lambda \ge 0$ are regularization parameters. Let $n\left(i\right)$ represent the set of outgoing edge weights of neuron i. Then, $\parallel W^i{\parallel }_0={\sum }_{j\in n\left(i\right)}{\parallel W_j^i\parallel }_0$ ($W_j^i$ is the jth outgoing edge weight of neuron i; $\parallel W_j^i{\parallel }_0=0$ when $W_j^i=0$ and $\parallel W_j^i{\parallel }_0=1$ otherwise) computes the number of the non-zero edges in $W^i$ and $\parallel W^i{\parallel }_g=\sqrt{{\sum }_j{\left(W_{j\in n\left(i\right)}^i\right)}^2}$ aggregates the weights associated with the ith neuron as a group. The core spirit of Equation (2) is illustrated in Figure 2a–c, the group lasso penalty $\parallel W^i{\parallel }_g$ tends to remove all the weights together with the ith neuron if it is less important. If a group is selected, the ${\ell }_0$ penalty further removes the weights with small magnitudes within the group. Group sparsity $\parallel W^i{\parallel }_g$ can help remove neurons in the neural network, which reduces the size of the neural network and further improve efficiency. Individual sparsity $\parallel W^i{\parallel }_0$ helps to achieve additional sparsity within the remaining neurons. Such structured sparsity can be used by the systolic tensor array [10,11]. The other regularization term ${\Gamma }_{\beta ,\gamma }^{\alpha }\left(T\right)$ is defined as following, where $\alpha ,\beta$, and $\gamma$ are non-negative regularization parameters. Equation (3) defines a hierarchical-structured sparse penalty in which structure is guided by the memory organization of GEMM using cuDNN [5]. As demonstrated in Figure 2d, the pruning strategy encoded in Equation (3) explicitly takes advantage of the GEMM used in cuDNN. $\parallel T^j{\parallel }_g$ enforces the group sparsity of all the filters applied to the jth channel and $\parallel T_{:,h,w}^j{\parallel }_g$ enforces the group sparsity across the filters on the same channel and $\parallel T^j{\parallel }_0={\sum }_f{\sum }_h{\sum }_w{\parallel T_{f,h,w}\parallel }_0$ prunes the small weights within the remaining channels and filters. Equation (3) can help to achieve an extremely compact model as Figure 1(c.4). Therefore, the computation can be accelerated at both software and hardware levels.

In this subsection, we first briefly review the general PALM (Proximal Alternating Linearized Minimization) framework used in our MobilePrune algorithm. Then, we introduce how we modify PALM to efficiently optimize the ${\ell }_0$ sparse group lasso for neural network compression. PALM is designed to optimize a general optimization problem formulated as: where $F(W,T)$ is a smooth function and ${\Phi }_1\left(W\right)$ and ${\Phi }_2\left(T\right)$ do not need to be convex or smooth, but are required to be lower semi-continuous. The PALM algorithm applies proximal forward–backward algorithm [37] to optimize Equation (4) with respect to W and T in an alternative manner. Specifically, at iteration $k+1$, the temporal values of $W^{(k+1)}$ and $T^{(k+1)}$ for the proximal forward–backward mapping are derived by solving the following sub-problems, where $U_i^{(k+1)}=W_i^{\left(k\right)}-\frac{1}{c^k}{\nabla }_{W}F(W^{\left(k\right)},T^{\left(k\right)})$ and $V_i^{(k+1)}=T_i^{\left(k\right)}-\frac{1}{d^k}{\nabla }_{T}F(W^{(k+1)},T^{\left(k\right)})$. Additionally, $c^k$ and $d^k$ are positive constants. This optimization process has been proven to converge to a critical point when functions F, ${\Phi }_1$, and ${\Phi }_2$ are bounded [37]. We further extend the convergence proof in [37] and prove that the global convergence of PALM holds for training deep learning models under mild conditions. The detailed proof can be found in Appendix A.

To optimize Equation (1), we define two proximal operators for the two penalty terms as the following,

Here, functions ${\Omega }_{\lambda }^{\eta }(·)$ and ${\Gamma }_{\beta ,\lambda }^{\alpha }(·)$ take vectors as inputs, which are equivalent to Equations (2) and (3) once we vectorize $W^i$ and $T^j$. The overall optimization process of MobilePrune is described in Algorithm 1. Once we can efficiently compute the optimal solution of ${\pi }_{\lambda }^{\eta }(·)$ and ${\theta }_{\beta ,\lambda }^{\alpha }(·)$, the computational burden mainly concentrates on the partial derivative calculation of functions $H(·)$ and $F(·)$, which is the same as training a normal DNN model.

Algorithm 1 The framework of MobilePrune Algorithm.

Initialize ${\left(W^i\right)}^0,\forall i$, ${\left(T^j\right)}^0,\forall j$ and $L_r$. for$k=0,2,\dots$do    for $i=1,2,\dots ,N$ do      $H_k^i={\nabla }_{W_k^i}\mathcal{L}\left(W_k^i,W_k^{j\ne i},T_k\right)$.      $W_{k+1}^i\in {\pi }_{\lambda /L_r}^{\eta /L_r}\left(W_k^i-\frac{1}{L_r}{H}_k^i\right)$ by solving Equation (7).    end for    for $j=1,2,\dots ,M$ do      $F_k^j={\nabla }_{T_k^j}\mathcal{L}(W_{k+1},T_k^j,T_k^{l\ne j})$.      $T_{k+1}^j\in {\theta }_{\beta /L_r,\gamma /L_r}^{\alpha /L_r}\left(T_k^j-\frac{1}{L_r}{F}_k^j\right)$ by solving Equation (8).    end for end for

To the best of our knowledge, ${\pi }_{\lambda }^{\eta }\left(y\right)$ and ${\theta }_{\beta ,\gamma }^{\alpha }\left(y\right)$ defined in Equations (7) and (8) are novel proximal operators that have not been attempted before. Solving ${\pi }_{\lambda }^{\eta }\left(y\right)$ and ${\theta }_{\beta ,\gamma }^{\alpha }\left(y\right)$ is the key to apply MobilePrune in Algorithm 1. Therefore, this subsection elaborates the algorithmic contributions we made to efficiently calculate ${\pi }_{\lambda }^{\eta }\left(y\right)$ and ${\theta }_{\beta ,\gamma }^{\alpha }\left(y\right)$.

In this subsection, we compared our proposed MobilePrune approach with other state-of-the-art pruning methods in terms of prune rate, computational costs, and test accuracy. We mainly compared our methods with structured pruning methods because DNN models pruned by non-structure pruning methods could not obtain practical speedup as shown in Figure 1. Notably, we only compared the results that can be reproduced by the source codes provided by the competing methods. First, we briefly summarized their methodology. PF [32] and NN slimming [33] were simple magnitude-based pruning methods based on $l1$ norm. BC [9], SBP [7], and VIBNet [39] cast the DNN pruning into probabilistic Bayesian models. C-OBD [40], C-OBS [2], Kron-OBD [40,41], Kron-OBS [2,41], and EigenDamage [42] are Hessian matrix-based methods. ${\ell }_0$ norm penalized method [8] and group lasso penalized method [24] are also well-known methods.

In our experiments, we use NVIDIA Corporation as the GPU and the number of cores of the CPU is 12. All the baseline models were trained from scratch via stochastic gradient decent(SGD) with a momentum of 0.9. We trained the networks for 150 epochs on MNIST and 400 epochs on CIFAR-10 and Tiny-ImageNet with an initial learning rate of 0.1 and weight decay of 5 × 10⁻⁴. The learning rate is decayed by a factor of 10 at 50, 100 on MNIST and at 100, 200 on CIFAR-10 and Tiny-ImageNet, respectively. The details of hyper-parameters for all experiments are summarized in Appendix B. We also provide the computational efficiency of our methods in Appendix C.

To demonstrate the efficacy and effectiveness of our proposed MobilePrune method, we perform a series of compassion studies with other state-of-the-art pruning methods such as ${\ell }_0$ norm, ${\ell }_1$ norm, ${\ell }_2$ norm, group lasso, and ${\ell }_1$ sparse group lasso for all three datasets—WISDM [49,50], UCI-HAR [51,52], and PAMAP2 [53,54,55]. We evaluate the pruning accuracy and pruning rate of weights (parameters) and nodes for our proposed MobilePrune approach and all other state-of-the-art pruning methods using the same learning rate (0.0001) and the same number of epochs (150) for all three datasets. The pruning thresholds are 0.015, 0.005, 0.01 for the pruning methods in the WISDM, HCI-HAR, and PAMAP2 datasets, respectively. In addition, we evaluate the computational cost and battery consumption for our proposed method with all other state-of-the-art pruning methods as well. The details of the dataset descriptions and the hyper-parameters for all experiments are summarized in Appendix D.

To demonstrate the efficacy and effectiveness of the ${\ell }_0$ sparse group lasso penalty, we performed a series of ablation studies on various DNN models. As shown in Table 4, if we only use ${\ell }_0$ norm penalty, there is no effect on a neuron or channel pruning as expected. Similarly, if we only employ the group lasso penalty, the pruned model still has more weights left. However, if we apply ${\ell }_0$ sparse group lasso, we can achieve pruned DNN models that are sparse in terms of both neurons (or channel) and weights. In addition, we compare our ${\ell }_0$ sparse group lasso model with ${\ell }_1$ sparse group lasso [59] on pruning DNN models. Table 4 shows their comparison on pruning various DNN models. More details can be found in Appendix A.3 and Appendix C. As shown in Table 4 and the results in supplementary, the ${\ell }_0$ sparse group lasso model significantly outperforms the ${\ell }_1$ sparse group lasso model in all aspects, which demonstrates its superiority in pruning DNN models.