O NE -S HOT P RUNING OF R ECURRENT N EURAL N ETWORKS BY J ACOBIAN S PECTRUM E VALUATION
A BSTRACT
Recent advances in the sparse neural network literature have made it possible to prune many large feed forward and convolutional networks with only a small quantity of data.
Yet, these same techniques often falter when applied to the problem of recovering sparse recurrent networks.
These failures are quantitative: when pruned with recent techniques, RNNs typically obtain worse performance than they do under a simple random pruning scheme.
The failures are also qualitative: the distribution of active weights in a pruned LSTM or GRU network tend to be concentrated in speciﬁc neurons and gates and not well dispersed across the entire architecture.
We seek to rectify both the quantitative and qualitative issues with recurrent network pruning by introducing a new recurrent pruning objective derived from the spectrum of the recurrent Jacobian.
Our objective is data efﬁcient (requiring only 64 data points to prune the network), easy to implement and produces% sparse GRUs that signiﬁcantly improve on existing baselines.
We evaluate on sequential MNIST, Billion Words, and Wikitext
I NTRODUCTION
Within the neural network community, network pruning has been something of an evergreen problem.
There are several motivations for pruning a neural network.
Theoretically, overparameterization is a well known but poorly understood quality of many networks.
Pruning algorithms provide a link between overparameterized models appropriately parameterized models.
Thus, these algorithms may provide insights into exactly why overparameterized models have so much success.
Indeed, recent work has closely linked the efﬁcient utilization of model capacity with generalization results (Arora et al., 2018).
From a more practical perspective, overparameterized networks require more storage capacity and are computationally more expensive than their pruned counterparts.
Hence, there is an incentive to use pruned networks rather than fully dense networks during deployment.
For years, many of the most successful network pruning techniques were iterative — relying on cycle of pruning and retraining weights to induce sparsity in the network.
As identiﬁed in Lee et al.
(2018), these methods usually either enforce a sparsity-based penalty on the weights (Han et al., 2015; LeCun et al., 1990), or else prune based on some ﬁtness criteria (Carreira-Perpinán & Idelbayev, 2018;
Chauvin, 1989).
Recent advances in pruning literature suggest that such costly cycles of pruning and retraining might not always be necessary.
For some problems, there exists a small subnetwork within the original larger network such that training against this smaller network produces comparable performance to training against the original fully dense network.
The Lottery Ticket Hypothesis (Frankle & Carbin, 2019) provides a method for recovering these networks, but only after training is complete.
SNIP (Lee et al., 2018) and GraSP (Wang et al., 2020) provide a saliency criterion for identifying this small subnetwork using less than 100 data points, no training, and no iterative pruning.
Our present work began by asking the question: “How well do these newly discovered pruning techniques, which optimize a network sensitivity objective, work on recurrent neural networks?”
Although Lee et al.
(2018) does evaluate the SNIP pruning criterion on both GRU and LSTM networks, we found these results to be somewhat incomplete.
They did not provide a comparison to random pruning, and the chosen tasks were not extensive enough to draw deﬁnitive conclusions.
When compared against random pruning, we found that the SNIP and GraSP pruning objective performed similarly to or worse than random pruning.
This left us wondering where those techniques were falling short, and if a better pruning objective could be developed that takes the temporal structure of recurrent networks into account.
In this paper, we propose a new pruning objective for recurrent neural networks.
This objective is based on recent advances in mean ﬁeld theory (Gilboa et al., 2019; Chen et al., 2018a), and can be interpreted as forcing the network to preserve weights that propagate information through its temporal depths.
Practically, this constraint is imposed by forcing the singular values of the temporal
Jacobian with respect to the network weights to be non-degenerate.
We provide a discussion about the similarities and differences between our objective and the SNIP and GraSP pruning objectives.
It can be shown that these prior objectives fail to ensure that the temporal Jacobian of the recurrent weights is well conditioned.
Our method is evaluated with a GRU network on sequential MNIST,
Wikitext, and Billion Words.
At 95% sparsity, our network achieves better results than fully dense networks, randomly pruned networks, SNIP (Lee et al., 2018) pruned networks, and GraSP (Wang et al., 2020) pruned networks.
P RUNING R ECURRENT N ETWORKS BY J ACOBIAN S PECTRUM E VALUATION 2. N OTATION
We denote matrices and vectors by upper- and lower-case bold letters respectively.
Vector-valued functions are bolded, whereas scalar valued functions are not.
Distributions over variables are with the following script: D , P .
We denote the standard (cid:96) p norm of a vector by (cid:107)·(cid:107) p .
Let [ · ] ij be the (i,j)-th element of a matrix, and [ · ] i the i-th element of a vector.
(cid:126) 1 ,(cid:126) 0 , denotes a vector of 1 s or 0 s of appropriate length, and use (cid:12) denotes a Hadamard product.
I A represents the standard indicator function.
For vectors, superscripts are always used for sequence lengths while subscripts are reserved for indexing
vector elements. 2. P RELIMINARIES 2.2. R ECURRENT M ODELS
Let X = (cid:8) x
( t ) (cid:9) S t =1 ; with each x ( t ) ∈ R D .
Similarly, let ˆY = (cid:8) ˆ y ( t ) (cid:9) S t
=1 , where each y ( t ) ∈
R O is an associated set of outputs, such that each tuple ( X , Y ) i.i.d.
∼ D .
Let M ( x ; θ ) : R D (cid:55)→ R O be a generic model, parameterized by θ ∈ R N , that maps X onto an output sequence.
Deﬁne a recurrent model as a mapping done through iterative computation, such that each ( ˆy ( t ) , h ( t ) )
=
M ( x ( t ) , ˆh ( t − 1) ; θ ) depends explicitly only on the current input and some latent state of the model, h .
We deﬁne a loss over the entire sequence of outputs as the sum of a non-sequential loss function, ˜ L , over an entire sequence: L ( M , X , Y )
= (cid:80)
Si =1 ˜ L ( ˆy ( t ) , y ( t ) ) .
We deﬁne a sparse model as one where the parameters factorize into θ =
c (cid:12) w , with c ∈ { 0 , 1 } N a binary mask and w ∈ R N the free values, typically trained by gradient descent.
We deﬁne a K -sparse condition on a sparse model M as the restriction (cid:107) c (cid:107)
=
K during the entire training trajectory.
A model is optimally K -sparse if it minimizes the expected loss, E D
[ L ( M , X , Y )]
after training while also being subject to a K -sparse condition.
2.2. M EMORY H ORIZON
We introduce the following terms: N is the size of the network hidden state h , and J t ∈ R N × N is the temporal Jacobian, of the hidden state at time t
+ 1 with respect to the previous hidden state, ∂ h ( t +1)
∂ h ( t ) , and σ ( t ) i the singular values of said matrix.
To arrive at a one-shot pruning criteria for recurrent neural networks, we consider the impact of the temporal Jacobian on both forward- and backward-propagation.
• (Backpropagation) The formula for backpropagation through time (BPTT), from the loss at time s can be given as:
[Equation]
Where ˜G h ( t ) ; θ is the Jacobian of h ( t ) considering only the explicit dependence on θ .
• (Forward Propagation)
A single time step of the network under small perturbations yields the following
[Equation]
With additional powers of the Jacobian appearing as we observe the entire sequence.
From Equation 1, it can easily be seen that increasing the normed singular values of each J t will on average exponentially increase the gradient signal from later sequence elements, which will expedite convergence by reducing the vanishing gradient problem.
From Equation 2, we additionally note that a well-conditioned Jacobian would enable the network to preserve separation of distinct input vectors, by preventing the additive perturbation from vanishing or exploding.
Prior works in mean-ﬁeld theory
Gilboa et al.
; Chen et al.
(2018a) provide an extensive analysis of a similar objective on the performance of a wide range of recurrent networks.
The Froebenius norm of the temporal Jacobian, deﬁned below, is thus key to both forward and backpropagation.
Both processes are signiﬁcantly expedited when the norm is close to 1.
[Equation]
(cid:12)(cid:12)(cid:12)
[Equation]
(cid:12)(cid:12)(cid:12)
[Equation]
(cid:33)
[Equation]
2. P RUNING C RITERIA
Under typical recurrent model initializations, where θ ∼ N ( µ θ , s 2 θ I ) or a similar distribution, with µ θ ∼ 0 , s 2 θ (cid:28) 1 , Gilboa et al.
(2019) has empirically observed that χ is < 1 , and that the singular values concentrate towards 0 (see Figure 2 for further evidence).
Therefore, we hypothesize that the fastest converging and best performing sparse models are those which simply maximize χ .
We would like to determine the effect of removing one parameter on the Jacobian during the training trajectory.
However, as we restrict ourselves only to information available at initialization, we approximate the effect of each parameter on the Jacobian by a ﬁrst-order Taylor expansion.
This is analogous to the derivations given in Lee et al.
; Wang et al.
[Equation]
We call d n the sensitivity score of parameter θ n .
This criterion will not be well-normed across different types of parameters.
This is due to numerous factors, including differing activation functions used for each gate, and differing distributions between the input and recurrent state.
Consequently, the variance of our objective is not uniform between groups of parameters (see Section 3.3 for empirical conﬁrmation).
We compensate for this by dividing our criterion by the expected magnitude of the gradient for each parameter.
The normalized sensitivity score becomes:
[Equation]
where ˜ D is either the data distribution or an approximate distribution (since we are only trying to estimate the approximate variance of the gradient distribution), and the sequence { ˜h ( t ) } is computed on inputs from that distribution.
This normalization scheme is similar in motivation to the normalization proposed in Pascanu et al.
(2013), and
allows us to consider all recurrent models with only one additional computation.
For our pruning objective, we simply take the K weights with largest sensitivity scores, as those represent the parameters which most affect the Jacobian objective near the initialization.
Formally, we ﬁnd the K -th largest sensitivity, ˜ d K , and set c n = I A ( d n ≥ ˜ d K ) .
Empirically, we ﬁnd that the sensitivity score remains an effective metric even if the weights are not restricted to a neighborhood where the Taylor expansion is valid (see Figure 2 for details).
This objective is simple to compute, requiring only two backward passes using auto-differentiation.
Furthermore, as we only depend on the Jacobian-vector product, it has a memory cost linear in the parameters.
[Table]
2. C OMPARISON T O E XTANT M ETHODS
There are two recently proposed criteria for pruning at initialization: GraSP (Wang et al., 2020), and
SNIP (Lee et al., 2018).
They are given by:
[Equation]
where [ H ] ij = E [ ∂ 2 L ∂θ i ∂θ j ] , g = E [ ∇ θ L ] are the expected gradient and Hessian respectively.
Both methods rely on the gradient of the loss with respect to the weights, with SNIP being more dependent on this gradient than GraSP.
Thus, the main term of interest is g , which can be decomposed as:
[Equation]
With ˜ G t , the Jacobian of h ( t ) with respect to θ , as deﬁned in Equation 1.
A consequence of the smaller singular values of J is that the successive terms of ˜ G t tend to vanish over time.
Thus, loss-based gradient objectives tend to be biased toward explicit dependency between h ( t ) on θ , thus neglecting long-term dependence between h ( t ) and h ( t − 1) .
In certain cases, (ex. when the hidden state is small relative to the input) SNIP and GraSP prune many recurrent connections while leaving the input connections largely untouched (see section 3).
In contrast, our algorithm considers the J matrix explicitly, which mitigates the problem of pruning too many recurrent connections.
[Table]
E VALUATION
For the following experiments
[Table]
we compute the (cid:96) 2 norm of J (cid:126) 1 using a single minibatch of 64 data samples
[Table]
and using only the last 4 steps of the sequence.
3. S EQUENTIAL MNIST B ENCHMARK
We ﬁrst test our method on the sequential MNIST benchmark (Lee et al., 2018), a relatively small dataset which contains long term dependencies.
We begin by verifying that our algorithm is robust across several common recurrent architectures.
The results in Table 1 conﬁrm that our method is not dependent on any speciﬁc recurrent neural architecture choice.
Our principal results for the Sequential MNIST benchmark are presented in Table 2.
Again, we see that our network’s performance improves with network size, with the largest gap between our method and others coming when the network grows to 1600 units.
We observe that SNIP and GraSP are surprisingly effective at small scales with good initialization, but fail when scaled to larger network sizes.
Of the baselines, only random pruning is competitive when scaled, a fact we found quite interesting.
For reference, we also provide results on standard L2 pruning (Reed, 1993)
(for which the schedule can be found in the appendix) and random pruning.
The reader should be careful to note that L2 pruning requires an order of magnitude more resources than other methods due to it’s prune – retrain cycle; it is only considered here as a lower bound for network compression.
Furthermore, while GraSP requires computing the Hessian gradient across the entire dataset, this is computationally infeasible in our case and we instead compute it with a single minibatch, for fairness.
Pruning Scheme
[Table]
GraSP and Random pruning are competitive for smaller models, but the results tend to diminish as the network size increases.
Our method obtains strong results even as the network size is large.
Further experimental details can be found in the appendix.
In the preceding section, we postulated that normalization across the objective was necessary for strong performance (see Equation 5).
This intuition is conﬁrmed in Table 2, where we present both the normalized results (with Glorot Glorot & Bengio (2010) and γ normalization) and the unnormalized results (without both).
Indeed, we see that this normalization is crucial for recurrent architectures, with unnormalized architectures having all of the retained network weights concentrated in a single gate.
This proved to be prohibitive to training.
Finally, in Table 3, we examine the performance of our algorithm at various sparsity levels.
Our algorithm continues to outperform random pruning, even at high sparsity levels.
[Table]
3. L INGUISTIC S EQUENCE P REDICTION
We assess our models on 3 sequence prediction benchmarks: 1) WikiText-2 (wiki2).
2) WikiText [Table] (wiki103), an expanded version of (1) with 10 times more tokens.
A truncated version of the
One Billion Words (1b) benchmark (Chelba et al., 2013), where only the top 100,000 vocabulary tokens are used.
The full experiment parameters are given in the appendix.
We report the training and validation perplexities on a random 1% sample of the training set in Table
[Table]
From the results
[Table]
it is clear that our algorithm succeeds in decreasing perplexity across all language tasks.
Despite their varying difﬁculties, our algorithm speeds up initial convergence on all tasks and maintains an advantage throughout training.
Finally, we perform an ablation experiment on the Penn Treebank Dataset (PTB) with an 800 unit
GRU at different sparsity levels.
The results are reported in Table 5.
[Table]
The loss from sparsity increases dramatically as the percentage of parameters remaining approaches zero.
This trend is similar to that reported in Gale et al.
(2019) and other prior works.
For reference, a dense 200 unit GRU (360k parameters) achieves 196.31 perplexity while a 100 unit GRU (150k parameters) achieves 202.97 perplexity.
3. Q UALITATIVE A NALYSIS
The success of our algorithm can be partially attributed to effective distributions across hidden units.
Whereas many of the other algorithms are overly concentrated in certain gates and biased towards the input weights, our algorithm effectively distributes sparse connections across the entire weight matrix.
We discuss the distribution of remaining connections on a 400 unit GRUs in Figure 1.
We also give a set of sample connections under each algorithm in Figure 3.
Finally, we perform an empirical study of the evolution of the Jacobian spectrum to verify our hypothesis on recurrency preservation.
We show a 400-unit GRU trained on sequential MNIST, with a dense network, our pruning scheme, and random pruning respectively.
It can be observed from
Figure 2 after 50000 training steps that our Jacobian has both higher mean and much fewer near-zero singular values, which helps to explain our performance and justify the intuition behind our algorithm.
The spectra at initialization also further conﬁrms that the initial singular values of J are small.
[Figure]
[Figure]
O THER R ELATED W ORK
Methods for Pruning Recurrent Networks: Our method is the latest in a series of attempts to generate sparse RNNs.
Perhaps the most well-known algorithm for sparse network pruning is Narang et al.
(2017a).
It is a modiﬁcation to magnitude based pruning wherein the pruning threshold evolves
according to several hyperparameters that have to be tuned by the user.
Kliegl et al.
(2017) uses iterative trace norm regularization to prune RNNs used for speech recognition.
This effectively reduces the sum of the singular values of the weight matrices.
But we found in our experiments that these values were often degenerate near 0.
Furthermore, this technique is iterative.
Narang et al.
(2017b) uses iterative ground lasso regularization to induce block sparsity in recurrent neural networks.
Wen et al.
(2017) alters the structure of LSTMs to decrease their memory requirements.
Their intrinsic sparse structures make structural assumptions about the sparsity distribution across the network.
Dai et al.
(2018) uses magnitude based pruning coupled with a special RNN structure
[Figure]
to make RNNs more efﬁcient.
The pruning algorithm itself is magnitude based.
See et al.
(2016) uses iterative pruning and retraining to prune a recurrent model for neural translation.
The underlying technique is simple iterative pruning, and the ﬁnal pruning percentage is only 80%.
While ﬁne for their application, we are interested in novel pruning techniques and higher levels of sparsity.
In summation, all the the methods discussed above utilized some variant of L1 or L2 pruning to actually sparsify the network.
The novel advances are all related to pruning schedules, modiﬁcations to recurrent architectures, or small transformations of the L1 or L2 objective.
Other Pruning Techniques: Many extant pruning techniques are applicable to recurrent network architectures, even if these methods were not designed from the ground up to work in the recurrent case.
Lee et al.
(2018) and Wang et al.
(2020) both provide a pruning objective that can be used to prune networks before training begins.
They are considered extensively in this work.
In Frankle &
Carbin (2019), it is shown that at initialization networks contain a small sparse set of connections that can achieve similar results to fully dense networks.
However, no known method yet exists to recover these sparse networks to the full extent demonstrated in that work.
Han et al.
(2015) showed impressive results with magnitude based pruning.
Follow up work made further use of magnitudebased pruning techniques (Carreira-Perpinán & Idelbayev, 2018; Guo et al., 2016); however, these techniques are primarily iterative.
Mean Replacement Pruning (Evci et al., 2018) uses the absolute-value of the Taylor expansion of the loss to as a criterion for which units in a network should be pruned.
This method can not be used with BatchNorm and achieves results comparable to magnitude based pruning.
Bayesian methods have recently seen some success in pruning neural networks.
(Ullrich et al., 2017), which is itself an extension of Nowlan & Hinton (1992), is the standard citation here.
In essence, this method works by re-training a network while also ﬁtting the weights to a GMM prior via a KL penalty.
Molchanov et al.
(2017) is another Bayesian pruning technique that learns a dropout rate via variational inference that can subsequently be used to prune the network.
Finally, there exists several classical pruning techniques.
Ishikawa (1996); Chauvin (1989) enforced sparsity penalties during the training process.
LeCun et al.
; Hassibi et al.
(1993) perform Hessian-based pruning, using the Hessian to get a sensitivity metric for the network’s weights.
While many of the above methods are effective in general, they do not explicitly consider the speciﬁcs of RNNs and sequential prediction.
Other Related Work Several interesting papers have recently taken a critical look at the problem of network pruning (Liu et al., 2018; Crowley et al., 2018).
The problem of network compression is closely related to network pruning.
It would be impossible to cite all of the relevant papers here, and no good literature survey exists.
Some worthwhile references are Gupta et al.
; Gong et al.
(2014); Courbariaux et al.
; Chen et al.
(2018b)
; Howard et al.
Both problems often share a common goal of reducing the size of a network.
Some notable papers explicitly consider the problem of recurrent network compression (Ye et al., 2018; Lobacheva et al., 2017; Wang et al., 2018).
In the context of the above work, our method is not iterative and can be fully completed before training even begins.
The tradeoffs in accuracy can be remedied by scaling up the network, since there is no longer a need to store fully dense weights during training.
Furthermore, our objective is speciﬁcally adapted to the sequential prediction context in which RNNs are deployed.
We are the ﬁrst pruning algorithm to consider the temporal Jacobian spectrum as a key to generating faster converging and better performance sparse RNNs.
Our method not only performs better in practice compared to other zero-shot methods, but also yields key insight into the factors behind RNN performance.
This may aid the development of new architectures and training schemes for sequential prediction.
C LOSING R EMARKS
In this work, we presented an effective and cheap single-shot pruning algorithm adapted toward recurrent models.
Throughout the work, we continually found the importance of the Jacobian spectrum surprising and interesting.
Future work could further examine the relationship between network width, the Jacobian spectrum, and generalization.
A PPENDIX A - E XPERIMENT H YPERPARAMETERS
Unless otherwise speciﬁed, our model consists of a single-layered RNN, followed by an appropriately sized softmax layer with sigmoidal activation.
The softmax layer is initialized with standard Xavier.
We use a minibatch size of 64 samples during training, and optimize using the AdaM optimizer (Kingma & Ba, 2014) with a learning rate of 1e-3.
We use an initial hidden state of zeros for all experiments.
For all networks, we only prune the recurrent layer while leaving prior and subsequent layers untouched, since we are primarily interested in performance of recurrent layers.
We trained all networks with a single Nvidia P100 GPU.
6. S EQUENTIAL MNIST
For seq.
MNIST, we follow the same process as SNiP, feeding in row-by-row.
We used N (0 , 0 . 1) for our own method, and Glorot initialization for SNiP and GraSP.
γ is computed from data sampled from a N (0 , 0 . 1) distribution.
We use only the activations from the last time step.
For L2, the density was annealed according to the schedule { 0 . 8 , 0 . 6 , 0 . 4 , 0 . 2 , 0 . 1 , 0 .
01 } every 10k training steps.
6. L ANGUAGE B ENCHMARKS
We use 2000-unit LSTMs for all language benchmarks.
To reduce the variance of our comparison, we freeze the embedding layer before training.
We use sampled sequential cross-entropy loss with 1000 tokens for wiki103 and 1b, and standard cross-entropy for wiki2.
We use He initialization for all papers.
Wiki2 was trained for 20k training steps (13 epochs), while wiki103 was trained for 12k training steps, and 1b was trained for 30k training steps.
A PPENDIX B - A DDITIONAL S TUDIES 7. I NITIALIZATIONS
We benchmark the performance of our algorithm against random pruning using 3 additional initializations, seen in Table 6.
With high variance, the ﬁrst-order expansion we use to estimate our objective fails to hold, so we do signiﬁcantly worse than the random benchmark.
[Table]
7. R UNTIME
We benchmark the runtimes of SNiP, GraSP and our own algorithm, using only a single batch and iteration for fairness, seen in Table 7.
A PPENDIX C - T RAINING C URVES
We present a sample training curve of a 400 unit GRU for sequential MNIST below.
As can be seen, random pruning is only competitive algorithm in this instance.
[Table]
Subsequently we present a sample training curve in Figure 5 for the 1b words experiment, detailed in
Table 4.
Our algorithm provides signiﬁcant beneﬁt over random pruning, but still lags behind the dense model.