Neural Network Pruning as Spectrum Preserving Process

Early attempts in pruning neural networks dated back to optimal brain damage (LeCun et al., 1990) and optimal brain surgeon (Hassibi and Stork, 1993) where the authors employed the second-order partial derivative matrix to decide the importance of connections and remove those unimportant ones. The simple yet very effective magnitude-based approach was examined in (Han et al., 2015b) for both dense layers and convolutional layers, where small entries in terms of absolute value were removed from the network. The work was further expanded (Han et al., 2015a) with the quantization techniques,which was previously introduced by (Gong et al., 2014). Channel-wise pruning (Li et al., 2016) for convolution was also examined based on the $l_{1}$-norm magnitude of the filters. The lottery ticket hypothesis (Frankle and Carbin, 2018) and a consequent work (Zhou et al., 2019) have again attracted the attention of the community on the magnitude-based pruning.

There is another line of efforts (Denton et al., 2014) (Vasilescu and Terzopoulos, 2002) (Lebedev et al., 2014) (Kim et al., 2015) (Liu et al., 2015) based on low-rank approximation, which has a similar taste to pruning but with a more clear theoretical support, falling within a larger scope of work namely neural network compression (Cheng et al., 2017). The low rank approximation was applied to the weight matrix or tensors (Kolda and Bader, 2009) in order to reduce the storage requirements or inference time for large pre-trained models. Some other works on neural network compression also include weight sharing (Chen et al., 2015) and hash trick (Chen et al., 2016), where they also look at the problem in the frequency domain.

Matrix sparsification is important in many numerical problems, e.g. low-rank approximation, semi-definite programming and matrix completion, which widely exist in data mining and machine learning problems. Matrix sparsification is to reduce the number of nonzero entries in a matrix without altering its spectrum. The original problem is NP-hard (McCormick, 1983)(Gottlieb and Neylon, 2010). The study of approximation solutions to this problem was pioneered by (Achlioptas and McSherry, 2007), and further expanded in (Achlioptas et al., 2013b) (Arora et al., 2006) (Achlioptas et al., 2013a) (Nguyen et al., 2009) (Drineas and Zouzias, 2011). An extensive study on the error bound was done in (Gittens and Tropp, 2009).

Since the spectrum of the sparsified matrix does not deviate significantly from that of the original matrix, serving as a linear operator the matrix retains its functionality, i.e. $A\in\mathbb{R}^{m\times n}$ is a mapping $\mathbb{R}^{m}\rightarrow\mathbb{R}^{n}$. We can define the matrix sparsification process as the following optimization problem:

where $A$ is the original matrix, $\tilde{A}$ is the sparsified matrix, $\|\cdot\|_{0}$ is the $0$-norm that equals the number of non-zero entries in a matrix, $\|\cdot\|$ denotes matrix norm, $\epsilon\geq 0$ is the error tolerance.

In matrix sparsification, we often use the spectral norm (2-norm) $\|\cdot\|_{2}$ and the Frobenius norm (F-norm) $\|\cdot\|_{F}$ to measure the deviation of the sparsified matrix from the original one.

In a dense layer, we focus on the $\textbf{z}^{T}A$ part since it contains most parameters. Neural network is essentially a function simulator that learns some artificial features, which is achieved by linear mappings, nonlinear activations, and some other customized units (e.g. recurrent unit). For the linear mapping, the analysis is usually done on the spectral domain.

Recall that Singular Value Decomposition (SVD) is optimal under both spectral norm (Eckart and Young, 1936) and Frobenius norm (Mirsky, 1960). The weight matrix $A\in\mathbb{R}^{m\times n}$ as a linear operator can be decomposed as

where $U=[\textbf{u}_{1},\textbf{u}_{2},...,\textbf{u}_{m}]\in\mathbb{R}^{m\times m}$ is the left singular matrix, $V=[\textbf{v}_{1},\textbf{v}_{2},...,\textbf{v}_{n}]\in\mathbb{R}^{n\times n}$ is the right singular matrix and and $\Sigma$ contains the singular values $diag(\sigma_{1},\sigma_{2},...,\sigma_{n})$ in a non-increasing order.

