Spectral Pruning for Recurrent Neural Networks

Recurrent neural networks (RNNs) are a class of neural networks used in sequential tasks. However, in general, RNNs have a large number of parameters and involve enormous computational costs by repeating the recurrent structures in many time steps. These make their application difficult in edge-computing devices. To overcome this difficulty, RNN compression has attracted increasing attention in recent years. It brings us more benefits in terms of the reduction of computational costs as the time step progresses, compared to deep neural networks (DNNs) without any recurrent structure. There are many RNN compression methods such as pruning [(Narang; tang2015pruning; Zhang; Lobacheva; wang2019acceleration; Wen2020StructuredPO; lobacheva2020structured), low rank factorization [(Kliegl; Tjandra), quantization [(Alom; Liu), distillation [(Shi; Tang), and sparse training [(liu2021selfish; liu2021efficient; dodge2019rnn; wen2017learning). This paper is devoted to the pruning for RNNs, and its purpose is to provide an RNN pruning method with the theoretical background.

Recently, Suzuki et al. [(Suzuki) proposed a novel pruning method with the theoretical background, called spectral pruning, for DNNs such as the fully connected and convolutional neural network architectures. The idea of the proposed method is to select important nodes for each layer by minimizing the information losses (see (2) in [(Suzuki)), which can be represented by the layerwise covariance matrix. The minimization only requires linear algebraic operations. Suzuki et al. [(Suzuki) also evaluated generalization error bounds for networks compressed using spectral pruning (see Theorems 1 and 2 in [(Suzuki)). It was shown that generalization error bounds are controlled by the degrees of freedom, which are defined based on the eigenvalues of the covariance matrix. Hence, the characteristics of the eigenvalue distribution have an influence on the error bounds. We can also observe that in the generalization error bounds, there is a bias-variance tradeoff corresponding to compressibility. Numerical experiments have also demonstrated the effectiveness of spectral pruning.

In this paper, we extend the theoretical scheme of spectral pruning to RNNs. Our pruning algorithm involves the selection of hidden nodes by minimizing the information losses, which can be represented by the time mean of covariance matrices instead of the layerwise covariance matrix which appears in spectral pruning of DNNs. We emphasize that our information losses are derived from the generalization error bound. More precisely, we show that choosing compressed weight matrices which minimize the information losses reduces the generalization error bound we evaluated in Section 4.1 (see sentences after Theorem 4.5). We also remark that Suzuki et al. [(Suzuki) has not clearly mentioned anything about how the information losses are derived from. As in DNNs [(Suzuki), we can provide the generalization error bounds for RNNs compressed with our pruning and interpret the degrees of freedom and the bias-variance tradeoff.

We also provide numerical experiments to compare our method with existing methods. We observed that our method outperforms existing methods, and gets benefits from over-parameterization [(chang2020provable; zhang2021understanding) (see Sections 5.2 and 5.3). In particular, our method can compress models with small degradation (see Remark 3.2) when we employ IRNN, which is an RNN that uses the ReLU as the activation function and initializes weights as the identity matrix and biases to zero (see [(Le)).

The summary of our contributions is the following:

• •

  A pruning algorithm for RNNs (Section 3) is proposed by the analysis of generalization error (Remark 4.3 and Theorem 4.8).

• •

  The generalization error bounds for RNNs compressed with our pruning algorithm are provided (Theorem 4.8).

One of the popular compression methods for RNNs is pruning that removes redundant weights based on certain criteria. For example, magnitude-based weight pruning [(Narang; narang2017block; tang2015pruning) involves pruning trained weights that are less than the threshold value decided by the user. This method has to gradually repeat pruning and retraining weights to ensure that a certain accuracy is maintained. However, based on recent developments, the costly repetitions might not always be necessary. In one-shot pruning [(Zhang; lee2018snip), weights are pruned once prior to training from the spectrum of the recurrent Jacobian. Bayesian sparsification [(Lobacheva; molchanov2017variational) induce sparse weight matrix by choosing the prior as log-uniform distribution, and weights are also once pruned if the variance of the posterior over weight is large.

While the above methods are referred to as weight pruning, our spectral pruning is a structured pruning where redundant nodes are removed. The advantage of the structured pruning over the weight pruning is that it more simply reduces computational costs. The implementation advantages of structured pruning are illustrated in [(wang2019acceleration). Although weight pruning from large networks to small networks is less likely to degrade accuracy, it usually requires an accelerator for addressing sparsification (see [(Parashar)). The structured pruning methods discussed in [(wang2019acceleration; Wen2020StructuredPO; lobacheva2020structured) induce sparse weight matrices in the training process, and prune weights close to zero, and does not repeat fine-tuning. In our pruning, weight matrices are trained by the usual way, and compressed weight matrices consist of the multiplication of the trained weight matrix and the reconstruction matrix, and no need to repeat pruning and fine-tuning. The idea of the multiplication of the trained weight matrix and the reconstruction matrix is a similar idea to low rank factorization [(Kliegl; Tjandra; prabhavalkar2016compression; grachev2019compression; denil2013predicting). In particular, the work [(denil2013predicting) is most related to spectral pruning, and it employs the reconstruction matrix replacing the empirical covariance matrix with kernel matrix (see Section 3.1 in [(denil2013predicting)).

In general, RNN pruning is more difficult than DNN pruning, because recurrent architectures are not robust to pruning, that is, even a little pruning causes accumulated errors and total errors increase significantly for many time steps. Such a peculiar problem for recurrent feature is also observed in dropout (see Introduction in [(Gal; Zaremba)).

Our motivation is to theoretically propose the RNN pruning algorithm. Inspired by [(Suzuki), we focus on the generalization error bound, and we provide the algorithm so that the generalization error bound becomes smaller. Thus, the derivation of our pruning method would be theoretical, while that of existing methods such as the magnitude-based pruning [(Narang; narang2017block; tang2015pruning; wang2019acceleration; Wen2020StructuredPO) would be heuristic. For the study of the generalization error bounds for RNNs, we refer to [(tu2019understanding; Chen; akpinar2019sample; joukovsky2021generalization).

We propose a pruning algorithm for RNNs inspired by [(Suzuki). See Appendix A for a review of spectral pruning for DNNs. Let $D=\{(X^{i}_{T},Y^{i}_{T})\}_{i=1}^{n}$ be the training data with time series sequences $X^{i}_{T}=(x^{i}_{t})_{t=1}^{T}$ and $Y^{i}_{T}=(y^{i}_{t})_{t=1}^{T}$, where $x^{i}_{t}\in\mathbb{R}^{d_{x}}$ is an input and $y^{i}_{t}\in\mathbb{R}^{d_{y}}$ is an output at time $t$. The training data are independently identically distributed. To train the appropriate relationship between input $X_{T}=(x_{t})_{t=1}^{T}$ and output $Y_{T}=(y_{t})_{t=1}^{T}$, we consider RNNs $f=(f_{t})_{t=1}^{T}$ as

for $t=1,\ldots,T$, where $\sigma:\mathbb{R}\to\mathbb{R}$ is an activation function, $h_{t}\in\mathbb{R}^{m}$ is the hidden state with the initial state $h_{0}=0$, $W^{o}\in\mathbb{R}^{d_{y}\times m}$, $W^{h}\in\mathbb{R}^{m\times m}$, and $W^{i}\in\mathbb{R}^{m\times d_{x}}$ are weight matrices, and $b^{o}\in\mathbb{R}^{d_{y}}$ and $b^{hi}\in\mathbb{R}^{m}$ are biases. Here, an element-wise activation operator is employed, i.e., we define $\sigma(x):=(\sigma(x_{1}),\ldots,\sigma(x_{m}))^{T}$ for $x=(x_{1},\ldots,x_{m})\in\mathbb{R}^{m}$.

Let $\widehat{f}=(\widehat{f}_{t})_{t=1}^{T}$ be a trained RNN obtained from the training data $D$ with weight matrices $\widehat{W}^{o}\in\mathbb{R}^{d_{y}\times m}$, $\widehat{W}^{h}\in\mathbb{R}^{m\times m}$, and $\widehat{W}^{i}\in\mathbb{R}^{m\times d_{x}}$, and biases $\widehat{b}^{o}\in\mathbb{R}^{d_{y}}$ and $\widehat{b}^{hi}\in\mathbb{R}^{m}$, i.e., $\widehat{f}_{t}=\widehat{W}^{o}\widehat{h}_{t}+\widehat{b}^{o}$, $\widehat{h}_{t}=\sigma(\widehat{W}^{h}\widehat{h}_{t-1}+\widehat{W}^{i}x_{t}+\widehat{b}^{hi})$ for $t=1,\ldots,T$. We denote the hidden state $\widehat{h}_{t}$ by

as a function with inputs $x_{t}$ and $\widehat{h}_{t-1}$. Our aim is to compress the trained network $\widehat{f}$ to the smaller network $f^{\sharp}$ without loss of performance to the extent possible.

Let $J\subset[m]$ be an index set with $|J|=m^{\sharp}$, where $[m]:=\{1,\ldots,m\}$, and let $m^{\sharp}\in\mathbb{N}$ be the number of hidden nodes for a compressed RNN $f^{\sharp}$ with $m^{\sharp}\leq m$. We denote by $\phi_{J}(x_{t},\widehat{h}_{t-1})=(\phi_{j}(x_{t},\widehat{h}_{t-1}))_{j\in J}$ the subvector of $\phi(x_{t},h_{t-1})$ corresponding to the index set $J$, where $\phi_{j}(x_{t},\widehat{h}_{t-1})$ represents the $j$-th components of the vector $\phi(x_{t},\widehat{h}_{t-1})$.

(i) Input information loss. The input information loss is defined by

where $\|\cdot\|_{n,T}$ is the empirical $L^{2}$-norm with respect to $n$ and $t$, i.e.,

where $\|\cdot\|_{2}$ is the Euclidean norm, $\|A\|^{2}_{\tau}:=\mathrm{Tr}[AI_{\tau}A^{T}]$ for the regularization parameter $\tau\in\mathbb{R}^{m^{\sharp}}_{+}:=\{x\in\mathbb{R}^{m^{\sharp}_{l}}\,|\,x_{j}>0,\ j=1,\ldots,m^{\sharp}_{l}\}$, and $I_{\tau}:=\text{diag}(\tau)$. Here, $\widehat{\Sigma}_{I,I^{\prime}}\in\mathbb{R}^{K\times H}$ denotes the submatrix of $\widehat{\Sigma}$ corresponding to the index sets $I,I^{\prime}\subset[m]$ with $|I|=K$, $|I^{\prime}|=H$, i.e., $\widehat{\Sigma}_{I,I^{\prime}}=(\widehat{\Sigma}_{i,i^{\prime}})_{i\in I,i^{\prime}\in I^{\prime}}$. Based on the linear regularization theory (see e.g., [(gockenbach2016linear)), there exists a unique solution $\widehat{A}_{J}\in\mathbb{R}^{m\times m^{\sharp}}$ of the minimization problem of $\|\phi-A\phi_{J}\|^{2}_{n,T}+\|A\|^{2}_{\tau}$, which has the form

where $\widehat{\Sigma}$ is the (noncentered) empirical covariance matrix of the hidden state $\phi(x_{t},\widehat{h}_{t-1})$ with respect to $n$ and $t$, i.e.,

We term the unique solution $\widehat{A}_{J}$ as the reconstruction matrix. Here, we would like to emphasize that the mean of the covariance matrix with respect to time $t$ is employed in RNNs, while the layerwise covariance matrix is employed in DNNs (see Appendix A). By substituting the explicit formula of the reconstruction matrix $\widehat{A}_{J}$ into (3.1), the input information loss is reformulated as:

(ii) Output information loss. The hidden state of a RNN is forwardly propagated to the next hidden state or output, and hence, the two output information losses are defined by

where $\widehat{W}^{o}_{j,:}$ and $\widehat{W}^{h}_{j,:}$ denote the $j$-th rows of the matrix $\widehat{W}^{h}$ and $\widehat{W}^{o}$, respectively. Then, the unique solutions of the minimization problems of $\|\widehat{W}^{o}_{j,:}\phi-\beta^{T}\phi_{J}\|^{2}_{n,T}$ $+$ $\|\beta^{T}\|^{2}_{\tau}$ and $\|\widehat{W}^{h}_{j,:}\phi-\beta^{T}\phi_{J}\|^{2}_{n,T}$ $+$ $\|\beta^{T}\|^{2}_{\tau}$ are $\widehat{\beta}^{o}=(\widehat{W}^{o}_{j,:}\widehat{A}_{J})^{T}$ and $\widehat{\beta}^{h}_{j}=(\widehat{W}^{h}_{j,:}\widehat{A}_{J})^{T}$, respectively. By substituting them into (3.5) and (3.6), the output information losses are reformulated as

Here, we remark that the output information losses $L^{(B,o)}_{\tau}(J)$ and $L^{(B,h)}_{\tau}(J)$ are bounded above by the input information loss $L^{(A)}_{\tau}(J)$ (see Remark 4.3).

(iii) Compressed RNNs. We construct the compressed RNN $f^{\sharp}_{J}$ by $f^{\sharp}_{J,t}=W^{\sharp o}_{J}h^{\sharp}_{J,t}+b^{\sharp o}_{J}$ and $h^{\sharp}_{J,t}=\sigma(W^{\sharp h}_{J}h^{\sharp}_{J,t-1}+W^{\sharp i}_{J}x_{t}+b^{\sharp hi}_{J})$ for $t=1,\ldots,T$, where $W^{\sharp o}_{J}:=\widehat{W}^{o}\widehat{A}_{J}$, $W^{\sharp h}_{J}:=\widehat{W}^{h}_{J,[m]}\widehat{A}_{J}$, $W^{\sharp i}_{J}:=\widehat{W}^{i}_{J,[d_{x}]}$, $b^{\sharp hi}_{J}:=\widehat{b}^{hi}_{J}$, and $b^{\sharp o}_{J}:=\widehat{b}^{o}$.

(iv) Optimization. To select an appropriate index set $J$, we consider the following optimization problem that minimizes the convex combination of the input and two output information losses:

for $\theta_{1},\theta_{2},\theta_{3}\in[0,1]$ with $\theta_{1}+\theta_{2}+\theta_{3}=1$, where $m^{\sharp}_{l}\in[m]$ is a prespecified number. The optimal index $J^{\sharp}$ is obtained by the greedy algorithm. We term this method as spectral pruning (for a schematic diagram of spectral pruning, see Figure 1). The reason why information losses are employed in the objective will be theoretically explained later, when the error bounds in Remark 4.3 and Theorem 4.5 are provided. We summarize our pruning algorithm in the following.

In this section, we discuss the generalization error bounds for compressed RNNs. In Subsection 4.1, the error bounds for general compressed RNNs are evaluated to explain the reason for deriving spectral pruning discussed in Section 3 in the error bound term. In Subsection 4.2, the error bounds for RNNs compressed with spectral pruning are evaluated.

In this section, numerical experiments are detailed to demonstrate our theoretical results and show the effectiveness of spectral pruning compared with existing methods. In Sections 5.1 and 5.2, we select the pixel-MNIST as our task and employ the IRNN, which is an RNN that uses the ReLU as the activation function and initializes weights as the identity matrix and biases to zero (see [(Le)). In Section 5.3, we select the PTB [(marcus1993building) and employ the RNNLM whose RNN layer is orthodox Elman-type. For RNN training details, see Appendix H. We choose parameters $\theta_{1}=1$, $\theta_{2}=\theta_{3}=0$ in (iv) of Section 3, i.e., we minimize only the input information loss. This choice is not so problematic because the bound of output information loss automatically becomes smaller with minimizing the input one (see Remark 4.3). We choose the regularization parameter $\tau=0$, where this choice regards $\widehat{f}$ as a well-trained network and gives priority to minimizing the approximation error between $\widehat{f}$ and $f^{\sharp}_{J}$ (see below Theorem4.5).