Joint Time-Frequency and Time Domain Learning for Speech Enhancement

1 Introduction

Single-channel speech enhancement addresses the problem of recovering clean speech from a noise-corrupted speech. According to the signal domain they work in, existing methods can be classified into two categories: time-frequency (T-F) domain methods and time-domain methods.

The T-F domain methods operate on the two-dimensional time-frequency spectrogram. As shown in Fig.1a, they often use short-time Fourier transform (STFT) to convert a raw audio signal into a T-F spectrogram. Then a T-F mask is predicted by the separation network. Finally, the output spectrogram is converted back to a time domain signal by the inverse STFT (ISTFT). T-F masking methods have been a great success for computational auditory scene analysis (CASA). They are also the mainstream methods for speech enhancement in the deep learning era. The theoretical ground for T-F analysis is that auditory patterns, such as proximity in frequency and time, harmonicity, and common amplitude and frequency modulation, are revealed on a T-F spectrogram [Wang and Chen, 2018]. However, the T-F domain methods are facing

• (c) Proposed cross-domain speech enhancement framework

Figure 1: Illustration of different speech enhancement frameworks. The color-filled modules are learnable neural networks. The separation network, either in T-F domain or in time domain, separates the features of clean speech from that of noise.

some challenges. First, STFT transforms the input signal into a complex domain with magnitude and phase. The phase of a signal is difficult to modify, but ignoring it will create a performance upper bound. Second, in supervised deep learning, a loss is computed at the output of the learnable network. It is a natural choice for T-F methods to compute the mean square error (MSE) loss for the T-F spectrogram. However, minimizing the MSE of spectrogram does not necessarily results in maximized signal-to-distortion ratio (SDR) of the output speech. This is known as the metric mismatch problem.

Time-domain methods try to avoid or solve these two problems by directly processing the raw waveform. Fig.1b shows the standard time-domain speech enhancement framework. Time-domain methods directly model the mixture waveform using an encoder-decoder framework and perform the separation on the output of the encoder. As phase no longer exists, the phase prediction problem is avoided. Besides, the loss is computed at the output of the learnable decoder. One is free to pick whatever metric that is important to the speech quality. As such, the metric mismatch problem is solved. However, the separation network in a time-domain method cannot leverage the known auditory patterns on a T-F spectrogram. This will potentially bring a performance loss.

A natural question to ask is whether it is possible to take advantages of methods in the two domains and avoid their respective drawbacks. Luckily, we find that the advantages of the T-F domain methods are mainly reflected in the front of the network and those of the time-domain methods are mainly in the latter part of the network. But an obstacle to connect these two parts is that the encoder and the decoder have to be paired in the conventional wisdom. Thanks to the powerful learning capability of deep neural network, this is no longer an issue with a learnable decoder. Fig.1c shows the flow chart of the proposed cross-domain framework, dubbed TFT-Net.

In TFT-Net, the input waveform is converted to T-F spectrogram through the predefined STFT. The subsequent T-F domain separation network can make full use of the auditory patterns revealed on the spectrogram. After the masked spectrogram is obtained, a learnable decoder can fully recover any spectrum modification after STFT. Note that supervision signal is given only after the decoder, and we do not intend to recover the ground-truth T-F spectrogram before the decoder. In fact, the output spectrogram is in a latent domain defined by the learned decoder. In the design of the separation network, we find that long-range correlations on the T-F spectrogram are important to the denoising performance. To tackle the high complexity, we propose dual-path attention blocks (DAB) to exploit the correlations along time and frequency dimensions in parallel. We further discover that the attention maps for different samples resemble each other, indicating sample-independent correlation is sufficient. Therefore, we employ sample-independent DABs (SDABs) in TFT-Net to balance the performance and computational cost.

In a nutshell, the contributions of our work are three-fold:

• We propose a cross domain learning framework TFTNet for speech enhancement. It takes advantages of both time domain and T-F domain approaches.
• We propose a novel sample-independent dual-path attention block (SDAB) to capture long-range temporal and frequency correlations with low computational cost.
• We perform comprehensive ablation studies and evaluate the proposed system TFT-Net on two major speech enhancement benchmarks. Results show that TFT-Net outperforms the state-of-the-art methods.

2 RelatedWork

This section reviews single-channel speech enhancement methods in different domains. Focuses are put on maskingbased methods in which the separation network generates a mask that describes the relationship between clean speech features and noisy features.

2.1 T-F DomainMethods

In a typical T-F domain masking method, as shown in Fig.1a, a pair of predefined encoder and decoder are used to convert the signal between time-domain waveform and T-F domain spectrogram. STFT and ISTFT are the most widely used transforms for the encoder and the decoder, respectively. A complete T-F spectrogram is composed of magnitude and phase. Early T-F masking methods only try to estimate the magnitude of the spectrogram due to the difficulty in modifying the phase. Later, research [Paliwal et al., 2011] reveals that phase also plays an important role in speech enhancement. There is a performance upper bound if only the magnitude is enhanced. PSM [Erdogan et al., 2015] and cIRM [Williamson et al., 2016] are then proposed to recover the phase in the spectrogram. While PSM is still a real-valued mask, cIRM is a complex-valued mask which can potentially recover both amplitude and phase of the clean speech.

In order to predict cIRM, the authors [Williamson et al., 2016] propose a DNN-based approach to estimate the real and imaginary components of the cIRM. However, the experimental results show that using cIRM does not achieve significantly better results than using PSM. We believe that the potential of a complex mask is not fully exploited in this work. In [Ephrat et al., 2018], a much deeper neural network with dilated convolution and bi-LSTM is designed. Dilated convolution enlarges the receptive field and bi-LSTM learns long-range correlations along the time axis. More recently, Yin et al. [Yin et al., 2019] propose a two-stream network for both magnitude and phase estimation. A frequency transformation block (FTB) is used in the front of the network to capture harmonic correlations along the frequency axis. A biLSTM is used at the latter part of the network to capture temporal dependencies. In [Kim et al., 2019], the spectrogram is treated as a time sequence and a Transformer is employed to capture long-range temporal correlations.

All these T-F domain methods benefit from the rich auditory patterns in the T-F spectrogram. However, we notice that when the long-range correlations along the time and the frequency axes are both considered, they are always learned in separate steps in prior works. We believe that learning longrange correlations along both axes is important, as harmonics exist in speeches [Plapous et al., 2005], [Krawczyk and Gerkmann, 2014], [Wakabayashi et al., 2018] and distinguishing noise needs long-term statistics. We also believe that learning the two types of correlations should be performed in parallel and the learned information needs to be fully fused. This becomes the motivation of our DAB design.

2.2 Time DomainMethods

Time-domain methods have emerged recently to tackle the difficulty in phase estimation and the metric mismatch problem in T-F domain methods. The most influential work in this category is Conv-TasNet [Luo and Mesgarani, 2019] which uses a pair of learnable encoder-decoder in time domain as an alternative to the predefined STFT-ISTFT. Before ConvTasNet, SEGAN [Pascual et al., 2017] uses generative adversarial networks (GANs) to directly predict the 1D waveform of the clean speech. In [Rethage et al., 2018], Wavenet is modified for speech enhancement. TCNN [Pandey and Wang, 2019] adopts a similar approach as Conv-TasNet, but it uses nonlinear encoder-decoder and longer frame length than Conv-TasNet.

These methods achieve their design goal in the sense that they avoid the problems of T-F domain methods. However, the most impressive performance is achieved by Conv-TasNet in the speech separation task. When applied to speech enhancement, its performance is inferior to the top-performing T-F domain method [Yin et al., 2019]. A possible reason for this phenomenon is that speech and noise patterns are easily distinguishable on T-F domain representations, and timedomain methods cannot make use of such prior knowledge.

3 The Proposed Scheme

We design a speech enhancement system with two objectives in mind. First, it should make full use of the prior knowledge about the T-F spectrogram. To realize this, the pipeline starts from the spectrogram and we design dual-path attention blocks to exploit the long-range correlations along both the time and frequency axes. Second, it should avoid the drawbacks of T-F domain methods. To this end, we propose an innovative cross-domain framework which directly uses timedomain metric to supervise the network learning.

3.1 The TFT-Net Framework

The most prominent feature of the proposed TFT-Net framework is that it starts with T-F spectrogram and directly outputs time-domain waveform. Fig.2 gives detailed illustration of the TFT-Net framework. Compared to Fig.1c, we omit the predefined STFT module so that we can concentrate on the end-to-end fully learnable parts.

Formally, the input noisy waveform Rin ∈ R is first transformed into STFT spectrogram, denoted by Sin in Fig.2. Here Sin ∈ RT×F×2 is a complex-valued spectrogram, where T represents the number of time steps and F represents the number of frequency bands. Sin is fed into two two-dimensional convolution layers to produce feature S0 ∈ RT×F×C , where C is the number of channels. S0 is then fed into N stacked learning blocks.

