ShrinkML: End-to-End ASR Model Compression Using Reinforcement Learning

TL;DR

This paper introduces ShrinkML, an AutoML system using reinforcement learning to optimize layer-wise compression of end-to-end ASR models with SVD, achieving effective compression without retraining and better results with retraining.

Contribution

The paper presents a novel RL-based AutoML approach for layer-wise ASR model compression using SVD, with practical considerations and improved compression results.

Findings

RL-based AutoML effectively compresses ASR models without retraining.

AutoML outperforms manual rank selection in model compression.

Compressed models maintain high accuracy on LibriSpeech.

Abstract

End-to-end automatic speech recognition (ASR) models are increasingly large and complex to achieve the best possible accuracy. In this paper, we build an AutoML system that uses reinforcement learning (RL) to optimize the per-layer compression ratios when applied to a state-of-the-art attention based end-to-end ASR model composed of several LSTM layers. We use singular value decomposition (SVD) low-rank matrix factorization as the compression method. For our RL-based AutoML system, we focus on practical considerations such as the choice of the reward/punishment functions, the formation of an effective search space, and the creation of a representative but small data set for quick evaluation between search steps. Finally, we present accuracy results on LibriSpeech of the model compressed by our AutoML system, and we compare it to manually-compressed models. Our results show that in the absence of retraining our RL-based search is an effective and practical method to compress a production-grade ASR system. When retraining is possible, we show that our AutoML system can select better highly-compressed seed models compared to manually hand-crafted rank selection, thus allowing for more compression than previously possible.

Tables

Table 1. Table 1: Evaluation of the best models found by both searches. Steps mean at what step during the search a model was found.

	AutoML		Manual	Baseline
	Slow	Fast	Guided
test-clean	8.35	8.34	8.67	8.29
test-other	21.61	21.33	22.04	21.13
step	1626	2363	NA
GPU hours	635	74

Equations

$$M=U\Sigma V^T$$

$$M\approx U^{\prime }{V}^{\ast }=U^{\prime }{\Sigma }^{\prime }{V^{\prime }}^T$$

$$Speedup=\frac{m\times n}{k\times (m+n)}$$

$$\mathcal{R}(w)=-\exp (w-w_b)$$

$$\mathcal{R}(w)=-\exp (\sqrt{\frac{w}{w_b}})$$

$${\mathcal{R}}_v({\Delta }_a)=-100{\Delta }_a-10$$

Full text

1 Introduction

End-to-end automatic speech recognition (ASR) models have outperformed traditional ASR systems. In particular, models with attention mechanism [1] have produced state-of-the-art results in speech recognition and translation application domains [2, 3, 4]. These models improve upon traditional systems, partially because they include acoustic models, language models and pronunciation dictionary jointly learned together. Furthermore, models with character/subword output units are especially useful in low latency deployments such as mobile phones and can produce good results with a beam-search decoding mechanism [5].

One challenge when encompassing all components of ASR into a single model is the need for larger model capacity and thereby deeper neural networks with several recurrent layers. This makes it harder to achieve real-time inference, especially with on-device deployments. One approach to reduce model size is to apply model compression techniques [6]. This is often done manually and may include many trial-and-error experiments until a reasonable compression scheme is found. However, more recently, automated search-based techniques have been employed to optimize the compression of deep neural networks [7]. More specifically, AutoML has been used to select the per-layer compression ratios to optimize a model globally.

In this paper, we use reinforcement learning to select the per-layer compression ratios based on matrix approximation using singular value decomposition (SVD) [8]. We base our work on a recently proposed end-to-end ASR model with attention mechanism [9, 10] that has obtained state-of-the-art results on LibriSpeech. We focus on practical aspects when designing the AutoML compression system in the specific context of ASR models. We show experimentally that our approach is superior to manual model compression for both modest and aggressive compression targets, in the absence or presence of retraining. In the following sections, we outline the ASR model and present insights into compressibility of layers in such encoder-attention-decoder models, describe the AutoML framework and present a data subsampling technique to speed up the search process, and finally delve into experimental results that prove the efficacy of our system, corroborated by on-device measurements on a mobile phone.

2 ASR Model Compression

2.1 ASR Model

We use an end-to-end attention based ASR model [11, 12] with an architecture similar to the one proposed in [10] as depicted in Fig. 1. The main differences in our model are the use of 1) unidirectional LSTMs in place of bidirectional LSTMs and 2) monotonic chunk-wise attention [13]. These changes were made to allow the model to decode speech in streaming mode, which is required for mobile applications. As a result, the model also has a higher baseline word error rate (WER) as compared to the results published in [10]. The unit size for all LSTM layers in the network is 1024 and the output layer comprises 10000 sized byte-pair encoded sub-word units [14, 15]. The model was trained on LibriSpeech training set [16] with MFCC front-end using Tensorflow [17] and the RETURNN framework [18]. A time reduction factor of 8 was used via the maxpool layers interleaved between the first four encoder LSTM layers. In this work we have not used a separate language model for further reduction of WER.

2.2 Low-Rank Matrix Approximation