Note that a input signal $\textbf{z}\in\mathbb{R}^{m}$ can be written into a linear combination of $\textbf{u}_{i}$, i.e. $\textbf{z}_{i}=\sum_{i}^{m}c_{i}\textbf{u}_{i}$ where $c_{i}$ are the coefficients. Thus, the mapping from $\textbf{z}\in\mathbb{R}^{m}$ to $\textbf{z}^{\prime}\in\mathbb{R}^{n}$ is

where $\sigma_{j}$ are non-increasingly ordered, since $\textbf{u}_{i}^{T}\textbf{u}_{j}=0,\forall i\neq j$ and $\textbf{u}_{i}^{T}\textbf{u}_{i}=1,\forall i$.

We are especially interested in finding out how the spectrum $\Sigma$, i.e., singular values, of the linear mapping evolve during neural network training. As a matter of fact, our empirical study shows that the neural network training process is a spectrum learning process. To investigate the spectral structure, we design the following task:

Task 1: check the spectra and norms of dense layer weight matrices and how they change during neural network training.

Following the logic, if the neural network training is a spectrum learning process, when we practise neural network pruning, we would like to preserve the spectrum in order to preserve the neural network performance. In other words, we want to obtain a sparse $\tilde{A}$ that has similar singular values to $A$. How to measure the wellness of spectrum preservation? We can use the spectral norm (2-norm) $\|A\|_{2}=\sigma_{1}$ which is the largest singular value since we care about the dominant principle component, and the Frobenius norm (F-norm)

which is usually considered an aggregation of the whole spectrum. Note that $\|A\|_{2}\leq\|A\|_{F}$ and $\|A\|_{F}\leq\sqrt{min(m,n)}\|A\|_{2}$.

Therefore the goal is to find a sparse $\tilde{A}$ such that $\|A-\tilde{A}\|_{2}\leq\epsilon$ or $\|A-\tilde{A}\|_{F}\leq\epsilon$. To show that the pruning process is a spectrum preserving process, we design the following task.

Task2: apply magnitude-based thresholding at different sparsity levels and check the relationship between the resulting weight matrix spectra and the performance of the pruned neural networks.

From the optimization perspective, the pretrained neural network achieved satisfying local optimum. Once the weight matrix is sparsified, the spectrum to some extent deviates and the optimality no longer holds. We wish to retrain the neural network such that it returns to the satisfying local optimum and the performance is preserved. Apparently we don’t want the neural net deviate from the local optimum too far, otherwise it would be difficult to return to the optimum due to the nonconvex optimization process of neural network. Hence, iterative pruning and retraining has been a common practice for neural network pruning. We also design a task to elaborate from the spectrum perspective.

Task3: Inspect the weight matrix spectra in each pruning iteration, i.e. the weight matrix spectra after pruning and after retraining respectively, and compare the corresponding neural network performances. We also include a comparison to one-shot pruning.

Magnitude-based neural network pruning have attracted a lot of attention and show supprisingly simplicity and superior efficacy. In the context of matrix sparsification, this is a straightforward approach, namely magnitude-based matrix sparsification or hard thresholding. Given a matrix $A$, let $\tilde{A}$ denote its sparsifier. Entry-wise we have

.

The fact can be trivially verified since using $|A_{ij}|$ and using $A_{ij}^{2}$ (on which F-norm is based) are equivalent in terms of deciding small entries in a matrix. However throwing away small entries does not always guarantee the optimal sparsification result in terms of 2-norm. And in many situations, we care more about the dominant singular value instead of the whole spectrum.

In randomized matrix sparsification, each entry is sampled according to some distribution independently and then rescaled. E.g. each entry is sampled according to a Berboulli distribution, and we either set it to zero or rescale it.

where $p_{ij}$ can be a constant or positively correlated to the magnitude of the entry. The following theorem provides the justification to this type of matrix sparsification.

By weak spectrum, it means small matrix norm. To be more concrete, since matrix norm is a metric and triangle inequality applies, we have

