CATRO: Channel Pruning via Class-Aware
Trace Ratio Optimization

The recent advances of weight pruning can be mainly categorized into unstructured pruning and channel pruning. Unstructured pruning can effectively prune the over-parameterized deep neural networks (DNNs) but results in irregular sparse matrices, which require customized hardwares and libraries to support the practical speedup, leading to a limited inference acceleration in most practical cases [14, 15]. Recent work concentrates on channel pruning [16, 4, 17, 18, 19, 20, 21, 22, 23, 24, 25], which prunes the whole channel and results in structured sparsity. The compacted DNN after channel pruning can be directly deployed on the existing prevalent CNN acceleration frameworks without dedicated hardware/libraries implementation. Early work [26] proposes a straightforward pruning heuristic, penalizing weights with $l_{1}$-norm group lasso regularization and then eliminate the least $l_{1}$-norm of channels. Network Slimming [16] introduces a factor in the batch normalization layer with respect to each channel and evaluates the importance of channels by the magnitude of the factor. Then, many criteria are proposed to evaluate the importance of neuron. SCP [6] measures the relative importance of a filter in each layer by calculating the $l_{1}$-norm and $l_{2}$-norm, respectively. FPGM [4] removes the filters minimizing the sum of distances to others. Variational CNN pruning [27] provides a Bayesian model compression technique, which approximates batch-norm scaling parameters to a factorized normal distribution using stochastic variational inference. GBN [28] introduces a gate for each channel and performs a backward method to optimize the gates. Then it independently prunes channels by the rank of gates. Similarily, DCPH [23] designs a channelwise gate to enable or disable each channel. Greedy algorithm is also applied to identify the subnetwork inside the original network [29]. Recently, MFP [22] introduces a new set of criteria to consider the geometric distance of filters. These works evaluate and prune the importance of channels individually. Another work [13] leverage Fisher’s linear discriminant analysis (LDA) to prune the last layer of the DNN, which, however is empirical heuristic and suffers from the expensive cost of cross-layer tracing. CCP [12] exploits the inter-channel dependency but stops with only pairwise correlations. It prunes channels without global consideration either. SCOP [30] utilize samples to set up a scientific control and then prunes the neural network in the group by introducing generated input features. Recent prevailing AutoML-based pruning approaches [31, 32, 33, 34, 35, 36] can achieve a high FLOPs reduction of DNNs. For instance, AMC [31] leverages deep reinforcement learning (DRL) to let the agent optimizes the policy from extensive experiences by repeatedly pruning and evaluating for DNNs. DNCP [21] uses a set of learnable architecture parameters, and calculates the probability of each channel being retained based on these parameters. However, they usually requires tremendous computation budgets due to the large search space and trial-and-error manner.

In recent years, the combinatorial selection based on submodular functions has been one of the most promising methods in machine learning and data mining. It has been a surge of interest in lots of computer tasks, such as visual recognition [37], segmentation [38], clustering [39, 40], active learning [41] and user recommendation [42]. The submodular set function maximization [43], which is to maximize a submodular function, is one of the most basic and important problem. There exist several algorithms for solving nonnegative and monotone submodular function maximization. Under the uniform matroid constraint [43], a standard greedy algorithm gives an $1-e^{-1}$ approximation [44, 45, 46] in polynomial time, and $1-e^{-1}$ is shown to be the best possible approximation unless $P=NP$ [44]. There are also many works investigating more effective algorithms [47, 48] and many other works investigating the maximization of a submodular set function under different constraints [49, 50, 51]. Although submodularity has been widely used in many tasks and widely studied, it has not been explored in model compression as far as we known.

For a CNN ${F}(\cdot;\mathbf{\Theta})$ with a stack of $L$ convolution (CONV) layers and weight parameters $\mathbf{\Theta}$, given a set of $\mathcal{N}$ image samples $\{{\mathbf{x}}^{[i]}\}_{i=1}^{\mathcal{N}}$ from $K$ classes and the corresponding labels $\{{\mathrm{y}}^{[i]}\in\{1,\cdots,K\}\}_{i=1}^{\mathcal{N}}$, channel pruning aims to find binary channel masks $\mathbf{m}_{l}$ with the following objective:

where $\mathcal{L}(\cdot)$ is the loss function (e.g., cross-entropy for classification), $c_{l}$ is the channel number of the $l$-th layer, $\mathcal{T}(\cdot)$ is the FLOPs function, and $\mathcal{T_{B}}$ is a given FLOPs constraint.

