Sub-band Knowledge Distillation Framework for Speech Enhancement

In recent years, single-channel speech enhancement methods based on deep learning have been proposed in large numbers, and they can be divided into time domain and frequency domain methods. Time-domain methods [1] [2] [3] [4] typically use the time-domain signal of noisy speech as a direct input to a neural network, and the learning target is the clean speech signal. Frequency-domain methods [5] [6] [7] [8] typically use the spectral features of noisy speech (magnitude spectrum, phase spectrum, complex spectrum, etc.) as inputs to a neural network, and learning targets are the spectral features of clean speech or some masks (e.g., ideal binary mask [9], ideal ratio mask [10], complex ideal ratio mask [11]). The frequency-domain methods still dominate the vast majority of current single-channel speech enhancement methods due to the high dimension and the lack of apparent geometric structure for the time domain signal.

In the frequency domain, most speech enhancement methods focus on the full-band spectral representation. They take the full-band feature sequence directly as the input of the neural network. Unlike them, this paper proposes a sub-band speech enhancement method. This method uses the magnitude feature sequence within a certain sub-band on the noisy spectrogram as the input of the model. The learning target is the magnitude feature sequence of the corresponding sub-band of the clean speech. This method has some advantages over the full-band methods. (1) For neural networks, the simpler the features to be learned, the easier the optimization of the network. Spectral features are highly patterned features, with low, medium, and high frequencies having significantly different data distributions. The reasonable division of sub-bands is equivalent to introducing a priori knowledge to the model so that the model can focus on learning more stable features on each sub-band. This method is similar to the partition technology in face recognition. The human face has a fixed pattern. By only hard segmenting the human face (eg., mouth, eye.), and then independently identifying each segment of the face and integrating them, the recognition accuracy can be improved [12]. (2) Most speech enhancement methods use an L-norm ($\ell_{1},\ell_{2}$) loss function, for which the contribution of each frequency band to the total loss is identical. This may not be suitable for the actual situation that the speech energy is mainly concentrated in the low and medium frequencies. Correspondingly, the low and medium frequencies play a more important role in the loss. However, high-frequency information is essential for perceptual evaluation. The proposed method optimizes and calculates the loss for each sub-band separately to avoid the problem of applying the L-norm loss to the full-band spectra. Besides, a proper division of sub-bands also allows the model not to lose the ability to capture cross-band information too much. Some of the previous works have been exploring non-full-band methods. In [13], a narrow-band LSTM, is proposed to deal with multichannel speech enhancement. The LSTM takes as input a sequence of TF bins associated with a frequency bin of multichannel noisy speech signals. The output of the model is a sequence of magnitude ratio masks at the same frequency bin. Experimental results show that the narrow-band method has outstanding performance. In [14], the spectrogram is divided into different non-uniform sub-bands, and spectral subtraction is performed separately in each sub-band to obtain accurate noise estimates. In [15], the sub-band feature is used to classify speech or noise, which indicates that the local spectral pattern is informative for discriminating between speech and other signals.

Different sub-bands have very different data distributions, and it is undoubtedly challenging to learn them simultaneously through a single small-scale model. The integration using multiple sub-models is an effective method, but it is undoubtedly troublesome during deployment. In this paper, we extend the sub-band enhancement model to a knowledge distillation framework containing multiple elite-level teacher models and one general-purpose student model for speech enhancement. Knowledge distillation was first proposed by [16]. As a compression [17] or training technique [18] [19], it has been widely used in machine translation [20], image object detection [21], speech recognition [22] [23] and other fields. Knowledge distillation usually uses the category probability labels of a large teacher model and original labels as the learning goal of a small student model. It makes the small student model have the learning ability close to the large teacher model. In the proposed sub-band knowledge distillation framework, we first trained multiple sub-band enhancement models (teacher models). These models specialize in feature mapping in a certain sub-band, and they can well conduct the speech enhancement in their sub-bands. After that, we use these teacher models to guide a sub-band enhanced model (student model) to learn all sub-bands. Finally, without increasing the number of parameters and the computational cost of the student model, the student model’s performance is further improved.

