PARP: Prune, Adjust and Re-Prune for Self-Supervised Speech Recognition

1 Introduction . For many low-resource spoken languages in the world , collecting large-scale transcribed corpora is very costly and sometimes infeasible . Inspired by efforts such as the IARPA BABEL program , Automatic Speech Recognition ( ASR ) trained without sufficient transcribed speech data has been a critical yet challenging research agenda in speech processing [ 31 , 33 , 42 , 32 , 21 ] . Recently , SelfSupervised Speech Representation Learning ( speech SSL ) has emerged as a promising pathway toward solving low-resource ASR [ 84 , 25 , 110 , 6 , 29 , 127 , 55 , 27 ] . Speech SSL involves pre-training a speech representation module on large-scale unlabelled data with a self-supervised learning objective , followed by finetuning on a small amount of supervised transcriptions . Many recent studies have demonstrated the empirical successes of speech SSL on low-resource English and multi-lingual ASR , matching systems trained on fully-supervised settings [ 6 , 29 , 127 , 4 , 126 ] . Prior research attempts , however , focus on pre-training objectives [ 84 , 25 , 110 , 72 , 57 , 74 , 71 , 73 , 55 , 22 , 27 , 17 , 129 ] , scaling up speech representation modules [ 5 , 6 , 53 ] , pre-training data selections [ 108 , 54 , 107 , 111 , 78 ] , or ⇤Equal contribution . 1Project webpage : https : //people.csail.mit.edu/clai24/parp/ 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) , virtual . applications of pre-trained speech representations [ 26 , 62 , 94 , 28 , 63 , 29 , 76 , 120 , 64 , 117 , 112 , 44 , 4 , 86 , 59 , 66 , 2 , 56 , 102 , 15 , 30 , 20 ] . In this work , we aim to develop an orthogonal approach that is complementary to these existing speech SSL studies , that achieves 1 ) lower architectural complexity and 2 ) higher performance ( lower WER ) under the same low-resource ASR settings . Neural network pruning [ 65 , 51 , 49 , 69 ] , as well as the more recently proposed Lottery Ticket Hypothesis ( LTH ) [ 39 ] , provide a potential solution that accomplishes both objectives . According to LTH , there exists sparse subnetworks that can achieve the same or even better accuracy than the original dense network . Such phenomena have been successfully observed in various domains : Natural Language Processing ( NLP ) [ 123 , 19 , 88 , 80 ] , Computer Vision ( CV ) [ 18 , 45 ] , and many others . All finding sparse subnetworks with comparable or better performance than the dense network . Given the lack of similar studies on pruning self-supervised ASR , we intend to fill this gap by finding sparse subnetworks within a pre-trained speech SSL that can achieve superior performance to the full pre-trained model on downstream ASR tasks . However , directly applying widely-adopted pruning methods , such as One-Shot Magnitude Pruning ( OMP ) and Iterative Magnitude Pruning ( IMP ) [ 49 , 39 ] , to pre-trained speech SSL suffers from two challenges . First , adopting these methods in the conventional pruning framework is extremely time-consuming for SOTA speech SSL models . OMP and IMP involve more than one round of finetuning on downstream tasks ( c.f . Figure 1 ) , and finetuning for ASR is timeconsuming and computationally demanding2 . The second challenge is that we do not observe any performance improvement of the subnetworks over the original dense network with OMP or IMP . Figure 3 shows the WER under low-resource scenarios of the subnetworks identified by OMP ( purple line ) and IMP ( blue dashed line ) at different sparsity levels . None of the sparsity levels achieves a visible drop in WER compared to the zero sparsity case , corresponding to the original dense network . These two challenges have prompted us to ask – do there exist sparse subnetworks within pre-trained speech SSL with improved performance on low-resource ASR ? How can we discover them efficiently in a single downstream finetuning run ? We propose a magnitude-based unstructured pruning method [ 41 , 11 ] , termed Prune-Adjust-Re-Prune ( PARP ) , for discovering sparse subnetworks within pre-trained speech SSL . PARP consists of the following two steps : 1 . Directly prune the SSL pre-trained model at target sparsity , and obtain an initial subnetwork and an initial pruning mask . 2 . Finetune the initial subnetwork on target downstream task/language . During finetuning , zero out the pruned weights specified by the pruning mask , but allow the weights be updated by gradient descent during backpropogation . After a few number of model updates , re-prune the updated subnetwork at target sparsity again . Step 1 provides an initial subnetwork that is agnostic to the downstream task , and Step 2 makes learnable adjustments by reviving pruned out weights . A formal and generalized description and its extension are introduced in Section 3 . Different from pruning methods in [ 49 , 39 ] , PARP allows pruned-out weights to be revived during finetuning . Although such a high-level idea was introduced in [ 48 ] , we provide an alternative insight : despite its flexibility , Step 2 only makes minimal adjustment to the initial subnetwork , and obtaining a good initial subnetwork in Step 1 is the key . We empirically show in Section 3 that any task-agnostic subnetwork surprisingly provides a good basis for Step 2 , suggesting that the initial subnetwork can be cheaply obtained either from a readily available task/language or directly pruning the pre-trained SSL model itself . In addition , this observation allows us to perform cross-lingual pruning ( mask transfer ) experiments , where the initial subnetwork is obtained via a different language other than the target language . Our Contributions . We conduct extensive PARP and baseline ( OMP and IMP ) pruning experiments on low-resource ASR with mono-lingual ( pre-trained wav2vec 2.0 [ 6 ] ) and cross-lingual ( pre-trained XLSR-53 [ 29 ] ) transfer . PARP finds significantly superior speech SSL subnetworks for low-resource 2Standard wav2vec 2.0 finetuning setup [ 6 ] on any Librispeech/Libri-light splits requires at least 50⇠100 V100 hours , which is more than 50 times the computation cost for finetuning a pre-trained BERT on GLUE [ 106 ] . ASR , while only requiring a single pass of downstream ASR finetuning . Due to its simplicity , PARP adds minimal computation overhead to existing SSL downstream finetuning . • We show that sparse subnetworks exist in pre-trained speech SSL when finetuned for low-resource ASR . In addition , PARP achieves superior results to OMP and IMP across all sparsities , amount of finetuning supervision , pre-trained model scale , and downstream spoken languages . Specifically , on Librispeech 10min without LM decoding , PARP discovers subnetworks from wav2vec 2.0 with an absolute 10.9 % /12.6 % WER decrease compared to the full model , without modifying the finetuning hyper-parameters or objective ( Section 4.1 ) . • Ablation studies on demonstrating the importance of PARP ’ s initial subnetwork ( Section 4.2 ) . • PARP minimizes phone recognition error increases in cross-lingual mask transfer , where a subnetwork pruned for ASR in one spoken language is adapted for ASR in another language ( Section 4.3 ) . PARP can also be applied to efficient multi-lingual subnetwork discovery for 10 spoken languages ( Section 4.4 ) . • Last but not least , we demonstrate PARP ’ s effectiveness on pre-trained BERT/XLNet , mitigating the cross-task performance degradation reported in BERT-Ticket [ 19 ] ( Section 4.5 ) . Significance . Findings of this work not only complement and advance current and future speech SSL for low-resource ASR , but also provide new insights for the rich body of pruning work . 2 Preliminaries . 2.1 Problem Formulation . Consider the low-resource ASR problem , where there is only a small transcribed training set ( x , y ) 2 Dl . Here x represents input audio , and y represents output transcription . Subscript l 2 { 1 , 2 , · · · } represents the downstream spoken language identity . Because of the small dataset size , empirical risk minimization generally does not yield good results . Speech SSL instead assumes there is a much larger unannotated dataset x 2 D0 . SSL pre-trains a neural network f ( x ; ✓ ) , where ✓ 2 Rd represents the network parameters and d represents the number of parameters , on some self-supervised objective , and obtains the pre-trained weights ✓0 . f ( x ; ✓0 ) is then finetuned on downstream ASR tasks specified by a downstream loss Ll ( ✓ ) , such as CTC , and evaluated on target dataset Dl . Our goal is to discover a subnetwork that minimizes downstream ASR WER on Dl . Formally , denote m 2 { 0 , 1 } d , as a binary pruning mask for the pre-trained weights ✓0 , and ✓l as the finetuned weights on Dl . The ideal pruning method should learn ( m , ✓l ) , such that the subnetwork f ( x ; m ✓l ) ( where is element-wise product ) achieves minimal finetuning Ll ( ✓ ) loss on Dl . 2.2 Pruning Targets and Settings . We adopted pre-trained speech SSL wav2vec2 and xlsr for the pre-trained initialization ✓0 . wav2vec 2.0 We took wav2vec 2.0 base ( wav2vec2-base ) and large ( wav2vec2-large ) pre-trained on Librispeech 960 hours [ 6 ] . During finetuning , a task specific linear layer is added on top of wav2vec2 and jointly finetuned with CTC loss . More details can be found in Appendix 8 . XLSR-53 ( xlsr ) shares the same architecture , pre-training and finetuning objectives as wav2vec2-large . xlsr is pre-trained on 53 languages sampled from CommonVoice , BABEL , and Multilingual LibriSpeech , totaling for 56k hours of multi-lingual speech data . We consider three settings where wav2vec2 and xlsr are used as the basis for low-resource ASR : LSR : Low-Resource English ASR . Mono-lingual pre-training and finetuning – an English pretrained speech SSL such as wav2vec2 is finetuned for low-resource English ASR . H2L : High-to-Low Resource Transfer for Multi-lingual ASR . Mono-lingual pre-training and multi-lingual finetuning – a speech SSL pre-trained on a high-resource language such as English is finetuned for low-resource multi-lingual ASR . CSR : Cross-lingual Transfer for Multi-lingual ASR . Multi-lingual pre-training and finetuning – a cross-lingual pretrained speech SSL such as xlsr is finetuned for low-resource multi-lingual ASR . 2.3 Subnetwork Discovery in Pre-trained SSL . One obvious solution to the aforementioned problem in Section 2.1 is to directly apply pruning with rewinding to ✓0 , which has been successfully applied to pre-trained BERT [ 19 ] and SimCLR [ 18 ] . All pruning methods , including our proposed PARP , are based on Unstructured Magnitude Pruning ( UMP ) [ 39 , 41 ] , where weights of the lowest magnitudes are pruned out regardless of the network structure to meet the target sparsity level . We introduce four pruning baselines below , and we also provide results with Random Pruning ( RP ) [ 39 , 41 , 19 ] , where weights in ✓0 are randomly eliminated . Task-Aware Subnetwork Discovery is pruning with target dataset Dl seen in advance , including One-Shot Magnitude Pruning ( OMP ) and Iterative Magnitude Pruning ( IMP ) . OMP is summarized as : 1 . Finetune pretrained weights ✓0 on target dataset Dl to get the finetuned weights ✓l . 2 . Apply UMP on ✓l and retrieve pruning mask m. IMP breaks down the above subnetwork discovery phase into multiple iterations – in our case multiple downstream ASR finetunings . Each iteration itself is an OMP with a fraction of the target sparsity pruned . We follow the IMP implementation described in BERT-Ticket [ 19 ] , where each iteration prunes out 10 % of the remaining weights . The main bottleneck for OMP and IMP is the computational cost , since multiple rounds of finetunings are required for subnetwork discovery . Task-Agnostic Subnetwork Discovery refers to pruning without having seen Dl nor l in advance . One instance is applying UMP directly on ✓0 without any downstream finetuning to retrieve m , referred to as Magnitude Pruning at Pre-trained Initailizations ( MPI ) . Another case is pruning weights finetuned for a different language t , i.e . applying UMP on ✓t for the target language l ; in our study , we refer to this as cross-lingual mask transfer . While these approaches do not require target task finetuning , the discovered subnetworks generally have worse performance than those from OMP or IMP . The above methods are only for subnetwork discovery via applying pruning mask m on ✓0 . The discovered subnetwork f ( x ; m ✓0 ) needs another downstream finetuning to recover the pruning loss3 , i.e . finetune f ( x ; m ✓0 ) on Dl .

Summary: The goal of this contribution is to reduce the complexity of self-supervised speech representation learning systems. Following a description of the landscape of unstructured magnitude pruning methods (i.e. methods that sort weights by their magnitude and remove the lowest ones with no other consideration), authors propose PARP. PARP is based on the observation that "task-agnostic" masks (obtained on a pre-trained network) are very similar to ones that would be obtained after fine-tuning on a target task. Taking benefit from this observation, and unlike OMP that cascades finetune -> prune -> finetune, PARP simply perform prune -> finetune. Experiments on low resource ASR and cross-lingual training show a clear advantage of PARP over IMP, OMP and MPI within a language (less obvious for high-to-low resource transfer and cross-lingual training), but also a surprising gain in performance w.r.t. unpruned architectures. Moreover, authors show that PARP also works on another modality, by pruning BERT models for GLUE downstream tasks.