An Information Theory-Inspired Strategy for Automated Network Pruning

Abstract

Despite superior performance achieved on many computer vision tasks, deep neural networks demand high computing power and memory footprint. Most existing network pruning methods require laborious human efforts and prohibitive computation resources, especially when the constraints are changed. This practically limits the application of model compression when the model needs to be deployed on a wide range of devices. Besides, existing methods are still challenged by the missing theoretical guidance, which lacks influence on the generalization error. In this paper we propose an information theory-inspired strategy for automated network pruning. The principle behind our method is the information bottleneck theory. Concretely, we introduce a new theorem to illustrate that the hidden representation should compress information with each other to achieve a better generalization. In this way, we further introduce the normalized Hilbert-Schmidt Independence Criterion on network activations as a stable and generalized indicator to construct layer importance. When a certain resource constraint is given, we integrate the HSIC indicator with the constraint to transform the architecture search problem into a linear programming problem with quadratic constraints. Such a problem is easily solved by a convex optimization method within a few seconds. We also provide rigorous proof to reveal that optimizing the normalized HSIC simultaneously minimizes the mutual information between different layers. Without any search process, our method achieves better compression trade-offs compared to the state-of-the-art compression algorithms. For instance, on ResNet-50, we achieve a 45.3%-FLOPs reduction, with a 75.75 top-1 accuracy on ImageNet. Codes are available at https://github.com/MAC-AutoML/ITPruner.

Appendices

Appendix A Proof of Theorem 1

Theorem 6

Assuming X is the input random variable follows a Markov random field structure and the Markov random field is ergodic. For a network that have L hidden representations $X_1, X_2,..., X_L$, with a probability $1-\delta$, the generalization error $\epsilon$ is bounded by

$$\begin{aligned} \epsilon \le \sum _{i=1}^L\sqrt{\frac{\log \frac{2}{\delta } + \log 2\text {I}(X;X_i)}{2n}}, \end{aligned}$$

where n is the number of training examples.

Proof

Let $\mathcal {D}_n = {(x_1, y_1),...,(x_n, y_n)}$ be an i.i.d random sample from $P_{X,Y}$. We define the loss function as a mapping $l:\mathcal {Y}\times \mathcal {Y}\rightarrow \mathbb {R}^+$ and the performance of a predictor $h:\mathcal {X}\rightarrow \mathcal {Y}$ is measured by the expectation on $P_{X,Y}$, which is also known as generalization error and formally defined as $R(h):=\mathbb {E}_{X,Y}\left[ l(Y, h(X))\right] .$ In practice, $P_{X,Y}$ is usually unknown, we thus estimate $\hat{h}$ based on $\mathcal {D}_n$ and the corresponding empirical risk is defined as $\hat{R}_n(h):=\frac{1}{n}\sum _{i=1}^nl(y_i, h(x_i)).$ For simplicity, we further define the mapping function $l = I\left( h(X)\ne Y\right)$, where I(., .) is the indicator function. In this case, applying the Hoeffding inequality (Hoeffding, 1994) we have

$$\begin{aligned} P\left( \left| \hat{R}_n(h)-R(h)\right| \ge \epsilon \right) \le 2\exp \left( -2n\epsilon ^2\right) . \end{aligned}$$

(A1)

Together with PAC learning theory (Mohri et al., 2018), we obtain the following corollary: $\square$

Corollary 1

Let $\mathcal {H}$ be a finite set of hyothesis. Then for all $\epsilon >0$, we have

$$\begin{aligned} P\left( \exists h\in \mathcal {H} \Bigg | \left| \hat{R}_n(h)-R(h)\right| \ge \epsilon \right) \le 2\left| \mathcal {H} \right| \exp \left( -2n\epsilon ^2\right) .\nonumber \\ \end{aligned}$$

(A2)

We further control the probability in Corollary 1 with a confidence $\delta$. That is, we ask

$$\begin{aligned} 2\left| \mathcal {H} \right| \exp \left( -2n\epsilon ^2\right) \le \delta . \end{aligned}$$

In other words, with a probability $\delta$, we have

