TutorNet: Towards Flexible Knowledge Distillation for End-to-End Speech Recognition

Generally, an end-to-end speech recognition model directly maps a sequence of input acoustic features $x_{1:T}=\{x_{1},...,x_{T}\}$ into a sequence of target labels $y_{1:N}=\{y_{1},...,y_{N}\}$ where $y_{n}\in\mathcal{Y}$ with $\mathcal{Y}$ being the set of labels in texts. $T$ and $N$ are respectively the total number of frames and the length of the target label sequence. To deal with the sequence-to-sequence mapping problem when the two sequences have unequal lengths, the CTC framework [1] introduces “blank” as an additional label and allows the repetition of all labels across frames. An alignment sequence $\pi_{1:T}=\{\pi_{1},...,\pi_{T}\}$ is a sequence of initial output labels for each frame, as every input frame $x_{t}$ is mapped to a certain label $\pi_{t}\in\mathcal{Y^{\prime}}$ where $\mathcal{Y^{\prime}}=\mathcal{Y}\cup\{blank\}$. A mapping function $\mathcal{B}$, which is defined as $y=\mathcal{B}(\pi)$, converts the alignment sequence $\pi$ into the final output sequence $y$ after merging consecutive repeated characters and removing blank labels. For example, two alignment sequences $\{\varepsilon,c,c,c,\varepsilon,a,\varepsilon,\varepsilon,t,t,\varepsilon\}$ and $\{c,c,\varepsilon,\varepsilon,a,a,\varepsilon,\varepsilon,\varepsilon,\varepsilon,t\}$ (using  ‘$\varepsilon$’ to denote blank label) correspond to the same sequence $\{c,a,t\}$ through the mapping function $\mathcal{B}$. The alignment between the input $x$ and the target output $y$ is not explicitly required in CTC training. The conditional probability of the target label sequence $y$ given the input sequence $x$ is defined as

where $\mathcal{B}^{-1}$ denotes the inverse mapping and returns all possible alignment sequences compatible with $y$. The conditional probability of a path ${\pi}=\{\pi_{1},...,\pi_{T}\}$ can be calculated as follows:

Given the target $y$ and the input $x$, the loss function $\mathcal{L}_{CTC}$ is defined as

where $Z$ represents a training dataset.

An alternative approach to the end-to-end mapping between speech and label sequences is to use the attention-based model. Unlike the CTC framework, AED directly predicts each target without requiring any alignment sequence. The framework consists of two sub-modules $Encoder$ and $AttentionDecoder$, so that it can learn two different lengths of sequences based on the cross-entropy criterion. The model predicts the posterior probability of the output transcription given the input speech features as follows:

where $x$ and $h$ are sequences of input speech features and encoded vectors respectively, and $y$ is a sequence of output text units whose length is $U$. $Encoder$ extracts encoded vectors $h$ from the input speech features $x$ in (4). Then $AttentionDecoder$ network predicts the next output symbol conditioned on the full sequence of previous predictions and acoustics, which can be defined as $p(y_{u}|x,y_{1:u-1})$. The attention mechanism determines which encoded vectors should be attended in order to predict the next output symbol.

Neural network models usually perform well with a large number of parameters. However, as a model architecture gets deeper, it requires heavy computation for both training and testing. To mitigate this computational burden, there has been a long line of research on KD, which aims at distilling knowledge from a big teacher model to a small student model. With this additional transfer procedure, the student can perform better compared to naive training. Existing KD methods typically fall into two categories: (1) transferring class probability and (2) transferring the representation of the hidden layer.

Generally, the output layer for the classification tasks uses softmax as an activation function. The output of the teacher model is a probability distribution over the target classes, and the sum of the outputs equals 1. Hinton et al. [9] first introduced KD, which distills class probability by minimizing the KL-divergence between the softmax outputs of the teacher and student. Compared to the one-hot label, the teacher’s softmax prediction has a nonzero probability value for each target class. This soft label is normally considered more informative than the one-hot encoded ground truth, further improving the student model in KD. The KD technique mentioned above only considers the output of the teacher model. In the case of transferring the hidden representation, some KD methods [10, 11, 12, 13, 14], especially in image processing, proposed transferring the representation-level information of the hidden layers where the mean squared error (MSE) between the representation-level knowledge of both models is minimized.