Most methods prune channels with some surrogate layer-wise criteria $\mathcal{C}(\cdot)$ instead of directly solving Eq.(1). Following this setting and with the goal of keeping $d_{l}$ channels in layer $l$, we can express the pruning objective as follows:

where ${\mathbf{o}}^{[i]}_{l}\in\mathbb{R}^{c_{l}\times w_{l}\times h_{l}}$ is the feature map of input ${\mathbf{x}}^{[i]}$ in the $l$-th layer, and $\odot$ is elementwise multiplication with broadcasting. A variety of $\mathcal{C}(\cdot)$ have been proposed [7, 8, 4, 5], which may not be completely equivalent to solving Eq.(1) but effective in practice.

Alternatively, we can formulate Eq.(2) as a matrix projection problem. For clarity, we denote the $i$-th element in a vector (or an ordered set) $\mathbf{a}$ as $\mathbf{a}(i)$. We build an ordered set of the channel index $\mathcal{I}_{l}=\{i|\mathbf{m}_{l}(i)=1\}$ with $|\mathcal{I}_{l}|=d_{l}$ and an indicator matrix $\mathbf{V}_{\mathcal{I}_{l}}=[\mathbf{e}^{(l)}_{\mathcal{I}_{l}(1)},\cdots,\mathbf{e}^{(l)}_{\mathcal{I}_{l}(d_{l})}]\in\mathbb{R}^{c_{l}\times d_{l}}$, where $\mathbf{e}^{(l)}_{t}\in{\{0,1\}}^{c_{l}}$ denotes the one-hot vector with its $t$-th element as one. Therefore, Eq.(2) can been expressed as follows:

For an input ${\mathbf{x}}^{[i]}$, we denote the pruned feature map in layer $l$ with $\mathcal{I}_{l}$ by ${\mathbf{\tilde{o}}}^{[i]}_{l}=\mathbf{V}_{\mathcal{I}_{l}}^{\top}{\mathbf{o}}^{[i]}_{l}\in\mathbb{R}^{d_{l}\times w_{l}\times h_{l}}$.

In order to obtain effective criteria for model pruning, inspired by the classical Fisher Linear Discriminant Analysis (LDA) [52], we dig into the feature discrimination for measuring multi-channel joint impact. As a simple example, Fig. 2 visualizes two feature spaces from the $18$-th layer in a well-trained ResNet-20 model trained on CIFAR-10. The red and blue points represent features of images containing ships and airplanes, respectively. We have an intuitive feeling that the features produced by channels in the right figure are more discriminative, and a classifier built on that space is more likely to outperform that on the left space. For deep CNNs with more layers, we hypothesis the discrimination effects accumulate and the classifier performance gap enlarges. We describe the relation between space discrimination and classifier accuracy in Assumption 1.

The above assumption enables us to concretely measure and optimize the discriminations. Suppose $\mathbf{X}=[{\mathbf{x}}^{[1]},\cdots,{\mathbf{x}}^{[N]}]\in\mathbb{R}^{c\times w\times h\times N}$ denotes $N$ data samples from the training set with $K$ classes and $\mathbf{y}=[{\mathrm{y}}^{[1]},\cdots,{\mathrm{y}}^{[N]}]$ denotes the corresponding class label. We build two graph weight matrices $\mathbf{G}^{(w)}\in\mathbb{R}^{N\times N}$ and $\mathbf{G}^{(b)}\in\mathbb{R}^{N\times N}$ to depict the prior knowledge about data sample relationship. $\mathbf{G}^{(w)}$ reflects within-class relationship, where each element $\mathrm{G}^{(w)}_{i,j}$ has a larger value if samples ${\mathbf{x}}^{[i]}$ and ${\mathbf{x}}^{[j]}$ belong to the same class, and smaller otherwise. Analogically, $\mathbf{G}^{(b)}$ represents the between-class relationship, where $\mathrm{G}^{(b)}_{i,j}$ is larger if samples ${\mathbf{x}}^{[i]}$ and ${\mathbf{x}}^{[j]}$ belong to different classes. Without any other priors, we directly use the Fisher score [52] as the graph metric for its simplicity and efficiency:

where $n_{k}$ is the number of samples in the $k$-th class, and $\sum_{k=1}^{K}n_{k}=N$.

Given Assumption 1 and the prior matrices in Eq.(6) and Eq.(7), a natural criterion $\mathcal{C}(\cdot)$ for channel pruning is to keep channels that reduce the discrimination of samples within the same class ($\sum_{ij}||{\mathbf{\tilde{o}}_{l}}^{[i]}-{\mathbf{\tilde{o}}_{l}}^{[j]}||\mathrm{G}^{(w)}_{i,j}$) and enlarge that of samples from different classes ($\sum_{ij}||{\mathbf{\tilde{o}}_{l}}^{[i]}-{\mathbf{\tilde{o}}_{l}}^{[j]}||\mathrm{G}^{(b)}_{i,j}$). Therefore, we can replace Eq.(3) by

