Transformer Transducer: One Model Unifying Streaming and Non-streaming Speech Recognition

Transformer Transducer [4] is a model architecture that can be trained with end-to-end RNN-T loss [7] using transformer based audio encoder and label encoder. As shown in figure 1, T-T model predicts a probability distribution over the label space at every time step. The probability of an alignment $P({\mathbf{z}}|{\mathbf{x}})$ can be factorized as

$\displaystyle P({\mathbf{z}}|{\mathbf{x}})=\prod_{i}P(z_{i}|{\mathbf{x}},t_{i},\mathrm{Labels}(z_{1:(i-1)})),$

(1)

where $\mathrm{Labels}(z_{1:(i-1)})$ is the sequence of non-blank labels in $z_{1:(i-1)}$. In T-T architecture $P({\mathbf{z}}|{\mathbf{x}})$ is parameterized with an audio encoder, a label encoder, and a joint network. The model defines $P(z_{i}|{\mathbf{x}},t_{i},\mathrm{Labels}(z_{1:(i-1)}))$ as follows:

$$\begin{split}\mathrm{Joint}=&\mathrm{Linear}(\mathrm{AudioEncoder}_{t_{i}}({\mathbf{x}}))+\\ &\mathrm{Linear}(\mathrm{LabelEncoder}(\mathrm{Labels}(z_{1:(i-1)}))))\end{split}$$

(2)

$$\begin{split}P(z_{i}|{\mathbf{x}},t_{i},\mathrm{Labels}(&z_{1:(i-1))})=\\ &\mathrm{Softmax}(\mathrm{Linear}(\mathrm{tanh}(\mathrm{Joint}))),\end{split}$$

(3)

where each $\mathrm{Linear}$ function is a different single-layer feed-forward neural network, $\mathrm{AudioEncoder}_{t_{i}}({\mathbf{x}})$ is the audio encoder output at time $t_{i}$, and $\mathrm{LabelEncoder}(\mathrm{Labels}(z_{1:(i-1)}))$ is the label encoder output given the previous non-blank label sequence.

More details about each transformer layer is shown in figure 2. It contains normalization layer, masked multi-head attention layer with relative position encoding, residual connection, stacking/unstacking layer and feed forward layer. The residual connections are applied with normalized input to the output of the attention layer or feed forward layers. The stacking/unstacking layer can be used to change frame rate for each transformer layer which helps to speed up training and inference. For further optimization, the label encoder can also be a bigram label embedding model.

In self attention block of transformer architecture we compute the self attention over entire input sequence and then mask it based on the left and right context of the layer. In variable context training we keep the left context of the layers constant and sample a right context length from a given distribution. The sampled right context length decides the mask for the self attention. In our experiments we found that randomly sampling right context length for each layer independently led to unstable training, and we had to limit the number of right context configurations. During training we specify right context configuration as a list of right contexts for each transformer layer in the encoder ex: $[0]\times 15+[8]\times 5$ specifies a right context of 0 for first 15 layers and right context of 8 frames for last five layers. During training we uniformly sample a right context configuration from all these available configurations. We show in section 4.3.2 that models trained in this way can be used with any right context configurations used for training.

Variable context training allows us to train transformer layers that are capable of using different right context from input, effectively at inference. We apply this technique to train a model that has input followed by several initial layers trained with zero right context and final few layers trained as variable context layers. With this model we can do speech recognition with different future audio context (depending on the right context used at inference time). This can be useful if an application can afford a higher latency for better quality results. We can also use this model for low latency speech recognition by running two parallel decoders using the same model, one with no or very small right context and one with larger right context. Since we keep the right context of shared layers constant (at zero) we only need to recompute the activation for the final layers with variable context. We call this a Y-model because we have parallel decoding branches running with different right context. We call the decoding branch with smaller right context ’low latency branch’ and decoding branch with higher right context ’high latency branch’

Transformer-Transducer minimizes RNN-T loss which contains all the alignment paths for each label sequence. Although optimizing label likelihoods should provide advantage on improving ASR performance, a model can have high alignment delay because there’s no mechanism to restrict prediction delay in optimizing RNN-T loss. The high alignment delay happens especially for a streaming model which does not have a right context because the model tries to improve prediction accuracy by looking ahead future frames.

