Wav2vec-Switch: Contrastive Learning from Original-noisy Speech Pairs for Robust Speech Recognition

We recapitulate the wav2vec 2.0 model on which our method is based. Compared to its precedents [8, 14], wav2vec 2.0 combines masked prediction and contranstive learning into a unified model during pre-training. It has a feature encoder $f:\mathcal{X}\mapsto\mathcal{Z}$ with a raw audio waveform $\mathbf{x}\in\mathbb{R}^{T}$ as its input, and a latent representation $Z=[\mathbf{z}_{1},\ldots,\mathbf{z}_{T^{\prime}}]$ with a time-domain down-sampling through a set of convolutional blocks as its output; a context network $g:\mathcal{Z}\mapsto\mathcal{C}$ that takes the masked $Z$ and outputs a contextualized representation $\mathbf{c}_{t}$ at each masked position $t$ through several blocks of Transformers [26]; and a quantization module $h:\mathcal{Z}\mapsto\mathcal{Q}$ discretizing the unmasked $Z$ to $Q$ from a finite set of codebook via Gumbel Softmax [27] and product quantization [28]. The contrastive loss is applied at each masked position $t$, discriminating the true quantized representation $\mathbf{q}_{t}$ (the positive sample) from $K$ distractors $Q_{t}^{-}=\{\mathbf{q}^{-}_{1},\ldots,\mathbf{q}^{-}_{K}\}$ (the negative samples) drawn from other masked positions within the same training example:

	$\displaystyle\mathcal{L}^{\mathrm{C}}(C,Q)$	$\displaystyle=$	$\displaystyle\sum_{t=1}^{N}\mathcal{L}_{t}^{\mathrm{C}}(C,Q)/N$		(1)
	$\displaystyle\mathcal{L}_{t}^{\mathrm{C}}(C,Q)$	$\displaystyle=$	$\displaystyle-\log\frac{\exp(\mathrm{sim}(\mathbf{c}_{t},\mathbf{q}_{t}))}{\sum_{\mathbf{q}^{-}\in Q_{t}^{-}}\exp(\mathrm{sim}(\mathbf{c}_{t},\mathbf{q}^{-}))}$		(2)

where $N$ is the number of masked positions where the loss is computed, and $\mathrm{sim}(\cdot,\cdot)$ is implemented as the cosine similarity function. In addition, a diversity loss $\mathcal{L}^{\mathrm{D}}$ encouraging the codebook utilization is implemented as negative perplexity of the Gumbel Softmax output. The total loss to optimize is:

$$\mathcal{L}=\mathcal{L}^{\mathrm{C}}+\alpha\mathcal{L}^{\mathrm{D}}$$

(3)

where $\alpha$ is a coefficient. During fine-tuning, the quantization module is discarded and the parameters inside the feature encoder are frozen. All the other network parameters are updated with the CTC loss [29].

In the pre-training stage of wav2vec 2.0, the decision that distractors are only sampled from masked positions within the same training example, rather than from any examples, is critical, or it will hurt the ASR performance [15]. The reason is that sampling only within the same training example can avoid learning features irrelevant to ASR, e.g., speaker or environmental characteristics. However, no mechanisms are designed to achieve noise robustness during the pre-training: when a noisy utterance is given, both the positive and negative samples that the contrastive loss is using contain noise, and there is no explicit way to differentiate the noise from the speech or to learn a contextualized representation invariant to noise. Our intuition is that, if the contextualized representation is robust to noise, the representation of an original/noisy speech should be capable of predicting the target of the noisy/original speech as well. Motivated by this, we propose wav2vec-Switch as follows.

For each mini-batch of the original waveform $X\in\mathbb{R}^{B\times T}$ where $B$ is the batch size, we duplicate $X$, apply an independently sampled noise to each row (example), and thus have a noisy version of $X$ as $\widetilde{X}$. Then both $X$ and $\widetilde{X}$ are fed into the wav2vec 2.0 network in parallel and forwarded through the feature encoder $f$, the context network $g$, and the quantization module $h$. At this point we have 4 quantities²²2Here we slightly abuse the use of the notation $C,Q,\widetilde{C},\widetilde{Q}$ to denote the batched versions of their respective quantities described in Section 2.1.:

	$\displaystyle C=g(f(X)),\qquad Q=h(f(X))$		(4)
	$\displaystyle\widetilde{C}=g(f(\widetilde{X})),\qquad\widetilde{Q}=h(f(\widetilde{X}))$

