A two-stage framework in cross-spectrum domain for real-time speech enhancement

Our FDFNet adopts the two-stage framework, and it generally consists of a STFT-based magnitude enhancement sub-network (dubbed FME-Net) and a STDCT-based spectrum refinement sub-network (dubbed DSR-Net). The overall architecture of FDFNet is illustrated in Fig. 1. The model input is the noisy speech $x$, which can be expressed as:

$$x=s+n$$

(1)

where $s$ and $n$ represent the clean speech and noise, respectively.

In the first stage, we adopt the typical magnitude-only convolutional recurrent network (CRN) [6] as FME-Net, which includes the convolutional encoder, recurrent neural network (RNN), and deconvolutional decoder. In detail, the convolutional encoder consists of several 2D convolutional (Conv2d) blocks, each of which is composed of a 2D convolutional layer followed by batch normalization and PReLU. For the deconvolutional decoder, it is the symmetric version of the encoder, where each 2D convolutional layer is replaced by the 2D deconvolutional layer. Between the encoder and decoder, several gated recurrent unit (GRU) layers are inserted to model the temporal correlations. In addition, skip connection is utilized to connect each encoder block to its corresponding decoder block for better performance. FME-Net receives the noisy speech’s magnitude spectrum $|X_{F}|$, which is derived by STFT, and estimates the enhanced magnitude spectrum $|\hat{X}^{1}_{F}|$. Then, we couple this enhanced magnitude $|\hat{X}^{1}_{F}|$ with the original noisy phase $\theta_{X_{F}}$ to obtain the enhanced STFT spectrum $\hat{X}^{1}_{F}$ of the first stage. This stage aims at suppressing the noise components coarsely.

After the FME-Net, we sequentially apply inverse STFT (ISTFT) and STDCT to $\hat{X}^{1}_{F}$, so as to obtain the pre-enhanced STDCT spectrum $\hat{X}^{1}_{D}$. Since STDCT is a real-valued TF transformation, the spectrum $\hat{X}^{1}_{D}$ contains both the magnitude and phase information. Therefore, the following DSR-Net which directly operates on the STDCT spectrum can simultaneously recover the target magnitude and phase.

In the second stage, DSR-Net takes both the pre-enhanced STDCT spectrum $\hat{X}^{1}_{D}$ and the original noisy STDCT spectrum $X_{D}$ as the input. Similar to previous works [5, 10], the output of the second stage is only used to make a further refinement to the first stage’s result. The details of our DSR-Net will be described in the following section. Overall, DSR-Net predicts a STDCT spectrum mask $\hat{M}_{D}$, which is subsequently multiplied with the pre-enhanced STDCT spectrum $\hat{X}^{1}_{D}$ to derive the final estimation $\hat{S}_{D}$ for the target STDCT spectrum.

In a nutshell, the whole procedure of our FDFNet can be formulated as:

$$X_{F}=|X_{F}|\cdot{e^{j\theta_{X_{F}}}}=\mathrm{STFT}(x)$$

(2)

$$|\hat{X}^{1}_{F}|=\mathcal{F}_{1}(|X_{F}|;\Phi_{1})$$

(3)

$$\hat{X}^{1}_{F}=|\hat{X}^{1}_{F}|e^{j\theta_{X_{F}}}$$

(4)

$$\hat{X}^{1}_{D}=\mathrm{STDCT}(\mathrm{ISTFT}(\hat{X}^{1}_{F}))$$

(5)

$$X_{D}=\mathrm{STDCT}(x)$$

(6)

$$\hat{M}_{D}=\mathcal{F}_{2}(X_{D},\hat{X}^{1}_{D};\Phi_{2})$$

(7)

$$\hat{s}=\mathrm{ISTDCT}(\hat{S}_{D})=\mathrm{ISTDCT}(\hat{M}_{D}\cdot{\hat{X}^{1}_{D}})$$

(8)

where $\mathcal{F}_{1}$ and $\mathcal{F}_{2}$ represent the functions of FME-Net and DSR-Net with parameter sets $\Phi_{1}$ and $\Phi_{2}$, respectively.