For speech, Li et al. [15] first attempted to apply the teacher-student learning. In previous studies for speech recognition, KD typically has been applied to DNN-HMM hybrid systems by minimizing the frame-level cross-entropy loss between the output distributions of the teacher and student [16, 17, 18, 19]. Geras et al. [20] proposed KD method that distills the softmax-level knowledge from RNN-based model to CNN-based model in the DNN-HMM framework. Wong et al. [21, 22] investigated sequence-level knowledge distillation to DNN-HMM trained by sequence discriminative criteria. In the case of KD under the cross-entropy criteria, the KD loss can be calculated as

where $p_{tea}(y|x)$ represents the posterior probability of the target label $y$ given the input $x$ yielded by the teacher model, and $p_{stu}(y|x)$ is that of the student model.

The same frame-level KD method has also been applied to KD of CTC models [23]. The frame-level KD in CTC framework can be computed as follows:

where $p_{tea}(k|x_{t})$ and $p_{stu}(k|x_{t})$ denote the posterior probability of the teacher and student CTC models, respectively. However, as reported in previous studies [23, 24, 25], applying the frame-level KD approach to the CTC-based speech recognition system can worsen the word error rate (WER) performance compared with the CTC model, which is trained only with the ground truth.

To address this problem, Takashima et al. [24] proposed a KD method that distills a sequence-level knowledge in the CTC framework using the teacher model’s N-best hypotheses. Kurata and Audhkhasi [26, 27] also introduced a KD approach for long short-term memory (LSTM)-CTC-based speech recognition model where a student can be trained using the frame-wise alignment of the teacher.

As mentioned above, the most frequently employed KD approach for speech recognition is to train a student with the teacher’s softmax prediction as a target, besides the one-hot encoded ground truth. However, the hidden representations of the teacher model are also considered important to provide essential knowledge for training the other. Moreover, if we can transfer hidden representations regardless of the types of neural network architecture, more flexible and effective distillation will be possible.

Before we describe SKD for distilling frame-level posterior in the CTC framework, it might be beneficial to review some characteristics of the CTC model. The output of the CTC model has two notable characteristics to consider when transferring a posterior distribution to the other. First, as noted in prior studies [1, 28], the softmax prediction obtained from a CTC-trained model is very spiky. This indicates that the softmax output tends to be similar to the one-hot vector. Secondly, since CTC is an alignment-free framework, CTC models trained with the same training data can have different frame-level alignments. In other words, the frame-level alignment yielded by the student can be different from that of the teacher. Unfortunately, this characteristic makes KD difficult because the KL-divergence and cross-entropy can hardly converge. Suppose we have two probability distributions $p_{tea}$ and $p_{stu}$ obtained at a certain frame $t$. As shown in Fig. 4, $p_{tea}$ has the highest probability on label ‘a’ and $p_{stu}$ on ‘blank’. In this case, the KL-divergence becomes infinity and KD is hard to converge.

In previous studies [23, 24, 25], it has been confirmed that applying the conventional frame-level KD approach can worsen the performance of a student CTC model compared with a model trained without KD. Also, we tried to train a student CTC model with the interpolation between the original CTC loss in (3) and frame-level KD in (7), but it failed to converge as reported in [24].

To deal with this instability problem, we propose to use $l_{2}$ loss instead of the KL-divergence. Since $l_{2}$ distance between two distributions is always bounded unlike the KL-divergence, it improves the numerical stability of the distillation loss. This alternative approach, namely SKD, allows student to stably learn the alignment of all the output labels, including blank ones. The SKD loss $\mathcal{L}_{SKD}$ to train the student model is given as

where $f_{tea}$ and $f_{stu}$ are the logits obtained from the teacher and student models respectively, and $\tau$ represents a temperature parameter. During SKD training, the SKD loss $L_{SKD}$ and the standard CTC loss $L_{CTC}$ are combined as an integrated loss. Thus, the final objective function for SKD can be formulated as

where $\lambda_{SKD}$ is a tunable parameter.

