Improving non-autoregressive end-to-end speech recognition with pre-trained acoustic and language models

In our proposed NAR CTC/attention model, the CTC branch plays an important role. It is expected to accurately predict a target length for the modality conversion mechanism and consider all the time boundaries during cache-based CTC/attention joint decoding. Given the promising results recently achieved by wav2vec2.0 [9, 17] based on the CTC criterion [18], we select it as our acoustic encoder.

The wav2vec2.0 encoder consists of a convolutional neural network (CNN)-based feature encoder and a Transformer-based context network. As shown in Fig. 1 (a), we use it to extract an acoustic representation ${\rm\mathbf{H}}^{\rm AC}=({\boldsymbol{h}}_{1},{\boldsymbol{h}}_{2},\cdot\cdot\cdot,{\boldsymbol{h}}_{T})^{\top}\in\mathbb{R}^{T\times d}$ from the raw speech waveform $\boldsymbol{x}$, where $d$ is the representation dimension and $T$ denotes time steps. During training, we apply a CTC branch after the encoder with the CTC objective $\mathcal{L}_{\rm ctc}$.

Incorporating BERT into an ASR system requires solving the mismatch between speech and text modalities. As shown in Fig. 1 (b), to efficiently connect acoustic and linguistic representations, we design a novel structure named modality conversion mechanism (MCM), which consists of the following two parts:

First, to bridge the length gap between the acoustic and linguistic representations, we take advantage of the monotonic property of the CTC branch and use it with a greedy search to predict a target length $\hat{L}$. To achieve efficient training, we choose a ground truth $L$ as our target length during training. We implement position encoding over this target length $L$ and broadcast its dimension to be consistent with the acoustic representation ${\rm\mathbf{H}}^{\rm AC}$. This broadcast is then fed into a cross attention block together with the ${\rm\mathbf{H}}^{\rm AC}$ to capture an alignment between the acoustic representation and its corresponding text. This alignment converts the $T$-length ${\rm\mathbf{H}}^{\rm AC}$ to an $L$-length embedding:

where ${\rm\mathbf{H}}^{\rm PE}\in\mathbb{R}^{L\times d}$ denotes the positional embedding, and ${\rm\mathbf{H}}\in\mathbb{R}^{L\times d}$ refers to the obtained acoustic embedding.

Second, we attempt to solve the mismatch between the acoustic embedding ${\rm\mathbf{H}}$ and the linguistic token embedding of BERT. BERT uses very large character/BPE tokens compared with the ASR system, and its adjustment is difficult. Previous works [11, 12] employ the token embedding of BERT as the learning target of ${\rm\mathbf{H}}$, but these approaches require a significant increase in the training iterations. Further, BERT’s token embedding is available during training and decoding, making it unnecessary to regard it as a learning target. Therefore, instead of adopting these methods [11, 12], we directly perform a dot product operation on the token embedding of BERT. Before using BERT, we first generate a preliminary prediction based on the obtained acoustic embedding:

where matrix ${\rm\mathbf{W}}\in\mathbb{R}^{d\times V}$ are trainable and ${\rm\mathbf{L}}^{a}\in\mathbb{R}^{L\times V}$ denotes the logits of acoustic embedding. Letting $V$ denote the vocabulary size of the ASR system, we then have ${\rm\mathbf{M}_{\rm BERT}}\in\mathbb{R}^{V\times d}$, a matrix that contains the token embedding of BERT shared with the ASR system, arranged in the order of tokens list in the ASR system²²2It should be noted that this approach works well with character-based BERT for Mandarin, but it needs extensions for BPE tokens used in English as the ASR system and pre-trained LM system have different BPE models, which are related to the vocabulary size of BPE thus cannot be shared.. We can then achieve an efficient modality conversion through dot product:

where ${\rm\mathbf{W}}^{a}\in\mathbb{R}^{L\times V}$ and each element ${\rm\mathbf{W}}^{a}_{ij}$ denotes the weight assigned to the $j$-th BERT token embedding at the $i$-th step. We then obtain the linguistic representation ${\rm\mathbf{H}}^{\rm LM}\in\mathbb{R}^{L\times d}$, which is fed into the BERT encoder. We use a cross-entropy (CE) criterion $\mathcal{L}_{ce1}$ to encourage the ${\rm\mathbf{L}}^{a}$ after softmax to generate correct predictions before feeding it into the BERT.

We choose to incorporate BERT [19] into our ASR system due to its powerful text processing capabilities enabled by its embedding layer and a multi-layer Transformer encoder [19, 20]. As shown in Fig. 1 (a), we apply a FC layer on the top of the BERT output to obtain logits of the linguistic representation ${\rm\mathbf{L}}^{l}$.