Considering that the CRN structure has been proven to be effective in the SE tasks [4, 6], we still follow this network topology when designing our DSR-Net. The details of our DSR-Net is illustrated in Fig. 2(a).

Since STDCT spectrum is the real-valued spectrum like the STFT magnitude spectrum, the structure of the encoder and decoder in our DSR-Net is exactly the same as that of FME-Net. While between the encoder and decoder, we design a time-frequency sequence modeling (TFSM) block, as shown in Fig. 2(b), to replace the ordinary RNN layer. Similar to the dual-path RNN architecture [12], the TFSM block aims to respectively model the sequential dependencies along the time and frequency dimensions. Specifically, the TFSM block firstly captures the local and global context features among different frequency bins at each frame. This is achieved by a bidirectional GRU (BiGRU) layer, followed by an add layer, layer normalization, and PReLU. The add layer after BiGRU is used to sum up the bidirectional outputs of BiGRU. Then, another GRU layer is applied to the previously processed result to model the temporal correlations. The GRU layer is also followed by layer normalization and PReLU. Residual connection is added between the original input and the sequence modeling result to yield the final output of TFSM block. Multiple TFSM blocks are stacked to ensure the performance.

The prediction target of our DSR-Net is the DCT ideal ratio mask (DCTIRM), which is defined as:

$$M_{D}=\frac{S_{D}}{\hat{X}^{1}_{D}}$$

(9)

where $S_{D}$ is the clean speech’s STDCT spectrum. Eventually, the enhanced speech $\hat{s}$ can be obtained according to Eq. (8).

Similar to the previous works [5, 10], We also adopt the two-stage training scheme to optimize our FDFNet. First, we train FME-Net with the mean square error (MSE) loss toward magnitude spectrum estimation, which can be expressed as:

$$\mathcal{L}_{\mathrm{FME}}=\left\||\hat{X}^{1}_{F}|-|S_{F}|\right\|_{F}^{2}$$

(10)

where $|S_{F}|$ is the magnitude spectrum of the clean speech $s$.

Subsequently, we freeze the optimized FME-Net and train the DSR-Net by a hybrid loss:

	$\displaystyle\mathcal{L}_{\mathrm{DSR}}$	$\displaystyle=\mathcal{L}_{\mathrm{T}}(\hat{s},s)+\mathcal{L}_{\mathrm{TF}}(\hat{M}_{D},M_{D})$		(11)
		$\displaystyle=\left\|\hat{s}-s\right\|_{1}+\left\|\hat{M}_{D}-M_{D}\right\|_{F}^{2}$

where $\mathcal{L}_{\mathrm{T}}(\hat{s},s)$ represents the L1 loss between the enhanced speech and clean speech, and $\mathcal{L}_{\mathrm{TF}}(\hat{M}_{D},M_{D})$ denotes the MSE loss toward DCTIRM estimation.

To evaluate the performance of our FDFNet, we conduct the experiments on the VoiceBank+DEMAND dataset [13]. It includes 11,572 utterances from 28 speakers for training and 824 utterances from another 2 unseen speakers for testing. All utterances are resampled to 16KHz.

During the experiments, a Hamming window with 32ms window length and 8ms hop size (75% overlap) is employed for both STFT and STDCT. Meanwhile, both the STFT and STDCT points are 512. The FME-Net contains five encoder blocks, three GRU layers, and five decoder blocks. The output channel of each convolutional layer in encoder is {16,32,64,128,256}, and correspondingly, the output channel of each deconvolutional layer in decoder is {128,64,32,16,1}. The kernel size and stride of all the (de)convolutional layers are set as (3,2) and (2,1), respectively. The hidden units of each GRU layer is {128,64,32}, and a fully-connected layer with 2304 units is after the last GRU layer. As for the DSR-Net, it also has five encoder blocks and five decoder blocks, and the setting of the output channels of the (de)convolutional layers is exactly the same as that of FME-Net. The convolutional stride is still (2,1), but the kernel size is changed to (5,2). DSR-Net includes three TFSM blocks, and the hidden GRU&BiGRU units of each TFSM block is {128,64,32}. In order to ensure the real-time SE, all the (de)convolutional layers in our FDFNet is implemented causally by applying the asymmetric zero-padding. Furthermore, the FME-Net and DSR-Net have the same training strategy. In the training phase, RMSprop optimizer with the initial learning rate of 2e-4 is used. The learning rate will be halved if the model performance does not improve for five consecutive epochs. The batch size and total training epochs are 16 and 80.