Our approach includes two stages of training: RKD as the initialization step and SKD in conjunction with the CTC objective as the fine-tuning step. Firstly, by minimizing $L_{RKD}$ without the CTC loss function, the student can mimic the teacher’s behavior at the representation-level. To effectively transfer the representation-level knowledge of the teacher, we apply the frame weighting mask $M_{FW}$ as in (9). After the RKD initialization, the student model is trained with SKD. As the student is trained to minimize the combined loss $L_{CTC-SKD}$, it can frame-wisely track the softmax prediction of the teacher.

We evaluated the performance of the proposed method on two speech datasets: LibriSpeech [29] and AISHELL-2 [30]. LibriSpeech is a large-scale (about 1000 hours) English speech corpus derived from audiobooks, sampled at 16kHz. The dataset is divided into clean and other. In the experiment, “train-clean-100”, “train-clean-360”, and “train-other-500” were used in the training phase. For evaluation,  “dev-clean”, “dev-other”, “test-clean”, and  “test-other” were applied. We also conducted our experiments on AISHELL-2, the Mandarin read-speech dataset (around 1000 hours of training data). The model evaluated over AISHELL2-2018A-EVAL data, containing dev set (2500 utterances from 5 speakers) and test set (5000 utterances from 10 speakers).

For LibriSpeech, we measured two metrics: word error rate (WER) and relative error rate reduction (RERR). WER is commonly employed to quantify speech recognition performance. To calculate WER, the number of errors is obtained by counting the substitutions, insertions, and deletions that occur in the recognition result. Then, it is divided by the total number of words in the correct sentence. RERR shows how much the WER is reduced, in proportion, compared to the baseline. For the Mandarin dataset, we measured the character error rate (CER) instead of WER. This is because a single character often represents a word for the Mandarin writing system. CER follows the same formula of WER, but with characters as the unit.

For KD, we adopted the following different speech recognition models which were used as the teacher or student model:

• •

  CTC model

  • –

    Jasper Dense Residual (Jasper DR) [6]: Jasper DR is a deep time-delay neural network (TDNN) composed of blocks of 1D convolutional layers. This differs from the original Jasper in that it has dense residual connections. The output of the convolution block is fed as an input to all the blocks via dense residual connections. In our experiments, we applied pre-trained Jasper DR, which consists of 54 convolutional layers, as the CNN-based teacher model.

  • –

    DeepSpeech2 [4]: DeepSpeech2 has the architecture of a deep RNN with a combination of convolutional and fully connected layers. As for the RNN-based model, we applied DeepSpeech2, consisting of two 2D convolutional layers followed by three 512-dimensional bidirectional LSTM (BLSTM) layers and one fully connected layer.

  • –

    Jasper Mini: Jasper Mini is composed of blocks of depthwise separable 1D convolutional layers. Depthwise separable convolution reduces the number of parameters and the computation required in the convolutional operations. The main structural difference from the Jasper DR is that Jasper Mini consists of depthwise separable convolutions with no dense residual connection. As for the CNN-based student model, we adopted Jasper Mini with 33 depthwise separable convolutional layers.

  • –

    QuartzNet [7]: QuartzNet is composed of multiple blocks with residual connections between them. Each block consists of one or more modules with 1D time-channel separable convolutional layers, batch normalization, and ReLU. We adopted Quartznet, which consists of 54 1D time-channel separable convolutional layers.

• •

  Hybrid CTC/attention model [31]

  • –

    Attention-based encoder-decoder (AED): Motivated by the prior studies [32, 33], the encoder contained the initial layers of the VGG net architecture [34] followed by four 1024-dimensional BLSTM layers and a linear projection layer, which yields better performance than the pyramid BLSTM [35] in many cases. We used a location-based attention mechanism, and the decoder network was a 2-layer LSTM with 1024 cells. The model was trained by both CTC and attention model objectives simultaneously.

Table I compares WERs obtained from the three baseline models when greedy decoding was applied and the number of parameters. Since the DeepSpeech2 baseline with BPE units failed to converge with a single BPE CTC loss, we adopted the CTC model, initialized with RKD over the first 100 steps, as the student baseline.

Our experiments were conducted using three toolkits: OpenSeq2Seq [36], ESPnet [37], and NeMo [38]. When training the CTC-based model, we mainly used OpenSeq2Seq. In the case of hybrid CTC/attention model, such as AED and Transformer, we applied the ESPnet toolkit. Pre-trained Jasper DR for the Mandarin dataset was provided by the NeMo toolkit.

