Partitioned Gradient Matching-based Data Subset Selection for Compute-Efficient Robust ASR Training (2022)

Abstract: Training state-of-the-art ASR systems such as RNN-T often has a high associated financial and environmental cost. Training with a subset of training data could mitigate this problem if the subset selected could achieve on-par performance with training with the entire dataset. Although there are many data subset selection(DSS) algorithms, direct application to the RNN-T is difficult, especially the DSS algorithms that are adaptive and use learning dynamics such as gradients, as RNN-T tend to have gradients with a significantly larger memory footprint. In this paper, we propose Partitioned Gradient Matching (PGM) a novel distributable DSS algorithm, suitable for massive datasets like those used to train RNN-T. Through extensive experiments on Librispeech 100H and Librispeech 960H, we show that PGM achieves between 3x to 6x speedup with only a very small accuracy degradation (under 1% absolute WER difference). In addition, we demonstrate similar results for PGM even in settings where the training data is corrupted with noise.

## introduction owing to their simplicity in directly mapping an acoustic input sequence to a output sequence of characters, or words, or even word-pieces, neural end-to-end methods @xcite @xcite @xcite have become ubiquitous. the most common end-to-end architectures include (i) connectionist temporal classification (ctc) models @xcite , (ii) attentionbased encoder-decoder models (aed) @xcite and (iii) sequence transduction models @xcite such as rnn-ts @xcite . due to their streaming and low-latency properties, sequence transduction †equal contribution. correspondence to: arakeshk@in.ibm.com, durgas@cse.iitb.ac.in architectures such as rnn-t @xcite @xcite are becoming state-of-the-art for modeling the asr problem. these successes in the asr have come at a cost, as most of the practical rnn-t models are trained on thousands of hours of labeled datasets @xcite . model training on these massive datasets leads to significantly increased training time, energy requirements, and consequently the carbon footprint @xcite @xcite . as per parcollet et al. @xcite , training an rnn-t model on librispeech 960h @xcite emits more than 10kg co 2 if trained in france, which becomes much worse for developing countries. this is exacerbated due to the many more training runs required for hyper-parameter tuning. this warrants a need for greener training strategies that rely on significantly lower resources while still achieving state-of-the-art results. one way to make asr training more efficient is to train on a subset of the training data, which ensures minimum performance loss @xcite @xcite @xcite @xcite @xcite . since training on a subset reduces end-to-end time, the hyperparameter tuning time is also reduced. while greedy subset selection algorithms employ various criteria to identify the appropriate subset of training points, the process of forming the subsets remains sequential. however, for a large scale speech corpus such as librispeech @xcite this requirement may be difficult to meet. in this work, we propose a partitioned gradient matching (pgm) approach, which scales well with huge datasets used in asr and takes advantage of distributed setups. to the best of our knowledge, this is the first such study performed for asr systems. ## background: rnn transducer the rnn-t model @xcite maps an input acoustic signal (x 1 , x 2 , . . . , x t ) to an output sequence (y 1 , y 2 , . . . , y u ), where each output symbol y @xmath0, and m is the vocabulary. an rnn-t model consists of three components -(i) transcription network -which maps an acoustic signal (x 1 , x 2 , . . . , x t ) to an encoded representation (h 1 , h 2 , . . . , h t ), t being the length of the acoustic signal and x i being a w dimensional feature representation, (ii) prediction network -which is a language model that maps the previously emitted non-blank tokens @xmath1, y 2 , . . . , y u-1 to an output space g u for the next output token. (iii) joint network -that combines the transcription network representation h t and prediction network representation g u to produce z t,u using a feed-forward network j and ⊕ as a combination operator (typically a sum). during the training, the output probability p rnnt (y t,u ) over the output sequence y is marginalized over all possible alignments using an efficient forward-backward algorithm to compute the loglikelihood. the training objective is to minimize the negative log likelihood of the target sequence. for inference, the decoding algorithms @xcite attempt to find the best (t, u) and their corresponding output sequence y using a beam search. in this work, we use the gradients of the joint network layer (j) for pgm, since the linear layer helps in fusing the audio(h t ) and the text(g u ) representations. ## limitations of existing subset selection algorithms an approach to the selection of a subset of points from the entire dataset is to rank points based on their suitability. this ranking can be done either via a some static metric such as diversity or representation among features @xcite or via a dynamic metric using instancewise loss gradients foot_0 to construct the subset greedily @xcite . in the latter case, ranking and re-ranking happens using instance-wise loss gradients. specifically, during the selection process, loss gradients of the entire set of instances have to be available in the memory in order to perform greedy selection, since otherwise, subset selection time would be prohibitively large owing to disk reads, etc. as keeping all the loss gradients in the memory would be resource intensive, we employ the following approximations, which have been also previously employed by @xcite , viz., (i) only last layer gradients are used and (ii) subsets are constructed for each class. the latter technique is not relevant in asr systems since asr requires sequential decoding into a large size vocabulary. similar to the figure 1 : as pgm is a adaptive dss algorithm, pgm is invoked after for every r epochs training rnn-t using stochastic gradient descent. at every time step, using the latest set of parameters, pgm forms partial subsets via gradient matching (gm) across gpus. these partial subsets are combined and used for the next r epochs of rnn-t training. this is repeated until the final set of parameters is obtained. last layer approximation, for the rnn-t model, we use the gradients of the joint network layer (j) which performs the important task of fusing speech (h t ) and text (g u ) features for sequence transduction. in killamsetty et al. (2021a) propose another technique, viz., the perbatch version, wherein one se-dataset single total per gradient size batch size (mb) (gb) size (gb) cifar10 0.0215 1.049 0.0082 librispeech 100h 4.096 111 28 another problem with the existing subset selection algorithms is that they are sequential in nature. this doesn't allow the selection algorithm to enjoy the speedup achieved using state of the art techniques such as parallelizing across multiple gpus etc. this may cause the subset selection algorithm to be a bottleneck while training rnn-t with datasets of the scale of librispeech. therefore, there is a need to design an data subset algorithm that doesn't need all the loss gradient to form a subset and could be distributed across gpus. ## partitioned gradient matching algorithm let @xmath2, y i )} n @xmath3, and @xmath4, y j )} m @xmath5, the validation set. let θ denote the asr system's parameters with θ t as the asr system's parameters at the t th epoch. the training loss associated with the i th instance is denoted by l i t (θ) = l t (x i , y i , θ) = ln p r(y i |x i ). we denote the validation loss by l @xmath6. let the training data be divided into d partitions, i.e., @xmath7. let b be the batch size, b @xmath8. let l d p t be the training loss associated with a data partition d p and this selects a subset of batches x t d p and associated weights w t d p , such that the weighted sum of loss gradients associated with each instance in the subset are the best approximation of the loss gradient of the entire data partition d p while honoring the budget constraints. we perform gradient matching on mini-batch wise loss gradients only as it helps in reducing the memory needs. similarly, we can define gradient matching problem with loss associated with the validation set as, the optimization problem given in eq.( 5 ) is weakly submodular @xcite . hence, we can effectively solve it using a greedy algorithm with approximation guarantees -we use orthogonal matching pursuit (omp) algorithm @xcite to find the subset and their associated weights. we also add to eq.( 5 ) an l 2 regularization component to discourage large weight assignments to any of the instances selected in the subset, thereby preventing the model from overfitting on some samples. the complete algorithm is presented in algorithm 1. in the algorithm, 'val' is a boolean flag that indicates whether to match the subset loss gradient with validation set loss gradient like in noisy settings ('@xmath9. depending on the choice of the loss gradient, we perform gradient matching with l d p t , current model parameters θ t , budget b k d , and a stopping criterion ϵ. we describe gradient matching in details in algorithm 2. once the appropriate batch for selection is determined, we form x f adding all the samples constituting the selected mini-batch. the model is then trained using the mini-batch sgd. we randomly shuffle elements in the subset x t , divide them up into mini-batches of size b, and run mini-batch sgd with instance weights. the complete block diagram of pgm is presented in figure 1 . as the subset selection process is dependant on the model parameters, we repeat the subset selection every r epochs. for each data ## experiments datasets we perform all our experiments on the librispeech dataset @xcite . we present results on the medium-scale librispeech 100h as well as on the large-scale librispeech 960h datasets. along with the standard librispeech benchmark, we also perform experiments on noisy librispeech, where the speech is augmented with noise across varying signal-to-noise ratios (up to 15db) on a fraction of the training data. we refer to this dataset as librispeech-noise, where up to 30% examples in the original dataset are augmented with noise across varying signal-to-noise ratios. architecture. we perform all our experiments on the speechbrain's @xcite librispeech transducer recipe. the transcription network of the rnn-t consists of a crdnn encoder which has 2 cnn blocks followed by 4 layers of bi-lstms and subsequently followed by 2 dnn layers. the prediction network consists of an embedding layer followed by a single layer gru unit. a joint network is a single linear layer that projects 1024 dimensional representations to output a vocabulary of 1000 bpe. the decoding is done through a time-synchronous decoding algorithm @xcite with a beam size of 4. the decoding involves an external transformer language model trained on the librispeech corpus @xcite @xcite . training details. for the training, we employ a learning rate of 2.0 with an annealing factor of 0.8 for the relative improvement of 0.0025 on validation loss (sometimes referred to as newbob scheduler). the training on librispeech 100h is performed on two a100 40gb gpus with the effective batch size of 8, whereas for librispeech 960h, we employ two a100 80gb gpus with an effective batch size of 24. all the training is done for 30 epochs. in all our experiments, the pgm algorithm is invoked after every 5 th epoch (@xmath10. the results for each setting are averaged over 3 runs with different random seeds. pgm details. for doing the subset selection with pgm, we use the gradients of the joint network parameters, which we believe would have the maximum information concentrated for the sequence. we freeze the rest of the network while we compute the gradient of the joint network of the rnn-t. we use @xmath11. subset selection is performed using training set loss gradients in experiments performed using librispeech 100h (figures 2 , 3 ) and librispeech 960h ( baselines. we compare the results obtained using the pgm method against three intuitive baselines -(i) random-subset baseline, in which the subset of the dataset is obtained by choosing points with uniform probability. (ii) largeonly -for each subset, we employ only the largest utterances based on duration. (iii) largesmall -for each subset size, half of the subset is filled with smallest utterances and the other half with the largest utterances based on duration, to remove the length bias of the largeonly baseline. ## conclusion we propose pgm, a distributable data subset selection algorithm which avoids the need to load the entire dataset at a time, by constructing partial subsets from smaller data partitions. pgm is an adaptive subset selection algorithm that improves the training time of the asr models while maintaining low relative test error as compared to the asr model trained with the entire dataset. this speed-up improves the efficiency of the training process and subsequently reduces the carbon footprint of training such models. our approach performs consistently better than random-subset baseline whilst providing good speed up, and robustness in the presence of noise. although we test the method on the rnn-t model, we believe a connections between pgm and grad-matchpb lemma 1 (triangle inequality). let v 1 , ..., v τ be τ vectors in r d . then the following is true: