DUAL LEARNING FOR LARGE VOCABULARY ON-DEVICE ASR

In order to perform ASR, TTS, and reconstruction in both domains, our implementation must include encoders and decoders for both audio and text. Imitating [26], we implement streaming with an architecture that can emit a provisional hypothesis immediately and then revise it after a short delay. In order to improve the likelihood of success on the two more difficult tasks (ASR and TTS), we adapt these components from existing ASR and TTS architectures. Our audio encoders and text decoder are adapted from conformer [27, 26]. Our text encoder and audio decoder are adapted from Tacotron 2 [28].

Formally, we frame our problem around three datasets. First, ${\cal S}$ consists of paired text and audio examples $(x,y)$, where $x=(x_{0},...,x_{m})$ is an audio sample of length $m$ and $y=(y_{0},...,y_{n})$ is the corresponding text transcript of length $n$. Second, ${\cal U_{T}}$ gives unpaired text examples $y$, and finally ${\cal U_{A}}$ gives unpaired audio examples $x$. We then define the components of the model as functions. We define the audio encoders $E_{A}^{\text{streaming}}$, which has only left-context and $E_{A}^{\text{delay}}$, which has 900ms of right-context. We also define the text encoder $E_{T}$, the decoders $D_{A}$ and $D_{T}$, and the linear transformations $T_{A\rightarrow T}$ which maps an audio embedding to a text embedding and $T_{A\leftarrow T}$ which maps a text embedding to an audio embedding.

We may then proceed to define the model’s objectives.

We might naively seek to train the above tasks together by alternating tasks across sequences of batches. We find that such a training scheme fails to achieve convergence, as each task is forgotten during the training of the others. Instead, we combine all the above tasks in a single batch. Each batch is split in thirds, the first coming from ${\cal S}$, the second from ${\cal U_{A}}$, and the last from ${\cal U_{T}}$. For the first third, we jointly optimize the supervised tasks:

$$\mathcal{L}_{S}=\frac{\mathcal{L}_{\text{ASR}}^{\text{streaming}}+\mathcal{L}_{\text{ASR}}^{\text{delay}}}{2}+\mathcal{L}_{TTS}+\mathcal{L}_{\text{Text Recon}}+\mathcal{L}_{\text{Audio Recon}}$$

for the second third, we optimize the unsupervised audio tasks:

$$\mathcal{L}_{A}=\mathcal{L}_{\text{U-TTS}}+\mathcal{L}_{\text{U-Audio Recon}}$$

for the last third, we optimize the unsupervised text tasks:

$$\mathcal{L}_{S}=\frac{\mathcal{L}_{\text{U-ASR}}^{\text{streaming}}+\mathcal{L}_{\text{U-ASR}}^{\text{delay}}}{2}+\mathcal{L}_{\text{U-Text Recon}}$$

We find that this method achieves convergence, so long as we initialize the model’s components from an ASR and TTS system trained on ${\cal S}$. Otherwise, the model generates incorrect pseudo-labels early in training, preventing progress.

Since dual learning involves the incorporation of unpaired text at training time, we naturally want to compare our method to the incorporation of unpaired text at inference time. To this end we evaluate our models with shallow fusion [29] with a pretrained LM. We also use a Hybrid Autoregressive Transducer (HAT) [30] text decoder, which permits the factorization of our models’ internal LM. Ultimately, at inference time for audio sample $x$ we seek:

$$y^{*}=\operatorname*{arg\,max}_{y}\log P(y|x)+\alpha P_{ELM}(y)-\beta P_{ILM}(y)$$

where $P(y|x)$ gives our acoustic model posterior, $P_{ELM}(y)$ gives the likelihood of a transcript in the external LM, $P_{ILM}(y)$ gives the likelihood of the transcript in the internal LM (as formulated in HAT), and $\alpha$ and $\beta$ are hyperparameters.

The ASR branch of our model is a cascading conformer adapted from [26], specifically sized to be realistic for an on-device streaming application. The streaming encoder is small to ensure fast inference. It consists first of 3 convolutional layers followed by 7 conformer layers with a 512-dimensional representation for a total of 56M parameters. The delayed encoder is larger, and is parameterized by 10 conformer layers with a 640-dimensional representation for a total of 99M parameters. Following [31], the HAT decoder consists of an embedding network and joint network, contributing another 9M parameters.

