A DAPTIVE P RUNING OF N EURAL L ANGUAGE M ODELS FOR M OBILE D EVICES
Anonymous authors
Paper under double-blind review
A BSTRACT
Neural language models (NLMs) exist in an accuracy–efﬁciency tradeoff space where better perplexity typically comes at the cost of greater computation complexity.
In a software keyboard application on mobile devices, this translates into higher power consumption and shorter battery life.
This paper represents the ﬁrst attempt, to our knowledge, in exploring accuracy–efﬁciency tradeoffs for NLMs.
Building on quasi-recurrent neural networks (QRNNs), we apply pruning techniques to provide a “knob” to select different operating points.
In addition, we propose a simple technique to recover some perplexity using a negligible amount of memory.
Our empirical evaluations consider both perplexity as well as energy consumption on a Raspberry Pi, where we demonstrate which methods provide the best perplexity–power consumption operating point.
At one operating point, one of the techniques is able to provide energy savings of 40% over the state of the art with only a 17% relative increase in perplexity.
I NTRODUCTION
An emerging application of neural language models (NLMs) is smart software keyboards on such mobile devices as smartphones and tablets that provide next-word prediction, allowing users to input entire words with a single tap.
For example, the app SwiftKey 1 advertises the use of neural networks for predictions; 2 according to Google Play Store, it has more than 100 million downloads, demonstrating its popularity.
Based on standard metrics such as perplexity, neural techniques represent an advance in the state of the art in language modeling (Merity et al., 2018b).
Better models, however, come at a cost in computational complexity, which translates to higher power consumption (Tang & Lin, 2018).
In the context of mobile devices, energy efﬁciency is, of course, an important optimization objective.
A casual web search, for example, reveals numerous complaints from users of the above apps about battery drain, indicating that this is not a hypothetical concern.
In reality, neural language models exist in a accuracy–efﬁciency tradeoff space.
Although this fact has been recognized for applications such as image recognition (Canziani et al., 2016) and keyword spotting (Tang et al., 2018), to our knowledge no one in the NLP community has explored these tradeoffs.
All previous papers on NLMs simply report single-point perplexity ﬁgures.
In contrast, the high-level goal of our work is to understand the tradeoffs between neural modeling accuracy and real-world efﬁciency constraints: in addition to perplexity, NLMs should be evaluated in terms of
FLOPs, 3 milliJoule per query (mJ/q), and inference latency.
We conduct exactly such experiments, using the Raspberry Pi (which shares the same architecture as most mobile devices today) as a more convenient hardware platform.
Ideally, NLMs should provide a “knob” that allows developers to tune accuracy–efﬁciency tradeoffs.
In this paper, we explore pruning approaches that take a pre-trained quasi-recurrent neural network (QRNN; Bradbury et al., 2017), representing the state of the art in NLM today, and provides exactly such a knob.
Furthermore, our techniques allow these tradeoffs to be tuned at inference time ,
which
[Figure]
allows a mobile device to adaptively control its behavior, e.g., favor efﬁciency at the cost of accuracy when the battery is low.
Thus, this paper makes the following contributions: First, to our knowledge, we are the ﬁrst to comprehensively explore accuracy–efﬁciency tradeoffs for NLMs with experimental evaluation of energy consumption on a Raspberry Pi.
Second, we evaluate a number of inference-time pruning techniques that takes any pre-trained QRNN and provides a tunable accuracy–efﬁciency “knob”.
B ACKGROUND AND R ELATED W ORK 2. Q UASI - RECURRENT N EURAL N ETWORKS
Quasi-recurrent neural networks (QRNNs; Bradbury et al., 2017) achieve highly competitive perplexity on word-level language modeling datasets, including state-of-the-art perplexity on WikiText103 (Merity et al., 2018b).
Although applying such techniques as dynamic evaluation (Krause et al., 2017), Hebbian softmax (Rae et al., 2018), and mixture of softmaxes (Yang et al., 2017) can produce lower perplexity, our focus is on the recurrent architecture.
Thus, we explore the task of pruning
QRNNs without using any other extensions.
Each word is encoded as a one-hot vector and then fed into a linear layer, which produces lowerdimensional word embeddings for the QRNN layers.
A single QRNN layer consists of two distinct components—convolution and recurrent pooling—that alternate to imitate an LSTM (Hochreiter &
Schmidhuber, 1997).
Given a stacked sequence of inputs X = x 1 ⊕ · · · ⊕ x n ∈ R k × n (e.g., word embeddings in language modeling), the one-dimensional convolution layer is deﬁned as
[Equation]
where W z , W f , W o are the weights associated with the input, forget, and output gates, respectively, · represents a masked convolution along time, and σ denotes the sigmoid function.
For W { z,f,o } ∈
R m × ( k × r ) , m is the number of output channels, k is the number of input channels, and r the window size across time.
Without loss of generality, we henceforth represent W { z,f,o } as two-dimensional matrices ∈ R m × s , where s = k × r .
The outputs are fed into a recurrent pooling layer:
[Equation]
where (cid:12) denotes element-wise product.
Altogether, these two layers deﬁne a single QRNN layer (Bradbury et al., 2017; see Figure 1).
Multiple layers can be stacked for greater expressiveness, where the output h 1: n of the previous layer is the input X to the current layer.
We tie the weights between the input and output layers, as used by Merity et al.
(2018a)
and proposed by Inan et al.
In addition to improving perplexity, weight tying reduces the number of parameters and hence the memory footprint, which is beneﬁcial to our task.
2. P RUNING
Weight pruning is an effective strategy for reducing the computational footprint of a model.
An inﬂuential pioneering work, LeCun et al.
(1990) proposes to discard weights using a error-approximation approach based on Hessian diagonals.
More recent work suggests pruning weights with small magnitudes (Han et al., 2016), with quantization and Huffman coding as additional steps.
However, these approaches introduce irregular sparsity to the weights, and they assume that re-training the weights is feasible.
In this work, we take a different approach and focus on techniques that eliminate entire ﬁlters.
This is because modern implementations of feedforward evaluation (e.g., im2col and particularly
NEON instruction on ARM processors) take advantage of dense matrix multiplications.
Pruning individual weights without changing the dimensions of the weight matrices has minimal effect on power consumption—this is conﬁrmed by our initial exploratory studies on the Raspberry Pi.
Hence, we only examine pruning techniques that discard entire ﬁlters of the convolutional layers:
Random pruning.
A simple baseline (Mittal et al., 2018) is random ﬁlter pruning, where n % of the ﬁlters are randomly pruned, layer-by-layer.
Interestingly, Mittal et al.
ﬁnd that random pruning is competitive with more advanced methods.
Filter norm.
Li et al.
(2017) propose ranking ﬁlters by their L 1 -norms, and then dropping off n % of the smallest ﬁlters on a layer-by-layer basis.
Mittal et al.
(2018) have previously found that L 1 -norm ﬁlter pruning (Li et al., 2017) outperforms a multitude of competing approaches.
Mean activation norm.
Among other approaches, Molchanov et al.
(2016) suggest pruning ﬁlters whose mean activations are small.
This approach is especially effective on ReLU, which both creates sparse activations and forces them to be non-negative.
L 0 regularization.
Louizos et al.
(2018) apply L 0 regularization to neural networks in order to learn sparse, efﬁcient structures.
Formally, deﬁne an objective
[Equation]
where L is the original loss function and θ the weights.
The dependence on the hypothesis and training examples has been omitted for brevity.
The optimal solution entails a non-differentiable objective and iteration over all 2 | θ | possibilities to choose the best θ ∗ ; hence, Louizos et al.
(2018) propose the following relaxation of the objective:
[Equation]
where z ∼ p ( z | φ ) is a binary discrete random mask parameterized by φ , and Q is the CDF.
Intuitively, for some choice of φ , the number of active parameters (on average) is penalized.
Inspired by the Concrete distribution (Maddison et al., 2016), Louizos et al.
(2018) propose the hard concrete distribution for z , further relaxing the discrete random mask into a continuous one:
[Equation]
where u ∈ R | θ | is a continuous random vector such that u i ∼ Uniform [0 , 1] , φ = log α are the mask parameters, and γ = − 0 .
1 , ζ = 1 . 1 , β = 2 / 3 are scaling hyperparameters.
Note that β can also be included as part of the mask parameters φ ; we follow Louizos et al.
(2018) and ﬁx β = 2 / 3 .
Louizos et al.
(2018) then apply the reparameterization trick (Kingma & Welling, 2014; Rezende et al., 2014) and make a Monte Carlo approximation to the objective:
[Equation]
N (cid:88) i = (cid:16) L ( θ (cid:12) z
( i ) )
(cid:17) + λ | θ | (cid:88) i = (1 − Q ( z i ≤ 0; φ i ))
A closed form expression is derived for the penalty, (1 − Q ( z i ≤ 0; φ i ))
=
σ (log α
i − β log − γζ ) .
At test time, the following estimator is used:
[Equation]
I NFERENCE -T IME P RUNING
In this section, we explain how the various techniques in Section 2.2 can be adapted to QRNNs.
For the following methods, we assume that a pre-trained model is provided.
We denote the weights at QRNN layer l as W ( l ) .
In all methods, we tie the indices across W z , W f , W o .
For example, if ﬁlter i is selected for pruning at layer l , then W ( l ) { z,f,o } :
=
W ( l ) { z,f,o } [ − i, :] , where − i denotes exclusion of index i .
This allows the removal of the column [: , − i ] in the next layer as well.
Random pruning.
We apply random pruning to W z , W f , and W o .
That is, we randomly prune ﬁlters associated with the same indices across the three weights.
Filter norm.
We apply ﬁlter norm pruning (Li et al., 2017), with the ﬁlter norms of W z acting as the criteria.
We ﬁnd W
z most helpful, since small ﬁlter norms should result in small hidden outputs, which is not necessarily the case for W f and W o .
Mean activation norm.
The hidden output H = h 1 ⊕ · · · ⊕ h n is a natural candidate for collecting mean activation statistics.
Intuitively, if (cid:107) H : ,i (cid:107) 1 is small on average, then the i th ﬁlters for
W z , W f , W o are less important.
Statistics are collected using a single pass of the entire training set.
For inference-time pruning, we store the collected statistics.
L 0 regularization.
Since we are given a pre-trained model and are prohibited from altering the weights, we learn the mask parameters only: φ ∗ = arg min φ ˆ R ( θ , φ ) .
We also enforce the sparsity on entire rows of W z , which corresponds to “group sparsity” in Louizos et al.
Speciﬁcally, we formulate the regularization on a feature map level instead, with Z as the target:
[Equation]
Z is chosen for the property that the i th feature map for h is zero if Z i is zero for c 0 = 0 .
This approach entails training and storing extra mask parameters for each operating point.
However, we ﬁnd this to be a non-issue for our task, since there are few operating points—three or four at most, out of which we use two for L 0 regularization—so the extra storage is negligible.
3. W ITH S INGLE -R ANK U PDATE
At speciﬁc operating points (e.g., 40% and 80% FLOPs), pre-trained weight updates can be stored and applied at inference-time to recover some perplexity.
Suppose W ∈ R m × n is a weight matrix in a neural network, and
W ∗ ∈
R m × n is some known set of weights that results in a lower loss.
Clearly, ∆ W := W ∗ − W can be stored and added at inference-time to obtain a better neural network.
However, it is obvious that this scheme is wasteful, since W ∗ could have directly substituted
W in the ﬁrst place.
Sacriﬁcing a negligible amount of storage to recover some perplexity, we propose learning a singlerank weight matrix update
[Equation]
to each weight in the convolution layers.
Speciﬁcally, the process is as follows, beginning with a pre-trained model:
Prune a pre-determined set of ﬁlters for some operating point (e.g., 40% FLOPs).
Initialize the weight updates ∆ W l = u ( l ) v ( l ) (cid:124) , u ( l ) i , v ( l ) i ∼ p ( (cid:15) ) for each convolution layer l , in our case Normal (0 , 0 . 1) .
Fixing the existing weights W l for each convolution layer, train a single-rank update such that W ∗ l := W l + ∆ W l , where W ∗ l is used as the new weight.
Store ∆ W l for use at inference time on the same operating point.
E XPERIMENTAL S ETUP
We evaluate the aforementioned pruning techniques for word-level language modeling on Penn
Treebank (PTB) (Marcus et al., 1993; as preprocessed by Mikolov et al., 2010) and WikiText- (WT-103) (Merity et al., 2017).
We denote the models for PTB and WT-103 as ptb-qrnn and wt103-qrnn , respectively.
4. D ATASETS AND T ASKS
For each model, we report word-level perplexity and recall-at-three (R@3), deﬁned as the percentage of top three token–logit outputs that contain the true next token—this was chosen due to most mobile keyboards providing three predictions.
For example, if { “cat”, “dog”, “baby” } are the top three predicted tokens for, “I adopted a ,” with “dog” being the ground truth, then the prediction is correct, regardless of the rank of “dog”.
Penn Treebank.
Built from Wall Street Journal articles, Penn Treebank
(PTB) is a small yet popular word-level dataset for language modeling.
In the standard pre-processed version (Mikolov et al., 2010), the dataset contains roughly 887K, 70K, and 78K training, validation, and testing tokens, respectively.
The number of unique tokens is capped at 10,000, yielding a relatively large 4.8% out-of-vocabulary (OOV) rate.
WikiText-103.
Merity et al.
(2017) introduce WikiText-2 and WikiText-103, datasets based on freely available Wikipedia articles.
We use only WikiText-103, since WikiText-2 was designed to be similar to Penn Treebank.
With 103 million training tokens, WikiText-103 is 103 times as large as
PTB.
WikiText-103 contains around 217K tokens for validation, and 245K for testing.
The number of unique tokens is 267K, resulting in a 0.4% OOV rate, signiﬁcantly lower than that of PTB.
4. H YPERPARAMETERS AND T RAINING
In all of the models, we chose the hyperparameters as suggested in Merity et al.
’s codebase.
For ptb-qrnn , we used a four-layer QRNN with 1550 hidden units for each layer and a 400dimensional embedding.
For wt103-qrnn , we used a four-layer QRNN with 2500 hidden units and 400-dimensional embeddings, along with a tied adaptive softmax (Merity et al., 2018b).
In both models, the ﬁrst layer uses a window size of two, while the rest use a windows size of one.
Following Merity et al.
(2018a)
, we also adopted the regularization techniques randomized backpropagation through time, embedding dropout, temporal activation regularization (TAR), activation regularization (AR), and variational dropout.
We followed the same training process as well, with non-monotonically triggered ASGD (NT-ASGD) as the optimizer.
We use the same hyperparameters as Merity et al.
(2018a)
and Merity et al.
(2018b) for each model–dataset pair.
During the training of wt103-qrnn , we follow Merity et al.
(2018b), using a tied adaptive softmax (Grave et al., 2017; Merity et al., 2018b) layer.
At inference time, we use a regular softmax instead, since we require R@3.
Pruning.
We selected a number of distinct operating points that represent discrete points in the accuracy–efﬁciency tradeoff space.
Based on previous work (Tang et al., 2018), ﬂoating-point operations (FLOPs) is a good proxy of both energy usage and latency, and so we use FLOPs as a way of selecting our operating points.
In L 0 regularization, the λ decay strength was selected so that the resulting model corresponds to roughly the FLOPs targets: To achieve 80% and 60% FLOPs for the model on PTB, we used λ
= 5 .
5 × 10 − 4 , 8 . 5 × 10 − 4 , respectively.
To achieve about 70% FLOPs on WT-103, we chose λ = 6 × 10 − 4 .
We trained the hard concrete mask parameters for roughly 5000 steps using Adam with a learning rate of 5 × 10 − 3 .
Since the weight decay penalty is incompatible with the objective, we removed it while training the mask.
For mean activation pruning, which requires some training examples to collect statistics, we used the entire training set for ptb-qrnn .
Since WikiText-103 is large, we used roughly 10% of the ﬁrst training examples for collecting statistics on wt103-qrnn .
Single-rank update (SRU).
For the PTB model, the single-rank update was trained for 10 epochs using NT-ASGD (Merity et al., 2018a) with a non-monotonic interval of three.
For WikiText-103, the update was trained for 2000 steps using Adam with a learning rate of 5 × 10 − 3 .
All other hyperparameters were the same as those used during the training stage.
4. I NFRASTRUCTURE D ETAILS
We trained all of our models on a commodity machine with a Titan V GPU, i7-4790k CPU, and
GB of RAM.
We used PyTorch 0.4.0 (commit 1807bac ) for developing and running our models.
We deployed our models on a Raspberry Pi (RPi) 3 Model B (ARM Cortex-A53) running Raspbian
Stretch (4.9.41-v7+).
Speciﬁcally, we copied the trained models over to the RPi, and ran them at the same operating points accordingly.
We plugged the RPi into a Watts Up Pro meter, a wattmeter that reports power usage at the rate of 1 Hz via a USB cable, which is connected back to the RPi.
Evaluating on the test set, we collected power draw statistics on 350 next-word predictions, which were averaged to produce a millijoule per query (mJ/q) estimate.
We obtained latency estimates in a similar manner by averaging the milliseconds per query (ms/q).
Finally, we subtracted off the idle power usage of the RPi to obtain a better estimate of the actual power for each query.
Although our ﬁnal application is NLMs running on mobile devices such as smartphones and tablets, there are many challenges to directly evaluating on such hardware, not to mention great variability between different smartphones.
The Raspberry Pi is a standard, convenient stand-in since it uses exactly the same ARM processor architecture as nearly all mobile devices today, while being easy to develop on for practitioners.
Evaluation on the RPi is widely adopted for research on efﬁcient
NNs today (Amato et al., 2017; Tang et al., 2018).
R ESULTS AND D ISCUSSION
In our results for PTB and WT-103, we compare to state-of-the-art results in the past.
In general, we ﬁnd that QRNNs are strong competitors to LSTM approaches, and achieve state-of-the-art perplexity on WikiText-103 (Merity et al., 2018b)
[Table]
is the same model that we use, but we achieve better perplexity following the same methodology.
The best results of each category are bolded.
“w/SRU” denotes the results after applying an SRU.
For PTB and WT-103, to demonstrate that the original model is well-parameterized, we train models with less parameters from scratch; see Tables 1 and 2, row 4—the smaller models achieve higher perplexity and lower R@3 than the originals do.
We note that a 20-point increase in perplexity may only correspond to a few points decrease in R@3, showing that perplexity changes on a much different scale than accuracy does (see Table 1, rows 3 and 8).
Furthermore, lower perplexity does not necessarily imply higher accuracy (see rows 6 and 8), conﬁrming that perplexity alone cannot completely determine the recall.
In Table 1, we chose 75 as the cutoff-point for perplexity—further
[Table]
results are illustrated in Figure 2.
For WT-103, we observe trends similar to those of PTB; A large drop in perplexity corresponds to a much smaller decrease in R@3 (see Table 2, rows 3 and 5).
5. A CCURACY –E FFICIENCY T RADEOFFS
We illustrate the accuracy–efﬁciency tradeoff space of the PTB and WT-103 models in Figure 2.
For each model, we tabulate the results at ﬁxed intervals according to the approximated percentage of
FLOPs, relative to that of the unpruned model.
We omit results that exceed 100 in test perplexity, since they are insufﬁcient for language modeling in practice.
[Figure]
We illustrate the perplexity–efﬁciency tradeoff space on the test set obtained before applying the single-rank update.
Surprisingly, random ﬁlter pruning is a strong baseline, which supports the ﬁndings from Mittal et al.
Random pruning not only outperforms ﬁlter norm and mean activation pruning, but also regains perplexity more easily with a single-rank update.
From Table 1 (rows 7–9) and Table(rows 6–8), we see that random pruning displays equivalent or superior performance to ﬁlter norm and mean activation pruning.
Interestingly, random pruning achieves the lowest perplexity with a single-rank update (Table 2, rows 5–8), out of all the baseline approaches on WT-103.
On the other hand, ﬁlter norm pruning is relatively weak, doing worse than random pruning in all cases—with or without a single-rank update—suggesting that ﬁlter norm pruning has no practical beneﬁt over random pruning.
L 0 regularization (Louizos et al., 2018) works best, as shown in rows 5–6 in Table 1 and row 5 in Table 2.
In general, testing on Penn Treebank and WikiText-103—two very different datasets—gives us consistent results, thus demonstrating the robustness of L 0 regularization (Louizos et al., 2018), compared to the other pruning approaches.
5. P OWER U SAGE AND L ATENCY
On the Raspberry Pi, the PTB models are relatively fast, while the WT-103 models are high latency, taking over one second (Table 2, rows 2–3 and 8) for the full models.
For type-ahead prediction on a mobile device, the WT-103 models are unsuitable as-is; further steps (e.g., more pruning then re-training, vocabulary reduction, quantization) would be required to deploy the models for practical use.
Supporting the ﬁndings from Tang et al.
(2018), the number of FLOPs scales linearly with latency and power: Full experimental results from Figure 2 yield Pearson’s r 2 = 0 .
98 for both latency– and power–FLOPs measurements, suggesting a strong linear relationship between the number of FLOPs and both latency and power.
In terms of extra parameters, a single-rank update costs less than 74 KB for ptb-qrnn , and less than 120 KB for wt103-qrnn .
Mean activation statistics requires 20 KB for ptb-qrnn , and
KB for wt103-qrnn .
Mask parameters for L 0 regularization cost about 20 KB on each power level for ptb-qrnn , and 30 KB for wt103-qrnn .
Filter norm pruning and random pruning do not require any extra storage.
5. E FFECTS OF P RUNING ON I NDIVIDUAL W ORD - LEVEL R ECALL
According to Zipf’s law, a small set of common words disproportionately outnumbers a much larger set of rare words, suggesting that rare words require more capacity to model.
Since pruning reduces model expressiveness by removing parameters, we hypothesize that rare words are more affected than common words are.
To validate this hypothesis, we compare the recall of each individual word between our original and pruned models.
Speciﬁcally, we collect the test set R@3 statistics of each present word in the vocabulary, order the words by rank, and then remove all words with zero R@ in the original model.
The last step is done to limit our attention to successfully modeled words, and
[Figure]
to avoid division-by-zero.
Next, to determine the relative difference in R@3 between the models, we compute the per-word R@3 ratio between the pruned and original models.
Finally, we apply valid convolution of a ﬁxed-window, uniform averaging ﬁlter across the x -axis to construct a less noisy visualization.
As shown in Figure 4, ﬁlter norm and random pruning yield models with worse R@3 on rare words than common words (see top 100 words on both datasets).
Mean activation pruning is similar to ﬁlter norm and random pruning on PTB; however, on WT-103, words with frequency ranks roughly between 500 and 3000 are penalized less than the most common 200 words are.
The best method,
L 0 regularization, surprisingly improves R@3 for many words on PTB.
On WT-103, however, it yields curves similar to those from ﬁlter norm pruning.
Overall, the graphs support the hypothesis that rare words are more difﬁcult to model.
5. S TUDY L IMITATIONS
We use the ofﬁcial PyTorch QRNN implementation, 5 which provides a fused pooling kernel for
CUDA-enabled GPUs but lacks one for CPUs
; such a kernel would provide further performance gains on a CPU. Since we perform our evaluation with a non-fused kernel, our latency results could be improved further in practice. Nevertheless, we successfully capture the general trend in accuracy– efﬁciency tradeoffs for NLMs, and we demonstrate the impact of inference-time pruning methods.  C ONCLUSION AND F UTURE W ORK
Motivated by the mass adoption of smart software keyboards on mobile devices, we explore the task of inference-time pruning on QRNNs, state-of-the-art neural language models.
Starting with existing training-time pruning methods, we extend their usability to QRNNs at run-time, obtaining multiple operating points in the accuracy–efﬁciency tradeoff space.
To recover some perplexity using a negligible amount of memory, we propose to train and store single-rank weight updates at desired operating points.
In this paper, we limit our inference-time power usage and latency evaluation to the RPi, a convenient surrogate for smartphones.
A potential extension is to conduct a similar study on NLMs running on
GPUs—in practice, this has applications in large-scale vocabulary speech recognition (Cho & Kumar, 2018).
Clearly, the GPU architecture differs greatly from the RPi, favoring high concurrency and minimal host–device data transfer.
Other extensions include evaluating other compression techniques on NLMs, such as knowledge distillation (Hinton et al., 2015) and vector quantization (Han et al., 2016).