Wav2vec-S: Semi-Supervised Pre-Training for Low-Resource ASR

We denote the pre-training and fine-tuning dataset as the source domain $S$ and target domain $T$. Suppose the source domain consists of an unlabeled dataset $\mathbb{U}^{S}=\left\{\mathbf{x}_{1}^{S},\ldots,\mathbf{x}_{N}^{S}\right\}$ and a labeled dataset $\mathbb{L}^{S}=\left\{\left(\mathbf{x}_{1}^{S},\mathbf{y}_{1}^{S}\right),\ldots,\left(\mathbf{x}_{M}^{S},\mathbf{y}_{M}^{S}\right)\right\}$, where $M\leq N$. These two datasets are used for pre-training, where the self-supervised pre-training uses $\mathbb{U}^{S}$ and semi-supervised pre-training uses both $\mathbb{U}^{S}$ and $\mathbb{L}^{S}$. Fine-tuning is performed on an additional target domain labeled dataset $\mathbb{L}^{T}=\left\{\left(\mathbf{x}_{1}^{T},\mathbf{y}_{1}^{T}\right),\ldots,\left(\mathbf{x}_{O}^{T},\mathbf{y}_{O}^{T}\right)\right\}$

We adopt the model structure in wav2vec 2.0 [1], which is shown in the first step of Fig. 1. A convolutional feature encoder is used to map the input raw audio $\mathcal{X}$ to higher-level latent speech representations $\mathcal{Z}$, which is then fed to the transformer context network to build context representations $\mathcal{C}$. During pre-training or fine-tuning, the mask module masked a proportion of the feature encoder outputs $\mathcal{Z}$ to $\mathcal{Z}^{{}^{\prime}}$ before feeding them into the context network. Note that the masked dimension is only the time dimension during self-supervised pre-training, while it consists of both time and channel dimensions during semi-supervised pre-training and fine-tuning like SpecAugment [26]. The masked time steps are denoted as gray color in Fig. 1.

We illustrate the wav2vec-S procedure in Fig. 1. The pre-training consists of two steps. Firstly, self-supervised pre-training is performed on the unlabeled source dataset $\mathbb{U}^{S}$ using the self-supervised loss. Then, semi-supervised pre-training is applied on both labeled source dataset $\mathbb{L}^{S}$ and unlabeled source dataset $\mathbb{U}^{S}$. The total loss for semi-supervised learning is:

$$\mathcal{L}_{\text{semi}}=\mathcal{L}_{\text{label}}+\lambda\mathcal{L}_{\text{unlabel}},$$

(1)

where $\mathcal{L}_{\text{label}}$ and $\mathcal{L}_{\text{unlabel}}$ denotes the loss for labeled and unlabeled data, respectively. $\lambda$ is the hyperparameter to be tuned, which is fixed to 1 in this work.

Specifically, for ASR, during semi-supervised pre-training, both labeled and unlabeled losses are CTC [27]. For labeled data, it is straightforward to compute the CTC loss as:

$$\mathcal{L}_{\text{label}}=-\mathbb{E}_{\mathbf{x},\mathbf{y}\sim p(\mathbf{x},\mathbf{y})}\log p_{\mathbf{\theta}}(\mathbf{y}\mid a(\mathbf{x})),(\mathbf{x},\mathbf{y})\in\mathbb{L}^{S}$$

(2)

where $(\mathbf{x},\mathbf{y})$ is the sample-label pair, $p(\mathbf{x},\mathbf{y})$ is the distribution of samples from $\mathbb{L}^{S}$, $\mathbf{\theta}$ is the model parameter, and $a(\cdot)$ is the augmentation function.

However, for the unlabeled data, since the ground-truth labels are not available, pseudo labels are used instead:

$$\mathcal{L}_{\text{unlabel}}=-\mathbb{E}_{\mathbf{x}\sim p(\mathbf{x})}\log p_{\mathbf{\theta}}(\hat{\mathbf{y}}\mid a(\mathbf{x})),\mathbf{x}\in\mathbb{U}^{S}$$