. We need to show that $N=A-\tilde{A}$ falls within the category of matrices described in Theorem 4.1. Since

and

, as long as $A_{ij}^{2}/p_{ij}$ is upper-bounded, which is true most of time, $var(N_{ij})$ is bounded. Therefore the randomized matrix sparsification can guarantee the error bound.

In this section, we propose a customized matrix sparsification algorithm to show the potential of designing a better spectrum preservation process in neural network pruning. We do not intend to present a new state-of-the-art neural network pruning algorithm. There are two important points in our proposed algorithm: truncation and sampling based on the principal components of explicit truncated SVD. To the best of our knowledge, sampling based on probability proportional to principal components is employed for the first time in designing matrix sparsification for neural network pruning .

First, we adopt the truncation trick that is common in existing work. As clearly pointed out by (Achlioptas et al., 2013a), the spectrum of the random matrix $N=A-\tilde{A}$ is determined by its variance bound. Usually, the larger the variance, the stronger the spectrum of the random matrix. Existing works took advantage of the finding and proposed truncation (Arora et al., 2006)(Drineas and Zouzias, 2011) in sparsification, i.e. to set small entries to zero while leaving large entries as is and sampling on the remaining ones.

where $p_{ij}\propto|A_{ij}|$, $t$ is decided by the quantile (leave large entries as is), and $c$, the lower threshold for zeroing weights, as a constant could be set manually, and $Bern(\cdot)$ denotes Bernoulli distribution.

Second, instead of sampling based on the probability calculated from the magnitude of the original matrix entry, we do sampling based on the probability calculated from the principal component matrix entry magnitude with a little compromise on complexity, in order to better preserve the dominant singular values. Matrix sparsification was originally proposed for fast low-rank approximation on very large matrices, due to the fact that sparsity accelerates matrix-vector multiplication in power iteration. Essentially, we desire to find the sparse sketch $\tilde{A}$ of $A$ that preserves the dominant singular values well. This coincides with the goal of layer-wise neural network pruning from the spectrum preserving viewpoint – we desire to preserve the dominant singular values, based on the fact that we often consider information lies in the low-frequency domain while noises are in the high-frequency domain. The major difference is that weight matrices in neural network, either from dense layers or convolutional layers, are usually not too large, and therefore explicit SVD or truncated SVD on them is fairly affordable. Once we have access to the principle components of the weight matrices, we are able to preserve them better in the sparsification process. Note that preserving dominant singular values is a harmonic approach between preserving the 2-norm and the F-norm, since $\|A\|_{2}=\lim_{p\to\infty}(\sum_{i}\sigma_{i}^{p})^{1/p}$ and $\|A\|_{F}=\lim_{p\to 2}(\sum_{i}\sigma_{i}^{p})^{1/p}$.

The crutial part is to find the low-rank approximation $B$ to $A$, where $B=\sum_{i=1}^{K}\sigma_{i}\textbf{u}_{i}\textbf{v}_{i}^{T}$ and $\sigma_{i}\textbf{u}_{i}\textbf{v}_{i}^{T}$ are from SVD on A. We set the entry-wise sampling probability based on $|B_{ij}|$, i.e. $p_{ij}\propto|B_{ij}|$. Alg 1 presents the sparification algorithm. The partition function is the one used in quicksort.

We need to demonstrate that the proposed sparsification algorithm preserves dominant singular values better and improves the generalization performance of the pruned network. We propose the following task to check whether it makes improvement based on our analysis.

Task 4 Apply the above algorithm on VGG19 layer-wise weight matrix sparsification, compare the generalization performance of the pruned network and the pruned network given by thresholding at the same sparsity level.

Here we provide a high level proof on sparsification error being upper bounded. Let $D$ denote the sparse sketch generated by setting smallest entries in $A$ to 0 and $\tilde{A}$ as usual the final sparsfied result. From Fact Fact we know that $D$ is the optimal sketch of $A$ in terms of F-norm, i.e. $\|A-D\|_{F}=\epsilon_{*}$. Based on Theorem 4.1 and its illustration we know that $N=\tilde{A}-D$ satisfies the zero-mean and bounded-variance condition. Hence $\|\tilde{A}-D\|_{F}\leq\epsilon_{0}$. Therefore if we apply triangle equality given matrix norm is a metric,