Each learning block consists of two parts, as shown in Fig.3. The first part is a group of 3 × 3 convolution layers that focus on capturing local correlations. We use four convolution layers and each of them is followed by batch normalization [Ioffe and Szegedy, 2015]. Around each of the two convolution layers, we employ a skip connection [He et al., 2016]. The second is the DAB which focuses on capturing long-range correlations and we will describe it in Section 3.2. The output features of each learning block are denoted by Si for i ∈ {1, 2, ...,N} where N is an adjustable hyperparameter. They have the same dimension as S0 .

After SN is produced by the last learning block, it is fed into the head layer to predict mask M ∈ T ×F ×Cm. Here, Cm is the number of channels. The enhanced spectrogram Sout in the latent domain can be calculated by the following generic function:

In our framework, we use the training target in [Choi et al., 2019]. So the head layer is a 1 × 1 convolution layer with Cm = 2. f(M) = tanh(|M|)∗M/|M|, g(Sin) = Sin and ⊙ indicates complex multiplication. Finally, Sout is fed into the decoder layer to produce the enhanced time-domain signal Rout ∈ R. Logarithmic MSE loss is computed between the original clean waveform and Rout .

In conventional wisdom, the encoder and the decoder need to be paired. Although Conv-TasNet [Luo and Mesgarani, 2019] has shown that the decoder does not have to perform the exact inverse operation of the encoder, we use the following method to ensure the convergence of our network. In short, the network structure of the learnable decoder is designed to be exactly the same as the convolutional implementation of ISTFT. Specifically, the learnable decoder is a 1D transposed convolution layer whose kernel size and stride being the window length and hop size in the STFT. As such, TFT-Net becomes a fully end-to-end learnable system.

3.2 Dual-path Attention Block

Long-range correlations exist along both the time and the frequency axes in a T-F spectrogram. It is quite obvious that audio signals, as time series, contain global correlations along the time axis. Along the frequency axis, there are well-known harmonic correlations. Although it is possible to treat the spectrogram as a 2D image and learn the correlations between every two pixels in the 2-D image, this is computationally too costly. DAB is designed to be a light-weight solution to capture the long-range correlations exhibited in T-F spectrogram.

As Fig.4 shows, we decompose the 2D T-F spectrogram (with C channels) into two stacks of 1D signals, one along the frequency axis and the other along the time axis. When the input feature dimension is T ×F ×C, it can be reshaped into TC vectors with dimension 1 × F or FC vectors with dimension 1 × T . Then, attention is applied along each axis in parallel. The attended features and the original features are then combined to generate the final output.

Ideally, we can use self-attention [Vaswani et al., 2017] to learn the attention map for each sample. But this might not be necessary. In classic speech enhancement literature [Scalart and FILHO, 1996], a uniform non-linear function is applied to the frequency axis to regenerate harmonics. It also matches with the intuition that harmonic correlations are sample-independent. On the time axis, when calculating Signal-to-Noise Ratio (SNR), the same set of parameters in recursive relation are used, which suggests that temporal correlation is time-invariant. Based on this understanding, we propose a sample-independent DAB, or SDAB. Specifically, the attention layer on F and T are implemented by fully-connected (FC) layers. Along the time path, the input and output dimensions of FC layers are T. Along the frequency path, the input and output dimensions of FC layers are F. FC layer is a kind of attention operation that uses learned weights. In other words, the relationship between every two elements is not a function of the input data. So using a FC layer is equivalent to learning a sample-independent correlations for all the training data. SDAB further reduces the computational complexity of the proposed network. In the experiment section, we will show through ablation study that SDAB performs almost as good as a sample-dependent self-attention.

4 Experiments

4.1 Datasets

AVSpeech+AudioSet

Audios from the AVSpeech dataset are used as clean speech. It is a large dataset proposed by [Ephrat et al., 2018]. It is collected from YouTube, containing 4700 hours of video segments with approximately 150,000 distinct speakers, spanning a wide variety of people and languages. The noisy speech is a mixture of the above clean speech segments with AudioSet [Gemmeke et al., 2017] which contains a total of more than 1.7 million 10-second segments of 526 kinds of noise. The noisy speech is synthesized by a weighted linear combination of speech segments and noise segments: Mixi = Speechj + 0.3 × Noisek, where Speechj and Noisek are 3-second segments randomly sampled from the speech and noise dataset. In our experiments, 100k segments are randomly sampled from the AVSpeech dataset and the "Balanced Train" part of AudioSet are used to synthesize the training set, while the validation set is the same as the one used in [Ephrat et al., 2018], synthesized by the test part of AVSpeech dataset and the evaluation part of AudioSet.