As $\mathbf{\Lambda}^{(w)}$ is a diagonal matrix with $\mathrm{\Lambda}^{(w)}_{i,i}=\sum_{j}\mathrm{G}^{(w)}_{i,j}$ and $\mathbf{\Lambda}^{(b)}$ is a diagonal matrix with $\mathrm{\Lambda}^{(b)}_{i,i}=\sum_{j}\mathrm{G}^{(b)}_{i,j}$, $\mathbf{\Lambda}^{(w)}-\mathbf{G}^{(w)}$ and $\mathbf{\Lambda}^{(b)}-\mathbf{G}^{(b)}$ are two Laplace matrices. Based on the Laplace matrix properties, Eq.(8) is mathematically equivalent to the following trace optimization:

where $\mathbf{O}_{l}=[{\mathbf{o}}^{[1]},\cdots,{\mathbf{o}}^{[N]}]\in\mathbb{R}^{c_{l}\times w_{l}\times h_{l}\times N}$ is the aggregated feature map for $\mathbf{X}$ in layer $l$, and $tr(\cdot)$ denotes the trace. Equivalently, our goal is to find the optimal channel index set $\mathcal{I}_{l}^{*}$ with the maximum trace ratio criterion for discrimination:

where $\mathbf{G}^{(w)}_{l}=\mathbf{O}_{l}(\mathbf{\Lambda}^{(w)}-\mathbf{G}^{(w)})\mathbf{O}_{l}^{\top}$ and $\mathbf{G}^{(b)}_{l}=\mathbf{O}_{l}(\mathbf{\Lambda}^{(b)}-\mathbf{G}^{(b)})\mathbf{O}_{l}^{\top}$. Since $tr(\mathbf{V}_{\mathcal{I}_{l}^{*}}^{\top}(\mathbf{G}^{(b)}_{l}-\lambda_{l}^{*}\mathbf{G}^{(w)}_{l})\mathbf{V}_{\mathcal{I}_{l}^{*}})=0$, for all $\mathcal{I}_{l}$ satisfying $|\mathcal{I}_{l}|=d_{l}$ we have

Therefore, to optimize Eq.(9) is further equivalent to optimize the following optimization:

To be more clear, the objective in Eq.(9) is the left-hand side terms of the inequation in Eq.(11), and the optimal solution $I_{l}^{*}$ of Eq.(9) makes the equality of the inequation in Eq.(11) hold (by the definition of $\lambda_{l}^{*}$). And similarly, given $\lambda_{l}^{*}$, the objective in Eq.(13) is the left-hand side terms of the inequation in Eq.(12), and the solution $I_{l}$ of Eq.(13) makes the equality of the inequation in Eq.(12) hold. Lastly, the equality of the inequation in Eq.(11) holds if and only if the equality of the inequation in Eq.(12) holds. Thus, the optimal solution of Eq.(9) makes equality of the inequation in Eq.(11) hold, makes equality of the inequation in Eq.(12) hold, and is also the optimal solution of Eq.(13). In contrast, the optimal solution of Eq.(13) makes equality of the inequation in Eq.(12) hold, makes equality of the inequation in Eq.(11) hold, and is also the optimal solution of Eq.(9).

Noting that solving Eq.(13) naturally considers the multi-channel joint impacts inherited in $\mathbf{V}$.

To solve Eq.(13), we introduce a submodular set optimization, provide theoretical guarantees on the pruned model accuracy, and propose a layer-wise pruning procedure.

Firstly, we introduce $s_{l,i}=\exp({\mathbf{e}^{(l)}_{i}}^{\top}(\mathbf{G}^{(b)}_{l}-\lambda_{l}\mathbf{G}^{(w)}_{l})\mathbf{e}^{(l)}_{i})$ as the discrimination score for channel $i$ in layer $l$ given a specified ratio $\lambda_{l}$, and $H_{l}(\mathcal{I})=\log(\sum_{i\in\mathcal{I}}s_{l,i})$ as the index set function given an index set $\mathcal{I}\subseteq\{1,\cdots,c_{l}\}$. When $\lambda_{l}^{*}=\lambda_{l}$, we notice that Eq.(13) is equivalent to the following one:

We introduce two nice properties of set functions and show $H_{l}(\cdot)$ satisfies them.

