SOA: Reducing domain mismatch in SSL Pipeline by Speech Only Adaptation for low resource ASR

Let $x_{src}$ be unlabeled data available from the source domain, $x_{src}^{\prime},y_{src}$ be paired speech-text (labeled) data from the source domain and $x_{tar}$ be unlabeled data available from the specified target domain.

We model the requirements for the pretraining procedure as possessing a large quantity of speech audio data, $x$, typically on the order of hundreds of hours of data. For the purpose of finetuning, we require paired speech-text data, $\{x^{\prime},y\}$ on the order of tens of hours of data. The pretraining process is modeled as learning a function $f:x\to z$ for learning latent space representations from the raw audio waveforms, and finetuning as learning $f^{\prime}:x\to y$, by iterating on $f$ to obtain the mapping from the audio to the vocabulary for the given dataset. By pretraining and finetuning on the source domain in this manner, we obtain the functions $g$ and $g^{\prime}$ respectively.

The issue with a domain shift is that the different distribution of the target domain data from the source domain implies models trained on the source domain do not generalize enough to learn the distribution of the target domain, i.e., $g^{\prime}$ does not act as a good approximation of $h^{\prime}$ for mapping between $x_{tar}^{\prime}$ to $y_{tar}$. The goal of domain adaptation then is to ’adapt’ $g^{\prime}$ to closer approximate $h^{\prime}$.

The simplest way to perform this would be to continually train the model on the available $x_{tar}$, thus adapting $g$ to $h$. Even after continual training, the model obtained is a mapping between $x_{tar}$ and $z_{tar}$, and thus without training on a CTC loss cannot be used for ASR. Thus, this requires a way to combine the latter layers of a model finetuned on a CTC loss, which is provided by SOA, as through the combination of encoder model layers, we approximate the iterating process to obtain $h^{\prime}$.

In this work, we utilize the base Wav2vec 2.0 model [24] as the speech foundation model for its relative computational simplicity. The model consists of a convolutional feature encoder $\theta:x\to z$ to map raw audio x to latent representations $z_{1},\dots,z_{N}$. These representations are input to a transformer based contextual encoder $\phi:z\to c$ to output context representations $c_{1},\dots,c_{N}$. During pretraining, latent representations are discretized to $q_{1},\dots,q_{N}$ using a vector quantization module. The model is trained to identify the true quantized latent $q_{t}$ using $c_{t}$ for each masked time step within a set of distractors sampled from other masked time steps using a contrastive loss [24]. During finetuning, we shift to using a Connectionist Temporal Classification (CTC) loss [25] objective, and task the model with predicting the characters of the paired text data using the representations $c_{t}$.

Let the speech foundation model (Wav2vec 2.0) be represented by $M_{1}:\{\theta_{1},\phi_{1}\}$, where $\theta_{1}$ represents the parameters present in the feature encoder and $\phi_{1}$ represents the parameters of the contextual encoder. Finetuning of $M_{1}$ using paired data $\{x_{src}^{\prime},y_{src}\}$, while keeping the weights in the feature encoder frozen, results in model $M_{2}:\{\theta_{1},\phi_{2}\}$, where $M_{2}$ is now a model optimized for performance on the source domain. This model $M_{2}$ can be reused for multiple target domains.

SOA is a two-stage (continual pretraining, and combination) training paradigm as shown in Figure 1. SOA can be described as:

Stage 1: Continual pretraining of $M_{1}$ using the Wav2vec 2.0 loss with data $x_{src}\cup x_{tar}$ while keeping the parameters of $\phi_{1}$ frozen. This results in model $M_{3}:\{\theta_{2},\phi_{1}\}$.

Stage 2: Combining the contextual encoder $\phi_{2}$ from $M_{2}$ with the feature encoder $\theta_{2}$ from $M_{3}$, resulting in $M_{4}:\{\theta_{2},\phi_{2}\}$.

Note that by freezing the parameters of $\theta_{1}$ during finetuning on the source data, and $\phi_{1}$ during continual pretraining, catastrophic forgetting is effectively prevented.

