Ask2Mask: Guided Data Selection for Masked Speech Modeling

1 INTRODUCTION . Self-training and self-supervised training techniques rely on huge amounts of unlabeled speech or text data for better generalization . The self-training techniques such as pseudo-labeling ( Scudder , 1965 ; Kahn et al. , 2020 ) and student-teacher training ( Park et al. , 2020 ) have shown promising improvements by incorporating the data selection process . This data selection step removes pseudolabels with less confidence as denoted by the teacher model before feeding the input to a student model . Xu et al . ( 2021 ) shows that self-training and self-supervised training are complementary to each other and also show that self-supervised models act as good initialization for self-training techniques . Self-supervised training ( Hinton & Zemel , 1994 ) is a representation learning approach which implicitly learns patterns in the data without relying on explicit labels . Masked speech modeling ( MSM ) is the recent and successful self-supervised learning technique , thanks to the advent of BERT ( Devlin et al. , 2018 ) in NLP which inspired learning speech representations from masked inputs . MSM techniques such as wav2vec2 ( Baevski et al. , 2020 ) , HuBERT ( Hsu et al. , 2021b ) and w2v-BERT ( Chung et al. , 2021 ) have shown considerable gains across various down-stream speech tasks and have become the go-to models for ASR . Unfortunately , MSM does not have a data selection scheme to discard the irrelevant input samples and instead imposes burden on the training criterion to learn the relevance of the input samples in learning meaningful representations . Hsu et al . ( 2021a ) noticed the impact of not selecting relevant data from the huge amounts of unsupervised data during pre-training by showing degradation in ASR performance when fine-tuned to a target dataset with limited data . To mitigate this constraint , Chan et al . ( 2021 ) introduced substantially more fine-tuning data related to the target dataset but did not achieve satisfactory results . Hsu et al . ( 2021a ) attempted to solve this issue by heuristically selecting the data from a closed set of unsupervised speech databases or by pooling in data relevant to target dataset along with the existing pre-training dataset . However , this data selection approach is not done within the existing pre-training dataset and it is not completely empirically motivated . In this study , in order to break the above limitation of the MSM techniques , we propose a simple strategy named ask2mask ( ATM ) to incorporate data selection within a chosen pretraining dataset . • In ATM , the masking is done over the input samples or speech frames with higher confidence as determined by the scorer . This is contrary to the random selection of frames to be masked in conventional MSM models . We hypothesize that this guided selection of frames to be masked allows the model to focus on the frames which can provide meaningful representations . The scoring model used in this work is necessarily a speech recognition model trained on small amount of data and provides frame-level confidence for each input . • The ATM technique is further extended to exploit the confidence values provided by the scorer by directly using them to re-weight the MSM loss . We denote this approach as ATM with loss scaling ( ATM+S ) . It allows the model training to focus on certain utterances by down scaling the utterances with low-confidences . Similar to our work based on masking with external guidance , there is work in NLP that also benefit by incorporating masking with knowledge . In Sun et al . ( 2019 ) , masking is done at phrase-level segments in BERT and has shown to learn semantic dependencies . In Wang et al . ( 2019 ) , phonetic knowledge is injected to mask over phonetic segments to perform spectral augmentation . Our ATM approach is primarily motivated based on the recent work by Veselý et al . ( 2017 ) on semi-supervised learning of conventional ASR systems which shows that performing data selection at frame-level or token-level on unsupervised data provides better performance . The importance of pruning out the input samples at frame-level has been studied in Ferreira et al . ( 2021 ) to improve both classification and regression tasks . Few works on unsupervised learning also highlight the importance of weighting the data based on its confidence ( Wessel et al. , 2001 ; Ren et al. , 2020 ; Coleman et al. , 2019 ) . We hypothesize that ATM can leverage the effect of data selection within a particular training corpus to further enhance the recognition performance of MSM techniques . To summarize , our contributions are listed as follows : • Novelty : To the extent of our knowledge , ATM is the first approach to incorporate a withincorpus data selection strategy in MSM . We also show that data selection can be simply performed inside MSM by guided selection of frames to be masked using a scorer model . • Technical contributions : We provide two simple strategies to incorporate data selection into MSM pretraining by applying the confidence of the scorer : 1 ) choosing the data at framelevel by applying guided masking 2 ) soft weighting the data at utterance level by scaling the MSM loss of each utterance with its corresponding confidence score . ATM is designed to be compatible to all MSM based pre-training techniques . • Empirical study : Analysis is done to find an optimal masking percentage for ATM and we highlight the effectiveness of ATM across varying masking percentages . The importance of masking frames with high confidence is substantiated by empirically comparing it with masking low confident frames and random frames respectively . Experiments are performed on AMI data which is from a distinct condition compared to Libri-light corpus used for MSM based pretraining . The results confirm the importance of ATM by improving the recognition performance on evaluation sets of AMI by a significant margin . 2 PRELIMINARIES ON MASKED SPEECH MODELING ( MSM ) . In this section , we formally define the masked speech modeling ( MSM ) technique and brief primary instantiations including wav2vec2 and w2v-BERT . The technique can be formulated by defining input speech sequence X = [ x1 , x2 , ... , xT ′ ] , where xt is the log Mel-filterbank feature vector at time t. X is sent to the feature encoder Φ to obtain the encoded representations E = Φ ( X ) . The feature encoder contains convolutional layers performing subsampling at a factor of 4 and reducing the total number of frames of an utterance from T ′ to T to get E = [ e1 , e2 , ... , eT ] . E is then sent to two parallel modules : 1 ) masking component , and 2 ) quantizer . 2.1 MASKING . The idea behind masking input samples and predicting them was initially proposed in BERT ( Devlin et al. , 2018 ) and later adopted to speech ( Baevski et al. , 2020 ) with modifications to suit the characteristics of speech input . The masking is done over sets of frames or blocks b1 , b2 , ... , bK and accommodates overlap between blocks . Here K is the number of masked blocks in a randomly masked encoded sequence Ẽ . The importance of block masking is motivated by the improvements observed in Span-BERT by Joshi et al . ( 2020 ) and ERNIE ( Sun et al. , 2019 ) . The block bk = [ ik , c ] , where ik is the starting index of the masked block and c is the corresponding right context size denoting the number of consecutive speech frames . Here ik are randomly sampled from a uniform distribution . It has been empirically observed by Baevski et al . ( 2020 ) that 49 % of the frames are masked and c = 10 is chosen as the golden ratio to attain best representation during pre-training . 2.2 QUANTIZER . Gumbel-softmax quantizer component Ψ is used to get quantized representations which act as targets for wav2vec2 and w2v-BERT models . These quantized representations align to phonetic units as described in Baevski et al . ( 2020 ) . Each quantized vector is of L dimensions which denote the number of targets or codes used in a codebook . Each incoming input E is projected to L dimensions within the quantizer before applying the Gumbel-softmax . 2.3 CONTEXT NETWORK AND MSM LOSS . Wav2vec2-conformer : In this model type , the unmasked sequence E is sent to Ψ to get Q = Ψ ( E ) , where Q = [ q1 , q2 , .. , qT ] as described in Baevski et al . ( 2020 ) . The masked sequence Ẽ is fed to the context network Ω which contains conformer blocks to learn contextual representations from the input . C = Ω ( E ) denotes the output of the context network . The contrastive loss Lctr ( cj , qj ) objective is computed between the quantized representation qj and context network output cj ∈ C for all masked time instances j ∈ J . Diversity loss Ldiv is computed as an auxiliary objective in wav2vec2 to force the model to choose diverse codes in the quantization codebook . Detailed description of Ldiv is in Baevski et al . ( 2020 ) . The final training objective is denoted as : Lwv = Lctr + 0.1 · Ldiv , ( 1 ) where Lctr = ∑J j=1 Lctr ( cj , qj ) . HuBERT-conformer : This is another variant of wav2vec2-conformer model with two major differences : 1 ) Targets are k-means cluster ids which are computed over a small portion of input 2 ) Cross-entropy loss Lce ( ŷj , yj ) is computed between the prediction of the context network ŷj and the k-means cluster id target yj . W2V-BERT : This model marries the concept of wav2vec2 and BERT model by including an additional context network Λ containing conformer blocks in addition to Ω as in wav2vec2 . The Λ receives the output of the Ω and strives to further learn refined contextual information to get H = Λ ( C ) . The targets of w2v-BERT yj is computed by taking an argmax over the codebook dimensions L of quantized representations qj , l as : yj = arg max l qj , l , l ∈ L ( 2 ) Finally , the cross-entropy loss Lce ( ŷj , yj ) is computed between the prediction ŷj = softmax ( hj ) and the target yj over the masked time instances J . The final training objective Lwb = Lce + Lwv is a combination of cross-entropy loss and wav2vec2 loss . A block diagrammatic overview of the above MSM architectures are available in appendix A.7 . 3 ASK2MASK ( ATM ) . The primary reason to employ pre-training models is to exploit the abundantly available unsupervised data for improving ASR under limited availability of supervised data . While the MSM models such as wav2vec2 and w2v-BERT described in Section 2 exploit the unsupervised data , they treat each data with equal weight for computing the final objective.Instead , we generate a score st for each encoded frame et . This is used to select relevant data in a fine-grained manner during masking for computing the loss objective . 3.1 METHODOLOGY . For each encoded feature frame et ∈ E , the scorer emits probabilities p ( vt = l |E ) ; l ∈ L of the frame belonging to a particular label . The scorer model is a CTC based frame-synchronous ASR model separately trained with a limited amount of data . Our initial intuition was to chose the scorer ’ s training data to match the target data condition , however our empirical analysis in ( cf . Section 5.3 ) shows that the performance is agnostic to the scorer model ’ s training data . Finally , the confidence score st of the frame is defined as the maximum probability across all labels : st = max l p ( vt = l |E ) ( 3 ) We sample K masking start indices { i1 , .. , ik } with probabilities : p ( ik = t ) = st∑ sv v/∈ { i1 , .. , ik−1 } · δt/∈ { i1. , ,ik−1 } , ( 4 ) That is , we sample beginning frames with probability proportional to the scores of each frame . The indicator function δt/∈ { i1. , ,ik−1 } ensures that we sample without replacement . This is the key difference between ATM and the random masking in prior works as described in Section 2.1 . Prior works uniformly sample the start indices of each masking block b1 : K , while the ATM uses the probability distribution induced by the scorer . K is determined by the percentage of frames to be masked . We hypothesize that frames with maximum confidence from an external scoring model will be 1 ) easiest to learn using an MSM training criteria and 2 ) most informative in for pretraining to facilitate fine-tuning . Conversely , the lowest confidence frames , those more confusable to an external scoring model , will be the least reliably learned by MSM and least informative for pretraining . The resulting frames are sent as input to the MSM architecture and the final loss objective L is determined by either of the MSM objectives Lwv or Lwb described in Sections 2.3 . This modified training objective allows the model to focus on learning from gradients calculated from the frames with high confidences .

Summary: This study investigates the use of an external frame-synchronous CTC-based ASR system to get confidence scores for frames in masked speech modeling which is called AskToMask (ATM). These confidence scores are used in two ways:  1) as a distribution to determine the segment to mask (higher confidence means higher probability),  2) the utterance level average of the per-frame confidence scores is used as a scaling factor of the per-utterance loss.   The experiments perform unsupervised pre-training using these two proposed methods and then performs supervised fine-tuning for the final ASR task. Experiments on Librispeech and AMI show that confidence score-based selection of masking improves the final WER as compared to using uniform distribution in the mask selection (the first approach versus baseline). The second approach provides marginal gains on top of the first approach.