SNR-based teachers-student technique for speech enhancement

We show the standard knowledge distillation in Figure 1. The idea behind knowledge distillation is that the soft probability output $p^{\prime}$ of a pre-trained teacher model contains more information than the correct class label $p^{\prime\prime}$. If the teacher model gives higher probabilities for specific categories, then this indicates that the categories of the input image should be near these categories. Knowledge distillation forces the student model to minimize the difference between its probability output $p$ and the teacher’s probability output $p^{\prime}$ to extract the additional knowledge that the teacher model obtained when calculating the correct probability. We usually use knowledge distillation to distill the information learned by large networks and pass it to networks with small parameters and weak learning capabilities. In this paper, we extend the standard knowledge distillation. The workflow are shown in Figure 2. First, we train multiple teacher models who are proficient in speech enhancement at a specific small SNR range. Then we use the teacher models to guide the training of the student model so that the student model can perform speech enhancement under both high SNR and low SNR. Below we describe the entire process.

Train the teacher models. We use $f_{T}=\{f_{t_{1}},f_{t_{2}},...\}$ to represent the set of teacher models. We use $f_{t_{i}}$ represents the $i$-th teacher model. The dataset of each teacher model covers only a small range of SNRs, and the range of SNRs covered by each teacher’s dataset does not overlap each other. After sufficient training, we get multiple teacher models that are good at dealing with the SNR covered by their dataset. We fixed the weights of the teacher models.

Forward calculation. We use $x$ to represent a noisy speech, and $y$ to represent the clean speech corresponding to the $x$. We input the $x$ to the SNR-based teacher models. According to the SNR of $x$, the training algorithm will select a teacher model $f_{t_{i}}$ to obtain the enhanced speech $f_{t_{i}}(x)$. We also input $x$ to the student model $f_{s}$ to get the enhanced speech $f_{s}(x)$.

Distill knowledge from the teacher model. We use a $L_{2}$ loss function to minimize the difference between the output of the student model and the teacher model. This loss can extract additional knowledge from the teacher model during the student training process, and strengthen the student model’s ability to enhance speech at a specific SNR. We still need to preserve the loss between the enhanced speech of the student model $f_{s}(x)$ and the clean speech $y$. The overall loss $L$ is as follows:

$\displaystyle L=\alpha\frac{1}{2}{||f_{s}(x)-f_{t_{i}}(x)||}^{2}+(1-\alpha)\frac{1}{2}{||f_{s}(x)-y||}^{2}$

(1)

where $\alpha$ is the weight of the knowledge obtained from the teachers, which is determined based on the validation dataset.

As the SNR of noisy speech decreases, the correct phase becomes more and more critical for speech enhancement. However, since the phase spectrum mapping is complicated, we conduct speech enhancement in the time domain. Inspired by Wave-U-Net [14], we designed a time-domain speech enhancement model based on one-dimensional (1D) convolutional neural networks. The teacher models and student model both take such the network structure. The input of the model is a fixed-length noisy speech signal, and the output is an enhanced speech signal. The model contains three parts: Encoder, Bottleneck, and Decoder, and its structure is shown in Figure 3. Below we describe it in detail.

The Encoder part consists of 12 consecutive encoder blocks. The first seven encoder blocks use the same structure: “1D convolutional layer + batch normalization + leaky ReLU + downsampling layer”. 1D convolutional layers can be used for feature extraction and integration. The parameters of the convolutional layers are listed above each encoder block. These seven encoder blocks all contain a downsampling layer, which can select one sampling point from two adjacent sampling points as output. After their processing, the size of the input signal in the time dimension will gradually become smaller. The structure of the remaining five encoder blocks in the Encoder is similar to that of the previous seven encoder blocks, but they do not contain a downsampling layer. For the time-domain model, since the first convolution layer of the model plays a critical role [17], we set the number of convolution kernels in the first convolution layer to a larger value (48), so that the model can better extract the features of the speech waveform. The number of convolution kernels in the remaining convolutional layers increases with the step size of 24 as the number of layers of the network increases. The structure of the Bottleneck section is the same as that of the Encoder.

The Decoder part contains 12 consecutive decoder blocks. The beginning of each decoder block contains the concatenation operation, which is used to concatenate the underlying features passed through the skip-connection. In order to maintain symmetry, the last seven decoder blocks in the decoder part all include the upsampling layer. The upsampling layer doubles the size of the feature map in the time dimension by linear interpolation, and finally restores the size corresponding to the input speech waveform. The remaining structure of the decoder blocks are similar to the encoder blocks.

We use public datasets to evaluate the proposed method.

For SNR-based teacher models: In this paper, we set up four teacher models. In order to train the teacher models for different SNR ranges, we randomly select 950 clean speeches from the TIMIT [18] training dataset, and mixed speech shaped noise, babble, destroy engine, destroyer ops and factory floor 1 (from NoiseX-92 dataset [19]) under different SNRs to generate the noisy speech dataset. The SNRs of the four teacher models are {-20, -17, -13, -11}dB, {-10, -7, -3, 1}dB, {0, 3, 7, 9}dB, and {10, 13, 17, 20}dB. In the end, we generated 19,000 noisy speeches for each teacher model, of which 18,000 speeches were used as the training dataset and the rest were used as the validation dataset.