$$\begin{aligned} 2\left| \mathcal {H} \right| \exp \left( -2n\epsilon ^2\right)&\le \delta \\ \exp \left( -2n\epsilon ^2\right)&\le \frac{\delta }{2\left| \mathcal {H} \right| }\\ -2n\epsilon ^2&\le \log \frac{\delta }{2\left| \mathcal {H} \right| }\\ \epsilon ^2&\ge \frac{\log \frac{2}{\delta }+ \log \left| \mathcal {H} \right| }{2n}. \end{aligned}$$

Meanwhile, we can also conclude that $\epsilon ^2$ is also bounded

$$\begin{aligned} \epsilon ^2\le \frac{\log \frac{2}{\delta }+ \log \left| \mathcal {H} \right| }{2n}, \end{aligned}$$

(A3)

with probability $1-\delta$. To prove our theorem, we further introduce a lemma (Shwartz-Ziv et al., 2018) about Asymptotic Equipartition Property (AEP) (Cover, 1999) as follows:

Lemma 1

Assuming X is a random variable follows a Markov random field structure and the Markov random field is ergodic. Then, we have that

$$\begin{aligned} \left| \mathcal {H} \right| \le 2^{H(X)}, \end{aligned}$$

(A4)

Let $X_i$ be a mapping of X, the size of typical set $|\mathcal {T}|$ is bounded by

$$\begin{aligned} |\mathcal {T}| \le \frac{2^{H(X)}}{2^{H(X|X_i)}} = 2^{\text {I}(X;X_i)} \end{aligned}$$

(A5)

Therefore, we conclude that for a mapping T, the square error is bounded by

$$\begin{aligned} \epsilon ^2_T\le \frac{\log \frac{2}{\delta }+ \log |\mathcal {T}|}{2n} \le \frac{\log \frac{2}{\delta }+ \text {I}(X;X_i)\log 2}{2n}. \end{aligned}$$

(A6)

Considering that there are L mappings in the network, we thus conclude that the overall square error is bounded by

$$\begin{aligned} \epsilon \le \sum _{i=1}^L\sqrt{\frac{\log \frac{2}{\delta } + \text {I}(X;X_i)\log 2}{2n}}, \end{aligned}$$

(A7)

with a probability $1-\delta$.

Fig. 8

The relationship between mutual information (MI) and accuracy on the validation set. The x-axis measures the nHSIC of the sampled feature, and the y-axis measures the accuracies of the randomly sampled architectures. The correlation between MI and accuracy is $-0.955872$. In other words, nHSIC is negatively correlated to the finetuned accuracy and thus can be used as an accurate predictor in model compression

Appendix B The Debate on the Information Bottleneck

ITPruner is based on the theory of the information bottleneck in deep learning (Shwartz-Ziv & Tishby, 2017), which is challenged by Saxe et al. (2019). Saxe et al. (2019) argue that the representation compression phase empirically demonstrated in Shwartz-Ziv and Tishby (2017) only appears when double-sided saturating nonlinearities like tanh and sigmoid are deployed. In this case, the claimed causality between compression and generalization was also questioned, which is foundation of ITPruner.

To clear such debate, we first propose Theorem 1 to mathematically demonstrate that reducing the mutual information between input and the hidden representations indeed leads to a smaller generalization error. Secondly, we also provide empirical evidence in Fig. 8 to show the performance on the validation set is highly correlated with the layer-wise mutual information. At last, the critical argument raised in Saxe et al. (2019) may cause by the inaccurate estimation of mutual information. Recent work (Wang et al., 2021) propose a new method to estimate the mutual information and empirically identify the original experiments in Shwartz-Ziv and Tishby (2017). For example, Fig. 9 reports the change of mutual information (MI) between input and different layers in the training process. As we can see, there is a clear boundary between the initial fitting (increase of mutual information) and the compression phases (decrease of mutual information). Overall, we thus conclude that the generalized information bottleneck principle behind ITPruner is solid theoretically and empirically.

Fig. 9

The change of mutual information (MI) between input and different layers in training process. There is a clear boundary between the initial fitting (increase of mutual information) and the compression phases (decrease of mutual information) in different layers

Table 12 The effect of different regularization terms on ResNet-20 and Invertable-ResNet. “baseline” indicates that the model is trained without weight-decay

Appendix C The Debate on the Invertable-ResNet

In the work of Gomez et al. (2017) and Jacobsen et al. (2018), they propose reversible networks, which is capable of constructing hidden representation that encode the input. Here we argue that Invertable-ResNet (Gomez et al., 2017) is not an experimental counterexample for our Theorem 6.