(3)

where $\hat{\mathbf{y}}$ denotes the pseudo label which is generated through:

$$\hat{\mathbf{y}}=\underset{\mathbf{y}}{\operatorname{argmax}}\log p_{\mathbf{\theta}}(\mathbf{y}\mid\mathbf{x}),$$

(4)

where $\operatorname{argmax}$ denotes the greedy decoding, which first takes the maximum probability tokens in each frame and then removes repeated and blank tokens. Note that the pseudo labels are generated using the up-to-date model $\theta$ on-the-fly as in [15].

After pre-training, the task-specific labeled loss $\mathcal{L}_{\text{label}}$ is also used to fine-tune the pre-trained model on the target domain labeled dataset $\mathbb{L}^{T}$.

The pre-training (source) dataset is LibriSpeech [28], where the 100h clean subset is used as the labeled dataset $\mathbb{L}^{S}$ and the other 860h is the unlabeled dataset $\mathbb{U}^{S}$. As for fine-tuning (target) labeled datasets $\mathbb{L}^{T}$, Wall Street Journal (WSJ) and the US accented part of AESRC [29] is used as the in-domain dataset since they are all read datasets. To verify the generalization ability, conversational dataset SwitchBoard (SWBD) [30] and lecture dataset TED-LIUM3 (TED) [31] are used as cross-domain datasets. Moreover, Mandarin Chinese dataset AISHELL-1 [32] and French dataset from Common Voice (CV French) [33] are used as cross-lingual datasets. We concentrate on the low-resource scenario, thus randomly selecting 1h or 10h subset of the above datasets for fine-tuning.

All experiments are conducted with fairseq [34]. For self-supervised pre-training, we use the open-source wav2vec 2.0 or data2vec base model pre-trained on Librispeech 960h. For semi-supervised pre-training, we use gradient accumulation to achieve an effective batch size of 25.6m samples. The maximum learning rate is $3\times 10^{-5}$, and the tri-state learning rate schedule from [1] is used. And the convolutional feature encoder is fixed during training. For fine-tuning, batch size and learning rate are the same with semi-supervised pre-training. Apart from the convolutional feature encoder, the transformer context network is also fixed for the first 10k updates. The total training updates for 10h and 1h fine-tuning are 20k and 13k, respectively. Beam-search decoding with a dataset-specific 4-gram language model is used for evaluation.

As shown in Table 1, we perform 1h and 10h fine-tuning on different pre-trained models. All pre-trained models are trained on the Librispeech but with different amounts of labeled/unlabeled data. The wav2vec 2.0 model trained on unlabeled data and the supervised pre-trained model trained on labeled data from scratch are used for comparison.

Comparing the supervised pre-trained model and wav2vec 2.0 model, we find that for 1h fine-tuning, the supervised pre-trained model outperforms wav2vec 2.0 on 3 out of 6 datasets (WSJ, TED and AISHELL-1). However, when fine-tuning data increases to 10h, wav2vec 2.0 consistently outperforms the 960h supervised pre-trained model on all datasets. It illustrates the effectiveness and generalization of self-supervised pre-training. The proposed wav2vec-S model consistently outperforms the supervised and wav2vec 2.0 model on all datasets, demonstrating the effectiveness of the simple pipeline of wav2vec-S. Note that in Table 1, only 100h labeled data is used in wav2vec-S, although more labeled data could provide better results (shown in subsection 4.4). Moreover, pre-training and fine-tuning both used character-level supervision and training updates is 20k. We will further discuss the impact of supervision level and training updates in subsection 4.5 and subsection 4.6. The following experiments are conducted with 10h fine-tuning on three representative datasets (WSJ, SWBD, AISHELL-1).

