Supervision-Guided Codebooks for Masked Prediction in Speech Pre-training

HuBERT is an SSL method benefiting from an offline clustering step to provide target labels for a BERT-like prediction loss [8]. Its backbone is a Transformer encoder [12]. During pre-training, the Transformer consumes masked acoustic features and the model is optimized to predict the discrete target labels on the masked regions. This forces the model to learn a combined acoustic and language model over the continuous inputs.

HuBERT adopts an iterative re-clustering and re-training process: in the first iteration, the targets are determined by clustering on MFCC features; in the second iteration, a new generation of training targets are created by clustering the latent representations from the model in the first iteration.

Our method shares the similar idea with self-training in the sense that they both improve ASR performance with bootstrapping [13, 14, 15, 16, 17, 18]. In self-training, a teacher model is trained on the available labeled data and then the unlabeled set is labeled with this initial model (a.k.a. pseudo-labeling). Finally, a new student model is trained on the combined labeled and pseudo-labeled data. This pseudo-labeling step can be repeated for multiple rounds to refine the labels. Instead of training a supervised student model to generate the whole transcription, our method predicts frame-level phoneme labels only on the masked region. It provides two advantages compared to self-training. First, our method shows excellent transferability in non-ASR tasks (presumably due to its self-supervised nature), enabling solving various speech tasks with one pre-trained model. In addition, the proposed method, like other SSL methods, is complementary to self-training methods for speech recognition (see the section below): substantial gain is observed in the speech recognition task.

Previous work [19, 20] shows a synergy between SSL and self-training by employing the pre-trained models in a self-training loop. Specifically, given the labeled dataset $S$ and the unlabeled dataset $U$, it first pre-trains a model on dataset $U$, fine-tunes it on dataset $S$. Then this fine-tuned model is used as the initial teacher model for pseudo-labeling. After that, we can either train a new supervised model or rerun the fine-tuning with the pre-trained model on the mixture of pseudo-labeled and truly labeled data. The combination of SSL and self-training has been shown very effective for extremely low-resource settings where only 10 minutes, 1 hour or 10 hours labeled speech are available. However, on the 100 hours setting, it is less effective and achieves comparable results as an SSL-only baseline [20]. This approach only considers SSL as a better starting point for the self-training. Our method is complementary in that it uses a supervised teacher model to guide pre-training, and the pre-trained model can further be included in the self-training loop.

We train a hybrid ASR model on the small amount of labeled data to provide accurate phoneme-level alignments for the unlabeled data. Traditional HMM-GMM models are capable of such a task by providing reliable estimate of the HMM state occupancy for each frame [9]. Using neural networks in place of GMM can usually further boost the performance due to their stronger modeling power. The neural networks can be trained either with a frame-level (e.g., cross-entropy) or a sequence-level (e.g., MMI [21] and MPE [22]) criterion. We adopt a special MMI criterion named Lattice-free MMI (LF-MMI) [23] where the denominator graph contains competing hypotheses derived from a phone LM, and the numerator graph encodes all possible sequences of HMM states pertaining to the transcript. In order to yield better alignments, especially for long utterances, the numerator graph is generated in a way so that it is acyclic, directly encoding alternative alignments information from a previous HMM-GMM training stage. The neural acoustic model trained in this way is advantageous over GMM models or neural models trained with other sequence-level criteria, in terms of ASR performance and alignment accuracy [24].

Once the hybrid model training is done, a regular decoding process is conducted on the unlabeled data with the trained model, and the state-level alignments are transduced to the phoneme-level alignments based on the predefined HMM topology and context-dependent phone tying rules¹¹1We use a phoneme-based system for illustration purposes. Actually grapheme units, like Chenones [25], are also applicable to our method.. The phoneme-level alignments are to be used as targets in PBERT pre-training.

Aside from the hybrid model approach, we also propose to train an end-to-end model and use that for K-means clustering. Specifically, given the small labeled dataset, we train a Transformer CTC model. Then we extract the hidden representations for all the speech in the unlabeled set and apply K-means clustering. The cluster assignments are used as the training targets of a HuBERT model. Compared with latent representations in an unsupervised pre-trained model (i.e., the first-stage HuBERT), the CTC latent representations have a higher correlation with meaningful phonemes. So K-means clustering on such representations may yield clusters more helpful for ASR even when the labeled speech data is much less than unlabeled speech. In addition, the training cost of a CTC model is much lower than that of an SSL model (see Section 4.5).

Our self-supervised model architecture design is the same as HuBERT. It contains a convolutional feature encoder with a downsampling rate 320 and a 12-layer Transformer encoder. Besides the original convolutional position embedding, we also add a bucket relative position bias following the implementation in [26] with 320 embeddings shared across all layers. Each embedding is a scalar that is added to the attention logits and corresponds to a range of possible key-query offsets. The range increases logarithmically up to an offset of 800, beyond which we assign all relative offsets to the same embedding. In our experiments we will show the contribution of the relative bias.