Voice Bank+DEMAND

This is an open dataset proposed by [Valentini-Botinhao et al., 2016]. Speeches of 30 speakers selected from Voice Bank corpus [Veaux et al., 2013] are used as clean speech: 28 are included in the training set and 2 are in the validation set. The noisy speech is synthesized using a mixture of clean speech with the DEMAND dataset [Thiemann et al., 2013]. A total of 40 different noise conditions are considered in the training set and 20 different conditions are considered in the test set. Finally, the training and test set contain 11572 and 824 noisyclean speech pairs, respectively. Both speakers and noise conditions in the test set are totally unseen by the training set.

4.2 EvaluationMetrics

We mainly use SDR [Vincent et al., 2006] and PESQ (perceptual evaluation of speech quality) in the ablation study. We further adopt CSIG [Hu and Loizou, 2007], CBAK [Hu and Loizou, 2007], COVL [Hu and Loizou, 2007], and SSNR (Segmental SNR) in the system comparison. For all the metrics, a higher value means a better result.

4.3 Implementation Details

Our method is implemented in Pytorch. All audios are resampled to 16kHz. STFT is computed using a Hann window of length 25ms, hop length of 10ms, and FFT size of 512, resulting in an input audio feature of 301 × 257 × 2 scalars. Convolution operation with zero padding, dilation=1 and stride=1 is used, making sure the input and output of the features are the same size. The number of channels for each convolution layer is 96. ReLU activations follow all network layers except for head layer (mask). Batch normalization [Ioffe and Szegedy, 2015] is performed after all convolutional layers.

When training model with the T-F domain framework, we use the training target and loss function in [Yin et al., 2019]. When training model with the cross-domain framework (TFT-Net), the training target in [Choi et al., 2019] is used and the loss function is logarithmic MSE.

Figure 5: Illustrations of learned weights of SDAB along frequency path and time path. From left to right is the 1st, 4th and 6th learned weights of SDAB. There are six SDABs in total.

4.4 Ablation Study

In the ablation study, all the networks are trained with the same random seed. Adam optimizer with a fixed learning rate of 0.0002 is used and the batch size is 8. Mean SDR and PESQ are used on the test dataset as the evaluation metric.

One-path Attention versus Dual-path Attention

SDAB operates on the T-F domain spectrogram, so the T-F domain framework is used when we study SDAB. Besides, two architectures are used to demonstrate the effectiveness of SDAB including TF-2stream and TF-1stream. Here, TF2stream represents PHASEN without FTB [Yin et al., 2019]. TF-1stream represents our TFT-Net without decoder and six learning blocks are used. Training target and loss function are the same as PHASEN.

Table 1 aims to demonstrate the effectiveness of dual-path attention. Based on TF-2stream, when dual-path attention is applied, 0.57dB and 0.05 gain on SDR and PESQ respectively are observed, compared to one-path attention only along frequency by using FTB in the original PHASEN. Based on TF1stream, significant gains on SDR and PESQ are observed when dual-path attention is applied. This large gain indicates the necessity of attention along both axes.

Sample-dependent versus Sample-independent

Next, we explore the effectiveness of sample-independent dual-path attention. Sample-dependent attention is implemented by using the self-attention mechanism [Vaswani et al., 2017]. The results are shown in Table 2. Based on the TF-2stream architecture, we find that sample-independent attention achieves almost the same results as sample-dependent attention in terms of SDR and PESQ. Based on the TF1stream architecture, the results of sample-independent dualpath attention are a little better than the sample-dependent one. Therefore, we can conclude that sample-independent attention will not negatively influence the accuracy of the model. On the other hand, the sample-independent attention mechanism within the SDAB is highly computationally efficient. Specifically, given an input feature map with a size of 301 × 257 × 96, the FLOPs of SDAB is 6.28B which is significantly smaller than the DAB's 14.7B FLOPs. Besides, the memory cost of SDAB is also less than DAB, since no big matrix multiplication is needed.