Specifically, Invertable-ResNet (Gomez et al., 2017) actually need to optimize layer-wise mutual information mentioned in Theorem 5 for better generalization. Firstly, in their released source code,⁵ weight decay is necessary for a better generalization. We also conduct an ablation study to validate the effectiveness of weight decay. As shown in Tab. 12, both ResNet-20 and Invertable-ResNet show a significant performance promotion when the regularization term is adopted in the training. Meanwhile, as theoretically demonstrated in Theorem 5, the adopted weight decay regularization is highly correlated to mutual information. Moreover, directly replacing wight decay with mutual information still shows a performance promotion, which is reported in Tab. 12. We thus conclude that optimizing mutual information in Invertable-ResNet (Gomez et al., 2017) is necessary.

There are two types of hidden representations exist in Invertable-ResNet (Gomez et al., 2017) simultaneously. One is highly related to the input X and uncorrelated to the output Y. The other one is the opposite, which is related to Y and uncorrelated to X. To validate our argument, we conduct a new experiment that extracts features in Invertable-ResNet and reports the mutual information with input and output. As illustrated in Fig. 10, most of the features have a certain degree of correlation with the X and Y. However, the two types of the feature described before do exist in Invertable-ResNet, which is marked as a darker red dot in Fig. 10.

Fig. 10

The mutual information between the intermediate feature map and the input X, output Y, respectively. Each point represents a different channel in the feature map

Theorem 5

Assuming that $X \sim \mathcal {N} (\varvec{0}, \varvec{\Sigma }_{X})$ and $Y \sim \mathcal {N} (\varvec{0}, \varvec{\Sigma }_{Y}), Y=W^TX$, $\min I(X;Y) \propto \min ||W||_F^2$.

Proof

The definition of mutual information is $I(X;Y)=H(X)+H(Y)-H(X,Y)$. Meanwhile, the entropy of multivariate Gaussian distribution $H(X)=\frac{1}{2}\text {ln}({(2\pi e)}^D|\Sigma _X|)$ and the joint distribution $(X,Y) \sim N\left( 0,\Sigma _{(X,Y)}\right)$, where

$$\begin{aligned} \Sigma _{(X,Y)}=\left( \begin{array}{cc} \Sigma _{X} & \Sigma _{XY}\\ \Sigma _{YX} & \Sigma _{Y}\\ \end{array} \right) . \end{aligned}$$

(C8)

Table 13 Top-1 accuracy of compressed VGG using different local metrics and architectures on CIFAR-10

Therefore, the mutual information between Gaussian distributed random variables X and Y is represented as,

$$\begin{aligned} I(X;Y)=\text {ln}|\Sigma _X|+\text {ln}|\Sigma _Y|-\text {ln}|\Sigma _{(X,Y)}|. \end{aligned}$$

(C9)

According to Everitt inequality, $|\Sigma _{(X,Y)}| \le |\Sigma _X||\Sigma _Y|$, this inequality holds if and only if $\Sigma _{YX}=\Sigma _{XY^T}=X^TY$ is a zero matrix. In other words,

$$\begin{aligned} \begin{aligned}&\min I(X;Y) \\ \propto&\min ~||\varvec{X}^T\varvec{Y}||_F^2 \\ \propto&\min ~||\varvec{X}^T\varvec{W}^T\varvec{X}||_F^2. \end{aligned} \end{aligned}$$

(C10)

For simplicity, we take an orthogonalized $\varvec{X}$. According to the unitary invariance of Frobenius norm, we have,

$$\begin{aligned} \begin{aligned}&\min ~||\varvec{X}^T\varvec{W}^T\varvec{X}||_F^2 \\ =&\min ~||\varvec{W}||_F^2. \end{aligned} \end{aligned}$$

(C11)

$\square$

Appendix D Ablation Study

We conduct ablation study to evaluate our calim that selecting weights in a layer need to be pruned does not help as much as finding a better $\alpha$. As we can see in Tab. 13, ITPruner always shows a significant improvement compared to their original methods. Meanwhile, integrating different pruning methods show a very similar performance. We thus conclude that selecting which weights in a layer need to be pruned does not help as much as finding a better $\alpha$, which is also well verified in the previous survey (Blalock et al., 2020).