In this paper, we evaluate our method on the LibriSpeech benchmark [11]. It contains 960 hours of training data train-960 from read audiobooks with three subsets: train-clean-100, train-clean-360 and train-other-500. We consider train-clean-100 as the supervised data (leveraging its transcripts), and the rest as the unlabeled data (discarding the transcripts).

In CTC clustering, we train the 12-layer Transformer CTC model from scratch on train-clean-100, with 8 GPUs and a batch size of 200 seconds of audio per GPU. After training, we run K-means clustering with 500 clusters on the latent features extracted from the CTC model at the 7-th transformer layer, which has the highest mutual information with phoneme tokens according to the analysis in [7].

Our pre-training setting follows LibriSpeech Base in [7]. During training, each audio example is cropped to at most 15.6 seconds. The model is pre-trained on 32 GPUs with 87.5 seconds of audio per batch for 400k steps. We use AdamW optimizer [31] with weight decay 0.01 and $\beta=(0.9,0.98)$. The learning rate ramps up linearly for the first 32k steps and then decays linearly back to 0. The peak learning rates is 5e-4.

Table 1 presents the results compared against the self-supervised and self-training baselines. The hybrid model and the CTC model that we used to generate codebooks are listed as the supervised baselines.

Our proposed methods rely on supervised models to provide high quality training targets for pre-training and perform well on both test sets. Without an LM, PBERT obtains a relative 17.0/10.0% WER reduction over HuBERT. Relative positional embedding can further bring a 2.6% relative gain on test-other. The CTC clustering method is slightly worse than PBERT since the pre-training targets are obtained by K-means clustering and thus less interpretable than phonemes. Compared to other baselines, PBERT with pseudo phoneme labels is significantly better on test-clean, indicating the supervision-guided codebook with train-clean-100 could provide a better pre-training objective for in-domain (clean) data. Moreover, we train a PBERT on the force-aligned phoneme labels³³3The phoneme labels are obtained by force-aligning labeled train-960 with a hybrid model trained on train-960., which is denoted as “Ground-truth phones + rel bias” in the table and can be considered as the upper bound of our methods.

In addition, PBERT could perform iterative pre-training like HuBERT (denoted as “+ 2nd iteration” in Table 1). We use the pre-trained PBERT model as a phoneme predictor to generate new frame-level pseudo targets for all training data. Then we run another round of pre-training with the new phoneme targets. The new pre-trained model is also fine-tuned on train-clean-100. This method improves the WER on test-other from 11.2% to 10.7%, but has no gain on test-clean. This implies that the second iteration of pre-training can generate better phoneme labels for unlabeled noisy data, while for clean data label quality is not improved. One possible explanation is that the hybrid model, trained with in-domain data only, does not perform well for out-domain data label generation, and the first round pre-trained model improves the model robustness on the out-domain data (i.e., test-other).

Self-training is a strong baseline, which performs better than HuBERT on test-clean but worse on the noisy test-other. This is due to our labeled training data having the same domain as the clean set and thus the pseudo-labels generated by the teacher model for the clean set are of higher quality than the noisy set. Our proposed method is compatible with self-training. We can see that one round of self-training could further bring WER down to 4.2/9.5% without LM, and 3.1/7.2% with LM (better than the self-training-only baseline by 6.1/14.3% respectively). An interesting phenomenon is that the self-training only provides limited gain in the 4-gram LM fusion setting. Our explanation is the pseudo-labeled data are created by decoding with LM, so the LM knowledge has been distilled into the fine-tuned model. Therefore, the 4-gram LM cannot contribute a lot to the final results.

Existing works have demonstrated that SSL enjoys outstanding generalizability and re-usability across various downstream scenarios [35, 34]. An interesting question is whether our model maintains the advantage of SSL if the phoneme is the target in pre-training. To answer this question, we fine-tune our models on VoxCeleb2 [36], and evaluate its performance on the three official trial lists Vox1-O, Vox1-E and Vox1-H, following the pipeline proposed in [37].

Table 2 shows the results of different pre-trained encoders being combined with the downstream model ECAPA-TDNN. “ASR encoder” means we pre-train a CTC model with labeled train-960, and feed the ASR encoder outputs to the downstream model to obtain speaker embeddings. We can observe that the ASR encoder does not enjoy a good generalizability, while all the pre-trained models outperform ECAPA-TDNN + fbank feature significantly. PBERT is slightly worse than HuBERT on the speaker verification with less than $3\%$ degradation. One possible explanation is that the phoneme labels, which are obtained from an ASR model, might lose speaker dependent characteristics compared to HuBERT.

We compare the training costs of different codebook generation approaches. HuBERT and CTC clustering method are implemented with PyTorch and trained on Tesla V100 GPUs. They require 866.7 GPU hours and 53.3 GPU hours respectively. The training cost of CTC clustering is 93.8% less than HuBERT. The training of our hybrid system is with Kaldi on Tesla K10 GPUs, and only 37.9 GPU hours are required⁴⁴4The alignment model (HMM-GMM) training is still performed on CPUs. However, recent work [24] has shown that an HMM-DNN model can serve as a better alignment model that runs on GPUs..