To utilize both acoustic and linguistic knowledge, we add the ${\rm\mathbf{L}}^{a}$ and ${\rm\mathbf{L}}^{l}$ together to get the final logits ${\rm\mathbf{L}}^{f}$:

where $\alpha$ is a tunable hyper-parameter. We calculate the CE loss $\mathcal{L}_{\rm ce}$ over the final output ${\rm\mathbf{L}}^{f}$ to provide further supervision.

During joint training, the loss function is defined by:

where $\lambda_{1}$, $\lambda_{2}$, and $\beta$ are tunable hyper-parameters. $\mathcal{L}_{\rm ctc}$ and $\mathcal{L}_{\rm ce1}$ are defined in Sections 2.1 and 2.2, respectively.

We employ the CTC branch to generate a target length $\hat{L}$, which enables the BERT to predict tokens in parallel. Inspired by [21], we design a cache-based CTC/attention joint decoding method to further improve the recognition accuracies, as shown in Fig. 1 (c). Unlike the conventional CTC/attention joint decoding [21], we precompute the attention-based score up to $\hat{L}$ length at one step instead of incremental computation, thus enabling our ASR model to achieve high inference speed ³³3Strictly speaking, this is not a NAR decoding, but it is used to further improve the proposed model with efficient computation..

The details of our method are shown in Algorithm 1, where $\Omega_{l}$ and $\hat{\Omega}$ denote queues that accept partial hypotheses of length $l$ and complete ones, respectively. $\mathcal{V}$ is the token vocabulary with $V$ size. Each $c\in\mathcal{V}\cup{\left\langle eos\right\rangle}$ is appended to a former partial hypothesis $g$ given by $\Omega_{l-1}$. The attention-based scores of $\hat{L}$ length sequence are calculated at one step and then stored in the cache. When $l>\hat{L}$, we predict ${\left\langle eos\right\rangle}$ with a 0.9 probability. We score each extended hypothesis $h\mbox{=}g\cdot c$ by a CTC prefix score and attention-based scores stored in the cache, as shown in line 14. This hypothesis is stored in either $\Omega_{l}$ or $\hat{\Omega}$ based on the value of $c$. While $c\neq{\left\langle eos\right\rangle}$, $h$ is added to $\Omega_{l}$, and $\Omega_{l}$ ’s size is then compared with the beam width for pruning by REMOVEWORST($\cdot$). We also utilize an additional function ENDDETECT$(\hat{\Omega},l)$ [21] to determine whether to stop the procedure before $l$ reaches $L_{\max}$. If finding $h\in\hat{\Omega}$ with higher scores is almost impossible as $l$ increases, ENDDETECT$(\hat{\Omega},l)$ returns $true$.

We apply softmax to ${\rm\mathbf{L}}^{f}$ to get an attention-based probability $p_{f}(c|x)$. We also cumulatively calculate a prefix CTC probability $p_{\rm ctc}(h,\cdot\cdot\cdot|x)$. The scores are then computed in log domain:

We evaluate our proposed NAR CTC/attention model on the Mandarin AISHELL-1 [22] and English Switchboard [23] corpora. We also use an unlabeled speech training set from the AISHELL-2 corpus to pre-train a Mandarin wav2vec2.0 base model [24].

The ESPnet2 toolkit [25] is used to build our wav2vec2.0 CTC baseline and our proposed NAR CTC/attention model. We employ raw speech as the acoustic input. For Mandarin, we use 4230 Chinese characters with 3 non-verbal symbols: blank, unknown-character, and sos/eos as the modeling units. For English, we utilize 2000 modeling units, including 1997 BPE [16] units and 3 non-verbal symbols.