The TTS branch of our model is adapted from Tacotron 2 [28]. The text encoder first maps wordpieces into a 512-dimensional embedding space, followed by three convolutional layers and a single bidirectional LSTM layer, totalling 8M parameters. The audio decoder consumes the audio sample autoregressively through a \saypre-net, which consists of two fully-connected layers with 50% dropout at each layer. We find that this aggressive dropout is critical to convergence, since during multi-task training with teacher forcing the TTS decoder has a strong tendency to rely entirely on the autoregressive signal instead of the encoder representation. This yields poor performance at inference, which in turn creates poor pseudo-labels for unpaired text. Scheduled sampling [32] was investigated as an alternative, but was found to be less effective than simple dropout. The audio sample is then passed to two LSTM layers, which also consume the encoder representation via cross-attention. After the LSTM has generated an audio prediction, that result is further processed by a full-context \saypost-net consisting of five convolutional layers. In total, the TTS decoder has about 26M parameters.

We use 960 hours of supervised audio from Librispeech as ${\cal S}$, 60k hours of unsupervised audio from Librilight [33] as ${\cal U_{A}}$, and 80M transcripts from the Librispeech LM set as ${\cal U_{T}}$. We process the audio into a 128-dimensional log-mel feature per 10ms of audio. We stack every third such feature with the three features before it, yielding 512-dimensional features at 30ms intervals. We then apply SpecAugment [34] with mask parameter $F=27$ and ten time masks, as in [27]. This forms the inputs to the audio encoder.

We evaluate our models using a beam search with a beam size of 8. For fusion experiments, we use an external language model trained on ${\cal U_{T}}$. The LM is a causal transformer [35] with 8 layers, 16 attention heads, and a model dimension of 1024, totaling about 100M parameters.

We note with interest that while the application of shallow fusion preserves the gains yielded by our method, further subtracting out the internal language model via HAT only partially preserves those gains. That is, subtracting out the internal language model substantially closes the gap between the baseline and our method. Figure 2 illustrates this effect by plotting WER for a parameter sweep of external LM interpolation (shallow fusion) weights and internal LM interpolation (HAT) weights. Quantitatively, subtracting out the internal LM with a factor of $\beta=0.1$ from a model with shallow fusion improves BASELINE by $7.9\%/6.2\%$ while only improving E-ALL by $1.2\%/0.6\%$.

This result suggests that our method largely benefits the internal language representation of the ASR system’s decoder. This explanation is consistent with other methods of incorporating unsupervised text into a model at training time such as deep fusion [36] and cold fusion [37], but without adding any additional parameters to the ASR model at inference time. That is, unlike conventional language model fusion, our method bakes knowledge of that data into the parameters of the decoder, providing much the same effect as combining the ASR system with a pretrained language model with no modifications to the architecture.

This result also suggests that the improvements due to our method come mostly, but not entirely, from the unsupervised text data, as opposed to the unsupervised audio. This is consistent with our design; unsupervised text yields pseudo-labeled examples for the ASR task, while unsupervised audio yields pseudo-labeled examples for the TTS task.

Since unsupervised data is often used to address parts of the data distribution that are absent from the supervised training set, we seek to understand the effect of our method on \saytail words, which we define as words that are underrepresented in the training data relative to their frequency in the language as a whole. To this end, we draw inspiration from [38] and generate a test set of examples containing words that are not represented in ${\cal S}$ but are represented in ${\cal U_{T}}$. Specifically, for some parameter $0<\tau<1$ we sample transcripts from ${\cal U_{T}}$ containing at least one unigram that occurs with frequency less than $\tau$ in ${\cal S}$ but greater than $\tau$ in ${\cal U_{T}}$. We set $\tau=0.00001$ and generate a test set of 10k transcripts. We then synthesize audio transcripts using a Tacotron TTS system as in [28]. Results on this tail test set are given in Table 1, and measurements of internal and external LM integration are included in Figure 2.

As above, we find that our method yields improvements across the board. However, we note that the improvements are smaller than they are on Librispeech test sets. In particular, without an LM we improve by $5.0\%$, with shallow fusion we improve by $8.7\%$, and with internal language model subtraction we improve by $3.9\%$.