Exploring Effective Distillation of Self-Supervised Speech Models for Automatic Speech Recognition

Recent years have witnessed great strides in self-supervised learning (SSL) on the speech processing. The SSL model is normally pre-trained on a great variety of unlabelled data and a large model size is preferred to increase the modeling capacity. However, this might limit its potential applications due to the expensive computation and memory costs introduced by the oversize model. Miniaturization for SSL models has become an important research direction of practical value. To this end, we explore the effective distillation of HuBERT-based SSL models for automatic speech recognition (ASR). First, in order to establish a strong baseline, a comprehensive study on different student model structures is conducted. On top of this, as a supplement to the regression loss widely adopted in previous works, a discriminative loss is introduced for HuBERT to enhance the distillation performance, especially in low-resource scenarios. In addition, we design a simple and effective algorithm to distill the front-end input from waveform to Fbank feature, resulting in 17% parameter reduction and doubling inference speed, at marginal performance degradation.


INTRODUCTION
Recently, self-supervised learning (SSL) has made great progress in various areas including speech processing. The training of speechbased SSL is similar to that of text-based SSL, such as BERT [1], by reconstructing or predicting itself in an unsupervised way. Through effective SSL algorithms [2], the SSL model is able to encode rich hidden acoustic and semantic information implicitly from oceans of unlabeled data. A series of SSL algorithms have been developed and presented promising performance in a range of speech-related tasks [3,4,5,6]. However, the large memory consumption and expensive computation load derived from the oversize SSL model might hinder their wide applications and deployments.
Inspired by compression works [7,8] on BERT model in NLP domain, there are several previous studies on the distillation of SSL models in speech domain [9,10,11,12], which attempts to reduce the model size for a well-trained SSL model in an unsupervised fashion. Most of these existing works are investigated on the SUPERB benchmark [13], a generic testing framework for pre-trained models on a range of downstream tasks. The SSL models are evalu- † equal contribution, ‡ corresponding authors ated in the constrained track, where the whole upstream model is frozen, and the hidden layer outputs from the Transformer encoder are weighted for downstream tasks. The constrained configuration might not be able to fully reflect the potential of the distilled model for specific tasks. In addition, most SSL models adopt waveform as input, and several CNN layers are applied for feature extraction. This waveform-based front-end is retained after distillation in the student model, which is cumbersome in terms of efficiency and memory consumption. As a result, the existence of a waveform front-end might introduce potential issues for the deployment of this distilled model.
In this work, several aspects are investigated to facilitate the distillation of SSL models for speech recognition. The major contributions of this paper can be summarized as follows: • A comprehensive study of SSL model distillation is conducted to compare performances between different student model structures (e.g. deep & thin, shallow & wide) given the same amount of student model parameters, under unconstrained conditions.
• A discriminative loss is introduced to the distillation of Hu-BERT model in an unsupervised fashion, as a supplement and replacement of the layer-by-layer regression loss widely used in previous work. Significant performance improvements can be achieved by applying the proposed loss in low-resource scenarios.
• An effective distillation on front-end input is proposed and explored, which allows the student model to discard the troublesome waveform front-end and adopt a more practical Fbank frontend, with significant parameter reduction and inference speedup, at minor performance degradation.

HuBERT
HuBERT [4] is one of the most popular speech-based SSL models, attracting increasing research interest due to its simplicity and superior performance. During the training of HuBERT, the K-Means algorithm is applied to obtain the pseudo label for each frame, on top of the MFCC features or hidden layer outputs of a well-trained HuBERT model. The raw waveform is taken as input and fed into several CNN layers for feature extraction. After that, the feature is partially masked successively and encoded by Transformer layers, to predict K-Means pseudo labels c. The probability distribution for one masked frame in time t can be computed as: where ec is the embedding of the K-Means clustering label c, ht represents the last hidden state of the Transformer encoder at time t, τ is the temperature to scale the logits, A is the projection matrix to generate the embedding of a hidden state, and cos(·, ·) denotes the cosine similarity. The objective function of HuBERT shown below aims to minimize the cross entropy loss over pseudo labels of the masked frames: where M denotes the set of masked frames.