Our wav2vec2.0 CTC baseline consists of a base wav2vec2.0 [9] encoder and a FC layer as the classifier with a size of 4233 and 2000 for Mandarin and English tasks respectively. If not specified, wav2vec2.0 refers to Mandarin wav2vec2.0 in the Mandarin task. As for our proposed NAR CTC/attention model, the wav2vec2.0 encoder and CTC branch are the same as those in the baseline model, and its MCM contains one-layer multi-head cross attention with 768 model dimensions and 4 heads. The FCs in Fig. 1(a) and Fig. 1(b) are of the same size as the CTC branch. The pre-trained Mandarin BERT (i.e., bert-base-chinese) and the English RoBERTa ⁴⁴4RoBERTa is a variant of BERT, and we choose it for our English tasks because its modeling unit is BPE and thus more suitable for our English tasks. [26] (i.e., roberta-base) are both provided by Huggingface Transformer Library [27], and the English wav2vec2.0 base model is provided by Fairseq [28]. For the first 5000 and 25000 steps, the parameters of both the wav2vec2.0 and the BERT/RoBERTa are fixed. In our embedding layer and the first 3 and 6 Transformer layers of BERT and RoBERTa, the parameters are always fixed. We set both the $\lambda_{1}$ and $\lambda_{2}$ in Eq. (6) to 0.5. The $\beta$ in Eq. (6) and $\alpha$ in Eq. (5) are set to 0.3. Real time factor (RTF) is measured on the test set of AISHELL-1 using a P100 GPU.

Following the ESPnet2 recipe [25], we fine-tune the Mandarin BERT and English RoBERTa with subsequent masks as external AR-based LMs for the baseline. During inference, we set the beam size as 10 and the weight of BERT/RoBERTa LM is 0.3 when used.

We compare the performance of our proposed NAR CTC/attention model with the strong wav2vec2.0 CTC baseline and other AR/NAR systems. The results are shown in Table 1, where W2v2 denotes the Mandarin wav2vec2.0. It should be noted that when the CTC baseline model uses an external BERT LM during one-pass decoding with beam search, we classify it as an AR system because the BERT is set to a unidirectional AR structure during the fine-tuning stage. And our NAR CTC/attention model employs greedy search during inference. The results show that our proposed NAR CTC/attention model greatly outperforms the strong wav2vec2.0 CTC baseline, yielding $15.1\%$ relative CER reduction.

Furthermore, even when a strong BERT-based AR LM is employed for the baseline to relax its conditional independence assumption, the results show that our proposed NAR CTC/attention model still achieves better recognition accuracy in addition to maintaining the advantage of the NAR model’s fast inference speed, which is not achieved by previous works [12, 33]. Comparing with other AR/NAR systems, our NAR model greatly improves over previous NAR systems with much fewer training epochs and it even exceeds the AR-based systems. Therefore, it can be concluded that our NAR CTC/attention model achieves impressive performance with high training efficiency.

In the main experiments, we employ the greedy search for our NAR model. In this section, we conduct ablation studies to verify the effectiveness of our cache-based CTC/attention joint decoding. The results are shown in Table 2, where Baseline represents our wav2vec2.0 CTC baseline and Cache CTC/Att denotes the cache-based CTC/attention joint decoding (i.e. beam search).

It can be seen from the results that the conclusions are the same whether we use the Mandarin or the English wav2vec2.0 encoder: 1) Our proposed NAR model outperforms the baseline even with vanilla attention-based greedy search. 2) After using the cache-based CTC/attention joint decoding method, our model has been further improved, especially on the English wav2vec2.0 encoder. 3) Although the cache-based CTC/attention joint decoding is slower compared to greedy search, it is significantly faster than the beam search with an AR-based external BERT LM, as our model only needs to be run for one time during inference.

In addition, the improvement on the English wav2vec2.0 encoder indicates that our proposed method may also help assist ASR tasks in low-resource languages based on cross-lingual pre-training.

Previous experiments have verified the effectiveness of our proposed model on Mandarin (i.e. logographic language) tasks. However, if we consider applying our method to alphabetic languages (e.g., English), it raises a problem regarding the vocabulary size of BPE, as an ASR system and BERT are likely to have different output units. For example, given the sentence ”$good\ weather$”, the tokenized input for RoBERTa is ”$\_good\ \_weather$”, but the corresponding input for the ASR system is ”$\_good\ \_wea\ ther$”. The BPE vocabulary size of RoBERTa is often word-level, which makes it too large and sparse for the ASR system to train on, especially for the CTC branch. A simple solution is that pre-training RoBERTa with the same BPE size as the ASR system. However, it is not always feasible due to the requirement of large computational resources.

To solve this challenge, first, since BPE units like ”$\_wea$” and ”$ther$” also exist in the token list of RoBERTa, we can directly convert these units to their corresponding index in RoBERTa’s token list before entering into the RoBERTa model. Second, we can also use CTC greedy search results as input to the MCM during training and decoding to avoid the length mismatch.

We choose the English Switchboard [23] corpus to evaluate our NAR model on alphabetic language. The results in Table 3 show that our method outperforms the strong wav2vec2.0 CTC baseline and AR-based Transformer system. However, after using a RoBERTa-based AR LM, the baseline surpasses our NAR model.