We compare using different amounts of labeled and unlabeled data during semi-supervised pre-training. As shown in Table 2, the performance is the best when using all 960h labeled data and is the worst when using only 100h labeled data. Semi-supervised pre-training with 100h labeled and 860h unlabeled data effectively bridges the performance gap and achieves comparable performance with the 960h labeled one.

We discuss the optimal supervision level for semi-supervised pre-training. Specifically, We consider phone-level supervision and character-level supervision. The phoneme transcripts are generated using phonemizer¹¹1https://github.com/bootphon/phonemizer.

As shown in Table 3, when phone-level fine-tuning is used, phone-level and character-level pre-training perform similarly. On the other hand, when character-level fine-tuning is used, character-level pre-training clearly outperforms phone-level. It shows that the higher-level supervision (character) during pre-training can generalize well to the lower level (phone) but not vice versa. This conclusion also stands in the cross-lingual dataset (AISHELL-1), although the character supervision during pre-training and fine-tuning are in different languages. Therefore, we could conclude that character-level supervision is better for semi-supervised pre-training.

We illustrate the relation between the number of training updates and downstream performance in Table 4.

When training updates increase, the WERs on source language datasets decrease, including in-domain (validation, WSJ) and the cross-domain (SWBD) datasets. On the contrary, the cross-lingual (AISHELL-1) WER increases. This indicates the trade-off between performances of the source language and other languages: with more training updates, the wav2vec-S model becomes more language-specific and the cross-lingual generalization ability is thus degraded.

We conduct experiments using 8 V100 GPUs to show the training time for the two steps in wav2vec-S and fine-tuning in Fig. 2. The semi-supervised pre-training requires much less training time than the wav2vec 2.0 self-supervised training. The reason is that self-supervised pre-training can speed up the convergence of the followed semi-supervised pre-training. Moreover, since the self-supervised pre-training is task-agnostic, it can be reused by all downstream tasks. Therefore, for a new task, only semi-supervised pre-training is required to be conducted before fine-tuning.

We use another self-supervised pre-trained model data2vec to test the generalization of wav2vec-S on different self-supervised pre-trained models. As shown in Table 5, data2vec outperforms wav2vec 2.0 on all datasets, illustrating the effectiveness of data2vec. And the wav2vec-S approach still consistently improves the performance of data2vec with the additional semi-supervised pre-training step. Therefore, wav2vec-S could act as a universal refinement approach to enhance a given self-supervised pre-trained model.

We analyze the representation similarity before and after fine-tuning for different pre-trained models to show if a particular layer of the pre-trained model is suitable for the ASR task. Specifically, the fine-tuning is performed on the 10h subset. We follow the practice in [8] and compute the CCA similarity between representations from each layer of a pre-trained model and the same layer of the corresponding fine-tuned model, where the lower CCA similarity means the representation changes more significantly during fine-tuning.

As shown in Fig. 3, the last few layers of wav2vec 2.0 change significantly during fine-tuning, which illustrates the representations in its last few layers are less effective for ASR. This phenomenon is alleviated in data2vec on both in-domain and cross-domain datasets (WSJ and SWBD), which could be the reason why data2vec outperforms wav2vec 2.0. With the semi-supervised pre-training, both wav2vec 2.0 and data2vec have more similar representation to the fine-tuned model, illustrating semi-supervised pre-training effectively closes the representation gap between the task-agnostic self-supervised pre-trained model and task-specific fine-tuned model.

There are some different phenomenons on the cross-lingual dataset AISHELL-1. Firstly, although data2vec performs better than wav2vec 2.0, it has lower CCA similarities in all layers. It means the CCA similarity could not directly reflect the ASR performance on the cross-lingual dataset, which is a significantly out-of-distribution dataset for pre-trained models. Secondly, the semi-supervised pre-training can not consistently improve the similarity in the last few layers of wav2vec 2.0. The reason might be that the labels in semi-supervised pre-training and fine-tuning are in different languages, which calls for the exploration of multi-lingual semi-supervised pre-training [35].