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: Many self-supervised speech representation learning methods use masked prediction at their core. This paper proposes Ask2Mask (ATM) approach for informed masking during learning through a supervised teacher model that provides frame-level posteriors probabilities of linguistic output units. The frame-posterior probabilities extracted from the supervised teacher model are also used as utterance weights to favor ones with high confidence.