• •

  LibriSpeech

  • –

    CTC model: We extracted 64-dimensional log-Mel filterbank features as the input, and the character set has a total of 29 labels with all lowercased Latin alphabet letters (a-z), apostrophe, space, and blank label. In the case of Jasper DR, we used the checkpoint provided by the OpenSeq2Seq toolkit. For character-level CTC-based student models, all the training excluding QuartzNet was performed on three Titan V GPUs, each with 12GB of memory. While training the RNN-based DeepSpeech2, Adam algorithm [39] was employed as an optimizer with an initial learning rate of 0.001 that was reduced with polynomial decay. In the case of the CNN-based Jasper Mini, we used NovoGrad optimizer [40] based on stochastic gradient descent (SGD). The initial learning rate of the optimizer started from 0.02, and it also had the same decaying policy as above. For RKD, we selected the last layer¹¹1CTC is regarded as a frame-level classification. Based on the general use of the last layer in KD of classification task, we chose the last layer for RKD training. of the teacher and student models to transfer the hidden representation, and the student was trained for 5 epochs. When transferring the knowledge from CNN-based Jasper DR to RNN-based DeepSpeech2, the kernel size of $c_{\theta}$, $D_{t}$, and $D_{s}$ were 1, 1024, and 512, respectively. In distilling from RNN-based DeepSpeech2 to CNN-based Jasper Mini, the kernel size of $c_{\theta}$, $D_{t}$, and $D_{s}$ were 1, 512, and 1024, respectively. After the RKD initialization, 50 epochs were spent for CTC training with SKD. In the case of QuartzNet, we trained the student model with 10 epochs for RKD initialization and 100 epochs for SKD training. The training of Quartznet was performed on three Titan RTX GPUs, each with 24GB of memory. In the case of QuartzNet, $D_{t}$, and $D_{s}$ were equal to 1024. We experimentally set the tunable parameter $\lambda_{SKD}$ to 0.25, which showed the best performance in the dev-clean set. When applying beam-search decoding with language model (LM), we used KenLM [41] for 4-gram LM, where the LM weight, the word insertion weight, and the beam width were experimentally set to 2.0, 1.5, and 256, respectively. When training the CTC model distilled from Transformer, since we used the byte-pair encoding (BPE) [42] method to construct the subword units (around 5000 tokens) for Transformer, the CTC-based student model was also trained with the same BPE units of the teacher model. The training was performed on three Titan RTX GPUs, each with 24GB of memory. In distilling the knowledge from Transformer to CTC-based model, 5 epochs and 50 epochs were spent for RKD and SKD, respectively.

  • –

    Hybrid CTC/attention model: When training AED, we used one Titan V GPU. Adadelta [43] was employed as an optimizer with an initial learning rate of 1.0. During training and inference, we set the CTC weight to 0.5. For RKD, the student was trained for 1 epoch, and then it was trained with SKD for 10 epochs. In the case of Transformer, we used the pre-trained model provided by the ESPnet toolkit. We used BPE to construct the subword units for the label of both AED and Transformer (around 5000 tokens). For the external LM, we used pre-trained RNNLM from ESPnet toolkit. The beam size was experimentally set to 20. When transferring the knowledge from Transformer to CTC-based model, $D_{t}$, and $D_{s}$ were equal to 512. In KD from Transformer to AED model, $D_{t}$, and $D_{s}$ were 512 and 1024, respectively.

• •

  AISHELL-2: For the Mandarin dataset, the character set has a total of 5207 labels. In the case of the teacher model, we used pre-trained Jasper DR, which was provided by the NeMo toolkit. When training the student model, we used three Titan V GPUs. For the CNN-based student model, the NovoGrad optimizer was adopted. In the case of the RNN-based student model, the Adam algorithm was employed as an optimizer. For RKD and SKD, 2 epochs and 20 epochs were spent, respectively. When training with SKD, we set $\lambda_{SKD}$ to 1.

We applied the following conventional KD techniques for performance comparison²²2We tried to train a student CTC model with the interpolation between the original CTC loss in (3) and frame-level KD in (7), and it failed to converge as reported in [24]. Therefore, the frame-level KD was not considered for comparison.:

• •

  CTC baseline [44]: The student model is trained solely based on the ground truth as the target, i.e., $\lambda_{SKD}=0$ in (11).

