Macro-block dropout for improved regularization in training end-to-end speech recognition models

Deep learning technologies have significantly improved speech recognition accuracy recently [1, 2]. There have been series of remarkable changes in speech recognition algorithms during the past decade. These improvements have been obtained by the shift from Gaussian Mixture Model (GMM) to the Feed-Forward Deep Neural Networks (FF-DNNs), FF-DNNs to Recurrent Neural Network (RNN) such as the Long Short-Term Memory (LSTM) networks [3]. Thanks to these advances, voice assistant devices such as Google Home [4], Amazon Alexa and Samsung Bixby are widely used at home environments. Tremendous amount of research has been conducted in the process of switching from a conventional speech recognition system consisting of an Acoustic Model (AM), a Language Model (LM), and a decoder based on a Weighted Finite State Transducer (WFST) to a complete end-to-end all-neural speech recognition system [5].

Such notable shifts happened not only in research on model architectures, but also in research on model robustness as well. In conventional approaches, researchers have focused on cleaning or transforming speech signals [6, 7, 8, 9, 10, 11] and features [12, 13]. However, it has been recently observed that data augmentation [14, 15] is especially powerful in enhancing the model robustness. Data augmentation during the training may be helpful in reducing the environmental mismatch between the training and testing time. The Small Power Boosting (SPB) algorithm [16] may be considered as an extreme example of reducing environmental mismatch by intentionally transforming portions of inputs which are more susceptible to noise or reverberation.

A large number of end-to-end speech recognition systems are based on the Attention-based Encoder-Decoder (AED) [5] and the Recurrent Neural Network-Transducer (RNN-T) [17] algorithms. These complete end-to-end systems have started outperforming conventional WFST-based decoders for large vocabulary speech recognition tasks [2]. Further improvements in these end-to-end speech recognition systems have been possible thanks to a better choice of target units such as Byte Pair Encoded (BPE) and unigram language model [18] subword units, and an improved training methodologies such as Minimum Word Error Rate (MWER) training [19]. In training such all neural network structures, over-fitting has been a major issue. For improved regularization in training, various approaches have been proposed [20]. Data-augmentation has been also proved to be useful in improving model training [4, 21, 22, 23]. The dropout approach [24] has been applied to overcome this issue in which both the input and the hidden units are randomly dropped out for regularization. This dropout approach has inspired a number of related approaches [25, 26, 27]. In DropBlock [28], the authors claimed that dropping out at random is not effective in removing semantic information when training Convolutional Neural Networks (CNNs) because nearby activations contain closely related information. Note that in CNN, filters are applied to nearby elements, thus activation units are spatially correlated. Motivated by this, they apply a square mask centered around each zero position. However, since this kind of spatial correlation does not hold for fully connected feed-forward layers or RNNs, the application of DropBlock is limited to CNNs.

In this paper, we present a new regularization algorithm referred to as macro-block dropout. We define a macro-block that contains multiple input units to a neural network layer. Rather than applying dropout to each unit, we apply random dropout to each macro-block. In our experiments using an RNN-T [17] and an Attention-based Encoder Decoder (AED) in Sec. 5, this simple algorithm has shown a quite significant improvement over the conventional dropout approach. Our contribution in this paper may be summarized as follows. First, by using large macro-blocks, the ratio of dropped input units in each layer has large variation even if we keep a constant dropout rate. This variation in the ratio of dropped units leads to better regularization. Thus, unlike DropBlock that relies on spatial correlation in CNNs, we may apply macro-block dropout to any kinds of neural networks such as Feed-Forward (FF) networks and RNNs. To the best of our knowledge, our work is the first in using big chunks consisting of large number of input units to RNNs for masking. Second, we propose a new way of input scaling after applying a mask in macro-block dropout. This new scaling approach is related to the fact that the portion of dropped units significantly varies in macro-block dropout. Third, we proposed a very low-cost regularization algorithm. As will be seen in Sec 5, the best performance is achieved when the number of macro blocks for each layer is only four. This means that we only need to generate four random values for each layer. Thus, this macro-block dropout is very simple with very small computational requirement.