In addition to the standard contrastive loss described in Eq. (1) where the loss takes $(C,Q)$ or $(\widetilde{C},\widetilde{Q})$ as its input arguments, we also switch the quantized targets $Q$ and $\widetilde{Q}$, and form two more tuples for the loss: $(C,\widetilde{Q})$ and $(\widetilde{C},Q)$. Therefore we obtain 4 contrastive loss quantities: $\mathcal{L}^{\mathrm{C}}(C,Q)$, $\mathcal{L}^{\mathrm{C}}(\widetilde{C},\widetilde{Q})$, $\mathcal{L}^{\mathrm{C}}(C,\widetilde{Q})$, and $\mathcal{L}^{\mathrm{C}}(\widetilde{C},Q)$. The new loss would be:

	$\displaystyle\mathcal{L}_{\mathrm{switch}}^{\mathrm{C}}(C,Q,\widetilde{C},\widetilde{Q})$	$\displaystyle=$	$\displaystyle\mathcal{L}^{\mathrm{C}}(C,Q)+\mathcal{L}^{\mathrm{C}}(\widetilde{C},\widetilde{Q})+$
			$\displaystyle\lambda\left(\mathcal{L}^{\mathrm{C}}(C,\widetilde{Q})+\mathcal{L}^{\mathrm{C}}(\widetilde{C},Q)\right)$

where the coefficient $\lambda$ controls the weight of the term calculated from the switched targets, relative to the one obtained from the original targets. Fig. 1 also illustrates how the 4 contrastive loss quantities are calculated. Finally $\mathcal{L}^{\mathrm{D}}$ is added to the total loss the same way as in Eq. (3). Note that when $\lambda=0$, the method reduces to the “wav2vec 2.0 with data augmentation” case, which will be compared against as the baselines in Section 3.

We also need to keep in mind that the network’s internal states for any specific input pair $(X,\widetilde{X})$ should be identical. Identical means that not only the architecture and parameters but also anything inside the network that relies on random states, including masked positions for the context network and dropout masks of all the dropout layers, should be the same for $X$ and $\widetilde{X}$. Otherwise the representations of the original speech and its noisy version will not follow with each other, and our approach will not behave as what we expect or learn representations with a meaningful interpretation. The ablation study in Section 3.7 also verifies the importance of enforcing such a constraint.

In practice, we batch $X$ and $\widetilde{X}$ together, and feed the resulting “large” mini-batch to the network. Every time before a random function is invoked, we save the current random state. That random state is restored immediately after the random function is invoked for the first half of the mini-batch, and then the same function is executed for the second half. By doing in this way, we are able to ensure that the network’s internal states for $X$ and $\widetilde{X}$ are always identical.

In order to validate the noise robustness of our proposed approach for ASR, we evaluate it on both synthesized and real-world noisy data. In this section, we first introduce how our data was prepared for these experiments. For the SSL paradigm, there are normally 3 stages: 1) self-supervised pre-training with a large unlabeled dataset; 2) fine-tuning with a relatively small set of labeled data; and 3) testing data in the target domain. We will explain the corresponding setup for each of these 3 stages.

In the synthesized data experiments, both the pre-training and test sets were formed by mixing the LibriSpeech corpus [30] with noise randomly drawn from the MUSAN corpus [31], and the signal-to-noise ratio (SNR) was uniformly sampled from the range between 5 and 10 dB. Note that the MUSAN corpus contains 3 categories of noise: music, noise and speech. To avoid introducing potential confusion with actual speech when learning speech representations, we only used the first 2 categories of the noise from MUSAN and added them on-the-fly to the original 960-hour LibriSpeech training set for pre-training. For the test sets, we prepared the original test-clean and test-other, and several synthesized versions with different levels of SNR or different subsets of the MUSAN corpus being added, in order to evaluate the model on test sets with a variety of mismatched conditions. We used the original LibriSpeech train-clean-100 labeled data for fine-tuning, and the reason is that we would like the final model not to degrade its performance on the original sets.

In the real noisy data experiments, we used the data from the CHiME-4 challenge³³3http://spandh.dcs.shef.ac.uk/chime_challenge/CHiME4/index.html [19]. The data in CHiME-4 is based on the “WSJ0” portion of the Wall Streat Journal corpus [32] (WSJ), either obtained by recording the actual talker’s speech in real noisy environments (on a bus, cafe, pedestrian area, and street junction) (denoted as “real data”), or artificially mixing the clean WSJ speech with background noise and recording them using a 6-channel distant microphone array and a close-talk microphone (denoted as “simulated data”). Specifically, we chose the real 1-channel track for testing. All the channels but channel 2 from both the “real data” and “simulated data” were used as independent utterances for pre-training and fine-tuning. Channel 2 was excluded because it is the one corresponding to the “behind” microphone which is of low recording quality. Since the CHiME-4 corpus is too small (less than 100 hours) to pre-train from scratch, we continued the pre-training with the CHiME-4 corpus on top of a model already pre-trained with LibriSpeech. The detailed comparisons between with and without continual pre-training are available in Table 3.

Our implementation was made upon the official release of wav2vec 2.0 from fairseq [33], and the network architecture used throughout this paper was identical to the LibriSpeech Base configuration specified in [15]: 12 transformer blocks with hidden dimension 768 and 8 heads. Most of the training settings and strategies for the Base model were also carried over into ours, except that the batch sizes were doubled to accommodate the pairs of original-noisy examples. The coefficient $\lambda$ in Eq. (2.2) was set to 0.3 empirically for all the wav2vec-Switch experiments. We also found that a smaller learning rate (e.g., 1/5 of the one used in the pre-training with LibriSpeech) for the continual training led to better ASR performance.