Some other techniques, e.g. quantization(Gong et al., 2014)(Han et al., 2015a), can be used together with sparsification to further compress matrices and neural networks. Essentially they are also spectrum preservation techniques (Achlioptas and McSherry, 2007)(Arora et al., 2006).

In this section we state and illustrate the following fact.

To see this, suppose we have a convolutional layer with input signal size of $W\times H$ as width by height, and with $C$ input channels and $O$ output channels. Here we consider a 2-d convolution on the signal. The kernel is of size $C\times K\times K$ and there are $O$ such kernels. For the sake of simplicity in notations, suppose the striding step is 1, half-padding is applied and there is no dilation (for even $W$ and $H$ the above setting results in output signal of size $W\times H$ as width by height). 2-d convolution means that the kernel is moving in two directions. Fact Fact has been utilized to optimize lower-level implementation of CNN on hardware (Chellapilla et al., 2006)(Chetlur et al., 2014). Here we take advantage of the idea to unify neural network pruning on dense layers and convolutional layers with matrix sparsification.

Let us focus on one single output channel, one step of the convolution operation is the summation of element-wise product of two higher-order array, i.e. the kernel $G\in\mathbb{R}^{C\times K\times K}$ and the receptive field of the signal of the same size $X\in\mathbb{R}^{C\times K\times K}$. Note that taking the summation of element-wise product is equivalent to vector inner product. Therefore if we unfold the kernel for a single output channel to a vector and rearrange the receptive field of the signal accordingly to another vector, a single convolution step can be treated as two vector inner product, i.e. $G*X=\textbf{g}^{T}\textbf{x}$ where $\textbf{g},\textbf{x}\in\mathbb{R}^{CKK}$. Since we have $O$ output channels in total, there are $O$ such kernels of the same size. All of them being unfolded, we then can convert the convolution into a matrix product $Z^{T}A$, where $A\in\mathbb{R}^{CKK\times O}$ being the kernels and $Z\in\mathbb{R}^{CKK\times WH}$ being the rearranged input signals. And consequently the output signal $Y\in\mathbb{R}^{WH\times O}$ (as mentioned before, stride 1, half padding and no dilation result in input signal and output signal being in the same shape). Figure 8 visualizes convolution as matrix multiplication.

The matrix multiplication representation of convolution discussed above generalizes to any other convolution settings. Also note that the way we unfold the filters does not affect the spectrum of the resulting matrix, since row and column permutations do not change matrix spectrum. Therefore, all the analyses based on simple linear algebra we have discussed so far generalize to convolutional layer pruning. To verify our analyses,

Task1,2,3 will also be conducted on convolutional layers. When we say weight matrix in a context of convolution, it refers to as the matrix unfolded from the convolution higher-order array in the way described in Figure 8.

Entry-wise pruning almost always results in unstructured sparsity that requires specific data structure design in network deployment in order to realize the complexity reduction from pruning. Therefore it’s desirable to prune entire channels from convolutional layers to achieve higher efficiency. There is another important work on pruning channels (Li et al., 2016). The approach is to take small $\sum_{ijk}|T_{ijk}|$ where $T$ denotes the filter for a specific channel. This is equivalent to remove a column in $A$ we just discussed with small-magnitude values. It’s also a spectrum preserving process as $\sum_{ijk}|T_{ijk}|$ is fairly a proximity to $\sum_{ijk}(T_{ijk})^{2}$ on which the F-norm is based. Hence pruning the whole filter with small $\sum_{ijk}|T_{ijk}|$ is to preserve the F-norm of the convolution matrix we discussed in the previous subsection. To check the relation between $\sum_{ijk}|T_{ijk}|$ and the F-norm of $\tilde{A}$, we propose the following task.

Task5 Pruning different filters and check the relationship between $\sum_{ijk}|T_{ijk}|$ and the resulted convolution matrix F-norm $\|\tilde{A}\|_{F}$.