Dropout is a simple regularization technique to alleviate the overfitting problem by preventing co-adaptations [24]. When the dimension of the input to a neural network layer is $\mathbbm{d}_{\mathbf{x}}$, we create a random mask tensor $\mathbf{m}$ with the same dimension $\mathbbm{d}_{\mathbf{x}}$. Each element $m\in\mathbf{m}$ follows the Bernoulli distribution $m\sim{\text{\it Bernoulli}(1-q)}$, where $q$ is the dropout rate. Given an input $\mathbf{x}$, the dropout output $\mathbf{x}_{\text{out}}$ is obtained by the following equation:

where $\odot$ is a Hadamard product. The scaling by $\frac{1}{1-q}$ is applied to keep the sum of input units the same through this masking process. In (1), we adopt the “inverted dropout” approach rather than the original form of dropout in [24] where the scaling is performed during the inference time. Before introducing our macro-block dropout, let us consider the variance of the dropped ratio of an input layer. When the total number of input units to a single RNN is $N$, the expected ratio of kept units is given by $\frac{(1-q)N}{N}=1-q$, and the standard deviation is given by $\sqrt{\frac{q(1-q)}{N}}$. From the central limit theorem, we know that this distribution can be approximated by a Gaussian distribution. If we use typical values of $q=0.2$ and $N=10^{5}$, the standard deviation becomes $0.00126$, which is very small. As shown in Fig. 1, when there are a large number of units, the distribution is very similar to a delta function centered at $1-q$. Since the ratio of kept units is always very close to $1-q$, we conclude that $\frac{\sum\mathbf{x}\odot\mathbf{m}}{\sum\mathbf{x}}\approx 1-q$. Thus we can safely use the fixed scaling of $\frac{1}{1-q}$ to keep the sum of the input units the same after dropout in (1). Dropout has been turned out to be especially useful in improving the training of dense network models for image classification [29], speech recognition [30], and so on. This dropout approach inspired many other related approaches such as DropConnect [25], drop-path [26], shake-shake [27], ShakeDrop [31], and DropBlock [28] regularizations.

The authors of DropBlock [28] claim that dropping out inputs to CNNs at random is not effective in removing spatially correlated information. In DropBlock, the zero position, which is the center of a square mask, is randomly sampled in the same way as the dropout in [24]. It has been reported that DropBlock significantly outperforms the baseline dropout when applied to CNNs. However, unlike our macro-block dropout, Bernoulli random numbers are generated for each unit of inputs to neural networks. Thus, the ratio of dropped units has very small variation similar to that of the conventional dropout case, which is one of the biggest differences from the proposed Macro-block dropout.

For speech recognition experiments, we employed an RNN-T speech recognizer and an Attention-based Encoder Decoder (AED).Our speech recognition system is built in-house using Keras models [35] implemented for Tensorflow 2.3 [34]. The RNN-T structures have three major components: an encoder (also known as a transcription network), a prediction network, and a joint network. In our implementation, the encoder consists of six layers of bi-directional LSTMs interleaved with 2:1 max-pooling layers [36] in the bottom three layers. Thus, the overall temporal sub-sampling factor is 8:1 because of these three 2:1 max-pooling layers. The prediction network consists of two layers of uni-directional LSTMs. The unit size of all these LSTM layers is 1024. Macro-block dropout is applied to all the inputs of each LSTM layer of the encoder and the prediction network except the first layer of the encoder. From the transducer output, a linear embedding vector with a dimension of $621$ is obtained and fed back into the prediction network. The AED model has three components: an encoder, a decoder, and an attention block. In our implementation, the encoder of the AED model is identical to the encoder structure of the RNN-T model explained above. As a decoder, we use a single layer of uni-directional LSTM whose unit size is 1024.

The loss employed for training the RNN-T model is a combination of the Connectionist Temporal Classification (CTC) loss [37] applied to the encoder output and the RNN-T loss [17] applied to the full network, which is represented by the following:

We refer to this loss in (4) as the joint CTC-RNN-T loss. For the AED model, we employ the joint CTC-CE loss, which is given by:

as in [1]. For better stability during the training, we use the gradient clipping by global norm [38], which is implemented in Tensorflow as the tf.clip_by_global_norm API.

To obtain further improvement in speech recognition accuracy, we incorporate an improved shallow-fusion technique. In this approach, we subtract log prior probabilities of each label obtained from the transcript of the speech recognition training database. This idea is initially described in [39]. Our formulation is based on our earlier work in [40]:

where $y_{l}$ is the output at the output label index $l$, $\bm{\mathsf{x}}_{[0:m]}$ is the input feature vector sequence from the zero-th frame up to the (m-1)-st frame, and $\hat{y}_{0:l}$ is the estimated output label sequence from the output index zero up to $l-1$. $\lambda_{p}$ and $\lambda_{\text{lm}}$ are constant weighting coefficients for the log prior probability denoted by $\log p_{\text{prior}}\left(y_{l}\right)$ and the log probability from the Language Model (LM) denoted by $\log p_{\text{lm}}\left(y_{l}\big{|}\hat{y}_{0:l}\right)$, respectively. $\log p\left(y_{l}\big{|}\bm{\mathsf{x}}_{[0:m]},\hat{y}_{0:l}\right)$ is the log probability from the speech recognition model. In our experiments, we use $\lambda_{p}$ of 0.002 and $\lambda_{\text{lm}}$ of 0.48 respectively as in [40].

In this section, we explain experimental results using the macro-block dropout approach with the RNN-T and the AED model described in Sec. 4. For training, we used the entire 960 hours LibriSpeech [42] training set consisting of 281,241 utterances. For evaluation, we used the official 5.4 hours test-clean and 5.1 hours test-other sets consisting of 2,620 and 2,939 utterances respectively. The pre-training stage has some similarities to our previous work in [43]. In this pre-training stage, the number of LSTM layers in the encoder increased at every 10,000-steps starting from two LSTM layers up to six layers. We use an Adam optimizer [44] with the initial learning rate of 0.0003, which is maintained for the entire pre-training state and one full epoch after finishing the pre-training stage. After this step, this learning rate decreases exponentially with a decay rate of 0.5 for each epoch. $\mathbf{x}[m]$ and $\mathbf{y}_{l}$ are the input power-mel filterbank feature of size 40 and the output label, respectively. $m$ is the input frame index and $l$ is the decoder output step index. We use the power-mel filterbank feature instead of the more commonly used log-mel filterbank feature based on our previous experimental results [40, 45, 43]. In out power-mel filterbank feature, we employ the power-law nonlinearity with the power coefficient of $\frac{1}{15}$, which is suggested in [13, 46]. For better regularization in training, we apply the SpecAugment as a data-augmentation technique in all the experiments [22]. In our experiments with different dropout rates ranging from 0.0 to 0.5, for both the conventional dropout and the macro-block dropout approaches, the best WERs are obtained when the dropout rate $q$ is 0.2.

Table 3 summarizes the experimental results with conventional dropout and macro-block dropout approaches using the RNN-T and AED models. In the case of the macro-block dropout approach, we conducted experiments with four different partition sizes $\left(\mathbbm{d}_{\text{(par)}}=\{(1,3),(1,4),(1,5),(1,10)\}\right)$ with the one-dimensional masking pattern that is shown in Fig. 2(b). From this table, we observe that the best WER is obtained when the number of blocks is four ($\mathbbm{d}_{\text{(par)}}=(1,4)$). As shown in this table, when the RNN-T model is employed, the macro-block dropout algorithm shows relatively 4.30 % and 6.13 % WER improvements over the conventional dropout approach on the LibriSpeech test-clean and test-other sets, respectively. In the same table, we observe that the performance improvement using the AED model is also similar. We obtain 4.36 % and 5.85 % relative WER improvements on the same sets, respectively. Finally, we apply the improved shallow fusion in 6 to further improve the performance. Table 4 shows WERs obtained with the AED model using the improved shallow fusion in (6) with a Transformer LM [41]. As shown in this Table, macro-block dropout shows 2.86 % and 5.72 % relative WER improvement on the LibriSpeech test-clean and test-other sets.

In this paper, we describe a new regularization algorithm referred to as macro-block dropout. In this approach, rather than applying dropout to each input unit to RNN layers, we apply a random mask to a bigger chunk referred to as a macro-block. We propose an improved way of performing scaling for better performance with macro-block dropout in Sec. 3.2. In experiments using RNN-T and AED models, we obtain significantly better results with macro-block dropout compared to the conventional dropout approach. We compare the variance of the ratio of input units that are not dropped with the conventional dropout approach and with macro-block dropout approach. We observe that this variance is significantly larger with macro-block dropout. We believe this larger variance helps regularization during training.