All the models were trained with 32 NVIDIA Tesla V100 GPUs, each with 32GB memory (required for the double-sized batches). We picked the best checkpoint (in terms of the validation loss), rather than the last one, for continual pre-training or fine-tuning after that, so that we do not need to worry about the risk of over-training.

We followed the base_100h setup in the wav2vec 2.0 code except the specific data being used for the CHiME-4 experiments. All the fine-tuning models were trained with 2 GPUs. We chose the best checkpoint according to the validation WER for final evaluations.

As we employed the CTC loss for fine-tuning, incorporating a language model during decoding is crucial for the best performance. For the LibriSpeech experiments, we downloaded an existing Transformer word-level language model⁴⁴4https://dl.fbaipublicfiles.com/wav2letter/sota/2019/lm/lm_librispeech_word_transformer.pt and adopted the same decoding configurations as those corresponding to the Base 100-hour experiment introduced in the wav2vec 2.0 paper. For the CHiME-4 experiments, an LSTM-based word-level language model with a vocabulary size of 65,000 was trained on the text portion of the WSJ corpus (see [34] for model details), and then the same decoding strategy was performed. The LM weight was tuned separately for each model.

We first show the WER (%) results on both the original and synthesized noisy sets under the matched condition in Table 1. Besides wav2vec-Switch (the third row), we also include the results from the wav2vec 2.0 model pre-trained on the original 960-hour LibriSpeech corpus (the first row), and the model trained on the synthesized noisy data (the second row, a.k.a. the data augmentation baseline). It is not surprising that without “seeing” the noisy data in training, the performance on the noisy test sets is much worse than that on the original ones: WERs increase by 2.5–3.2 times. After adding noise to the training data, the performance on the noisy data gets greatly improved, while not changing significantly on the original sets. When further replacing the wav2vec 2.0 model with wav2vec-Switch, the performance on the noisy sets is improved relatively by 7.1–11.0% without LM, or 2.9–4.9% with LM. In addition, the WERs on the original sets are even slightly better than those in the baseline in the “no LM” case, and remain almost the same if the LM is used. These results are promising as it is shown clearly that the model trained with wav2vec-Switch achieves noise robustness on the synthesized noisy data without hurting the performance on the original data.

Next we present the results under mismatched conditions. Mismatched conditions refers to the cases where the noisy conditions for testing are different from those for training. Given that the noisy condition for training was music+noise (5–10 dB), we created 3 versions of noisy test sets with mismatched conditions: 1) music+noise (0–5 dB) was the version where the SNR range in the test sets has no overlap with that in the training set; 2) speech (5–10 dB) was the one where the noise type being added to the test sets is different; and 3) speech (0–5 dB) was when both the SNR range and noise type are different. Table 2 demonstrates the results without LM along with the gains obtained by using the proposed wav2vec-Switch instead of wav2vec 2.0. We can see that the improvements hold in all the mismatched conditions. However, the relative gain becomes smaller when the degree to which the test noise condition mismatch with the training noise condition gets larger.

We now evaluate wav2vec-Switch on the CHiME-4 corpus. The 1-channel track real noisy data was used for model validation and evaluation. We also include the best results reported in the challenge [35] and other more recent ones [36, 21]. All of them adopted a supervised training paradigm and had a dedicated speech enhancement module to preprocess the noisy speech input before an acoustic modeling module, and some of them even leveraged model ensemble. As presented in Table 3, the best results are from wav2vec-Switch with the continual pre-training. The relative improvement from the corresponding baseline is 7.8% without LM (16.5 vs. 17.9), or 5.7% with LM (6.6 vs. 7.0). It is also worth noting that, without any speech enhancement, our self-supervised approach followed by a simple CTC fine-tuning achieves better results than those using carefully designed enhancement algorithms. The additional data we were using was just the unlabeled 960-hour LibriSpeech audio and the MUSAN corpus.

To investigate the impact of not keeping the dropout masks or the masked positions the same within each original-noisy speech pair, we conducted one experiment on LibriSpeech where the dropout masks within the context network were different between the examples within each pair, and another experiment where the masked positions were different⁵⁵5We still keep the number of masked positions the same to ensure dimension match after switching the targets.. We report the results without LM in Table 4. For easy comparisons, we also copy 2 relevant rows from Table 1. It shows that despite a clear degradation from the one with identical dropout masks, the one with non-identical dropout masks is still better than the baseline on all the test sets except the original test-clean. However, having different masked positions resulted in a significant deterioration in WER (15.8–27.6% increase), which is expected since we cannot force a representation to still have consistent predictions when they are actually making predictions for different positions before and after the switch, and consequently it will end up with a sub-optimal solution if doing so. These two experiments attest the importance of maintaining identical dropout masks and masked positions for input pairs in wav2vec-Switch.