The advantage in the usage of SOA is twofold: 1) As only the contrastive loss during pretraining is employed for the target domain adaptation, we do not require any paired speech-text data $\{x_{tar}^{\prime},y_{tar}\}$, which is difficult to procure for low resource tasks. 2) Models $M_{1}$ and $M_{2}$ pretrained and finetuned on the source domain can be reused for subsequent adaptation to several target domains, thus reducing the computational cost.

Our experiments utilized a diverse array of datasets, each categorized as follows. The foundational base models were initially pretrained using the LibriSpeech corpus [26], consisting of 960 hours of read adult speech. We refer to the model optimized for clean, adult speech as $M_{2}$, achieved through finetuning on the 100-hour clean section of this dataset. In pursuit of enhancing performance for various low resource tasks, we utilized the following datasets for the continuous pretraining of $M_{3}$.

For the base foundation model $M_{1}$, we use the open-sourced base Wav2vec 2.0 model (95M parameters) from the fairseq toolkit [31]. During different stages of the framework, we obtain different models with varying numbers of updated parameters. Note that $M_{2}$ (with updated encoder layers) has close to 92M updated parameters, while $M_{3}$ (with only updated convolution layer weights) has close to 5M updated weights, thus leading to an efficient finetuning process for domain adaptation, as $M_{2}$ does not need to be retrained for every new domain task.

For the finetuning of $M_{2}$, we update using a noam scheduler [32] with warm up steps of 8k, and a peak learning rate of 3e-5. The peak learning rate holds for the next 32k steps, then exponentially decays to the ratio $\lambda$ of the initial learning rate, where $\lambda$ is set to 0.05.

We perform continual pretraining of $M_{3}$ using the Adam optimizer. During the first 8% of all training steps, the learning rate increases to 3e-5 and then decays polynomially. We experiment with different target domains $x_{tar}$, as well as varying the size of source $x_{src}$ and target $x_{tar}$ domain corpus used for SOA.

Training for all models is conducted on 2 Nvidia A4000 GPUs. The number of Floating Point Operations needed for training is computed by multiplying the training time of the model, the number of GPUs used during training, and an estimate of the single-precision floating-point capacity of the GPU (19.17 TFLOPS).

We also evaluate the effectiveness of SOA on adaptation for noisy speech in Table 5 through the variation in the number of training hours of the noisy (target) data used for SOA. Here, we refer to Baseline as model $M_{2}$ finetuned on only LibriSpeech (source) data without any SOA, and note that the performance for the baseline method is consistent with previously published results [34]. SOA was performed for all the listed models for a total of 50k updates. We note that by utilizing 100 hours of the target domain, we are able to achieve a 28.9% relative Word Error Rate reduction on the noisy test set. We also report the results from decoding using 4-gram LibriSpeech LM for the baseline and best performing SOA model.

To explore how the SOA model shifts the representations of the features vectors $z$, we perform the following experiment. We first feed both the baseline and SOA models a 1-second signal $x_{1}=\sin(2\pi f_{l}t)$ with $f_{l}$ ranging from 10Hz to 8kHz at 10Hz intervals, as in [35], to obtain the output representation $z_{1}$. We then feed the models with a speech signal $x_{2}$ from the MyST dataset, whose output representation is $z_{2}$. We proceed to compute the cosine similarity between $z_{1}$ and $z_{2}$ for different values of $f_{l}$ to demonstrate the ability of the feature encoder to ’capture’ the formant frequency.

Figure 2 demonstrates a plot of the cosine similarities of the feature space encodings as a function of frequency for the vowel in the word ’Good’ from the MyST dataset. We note that there is a spike in the values of the cosine similarities at the formant frequencies for the vowel ($F_{1}:568$Hz$,F_{2}:1559$Hz$,F_{3}:2944$Hz). We see that the SOA model shows sharper and more pronounced peaks, indicating that the SOA method leads to the feature encoder being more attuned to the formant frequencies of the target domain. While this plot is only for an individual audio sample, this trend holds up across a variety of utterances, indicating that this shift in the representation space could be the reason behind the improved performance of SOA.