For student model: We randomly selected 950 speeches from the TIMIT training dataset and mixed them with the five noises that appeared in the training dataset of the teacher models at {-20, -10, 0, 10, 20}dB. This produces 23,750 noisy speeches. Among them, 22,000 were used as the training dataset, and 1,750 were used the validation dataset. To evaluate the student model, we randomly selected 100 speeches from the TIMIT test dataset, mixing with five noises that appeared in the above two datasets and four noises from the CHiME-4 [20] dataset (pedestrian area noise, on the bus noise, cafe noise, and street noise) at {-20, -15, -10, -5, 0, 5, 10, 15, 20}dB. The resulting 8,100 noisy speeches act as the test dataset. It is worth mentioning that the test dataset of the student model contains noises, SNRs, and speakers that have not appeared in the training dataset, which makes speech enhancement very challenging.

We use PESQ [21] and STOI [22] to measure speech quality and intelligibility, respectively.

The sampling rate of all speeches is 16,000Hz. The input length of all models is fixed. Before the beginning of each training epoch, we will select consecutive 16,384 sampling points (1.024 seconds) from the random position of noisy speech as the input of the model. Except for the learning rate, other experimental parameters are the same for all models. We use Adam optimizer [23] (decay rates $\beta_{1}=0.9$, $\beta_{2}=0.999$) and set the batch size to 16. we set $\alpha$ of leaky ReLU to 0.1. We set the learning rate of the teacher models to a small constant of 0.0002. We set the initial learning rate of the student model to 0.002, which is reduced twice by every 300 epochs until the validation loss does not decrease. According to the validation dataset, we set $\alpha$ in Equation 1 to $0.5$.

We compared with the three methods, which are Wavenet-denoising [7], PSM-BiLSTM [24], and CRN [6]. The first one is to mapping features on the time domain, and the last two are mapping features on the frequency domain.

• •

  Wavenet-denoising: It is an end-to-end method for speech denoising based on wavenet [25]. It retains wavenet’s powerful acoustic modeling capabilities, while significantly reducing time complexity by eliminating its autoregressive nature. We used its official implementation.

• •

  PSM-BiLSTM: It is a bidirectional LSTM network for speech enhancement, using a phase-sensitive spectrum approximation (PSA) cost function. We implemented it and used the same hyperparameters as that in the paper.

• •

  CRN: It contains a convolutional encoder-decoder and a LSTM bottleneck. We implemented it and used the same hyperparameters as that in the original paper.

Table 1 presents the noisy speech and the enhanced speech from -20dB to 20dB in terms of PESQ and STOI. The “Noisy” lines in the table represent the noisy speech, and the “Enhanced” lines represent the enhanced speech using the proposed method. The “Seen” columns mean the SNRs exist in the training dataset of the student model while the “Unseen” columns indicate the SNRs that do not exist in the training dataset of the student model.

When comparing the noisy and enhanced speech, we can found that the PESQ and STOI are improved after speech enhancement with the proposed method even some noises do not appear in the training dataset (the last four noises). It is also clear that our method improves the PESQ and STOI for both “Seen” SNR and “Unseen” SNRs. The SNRs lower than -5dB are generally considered as extremely low SNRs, and the speech enhancement under such conditions is very challenging. However, our method works well, even if the range of the SNR of the dataset is so wide (-20dB to 20dB). The average improvement of the PESQ and STOI is 38.71% and 12.73%, respectively. This improvement indicates that the proposed method is effective.

To demonstrate the effectiveness of the SNR-based teacher models, Table 2 lists the performance of the four SNR-based teacher models (T1, T2, T3, T4), the student model trained without teacher supervision (S1), and the proposed method (S2). As shown in Table 2, we only test the teacher models on the trained small range of SNRs. For clarity, the PESQ and STOI of the noisy speech are also listed in the table.

Comparing with the noisy speech, there are notable improvements on the PESQ and STOI for the teacher models on their corresponding SNRs. This suggests that the teacher models can perform well on a small range of SNRs. The S1 and the S2 use the same training dataset and network structure. Compared to all the above models, the S1 performs the worst no matter which SNR. We guess that this is because the SNR range of the dataset is too large. With the learning ability of the S1, it is very difficult to consider all the SNRs from low to high.

The PESQ and STOI of the proposed method S2 approximate those of the teacher models at the corresponding SNR, and are far better than the S1 at all SNRs. These results indicate that the proposed SNR-based teachers-student technique can help the student model improve its ability to handle high SNR and low SNR at the same time.

We calculated the PESQ and STOI on the validation dataset for every ten epochs during the training process of the S1 and S2. The growth curves are shown in Figure 4. The horizontal axis indicates the epoch of training, and the vertical axis is the average metric. We can clearly notice that from about the 200th epoch, the PESQ and STOI of S2 exceed the PESQ and STOI of S1. In the subsequent training process, the performance gap between the two models is still increasing.

Table 3 reports the performances of the proposed method and the baseline methods on the test dataset (-20dB to 20dB). We can easily notice that the proposed method (S2) achieves the best performance, whether it is PESQ or STOI. Compared to the most robust method CRN in the baseline methods, we still improved 0.143 on PESQ and 0.021 on STOI. In the scenario that contains both high SNR and low SNR noisy speeches, our proposed method based on SNR-based teachers-student technique and time-domain U-Net is undoubtedly more advantageous.