Distillation of SSL Models
Knowledge distillation [14] provides an effective solution to compress the model size [15], which is also known as teacher-student learning [16] in literature. There are various ways for the student model to learn useful information from the teacher model. One typical approach is to minimize the Kullbac-Leibler divergence (KL-divergence) between the probability distributions of student and teacher models. Another alternative loss is to apply the L1 or L2 loss between the hidden layer outputs of student and teacher models [17]. In most previous works, the knowledge distillation is operated in a supervised way like [18], where the ground truth label is given, and the student model learns from the teacher and ground truth simultaneously for better performance.
In the realm of speech-based self-supervised learning, the distillation is normally conducted on unsupervised data. There are several successful attempts of applying knowledge distillation for speech based SSL models. The DistilHuBERT [9] model with shallow and wide (S&W) structure is used to learn from the hidden layer outputs of the teacher model. [10] proposes an efficient but complex two-stage distillation method. In [11], the FitHuBERT model, with a deep and thin (D&T) model structure, is able to achieve better results on content-related downstream tasks compared to [9]. [12] presents a comprehensive study by comparing the performances between two different model architectures, S&W and D&T. It is worth noting that both DistilHuBERT and FitHuBERT are investigated under the contrained track on the SUPERB benchmark, which might not be able to reflect the potential effect of the distilled model as it does not fully explore the powerful modeling capacity lied in the Transformer encoder and merely treats it as a frozen feature extractor during the whole fine-tuning stage, missing the opportunity to pursue strong but non-linear features [19].

METHODS
Unlike most previous works [9,11,12], where the SSL model distillation is investigated on a range of downstream tasks in the constrained track on the SUPERB benchmark, in this paper, we mainly focus on the ASR task, and the distillation model is fine-tuned under the unconstrained condition, i.e., the SSL model can also be updated to get the better performance. Given the superior performance of HuBERT and its popularity in previous distillation works, HuBERT is adopted for distillation in this work.

Exploration of Student Model Structures
Motivated by the investigation and findings in [12], we first compare the performance of two different student model structures with similar number of parameters under the unconstrained track, to determine the optimal model architecture for distillation. Figure 1 illustrates the distillation process for these two student model structures.

Discriminative Loss for SSL Model Distillation
As described in Section 2.2, two types of information sources from the teacher model can be used for distillation, one is the probability distribution and the other is the hidden layer output. In the existing SSL model distillation algorithms, the latter information is widely adopted and the regression loss Lreg is applied to minimize the L1 or L2 distance of the hidden layer outputs between the student and teacher models. Inspired by the fact that pseudo labels are used as discrete targets in HuBERT and the cross entropy loss is applied for optimization, it is straightforward to utilize the probability distribution calculated in the output layer of HuBERT for knowledge distillation. In order to make a distinction with the regression loss Lreg used in DistilHuBERT and FitHuBERT [9,11], we name this distillation loss based on probability distribution as the discriminative loss L disc .
These two types of distillation loss, Lreg and L disc , can be combined to derive the final distillation loss with different weights, which can be denoted as below: where λreg and λ disc are set to tune the impact of regression and discriminative losses. For the distillation with the regression loss Lreg, we follow the previous works [9,11]. σ denotes the sigmoid function, T denotes the set of teacher model's layers to learn. L1 and cos similarity are adopted to close the gap between the hidden layer outputs of student and teacher models. For the discriminative loss L disc , KL-divergence is applied to minimize the probability distribution distance between the student and teacher models. The probability distributions over the pseudo labels of teacher pt and student ps can be computed as follows: where h last t and h last s are the last hidden layer outputs of the teacher and student models. The computation of pt is the same as teacher HuBERT model, where τ is the temperature used in the teacher, We is the embedding layer to produce the teacher's output embeddings, and l is the embeddings for K-Means pseudo labels. For the computation of ps, we adopt a linear output layer Wp followed by a softmax operation to compute the probability distribution for the student model for simplicity.

Distillation of Front-end Features
Due to the widespread influence of the Fairseq toolkit [20] in the speech based SSL community, most publicly released SSL models follow the front-end configuration used in Fairseq, where the waveform is taken as input and a CNN module is applied for feature extraction. Conventional, previous distillation works on SSL models retain the waveform-based front-end in the student model, and mainly focus on the compression of the Transformer encoder. However, this waveform-based front-end takes a fair amount of parameters and is computationally expensive, especially when the backbone Transformer model shrinks largely in the student model. It turns out to be the potential bottleneck for the optimization of model size and efficiency. In order to address this issue, in addition to distilling the backbone Transformer model, the distillation of the front-end is  explored. Instead of maintaining the waveform front-end, a more friendly Fbank front-end is adopted in the student model, which is distilled from the teacher model with the waveform front-end. We propose the following Fbank front-end distillation pipeline inspired by the front-end adapter adopted in the fine-tuning stage [21]: where N is the specified number of steps. For the first N steps, only the front-end distillation is applied. Afterwards, the distillation is conducted as described in Section 3.2. L f rontend can be the L1 or L2 loss to minimize the distance between the outputs of the teacher's waveform front-end and the student's Fbank front-end. The pipeline of L f rontend calculation is illustrated in    Table 1. For the student with waveform front-end, we keep the same front-end settings as the teacher model with seven CNN layers, while for the student with Fbank front-end, we select Fbank features with a frameshift of 10ms and downsample them to a frameshift of 20ms using only one CNN layer with a stride size of 2. Distillation The audios from the 960-hour LibriSpeech dataset [22] are used for knowledge distillation in an unsupervised fashion. The default number of update is 25k, and other settings are kept the same as DistilHuBERT [9]. For the Fbank front-end student, we first apply the L f rontend loss for 5k steps, then the L distill loss for 25k steps. For the student with waveform front-end, the S&W student is initialized with the teacher's front-end and the first two layers of the transformer encoder, while the D&T student is only initialized with the teacher's front-end and the Transformer encoder is randomly initialized due to mismatch of dimension in Transformer between the teacher and student models. The S&W student uses 3 prediction heads to respectively learn the teacher's 4th, 8th and 12th layers, while the D&T student performs a layer-to-layer distillation over all layers. We change λ disc and λreg in Equation 3 to evaluate the distillation model with different settings. Fine-tuning The fine-tuning experiments are implemented on Fairseq [20]. The 1-hour & 10-hour Libri-light [23] and 100-hour Librispeech splits are used as fine-tuning data, with fine-tuning steps of 25k, 25k and 80k, and learning rates of 2e-5, 5e-4 and 5e-4 respectively.