Constrained alignment training was originally proposed in [8]. It puts constraints on RNN-T loss by masking out high delay alignment paths from reference alignments. To find reference alignments, we used full-attention non-streaming Transformer-Transducer model as a reference alignment model because it showed almost no delay. Based on the reference alignment, we put constrained window on word-boundaries and any word label path outside of the constrained window is masked out from the RNN-T loss. We applied constrained alignment training on Y-model and its evaluation results are in section 4.3.2.

Streaming audio encoding is the process of taking audio frame(s) and previous states as inputs, and outputting the corresponding encoded feature(s) and next states. Unlike RNN/LSTM based audio encoder, Transformer encoder has no time dependency while encoding audio. In other word, to encode the $i$-th frame, RNN/LSTM based encoders need to first encode from the $0$-th to the $i-1$-th frame before encoding the $i$-th frame, whereas, Transformer encoder can encode all the $i$ frames at the same time in parallel. We refer to the way of encoding a batch of time steps in parallel ”batch step” inference. In the later section, we will discuss the effectiveness of batching more time frames in terms of inference speed.

In non-streaming inference, ideally we can run the Transformer encoder exactly the same way it is run during training time. However, in practice, because of the $O(T^{2})$ memory consumption in attention matrix computation, the Transformer encoder can be very memory expensive. As showed in [4], limited context Transformer encoder provides about the same quality as infinite context encoder. Based on this model architecture, we can compute the attention matrix for a few $Query$ blocks at a time, and each $Query$ block only attends to a limited set of $Key$s, which makes the memory consumption constant. We refer to this inference method as ”query slicing”.

In the transformer transducer model, we use a label encoder as an auto-regressive model to encode predicted label history. The label encoder output is combined with the audio encoder output using a dense layer before the softmax function. The decoding speed is highly dependent on the computational complexity of label encoder network because it is run for every hypothesis. For RNN-T models, it has been shown that the label history context length can be reduced significantly without affecting the accuracy of the models [9]. In this paper, we experimented with two models for label encoding. The first one is a transformer model with limited label context length as described in [4]. The other one is an embedding model with a bigram label context. The embedding model learns a weight vector of $d$ dimension for each possible bigram label context, where $d$ is the dimension of audio and label encoder outputs. The total number of parameters is $N^{2}*d$ where $N$ is the vocabulary size for the labels. The learned weight vector is simply used as the embedding of bigram label context in the T-T model. Since this is a simple embedding lookup based on label bigrams, the runtime for label encoder is very fast.

To further speed up the inference, we maintain a cache to store and reuse the computed label encoder outputs for the limited label contexts seen so far in decoding since during decoding the model will encode the same label history multiple times for different hypothesis.

For our experiments we use 30K hours of speech data from voice-search application. The test set we use consists of 14K Voice Search utterances with duration less than 5.5 seconds long. The training and test sets are all anonymized and hand-transcribed. The input speech wave-forms are framed using a 32 msec window with 10 msec shift. We use 128 dimension logmel energy features and use as acoustic features after stacking 4 of them and subsampling by a factor of 3, resulting in 30msec features. During training we perform specaugment [10] on the acoustic features.

Tabel 1 shows WER and benchmarking results for different label encoder architectures. From the results we see that there is not much difference in WER when the context is reduce for the label encoder. Similar results have been reported before in [9]. For the case of limited context of 3 graphemes, we see a huge improvement in speed of decoder over the case when label encoder has 40 labels. This is because of the label encoder output caching as explained in section 3.3. For label context 2 it is even faster since the model itself is just a lookup table.

We present results on two Y-architecture models:

In  Figure3, we benchmark the time taken to encode 100 seconds audio with different modes on a single TPU [11] core with respect to the number of time steps the encoders run at a time. We can see that the inference speed is faster when we encode more time steps at a single inference run for all the modes because it allows better parallelization. We can also see that training mode is faster than query slicing mode, and batch step mode is the slowest among them because of extra overheads for handling the query block, and states. In high latency scenarios, where the encoder encodes 120 frames at a time, the encoding times are even faster. It takes 0.3 second, 0.6 second, and 1.8 seconds to encode 100 seconds audio with training mode, query slicing mode and batch step mode respectively.

In  Figure.4, we benchmark query slicing inference mode on desktop CPU for a 36 minutes long audio. We can see that this audio can be recognized in less than 3 minutes with 8 CPU cores (i.e. less than 8% real time factor). The recognition time can be further improved with a bigrapheme embedding lookup decoder with slight WER regression.