To demonstrate the effectiveness of our design, an ablation study is conducted as shown in Table 1. We quantitatively evaluate the performance of model by a set of commonly used metrics, including wide-band perceptual evaluation of speech quality (WB-PESQ) [14], and three MOS based metrics [15] for signal distortion (CSIG), intrusiveness of background noise (CBAK), and overall audio quality (COVL).

We respectively test the performance of two single-stage models, i.e., FME-Net and DSR-Net. They are actually the CRN [6] and the DCTCRN [16] combined with our TFSM block. In addition, we also test the performance after adding our TFSM block to the FME-Net (marked as FME-Net^†), as well as the performance after replacing the TFSM block in the DSR-Net with the GRU layer (marked as DSR-Net^⋆). From the results, we can obtain the following observations: (1) Going from FME-Net (i.e., CRN) to DSR-Net^⋆ (i.e., DCTCRN), the WB-PESQ and CBAK increase by 0.06 and 0.08, but the CSIG decreases by 0.05, which leads to the similar COVL. And same performance differences also exist between FME-Net^† (i.e., CRN+TFSM) and DSR-Net (i.e., DCTCRN+TFSM). This illustrates that although the STDCT-based model generally has superior performance, especially in noise suppression, it also causes more serious damage to the speech components. (2) The introduction of TFSM block improves the WB-PESQ, CSIG, CBAK, and COVL by 0.10, 0.01, 0.06, and 0.05 for FME-Net. As for DSR-Net, TFSM block also provides performance gain of 0.17 on WB-PESQ, 0.09 on CBAK, and 0.08 on COVL. This demonstrates the benefit of our TFSM block.

The two-stage model FDFNet is constructed by connecting FME-Net and DSR-Net together. In FDFNet, the input of DSR-Net includes the pre-enhanced spectrum of FME-Net, which has an obvious speech contour. Thus, this can guide the DSR-Net to better preserve the speech components. Meanwhile, the great noise reduction capability of the STDCT-based model ensures DSR-Net to effectively suppress the residual noise. In addition, the phase is also recovered by an implicit way in the STDCT spectrum domain. We find that the WB-PESQ, CSIG, CBAK, and COVL of FDFNet are respectively improved to 3.05, 4.23, 3.55, and 3.65. We have also tried to incorporate the TFSM block to FME-Net, which is marked as FDFNet^∗. But this cannot further improve the model performance, which illustrates that in the two-stage pipeline, over-optimization for the first stage model may be unnecessary.

We further compare our FDFNet with the previous advanced systems, and the results are presented in Table 2. In these benchmarks, CTS-Net [5] also adopts the two-stage pipeline, but its calculation process is only defined in the STFT spectrum domain. In CTS-Net, the second stage model adopts a dual-decoder network topology to capture a complex residual estimation, which is directly added to the first stage’s output to obtain the final estimation. It can be observed that with similar parameter size, our FDFNet outperforms the CTS-Net on the WB-PESQ, CBAK, and COVL scores. As for the CSIG, our FDFNet is slightly lower than that of CTS-Net, which may be due to the fact that the speech damage caused by STDCT has not been fully compensated. Furthermore, with the same CRN structure and slightly fewer parameters, the DCTCRN (3.1M), i.e., the DSR-Net^⋆ in Table 1, outperforms DCCRN [4] (3.7M). This further proves the advantage of STDCT over STFT. In addition, compared with other advanced systems, our FDFNet also illustrates superior performance. The enhanced audio clips can be found at https://github.com/Zhangyuewei98/FDFNet.git.