Secondly, to maximize $H_{l}(\mathcal{I})$ with a fixed $\lambda_{l}$, we show that sequentially selecting the indexes $i$ with the largest score $s_{l,i}$ obtains a model with a guaranteed accuracy according to Theorem 1.

Finally, we consider to iteratively optimize the discrimination ratio. Suppose the current index set is $\mathcal{I}_{l}$ and the discrimination ratio is $\lambda_{l}$. After maximizing $H_{l}(\mathcal{I}_{l};\lambda_{l})$ with the index sequential selection procedure, we get a new index set $\mathcal{I}^{\prime}_{l}$ and a new ratio $\lambda^{\prime}_{l}$. From Theorem 2, we have $\lambda^{\prime}_{l}\geq\lambda_{l}$.

The whole layer-wise pruning procedure, given the channel budgets for each layer, is shown in Algorithm 1.

Theorems 2 and 1 show that, given a well-trained unpruned network and channel numbers $d_{l}$, the pruning procedure in Algorithm 1 leads to monotonically increasing feature discrimination value until reaching the optimal solution (by Theorem 2) and obtains the pruned network with an accuracy lower boundary proportional to the best-pruned network with large probability (by Theorem 1).

Note that the objective in each layer is not a simply linear sum of each channel. Determining the value of $\lambda_{l}$ in $s_{l,i}$ incorporates the joint-channel impact and depends on other channels. With the update of $\lambda_{l}$, channels selected in the previous rounds may be unselected. In other words, the set of channels is not built by step-by-step greedy inclusion.

Compared to training-based pruning, the time complexity of our layer-wise pruning procedure is linear to the number of sampled data, which is usually much smaller than (and increased sub-linearly to) the number of original training data, and it only performs inference once and no back-propagation. Theoretically, the convergence of the pruning step is guaranteed; Empirically, it only uses three to four iterations to converge to the pruned network structures, taking a negligible amount of time.

We further propose an approximate algorithm to determine the architecture of the pruned model, i.e., to find the optimal channel numbers $\boldsymbol{d}$, as a prerequisite of the layer-wise pruning stage shown in Algorithm 1. Besides the discrimination criteria, we also consider the FLOPs constraint, aiming to quickly estimate the effects of local channel number changes on the whole model’s power and efficiency.

We first introduce a network discrimination approximation function with given $\{\lambda_{l}\}_{l=1}^{L}$, which is $\mathcal{D}(\boldsymbol{d})=\sum_{l=1}^{L}\max_{|\mathcal{I}_{l}|=d_{l}}H_{l}(\mathcal{I}_{l};\lambda_{l})$. Taking the suboptimal solution in Theorem 1 and ignoring the effects of $d_{l}$ in other layers, we have its derivative as follows:

where $[\tilde{s}_{l,1},\cdots,\tilde{s}_{l,c_{l}}]$ is the descending sorted elements of $\{s_{l,i}\}_{i=1}^{c_{l}}$. The two approximations come from the sequentially selection of element $d_{l}$ in Theorem 1 and Taylor expansion, respectively.

On the other hand, to measure the network computational efficiency, we define the FLOPs function of layer $l$ as $\mathcal{T}_{l}(\boldsymbol{d})=d_{l-1}d_{l}w_{l}h_{l}q_{l}^{2}$, where $q_{l}\times q_{l}$, $w_{l}$, and $h_{l}$ is the $l$-th layer’s kernel size, feature map width, and feature map height, respectively. The FLOPs function for the entire network is $\mathcal{T}(\boldsymbol{d})=\sum_{l=1}^{L}\mathcal{T}_{l}(\boldsymbol{d})$, and its derivative for $d_{i}$ is as follows:

As a tradeoff to jointly considering computations and discriminations, we define the following function to evaluate the effect of changing each channel number $d_{l}$:

For the global architecture optimization, we manipulate the channel numbers without selecting the channels. Therefore, we take a greedy coordinate ascent strategy to find the optimal channel numbers, which is describe in Algorithm 2. It is worth noting that, similar to the layer-wise pruning procedure, this algorithm only performs inference once without backpropagation. After that, it only needs to recalculate a small portion of variables in each iteration.

We take the two stages together to build CATRO, the proposed channel pruning method. As illustrated in Fig. 3, CATRO first employs Algorithm 2 (Greedy Selection for Channel Numbers) to select the channel numbers for each layer and then uses Algorithm 1 (Layer-wise Pruning with Trace Ratio Criterion) to prune each layer. It fine-tunes the pruned network as the last step on all available training data, similarly to other pruning methods.