In order to evaluate the effectiveness and performance of the proposed method, we conducted several experiments and analyses on an open-source data set. The final experimental results show that the teacher model’s guidance dramatically improves the student model’s performance. Moreover, the performance of the student model exceeds the corresponding full-band model.

We use the representation of the speech signal in the short-time fourier transform (STFT) domain:

$$X(t,f)=S(t,f)+N(t,f)$$

(1)

where $X(t,f)$, $S(t,f)$ and $N(t,f)$ respectively represent the time-frequency (T-F) bin of noisy speech, clean speech and interference noise at time frame $t$ and frequency bin $f$. $t\in Z^{+},1\leq t\leq T$. $f\in Z^{+},0\leq f\leq F-1$. $T$ and $F$ denote the total number of frames and the total frequency bins expressed in the frequency domain, respectively. For the closed interval $[0,F-1]$, we take an ordered sequence of equally spaced points:

$\displaystyle 0=f_{0}<f_{1}<f_{2}<...<f_{n}=F-1$

(2)

We refer to the closed interval $[f_{i},f_{(i+1)}]$ as the sub-band, where $0\leq i\leq n-1$. $[0,F-1]$ contains $n$ sub-bands. The frequency width of all sub-bands is fixed at $\lfloor\frac{F}{n}\rfloor$ except for the last one. In this paper, we use the magnitude component of the speech in the STFT domain as the input and training target. We use a single neural model to map all sub-bands:

	$$\displaystyle|\hat{S}(i)|=G(|X(i)|)$$		(3)
	$$\displaystyle|X(i)|=\left(\begin{smallmatrix}|X(1,f_{i})|&|X(2,f_{i})|&\cdots&|X(T,f_{i})|\\ |X(1,f_{i}+1)|&|X(2,f_{i}+1)|&\cdots&|X(T,f_{i}+1)|\\ \vdots&\vdots&\ddots&\vdots\\ |X(1,f_{i+1})|&|X(2,f_{i+1})|&\cdots&|X(T,f_{i+1})|\end{smallmatrix}\right)$$

where $|X(i)|,0\leq i\leq n-1$ is the noisy magnitude features in the $i$-th sub-band, $G(\cdot)$ is the sub-band enhancement model, and $|\hat{S}(i)|$ is the enhanced feature in the $i$-th sub-band with the same dimension as $|X(i)|$. We treat the divided sub-bands as independent units, and the model can be optimized for each sub-band individually. In an extreme case, we can divide $F$ sub-bands.

Considering that the model simultaneously handles sub-bands with different data distributions is very challenging, inspired by [16], we introduce the sub-band enhancement model into a knowledge distillation-based framework. The basic principle of knowledge distillation is that a soft label generated by the large model contains more information than the original target label. In general, if a large model gives a higher probability of the input features for specific categories, it means that those categories are similar to the real ones. We typically use knowledge distillation to distill the knowledge contained in a large-scale model (teacher model) into a smaller-scale model (student model) to compress model or to enhance performance.

In this paper, our knowledge distillation framework contains multiple teacher models for local sub-bands and a student model for all sub-bands. We train these teacher models so that they have excellent performance in their specific sub-band. During the student model training, we will select different teacher models according to the input sub-band to guide the student model, improving the student model’s performance without changing the original student model scale.

In this paper, we propose a sub-band knowledge distillation framework for single-channel speech enhancement. We divide the spectrogram into multiple sub-bands, and train elite-level sub-band enhancement models (teacher models) on each sub-band. These models are then used to guide the training of general sub-band enhancement models (student model). We evaluated our proposed method on an open-source data set. We found that compared to the full-band model, although the sub-band enhancement model loses the ability to capture global cross-band information, the performance does not show a significant drop as expected. We also found that the teacher models further improves the performance of the sub-band enhancement model. Moreover, the performance of the student model exceeds the corresponding full-band model.