In order to better understand the attention mechanism, we visualize the learned weights of SDABs. The network we used for analysis has six learning blocks, so there are six pairs of attention weights. Fig.5 shows three pairs of them, taken from the first, the fourth, and the sixth SDABs. The attention weights along both frequency and time axis show a local-global-local pattern. In the first SDAB, we observe large weights on the diagonal, showing that more attention is paid to local components. Large weights also appear in other locations, but in a more random way. In the fourth DAB, large weights are scattered around in both frequency and time attention maps. But divergent lines can be observed in the frequency attention map, which are consistent with the harmonic correlations observed across the frequencies. In the last SDAB, weights are more concentrated on the diagonal than in the first SDAB, showing that global correlations have been successfully utilized.

Cross-domain Framework

Table 3 shows the results of training the network with different frameworks. TF-1stream represents training the network with the T-F domain framework. TFT-Net represents training the network with the proposed cross domain framework. Compared to TF-1stream, TFT-Net obtains 0.84dB gain on SDR and a decrease of 0.05 on PESQ which are consistent with our expectation, since logarithmic MSE is used as the training criteria which is equivalent to optimizing the SDR to a certain extent. This verifies the effectiveness of the cross domain framework. We also try replacing the learnable decoder layer by ISTFT (TFT-Net*). When comparing it to TFT-Net, there are 0.26dB and 0.01 loss on SDR and PESQ respectively which also verifies the effectiveness of the proposed framework.

The Number of Learning Blocks

In order to trade-off accuracy and complexity, we also train the model with different number of learning blocks. Table 4 shows the results. It is consistent with our expectation that the larger the model is, the better the result is. Compared to the case of three learning blocks, when the block number is six, we can obtain 0.77dB and 0.11 gain on SDR and PESQ respectively. When compared to six learning blocks, nine learning blocks only obtain slight improvement. Therefore, for a good trade-off between speed and accuracy, we use six learning blocks as our final proposed model to do comparisions to other state-of-the-art works.

4.5 Comparision to State-of-the-Arts

AVSpeech + AudioSet

We compare our method with three other recent methods, Conv-TasNet [Luo and Mesgarani, 2019], "Google" [Ephrat et al., 2018], and PHASEN [Yin et al., 2019]. Our networks are trained with Adam optimizer. Learning rate is 0.0002 and batch size is 8. The results in Table 5 show that our method outperforms all these three methods. Note that we use the same training data as PHASEN which is only a small fraction (100k/2.4M) of that used by "Google". Such superior performance on a large dataset demonstrates that our method can be generalized to various speakers and various kinds of noisy environments.

Voice Bank + DEMAND

In order to fairly compare the proposed method with many other methods, we also train our model on this small but commonly-used dataset. In this experiment, our networks are trained for 40 epochs, with Adam optimizer. Learning rate is 0.0005 and batch size is 8.

Table 6 shows the comparision results. Our TFT-Net achieves the best results on SSNR among all the methods listed. This is consistent with our expectation since our method uses logarithmic MSE as the training criteria which is equivalent to optimizing the SSNR to a certain extent. Compared to time-domain methods like SEGAN [Pascual et al., 2017], Wavenet [Rethage et al., 2018], and DFL [Germain et al., 2018], our method outperforms them on all the five metrics. This demonstrates the advantages of using information of the T-F domain. Compared to T-F domain methods like MMSE-GAN [Soni et al., 2018], PHASEN [Yin et al., 2019], our method does not get best results on the other four metrics. This maybe because these metrics depend on the magnitude spectrogram of speech and TF-domain methods optimize the spectrogram directly while our method uses time-domain signal as supervision. Besides, authors of Conv-TasNet have given very detailed explanation of why time-domain methods get a lower PESQ score but may still excel in (real subjective) MOS evaluation. MDPhD [Kim et al., 2018] performs processing the two domains in a cascaded way and the input signal is separately processed by two networks. Our method outperforms it on all the metrics indicating the superiority of our fully end-to-end learnable system in taking advantages of both domains.

5 Conclusion

In this paper, we propose a cross-domain framework TFTNet which takes advantage of the knowledge we have about spectrogram and avoids some of the drawbacks that T-Fdomain methods have been suffering from. We also propose a novel long-range correlations learning module SDAB which is very lightweight and effective. Through experiments on two datasets, we show the superiority of the proposed method over prior arts.

In the near future, we plan to speed up our model and apply it to low-latency applications. Another important direction is to study how different training targets and loss functions will influence the proposed method. Finally, we plan to extend the proposed method to various audio-related tasks such as dereverberation.