Fast and accurate factorized neural transducer for text adaption of end-to-end speech recognition models

A neural transducer model [1] consists of encoder, prediction, and joint networks. The encoder network is analogous to the acoustic model in hybrid models, which converts the acoustic feature $x_{t}$ into a high-level representation $f_{t}$, where $t$ is the time index. The prediction network works like a neural LM, which produces a high-level representation $g_{u}$ by conditioning on the previous non-blank target $y_{u-1}$ predicted by the RNN-T model, where $u$ is output label index. The joint network combines the encoder network output $f_{t}$ and the prediction network output $g_{u}$ to compute the output probability with

	$\displaystyle z_{t,u}=W*\text{relu}(f_{t}+g_{u})+b$
	$\displaystyle P(\hat{y}_{t+1}|x_{1}^{t},y_{1}^{u})=\text{softmax}(z_{t,u})$		(1)

To address the length differences between the acoustic feature $\textbf{x}_{1}^{T}$ and label sequences $\textbf{y}_{1}^{U}$, a special blank symbol, $\phi$, is added to the output vocabulary. Therefore the output set is $\{{\phi}\cup\mathcal{V}\}$, where $\mathcal{V}$ is the vocabulary set.

Two prediction networks are used in FNT [31], as shown in figure 1. One ($Predictor\_b$) is for the prediction of the blank label $\phi$, and the other ($Predictor\_v$) is vocabulary prediction (rightmost orange part in figure 1). The vocabulary predictor could be considered as a standard LM. The combination methods with encoder output $f_{t}$ for these two prediction outputs are different. For the blank prediction, it is the same as in standard neural transducer models.

$\displaystyle z^{b}_{t,u}=W^{b}*\text{relu}(f_{t}+g_{u}^{b})+b^{b}$

(2)

For the vocabulary prediction, it is firstly projected to the vocabulary size and converted to the log probability domain by the operation of log softmax. After this, it is added with the encoder output.

	$\displaystyle d_{t}^{v}$	$\displaystyle=$	$\displaystyle W_{enc}^{v}*\text{relu}(f_{t})+b_{enc}^{v}$
	$\displaystyle d_{u}^{v}$	$\displaystyle=$	$\displaystyle W_{pred}^{v}*\text{relu}(g_{u}^{v})+b_{pred}^{v}$
	$\displaystyle z_{u}^{v}$	$\displaystyle=$	$\displaystyle\text{log\_softmax}(d_{u}^{v})$
	$\displaystyle z^{v}_{t,u}$	$\displaystyle=$	$\displaystyle d_{t}^{v}+z_{u}^{v}$		(3)

Two combination outputs are concatenated and softmax is applied to get the final label probability

$\displaystyle P(\hat{y}_{t+1}|x_{1}^{t},y_{1}^{u})=\text{softmax}([z^{b}_{t,u};z^{v}_{t,u}])$

(4)

The loss function of FNT is

$$\mathcal{J}_{f}=\mathcal{J}_{t}-\lambda\log P(\textbf{y}_{1}^{U})$$

(5)

where the first term is the standard neural transducer loss and the second term is the LM loss with cross entropy (CE). $\lambda$ is a hyper-parameter to tune the effect of LM loss.

As showed in section 2.2, the encoder output and the vocabulary predictor output are combined by sum operation. The predictor output is log probability, but the encoder output is not. According to Bayes’ theory, the acoustic and language model scores should be combined by weighted sum in log probability domain. Therefore we refine FNT by converting the encoder output to the log probability by adding log softmax

Furthermore, to force the encoder part act more like the acoustic model, CTC criterion is added for the encoder output as shown in the blue frame part in figure 2. The reason we choose CTC instead CE is that it’s not easy to get the sentence piece level alignment for training data, while sentence piece unit is commonly used as the output unit for neural transducer E2E model.

With such changes, the combination of encoder output and vocabulary predictor output is shown in below equations.

	$\displaystyle z_{t}^{v}$	$\displaystyle=$	$\displaystyle\text{log\_softmax}(d_{t}^{v})$
	$\displaystyle z^{v}_{t,u}$	$\displaystyle=$	$\displaystyle z_{t}^{v}[:-1]+\gamma*z_{u}^{v}$		(6)

where $\gamma$ is a trainable parameter, which could be taken as LM weight. One thing needs to be mentioned is after adding CTC, the dimension of $d_{t}^{v}$ and $z_{t}^{v}$ become vocabulary_size+1 because CTC needs one extra output “blank”. Here we put “blank” as the last dimension. and it’s excluded when $z_{t}^{v}$ is added with $z_{u}^{v}$.

The final loss function can be written as

$$\mathcal{J}_{f}=\mathcal{J}_{t}-\lambda\log P(\textbf{y}_{1}^{U})+\beta\mathcal{J}_{ctc}$$

(7)

where $\mathcal{J}_{ctc}$ is CTC loss and $\beta$ is a hyper-parameter to be tuned in the experiments.

To adapt FNT model with text data, the most straightforward way is to finetune the vocabulary predictor with the adaptation text based on cross-entropy loss. But this may degrade model performance on general domain. To avoid this, KL divergence between the vocabulary predictor outputs of adapted model and baseline model is added during the adaptation as shown in the orange frame part in figure 2.

The adaptation loss with KL divergence is

$$\mathcal{J}_{adapt}=\text{CE}(Z_{u}^{v},Y_{adapt})+\alpha\text{KL}(Z_{u}^{v},Z_{u}^{{}^{\prime}v})$$