Our primary target is to build a generic reinforcement learning (RL) based AutoML system that automatically optimizes the per-layer compression ratios for ASR models. We decide to choose low-rank matrix factorization, and specifically, singular-value decomposition (SVD) to compress the weight matrices, since it has proven to be an effective method for compressing similar models [19, 20, 21, 22, 23, 24]. However, we design our system such that it is generalizable to different compression methodologies. We refer the reader to prior work for further reading on SVD, but a summary is given below. We factorize a matrix $M$ into three matrices $U$, $\Sigma$ and $V$:

[TABLE]

where $\Sigma$ is a diagonal matrix which consists of the singular values (typically in descending order of magnitude) that uniquely identify $M$. We compress $M$ by removing the lower-magnitude values in $\Sigma$ – we refer to the number of remaining element as a factorization rank and denote it with $k$. We select $k$ in such a way that it preserves a given energy, where energy refers to the normalized summation of the remaining singular values. We then remove corresponding columns from matrices $U$ and $V$, creating a set of three new matrices $U^{\prime}$, $\Sigma^{\prime}$ and $V^{\prime}$ which approximate original matrix $M$. Finally, we combine $\Sigma^{\prime}$ with $V^{\prime}$ so that we can replace $M$ according to the following equation:

[TABLE]

If $M$ is an $m\times n$ matrix, then $U^{\prime}$ and $V^{*}$ have dimensions $m\times k$ and $k\times n$ respectively, and we get a theoretical speedup and model size reduction:

[TABLE]

3 AutoML for SVD Rank Selection

Fig. 2 plots the WER degradation when each of the layers are factorized separately – as shown, the model accuracy is much more sensitive to changes in the input and output layers, but less so to intermediate encoder LSTM cells, attention and decoder layers. Selecting the ranks per layer for the 18 matrices in the model is a hard problem because of the large number of combinations. For instance, if we discretize compression range for each layer to 5 rank options, we have $5^{18}=3.8\times 10^{12}$ possible factorizations. We use reinforcement learning to navigate that search space and find the best compression schemes.

3.1 AutoML System

We use an RL-based system similar to the one proposed by [25] after slightly adapting it for the model compression task. More specifically, in our case an agent is responsible for deciding the factorization scheme which would: a) guarantee a predefined speedup, and b) minimize WER of the factorized model. The scheme in this context is a series of decisions, with each decision representing a compression ratio for an individual layer. Both the list of layers and the set of discrete compression levels per layer are provided a priori to the search. We call the set of all possible factorization schemes a search space and the set of all parameters which modify it hyperparameters of the search space. More formally, our search space can be defined as a 2-dimensional matrix: $S\in\mathbb{Z}^{l\times d}$, where $S_{i,j}$ represents a rank $k$ to use when factorizing the $i$-th layer according to the $j$-th option. From this search space, the agent selects a factorization scheme $s\in\mathbb{Z}^{l}$ by choosing one value from each row of matrix $S$. The selection is done by sampling $l$ probability distributions over $d$ options which are produced by a trainable policy $\pi(\theta)$.

In our system, this policy is modeled with a single LSTM layer with 100 input units and 100 hidden units which takes a sequence of $l$ inputs (one for each layer) – each element of output sequence is then passed to its individual fully-connected layer with $d$ units, to match its length with the number of available decisions, followed by softmax. The final output is $D\in\mathbb{R}^{l\times d}$ where each vector $D_{i}\in\mathbb{R}^{d}$ represents probability distribution of selecting different factorization options for the $i$-th layer. We reward the agent depending on the WER of the model when compressed according to the proposed scheme, and use policy gradient to update $\theta$.

3.2 AutoML System Tuning for Compression

One of the main challenges for many AutoML systems is reducing the time required to produce useful results. RL-based search is slow because evaluation of a proposed action (in our case, running evaluation on a validation set) takes significantly more time than proposing a new action. In this work we address this problem and propose a number of techniques which we have successfully used to significantly reduce evaluation time. This way, we speed up the search, while still providing representative feedback to the RL agent.

To quickly reject unprofitable points, we estimate speedup of a proposed scheme ($a$) and compare it to the predefined target ($a_{t}$), as shown in Algorithm 1. If it falls below the threshold, we punish the agent using a dedicated reward function $\mathcal{R}_{v}$, otherwise the scheme is accepted for full evaluation and we reward the controller according to returned WER using reward function $\mathcal{R}$. We observed that combined usage of the two reward functions makes the search converge to a certain area of the speedup-WER plane – as presented in Fig. 4 – with $\mathcal{R}$ pushing points towards smaller WER and $\mathcal{R}_{v}$ pushing above $a_{t}$. However, both functions need to be tuned relative to each other, as well as independently, to obtain good results. In our work we empirically found that when targeting conservative $a_{t}$ values relative differences in WER between the original model ($w_{b}$) and its compressed versions ($w$) might be too small to make the agent discriminate between good models with similar WERs. Therefore, we decided to use exponential function to add additional emphasis to the difference:

[TABLE]

However, the range of this function is too big when targeting high speedups $a_{t}$, because the WER range is much larger, so we used a slightly modified version for such aggressive searches:

[TABLE]