Results
The first experiment is to investigate the performance of different student model structures and distillation losses, and the results are shown in Table 2. According to the results, we notice that the D&T student model consistently outperforms the S&W student model with similar amounts of parameters. This validates the previous findings [11,12] reported on SUPERB under the constrained condition. Therefore, the D&T model is more advantageous than the S&W model for ASR and will be used as the baseline in the following experiments. It is also interesting to note that the improvement on the unconstrained condition (35.2% relatively) is much larger than that on the constrained condition (9.6% relatively) reported in [12] with 100-hour fine-tuning data and regression loss Lreg. By introducing the discriminative loss, where λ disc = 1, the ASR performance can be further improved compared to the baseline with regression loss, in most fine-tuning data and student model configurations, especially for the low-resource scenarios with 1-hour and 10-hour fine-tuning data. Even in the worse case, with the 100hour fine-tuning data and the D&T student model, the word error rate (WER) degradation caused by the additional discriminative loss is minor, which is from 7.50% to 7.88%.
The next experiment is to explore the distillation of front-end. Student models with the waveform front-end and Fbank front-end are compared. The D&T student model is adopted for distillation because of its superior performance, and the statistics of student models can be found in Table 1. Table 3 gives the WER results of these two types of front-end. It can be seen that the student model with Fbank front-end can yield a decent WER performance after distillation, although it is slightly worse than the student model with waveform front-end. The degradation can be partially explained by the 17% reduction of parameters in the student with Fbank frontend. It is also noteworthy that the introduction of discriminative loss consistently improves the ASR performance for student models with Fbank front-end in various amount of fine-tuning data. Table 4 presents the elapsed inference time comparison for the student models with waveform and Fbank front-ends. A decomposition of consumed time on front-end and Transformer encoder is also given. According to the results, the inference speed of the student model with Fbank front-end is significantly faster than that of student with waveform front-end, in both 1 and 4 CPU threads, and a total inference speedup of 2 can be achieved by adopting the Fbank front-end in the student model.

Details of Replacing Front-end
In our approach, the student's Fbank features need to adapt to the teacher's waveform features first when distilling. Without this process, the performance of the student will be deteriorated a lot. Table 5 presents our results. Adaptation to the front-end proves to be effective, probably because the features of the teacher's Transformers are adapted to the original waveform features. Using Fbank di-rectly is more likely to result in a mismatch between the student's hidden states and the teacher's during distillation, reducing the distillation performance. Zhu et al. [24] identified the ASR capability of a pre-trained model by analysing the CCA similarity [25,26], reasoning that the higher the CCA similarity, the less the model changed during fine-tuning, indicating a more capable model. Based on this, we calculate the CCA similarity for the distillation models before and after finetuning. As Figure 3 shows, models distilled with L disc have more similar representation to the fine-tuned model at low-resource cases, which implements that L disc is able to build better linguistic representational ability under low-resource scenarios for distilled models. Fig. 3: CCA similarity of hidden states between models before and after fine-tuning on low-resource Libri-light fine-tuning datasets (1h & 10h). D&T models are analyzed here, and we choose 100 pieces of audio in Librispeech test-clean dataset as the test data.

CONCLUSION
The distillation of speech-based SSL models is of practical value and is critical to facilitate their widespread applications. In this work, a series of explorations are conducted on the SSL model distillation for the ASR task. We validate that the deep and thin structure is more advantageous to the ASR task given limited model parameters, which is consistent with the findings in the constrained track on SUPERB benchmark. As a supplement to the regression loss used in previous work, a discriminative loss is introduced into the distillation of HuBERT in an unsupervised fashion, which significantly improves the performance on ASR tasks, especially in low-resource situations. Furthermore, we propose a front-end distillation framework by replacing the waveform with the Fbank feature as the input of the student model. This can reduce 17% of the parameters and halve the inference time at the cost of minor performance degradation.