(8)

where $Y_{adapt}$ is the adaptation text data, $Z_{u}^{v}$ is the log softmax of adaptation text from the adapted model and $Z_{u}^{{}^{\prime}v}$ is the log softmax of adaptation text from the baseline model. $\alpha$ is the KL divergence weight to be tuned in the experiments.

Since the vocabulary predictor in FNT is designed to be an LM, we explore the possibility of training it independently on a much larger text corpus than the transcriptions in the FNT training data. The parameters of this pre-trained external LM could be further updated to potentially improve accuracy. In principle, we could choose a variety of architectures for the external LM. In this paper, we limit ourselves to an architecture that is very close to that of the standard prediction network for a fair comparison with the baseline system. The vocabulary of the external LM is the same as that of the FNT system, and is trained using the conventional cross-entropy loss. Experimental results in Section 4 show that external LM trained with more data improves model accuracy, and updating the external LM parameters during FNT model training further improves the results.

As noted before, fine-tuning the vocabulary predictor is one of the straightforward adaptation method for FNT. Although updating vocabulary predictor is much faster than updating the whole FNT network, it could not meet the immediate adaptation requirement for some applications. In this paper, we propose to use n-gram integration for fast adaptation of FNT.

In this method, a n-gram based language model is firstly trained with the adaptation text data. Then it is interpolated with the probability output from the vocabulary predictor during the decoding. The vocabulary log probability after interpolation with n-gram probability $P(y_{u}|y_{1}^{u-1})_{ngram}$ is calculated as

$\displaystyle z_{u}^{v}=\text{log}((1-w)*P(y_{u}|y_{1}^{u-1})_{pred}+w*P(y_{u}|y_{1}^{u-1})_{ngram})$

(9)

where $P(y_{u}|y_{1}^{u-1})_{pred}=\text{softmax}(d_{u}^{v})$ is the label probability output from vocabulary predictor. Then $z_{u}^{v}$ is plugged into Equation (6) to calculate $z^{v}_{t,u}$ which is used to generate the final output of FNT.

In this method, no neural network training is involved, and n-gram LM model training is super fast with the adaptation text data. Experimental results in section 4.3 show it has much lower computational cost compared to the fine-tuning based method.

All results on general testing set are given in table 2. The baseline model is a standard C-T model. The standard FNT model is with the same structure as in [31], and it is trained with 30k training data from scratch. Compared with the baseline C-T model, the standard FNT model got 1.29% relative WER increase. To reduce the accuracy degradation, FNT model is refined by adding the CTC criterion based on equation 6 and 7. The CTC loss weight $\beta$ is 0.1. We can see that adding CTC decreases WER from 11.01 to 10.97, but is still worse than the baseline C-T model. The degradation is reduced further by initializing vocabulary predictor from an external language model. Two recipes are examined: one is the external LM is fixed during the FNT model training, and the other is the external LM is updated together with other parts of FNT model. Updating the external LM got the best result, which is even better than the baseline C-T model.

In this section, We first evaluated the impact of KL divergence using Librispeech set based on standard FNT model. The results are given in table 3. The results showed the adapted model without KL divergence degraded the accuracy on general testing set largely. FNT adaptation with KL divergence helped to recovered the loss on general testing set obviously with very small WER increase on adaptation set compared to the standard FNT adaptation. In the following adaptation experiments, KL divergence weight $\alpha$ is always set to 0.1.

Table 4 shows the adaptation results for different FNT models on all adaptation sets. For each task, the results in “base” column are WER before adaptation, and the results in “adapt” column are WER after adaptation. Simple average WERs are also reported by averaging the WERs from all three tasks. Comparing “base” results for B0 and F0, we could find the accuracy gap between the baseline C-T model and the standard FNT model on these adaptation set is much larger than that on general set, especially for task1 and task2. The possible reason is that the domains in the training data may have some coverage for the general testing set, but they are totally irrelevant to these adaptation sets. With the proposed refinements for FNT, this gap is decreased step by step, the best FNT model F3, which is with full combination of the proposed methods, could get the similar accuracy as the baseline C-T model. The same trend could be also observed for the “adapt” results. Each method contributes accuracy improvement to the adapted FNT model. Compared with the adapted standard FNT model (F0), the adapted model F3 reduces WER for three adaptation sets by relatively 9.48% in average (from 12.80% to 11.59% ). And compared with the baseline C-T model, the adapted model F3 gets 29.21% relative WER reduction (from 16.37% to 11.59%).

In this section, we evaluated n-gram interpolation performance based on the best FNT model (F3) for three adaptation tasks. For each task, a 5-gram LM is trained with the adaptation text. To make the interpolation simple and efficient, sentence piece instead of word is used as the basic unit for the 5-gram LM. The interpolation weight $w$ is always set as 0.3. The results are shown in the last row of table 4. Compared with the adaptation of fine-tuning vocabulary predictor, n-gram interpolation based adaptation got a little higher WER, but the relative WER reduction over the baseline C-T model is still satisfying, which is 21.04% (from 16.37% to 12.92%). More importantly, the adaptation speed is improved hugely. Experiments show that for fine-tuning method, the adaptation process cost about 10 seconds per 1,000 words on GPU and the cost is formidable when adapting with CPU. In contrast, it only needs about 0.002 seconds per 1,000 words on CPU for n-gram LM training. This is very useful for those application scenarios which need immediate adaptation.