For both aggressive and conservative compression we observed that a simple linear function works well for $\mathcal{R}_{v}$, with parameters $-100$ and $-10$ tuned to make its values visibly worse than $\mathcal{R}$, even when $\Delta_{a}\rightarrow 0$.

[TABLE]

3.3 Small Dataset Creation

LibriSpeech validation sets consist of 5567 utterances split among dev-clean and dev-other. Using the validation sets to evaluate WER for each model proposed by AutoML was too slow and led us to investigate if a representative subset could be found. First, we tried randomly selecting utterances from the validation sets. However, we observed experimentally that this did not work well with small random sets – if we minimize WER on the random set, this does not always translate into a lower WER on the full-size validation or test sets. Instead, we used our “condense” algorithm in Listing 1 to select the utterances that were most representative of the validation sets.

As Listing 1 shows, we find the WER per utterance from different cohort models. These models are simple variations of our baseline model presented in Section 2.1 – we use 8 cohort models trained with different number of layers, layer sizes in the encoder, and few variations of SVD-based compression schemes. We compile the WERs of each utterance when decoded with each of these models, and correlate that to the WERs of the entire validation set with the same models. We then choose the utterances that correlate highly in WER with the entire set across all the cohort models. As Fig. 3 shows, our “condensed” datasets correlate much better than “random” sets, especially with very small dataset sizes. With ${correl}\_{min}=0.95$ we created an 83-sample dataset for use with AutoML – this is 67$\times$ smaller than the validation sets but was highly-correlated with the full set as Fig. 3 shows. We used it as a drop-in replacement for AutoML thus making the whole system $>$10$\times$ faster overall.

4 Results

We use AutoML for two experiments. First, we want to find a modestly-compressed model with minimal WER loss in the absence of retraining. We believe that full training data is not always available, especially in systems where all of the training data is not in one place [26], or when app developers are using pretrained models and do not have access to training sets. Second, we use AutoML to find the best aggressively-compressed seed model for retraining. In this case we maximize compression, and rely on retraining to recover the accuracy.

4.1 Compression without Retraining

We launched two searches, slow and fast, both targeting speedup $a_{t}=1.2$, and using Equations 4 and 6. The slow search evaluated models on the entire dev-clean dataset and selected compression ratios for each layer from the same range. The fast search used the condensed dataset described in the previous section and had per-layer compression ranges adjusted according to the information from Fig. 2. To extract the best model, for each search independently we gathered all explored configurations and selected top-5 (as evaluated by the search) which were then evaluated on the “test” datasets from LibriSpeech (test-clean and test-other). The best model reported here is the best of the 5 on test-clean. As shown in Table 1, the fast search was able to find an equally-good compression scheme much faster. Both searches were also able to find better compression schemes than the hand-crafted manual ones as highlighted in Fig 4.

4.2 Aggressive Compression With Retraining

Previous work has repeatedly proven that SVD-based compression is very effective in attaining speedup while maintaining model accuracy when retraining is available [19, 21, 23]. We use our “guided manual” method of compression to test the limits of compression with retraining. We train for 50 epochs over LibriSpeech with an epoch split of 20 – this means we go through the training data 2.5 times. As Fig. 5 shows, we achieved up to ~3$\times$ compression without degrading WER on test-clean. Can we push this any further using our AutoML system? To answer this question, we set our speedup threshold to 3.7$\times$ and launched our fast AutoML search for only a few hours, we then used the best-found model as the seed model for retraining and the results are shown in Fig. 5. The model found by AutoML achieved 3.75$\times$ speedup and had approximately half the WER before retraining when compared to the manually-compressed model at the same speedup. After retraining, we were able to recover almost full accuracy, as shown in Fig. 5, therefore, we believe this is an effective systematic method to selecting highly-compressed models for retraining. After 8-bit quantization, the model size is compressed to 24 MB from the 32-bit compressed model size of 72 MB.

4.3 On-Device Measurements and Considerations

So far, we have been using Equation 3 as a proxy for speedup, however, we measured the actual on-device runtime of our retrained model from Section 4.2 to validate our estimate. We evaluate on a Qualcomm Snapdragon 845 development board, using TFLite [27] on CPU. As depicted in Fig. 6, we acquire a significant speedup of 2.16$\times$ with our SVD-based compressed model over the baseline – especially in the encoder and decoder parts of the network. However, this is much lower than the estimated 3.74$\times$ shown in Fig 5. We shall investigate this discrepancy in future work but mention this here as a note that theoretical and measured speedups often vary greatly.

5 Conclusion

In this work, we presented an AutoML framework to push the boundaries of SVD-based ASR model compression beyond what is possible manually. We improved upon the WER attainable by manual compression when retraining is not possible. Even when we could retrain, we have shown that AutoML can improve the compression ratio, and therefore speedup, of our ASR model without any loss of accuracy. In the future we aim to use AutoML to optimize and mix different compression techniques, and we hope to make the estimated speedup/accuracy in such a system more faithful to the actual results.

6 Acknowledgements

MoChA within the experimental framework was implemented by Dr. Chanwoo Kim and his team in Samsung research. We would like to thank them for their support and feedback.