Shrinking Bigfoot: Reducing wav2vec 2.0 footprint (2021)

Abstract: Wav2vec 2.0 is a state-of-the-art speech recognition model which maps speech audio waveforms into latent representations. The largest version of wav2vec 2.0 contains 317 million parameters. Hence, the inference latency of wav2vec 2.0 will be a bottleneck in production, leading to high costs and a significant environmental footprint. To improve wav2vec’s applicability to a production setting, we explore multiple model compression methods borrowed from the domain of large language models. Using a teacher-student approach, we distilled the knowledge from the original wav2vec 2.0 model into a student model, which is 2 times faster, 4.8 times smaller than the original model. More importantly, the student model is 2 times more energy efficient than the original model in terms of CO2 emission. This increase in performance is accomplished with only a 7% degradation in word error rate (WER). Our quantized model is 3.6 times smaller than the original model, with only a 0.1% degradation in WER. To the best of our knowledge, this is the first work that compresses wav2vec 2.0.

## introduction speech recognition technologies enhance people's lives through a wide range of applications such as virtual assistants, home automation, and realtime transcription and captioning. neural networkbased speech recognition models achieve superior performances, but successfully training them requires a lot of labeled data @xcite . wav2vec 2.0 @xcite is a framework that combines self-supervised masking pretraining with a tuning step that achieves impressive results with little labeled data. with only 10 minutes of labeled data, wav2vec 2.0 achieves word error rates (wer) of 4.8% and 8.2 % on librispeech @xcite dev-clean and dev-other datasets, respectively. moreover, wav2vec 2.0 achieves state-of-the-art performances on the lib-rispeech benchmark while using a minimal amount of labeled data. self-supervised learning enables wav2vec 2.0 to learn speech representations efficiently without much labeled data. this allows easily training an end-to-end automatic speech recognition (asr) system for a low resource language foot_0 or for a specific domain. the largest version of wav2vec 2.0 has 317 million parameters, which makes real-time inference inefficient in production. language modeling is a central research question in natural language processing (nlp). in recent years massive lm architectures, based on novel pre-training approaches, revolutionized nlp @xcite . the bert @xcite architecture of masked pre-training and transformers is similar to wav2vec2.0. to improve wav2vec's applicability to a production setting, we explore multiple model compression methods borrowed from the domain of massive lms. a significant drawback of massive lm models is high resource consumption. the negative impact of such models, beyond the costs for organizations using this technology in production, is cause for concern @xcite . various compression methods have been developed to overcome the problem, such as @xcite . to compress wav2vec 2.0, we explore quantization and knowledge distillation methods. our quantized model is 3.6 times smaller than the original model, with a 0.1% loss in wer, making it easy to deploy in a production environment. our student model is at least 2 times faster on both cpu and gpu and 4.8 times smaller than the original model, with a 7% loss in wer through knowledge distillation. since the knowledge distilled model is smaller than the original model, it is more energy efficient. in section 7, we estimated that the distilled model could potentially save 21 tons of co 2 equivalents (comparing to the original model) if running inferences on 54 millions of hours of speech data. our work makes two key contributions: 1. this research is the first work that compresses the stateof-the-art wav2vec 2.0 speech recognition model to the best of our knowledge. 2. our distilled student model has a faster inference speed and makes wav2vec 2.0 more cost-efficient and environmental friendly. ## related work in this section, we relate research works done on model compression in asr systems. then we describe compression methods applied in language modeling, which we wish to explore for compressing wav2vec. ## wav2vec wav2vec 2.0 @xcite achieved a breakthrough in asr by adopting the masked pretraining method employed in the massive language model bert. bert masks a few words in each training sentence, and the model learns by attempting to fill the gaps. instead of masking words, wav2vec 2.0 masks parts of the audio representa-tion and lets the transformer network fill in the gaps. this self-supervised setup makes it possible to learn mainly from unlabeled data while only finetuning on a small labeled dataset. the wav2vec 2.0 model comprises two significant components that we consider for compression: the cnn layers and the transformer layers. in the following sections, we explore how compressing a fine-tuned version of wav2vec 2.0 foot_1 using knowledge distillation and quantization affects the size of the model and its inference speed. ## knowledge distillation for wav2vec 2.0 we based our distillation on the knowledge distillation used by @xcite to distilbert @xcite . they initialized their student models by taking alternating layers from a well-performing larger model and scheduled the learning rate by first increasing, then gradually decreasing it. since it is possible to reduce either the cnn layer or the transformer layers of wav2vec 2.0, we decided to reduce the transformer ones since wav2vec 2.0 has only 7 cnn layers, which does not constitute a substantial computational bottleneck compared to transformer layers, see ## quantization for wav2vec 2.0 in this section, we discuss how the two quantization techniques available in pytorch apply to our task of quantizing wav2vec 2.0. ## experiments the wav2vec 2.0 model we are trying to compress was trained with 960 hours of unlabeled data from the librispeech dataset @xcite , and then fine-tuned with the labeled version of the same 960 hours. we did not attempt to compress the model before it was fine-tuned. we used a viterbi decoder to convert the output of the student wav2vec 2.0 model into text. we conducted training on 4 nvidia v100 gpus and set the batch size to 1 per gpu. we trained with 16-bit floating-point arithmetic and fixed the seed at 42 in all our runs to make the results reproducible. for knowledge distillation, we used the adam optimizer with weight decay. we set @xmath0, @xmath1.9, 0.98), and the weight decay coefficient to 0.01. we scheduled the learning rate by the number of epochs. the learning rate starts from 2.5 × 10 -5 and linearly increases for the first two epochs. after two epochs, we linearly decrease the learning rate in every epoch. also, we assessed the effect of using a different number of transformer layers on the target metrics, see for quantization, we used pytorch's dynamic quantization to quantize linear layers of wav2vec 2.0. our implementation is available at https:// git.io/jmpjm . ## results and discussion considering the quantized model, we notice that even though the model is 3.6 times smaller than the original model, it is not significantly faster. this discrepancy can be attributed to the need for de-quantize tensors entering multi-head attention; see section 5.2 for more details. pytorch does not support gpu quantization; therefore, we report only cpu inference time for the quantized model. its significant reduction in size, together with its ability to maintain the original wer, albeit with a slight loss, make the quantized model easily deployable in a production environment. ## conclusion jointly training teacher and student model has become a popular technique for compressing asr systems @xcite . in contrast, we employed a traditional knowledge distillation approach and it will be an interesting future work direction to investigate why such ap-proach did not achieve great performances for the student model. despite performance limitations, we found that knowledge distillation is the most effective compression method for wav2vec 2.0 for saving energy by lowering cpu or gpu inference time. using knowledge distillation, we trained a student wav2vec 2.0 model, which is at least two times faster than the original model on both cpu and gpu. moreover, the trained student model is also two times more energy efficient than the original model. wav2vec models present a valuable technology for democratizing asr. we explored multiple methods for reducing memory usage and speeding up such models for deployment. more remarkably, our compressed model is environmental friendly during inference. maintaining low carbon footprint and improving model's performance will be our future research direction.