• •

  Sequence-level knowledge distillation [24]: N-best hypotheses of the teacher are used as a sequence-level knowledge for KD. In this experiment, the 5-best³³3The Sequence-level KD criteria can be summarized as the weighted mean of the original CTC loss regarding each hypothesis of the label sequence. The posterior probabilities of the hypotheses estimated by the teacher CTC model are used as the weights of each CTC loss. In our experiments, most probabilities were concentrated in the highest-scoring hypothesis, which was at the top of the beam. Even when we extracted more sentences than 5-best, the posterior probabilities were mostly concentrated on very few sentences. That is why we used 5-best hypotheses for the comparison. hypotheses were extracted using KenLM, where the LM weight, the word insertion weight, and the beam width were experimentally set to 2.0, 1.5, and 256, respectively.

• •

  Guided CTC training [27]: From the posterior distribution of the teacher, guided CTC training makes a mask that sets 1 only at the output label of the highest posterior at each frame. The student can be guided to align with the frame-level alignment of the teacher by using the guided mask.

For the convenience of notation and ease of comparison, we let CNN_DR, RNN_DS, and CNN_Mini denote Jasper DR, DeepSpeech2, and Jasper Mini, respectively. In the subsequent part of this paper, $A$ $\rightarrow$ $B$ means that model $A$ transfers knowledge to model $B$ where the former is a teacher and the latter is a student. To verify the effectiveness of the proposed method in various situations, we tried six different transfer scenarios on LibriSpeech dataset:

• •

  CNN_DR $\rightarrow$ RNN_DS: From CNN-based model (Jasper DR) to RNN-based model (DeepSpeech2)

• •

  CNN_DR $\rightarrow$ CNN_Mini: From CNN-based model (Jasper DR) to CNN-based model (Jasper Mini)

• •

  RNN_DS $\rightarrow$ CNN_Mini: From RNN-based model (DeepSpeech2) to CNN-based model (Jasper Mini)

• •

  RNN_DS $\rightarrow$ RNN_DS: From RNN-based model (DeepSpeech2) to RNN-based model (DeepSpeech2)

• •

  (CNN_DR & RNN_DS) $\rightarrow$ CNN_Mini: From two teachers (Jasper DR & DeepSpeech2) to CNN-based model (Jasper Mini)

• •

  (CNN_DR & RNN_DS) $\rightarrow$ QuartzNet: From two teachers (Jasper DR & DeepSpeech2) to CNN-based model (QuartzNet)

In order to show the versatility of the proposed method, we also conducted our experiment on AISHELL-2:

• •

  CNN_DR $\rightarrow$ RNN_DS: From CNN-based model (Jasper DR) to RNN-based model (DeepSpeech2)

• •

  CNN_DR $\rightarrow$ CNN_Mini: From CNN-based model (Jasper DR) to CNN-based model (Jasper Mini)

We experimented with TutorNet in both CNN_DR $\rightarrow$ RNN_DS and CNN_DR $\rightarrow$ CNN_Mini scenarios with AISHELL-2 dataset. For the teacher model, we used pre-trained Jasper DR (CER % on iOS dev dataset), which was provided by NeMo toolkit. When training the CNN-based and RNN-based student models, we utilized OpenSeq2Seq toolkit. Table VIII summarizes the results on AISHELL-2. The results show that TutorNet still works well with the Mandarin dataset. When distilling knowledge to RNN_DS, the student provided 12.30 % in dev iOS dataset. In the case of CNN_DR $\rightarrow$ CNN_Mini, TutorNet achieved WER 10.65 % and RERR 9.52 % in dev iOS dataset.

In the previous experiments, we mainly focus on KD between CTC-based models, especially across different neural networks. For speech recognition, besides the CTC-based model, there have been various types of end-to-end models such as Transformer and AED model. If we can transfer the knowledge regardless of the types of models, a more flexible KD will be possible. To check the applicability of the proposed method to the other types of models, we conducted two different scenarios with LibriSpeech dataset:

• •

  Transformer $\rightarrow$ CTC-based model: From Transformer to CTC-based model (DeepSpeech2)

• •

  Transformer $\rightarrow$ AED model: From Transformer to AED model (VGG-BLSTM)