Phase-aware Single-stage Speech Denoising and Dereverberation with U-Net

Designing a mask that is not limited to output the optimal value of ideal mask requires two conditions to satisfy. First, the range of magnitude mask should not be limited. Second, the mask has to be complex-valued so that it can correct both the magnitude part and phase part of the mixture signal. The proposed phase-aware $\beta$-sigmoid mask (PHM) is designed to handle both conditions while systemically restricting the sum of estimated complex values to be exactly the value of mixture, $X_{t,f}=Y^{(k)}_{t,f}+Y^{(\lnot k)}_{t,f}$. The PHM separates the mixture $X_{t,f}$ in STFT domain into two parts as one-vs-rest approach, that is, the signal $Y^{(k)}_{t,f}$ and the sum of the rest of the signals $Y^{(\lnot k)}_{t,f}=X_{t,f}-Y^{(k)}_{t,f}$, where index $k$ could be one of direct path source ($d$), reverberation ($r$), and noise ($n$) in our setting, $k\in\{d,r,n\}$. The complex-valued mask $M^{(k)}_{t,f}\in\mathbb{C}$ estimates the magnitude and phase value of the source of interest $k$. The mask is composed of two parts, (1) magnitude mask estimation, (2) phase estimation by reusing the magnitude estimation from (1) and two-class sign prediction.

First, the network outputs the magnitude part of two masks $\mathinner{\!\left\lvert M^{(k)}_{t,f}\right\rvert}$ and $\mathinner{\!\left\lvert M^{(\lnot k)}_{t,f}\right\rvert}$ with sigmoid function $\sigma^{(k)}(\bm{z}_{t,f})$ multiplied by coefficient $\beta_{t,f}$ as follows,

$$\mathinner{\!\left\lvert M^{(k)}_{t,f}\right\rvert}=\beta_{t,f}\cdot\sigma^{(k)}(\bm{z}_{t,f})=\beta_{t,f}\cdot\frac{1}{1+e^{-(z^{(k)}_{t,f}-z^{(\lnot k)}_{t,f})}}$$

(1)

where $z^{(k)}_{t,f}$ is the output located at $(t,f)$ from the last layer of neural-network function $\psi^{(k)}(\phi)$, and $\phi$ is a function composed of network layers before the last layer. $\mathinner{\!\left\lvert M^{(k)}_{t,f}\right\rvert}$ serves as a magnitude mask to estimate source $k$ and the value of it ranges from 0 to $\beta_{t,f}$. The role of $\beta_{t,f}$ is to design a mask that is close to the optimal mask with a flexible magnitude range so that the value is not bounded between 0 and 1 unlike the typically used sigmoid mask. In addition, because the sum of the complex valued masks $M^{(k)}_{t,f}$ and $M^{(\lnot k)}_{t,f}$ must compose a triangle, it is reasonable to design a mask that satisfies the triangle inequalities, that is, $\mathinner{\!\left\lvert M^{(k)}_{t,f}\right\rvert}+\mathinner{\!\left\lvert M^{(\lnot k)}_{t,f}\right\rvert}$ $\geq 1$ and $\mathinner{\!\left\lvert\,\mathinner{\!\left\lvert M^{(k)}_{t,f}\right\rvert}-\mathinner{\!\left\lvert M^{(\lnot k)}_{t,f}\right\rvert}\right\rvert}\leq 1$. To address the first inequality we designed the network to output $\beta_{t,f}$ from the last layer with a softplus activation function as follows, $\beta_{t,f}=1+\texttt{softplus}((\psi_{\beta}(\phi))_{t,f})$, where $\psi_{\beta}$ denotes an additional network layer to output $\beta_{t,f}$. The second inequality can be satisfied by clipping the upper bound of the $\beta_{t,f}$ by $1\mathbin{/}\mathinner{\!\left\lvert\,\sigma^{(k)}(\bm{z}_{t,f})-\sigma^{(\lnot k)}(\bm{z}_{t,f})\right\rvert}$.

Once the magnitude masks are decided, we can construct a phase mask $e^{j\theta_{t,f}^{(k)}}$. Given the magnitudes as three sides of a triangle, we can compute the cosine of absolute phase difference $\Delta\theta_{t,f}^{(k)}$ between the mixture and source $k$ as follows, $\cos(\Delta\theta_{t,f}^{(k)})=\nicefrac{{(1+\mathinner{\!\left\lvert M_{t,f}^{(k)}\right\rvert}^{2}-\mathinner{\!\left\lvert M_{t,f}^{(\lnot k)}\right\rvert}^{2})}}{{(2\,\mathinner{\!\left\lvert M_{t,f}^{(k)}\right\rvert})}}$. Next, the rotational direction (clockwise or counterclockwise) for phase correction can be decided by estimating sign value $\xi_{t,f}\in\{-1,1\}$ as follows,

$$e^{j\theta_{t,f}^{(k)}}=\cos(\Delta\theta_{t,f}^{(k)})+j\xi_{t,f}\sin(\Delta\theta_{t,f}^{(k)}).$$

(2)

Two-class straight-through Gumbel-softmax estimator was used to estimate $\xi_{t,f}$ [21]. It allows us to discretize the output of the Gumbel-softmax function $\gamma^{(i)}$ with $\operatorname*{arg\,max}$ and still be able to train the network in an end-to-end manner using a continuous approximation in the backward pass. $\xi_{t,f}$ is defined as follows,

$$\xi_{t,f}=\begin{cases}-1,&\gamma^{(0)}(\bm{q}_{t,f})>\gamma^{(1)}(\bm{q}_{t,f})\\ 1,&\text{otherwise}\end{cases}$$

(3)

where $\gamma^{(i)}(\bm{q}_{t,f})$ is defined as follows,

$$\gamma^{(i)}(\bm{q}_{t,f})=\frac{e^{q^{(i)}_{t,f}}}{\sum_{i}{e^{q^{(i)}_{t,f}}}}=\frac{e^{((\psi_{i}(\phi))_{t,f}+g_{i})\mathbin{/}\tau}}{\sum_{i}{e^{((\psi_{i}(\phi))_{t,f}+g_{i})\mathbin{/}\tau}}},$$

(4)

and $g_{0}$ and $g_{1}$ are samples from Gumbel$\left(0,1\right)$, $\psi_{i}$ is an additional network layer to output logit value $q^{(i)}_{t,f}$, and $\tau$ is a temperature parameter for Gumbel-softmax. Finally, $M^{(k)}_{t,f}$ is defined as follows,

$$M^{(k)}_{t,f}=\mathinner{\!\left\lvert M^{(k)}_{t,f}\right\rvert}e^{j\theta_{t,f}^{(k)}}.$$

(5)

As we desire to extract both direct source and reverberant source, two pairs of PHMs are used for each of them. The first pair of masks separates direct source and the rest of the component, denoted as $M^{(d)}_{t,f}$ and $M^{(\lnot d)}_{t,f}$. The second pair of masks separates noise and reverberant source component, denoted as $M^{(n)}_{t,f}$ and $M^{(\lnot n)}_{t,f}$. Since PHM guarantees the mixture and separated components to construct a triangle in the complex STFT domain, the outcome of the separation can be seen from the perspective of quadrangle as in Fig 1. In this setting, as the three sides and two side angles are already determined by the two pairs of PHMs, the last fourth side of quadrangle, the reverberation component $\hat{Y}^{(r)}_{t,f}$, is uniquely decided.

Learning to maximize cosine similarity can be regarded as maximizing the signal-to-distortion ratio (SDR) [6]. Cosine similarity loss $C$ between estimated signal $\hat{\bm{y}}^{(k)}\in\mathbb{R}^{N}$ and ground truth signal $\bm{y}^{(k)}\in\mathbb{R}^{N}$ is defined as follows,

$$C(\bm{y}^{(k)},\hat{\bm{y}}^{(k)})=-\frac{<\bm{y}^{(k)},\hat{\bm{y}}^{(k)}>}{\mathinner{\!\left\lVert\bm{y}^{(k)}\right\rVert}\mathinner{\!\left\lVert\hat{\bm{y}}^{(k)}\right\rVert}},$$

(6)

where $N$ denotes the temporal dimensionality of a signal and $k$ denotes the type of signal ($k\in\{d,r,n\}$). Consider a sliced signal $\bm{y}^{(k)}_{[\frac{N}{M}(i-1)\mathrel{\mathop{\ordinarycolon}}\frac{N}{M}i]}$, where $i$ denotes the segment index and $M$ denotes the number of segment. By slicing the signal and normalize it by its norm, each sliced segment is considered as an unit for computing $C$. Therefore, we hypothesize that it is important to choose a proper segment length unit $\frac{N}{M}$ when computing $C$. In our case, we used multiple settings of segment lengths $g_{j}=\frac{N}{M_{j}}$ as follows,

$$\mathcal{L}(\bm{y}^{(k)},\hat{\bm{y}}^{(k)})=\sum_{j}\frac{1}{M_{j}}\sum_{i=1}^{M_{j}}C(\bm{y}^{(k)}_{[g_{j}(i-1)\mathrel{\mathop{\ordinarycolon}}g_{j}i]},\hat{\bm{y}}^{(k)}_{[g_{j}(i-1)\mathrel{\mathop{\ordinarycolon}}g_{j}i]}),$$

(7)

where $M_{j}$ denotes the number of sliced segments. In our case the set of $g_{j}$'s was chosen as follows, $g_{j}\in\{4064,2032,1016,508\}$, assuming they moderately cover the range of duration of phonemes in speech.

To further improve the design of the loss function, we applied two simple techniques — 1. pre-emphasis ($\pi$) and 2. $\mu$-law encoding ($\mu$) — on signals. As most of the speech signal components are concentrated in the lower frequency bands, we found that applying pre-emphasis on loss function can help penalize high frequency components. In addition, since the samples of speech signals are usually centered around zero, we found that it is helpful to use 16-bit $\mu$-law encoding as it distributes samples more uniformly by the nature of continuous logarithmic transform. The proposed loss function $\mathcal{L}^{+}$ is defined as follows,

$$\begin{split}\mathcal{L}^{+}(\bm{y}^{(k)},\hat{\bm{y}}^{(k)})&=\mathcal{L}(\bm{y}^{(k)},\hat{\bm{y}}^{(k)})+\mathcal{L}(\pi(\bm{y}^{(k)}),\pi(\hat{\bm{y}}^{(k)}))\\ &+\mathcal{L}(\mu(\pi(\bm{y}^{(k)})),\mu(\pi(\hat{\bm{y}}^{(k)})))\end{split}$$

(8)

Finally, we used the proposed loss function for every $k$ and $\lnot k$ combinations as follows,

$$\mathcal{L}_{\text{final}}=\sum_{k}(\mathcal{L}^{+}(\bm{y}^{(k)},\hat{\bm{y}}^{(k)})+\mathcal{L}^{+}(\bm{y}^{(\lnot k)},\hat{\bm{y}}^{(\lnot k)})).$$

(9)

We used the DNS challenge dataset [22] and internally collected dataset for training. The former is a large scale dataset where the speech samples were collected from Librivox [23], and noise samples from Audioset [24] and Freesound [25]. Note that we did not use the provided noisy speech data from the DNS dataset but used on-the-fly augmentation with the clean speech and noise in the two datasets during the training phase. Since our goal is to perform both denoising and dereverberation, we used pyroomacoustics [26] to simulate an artificial reverberation with randomly sampled absorption, room size, location of source and microphone distance. We also trimmed random segments of 2 seconds from speech and noise data, and mixed them with uniformly distributed source-to-noise ratio (SNR) ranging from -10 dB to 30 dB.

For test, we used two datasets such as the synthesized testset in the DNS challenge (DNS) and WHAMR [20]. The DNS synthesized testset provides noisy-reverberant mixtures and noisy mixtures without reverb. DNS was used only to test the denoising performance since it does not provide the direct source signal of synthesized mixture samples. Therefore, a reverberant source was given as ground truth when the model is tested on noisy-reverberant mixtures. Both the denoising and dereverberation performance were tested on the min subset of WHAMR dataset which contains 3,000 audio files. To test the denoising and dereverberation performances both simultaneously and separately, we tested our models on four scenarios: 1) nr2d: noisy-reverberant mixture to direct source 2) nr2r: noisy-reverberant mixture to reverberant source 3) n2d: noisy mixture to direct source 4) r2d: reverberant source to direct source. The corresponding four pair of test subsets, denoted in a following way (mixture, ground_truth), were used as follows, 1. nr2d: (mix_single_reverb, s1_anechoic), 2. nr2r: (mix_single_reverb, s1_reverb), 3. n2d: (mix_single_anechoic, s1_anechoic), 4. r2d: (s1_reverb, s1_anechoic).

Input features were used as a channel-wise concatenation of log-magnitude spectrogram, real and imaginary part of demodulated phase [13], group delay, and delta-phase [27]. The window size of model20 was 1024 with 256 hop size and the window size of model10 was 512 with 128 hop size. All models were trained for 125k iterations with AdamW optimizer [28]. The learning rate was set to 0.0004 and halved at 62.5k iteration. Every test was done with a non-causal inference using model20 except the experiments in subsection 5.5.

To show the effect of loss functions we observed SI-SDR [10] and PESQ [29] while changing four different loss functions. Complex MSE ($\mathbb{C}$MSE) and three different cosine similarity based loss functions — SingleScale, MultiScale, and MultiScale+ — were compared each of which corresponds to Eq. 6, Eq. 7, and Eq. 8, respectively. The quantitative results in Table 1 show that the proposed multi-scale and emphasis functions are beneficial for both denoising and dereverberation tasks in most of the cases.

Here, we used the phase distance defined in [6] to quantitatively measure the phase enhancement performance. Phase distance ($PD$) between spectrogram $A$ and $B$ is formulated as follows,

$\displaystyle PD(\bm{A},\bm{B})=\sum_{t,f}\frac{\mathinner{\!\left\lvert A_{t,f}\right\rvert}}{\sum_{t^{\prime},f^{\prime}}\mathinner{\!\left\lvert A_{t^{\prime},f^{\prime}}\right\rvert}}\angle(A_{t,f},B_{t,f}),$

(10)

where $\angle(A_{t,f},B_{t,f})$ is the angle between $A_{t,f}$ and $B_{t,f}$, ranging from 0°to 180°. We measured $PD$ between ground truth $\bm{Y}$ and mixture $\bm{X}$, and $PD$ between ground truth $\bm{Y}$ and estimation $\bm{\hat{Y}}$, and checked how much $PhaseGain(\%)$ was obtained. This was tested on all four scenarios of WHAMR testset and shown in Table 2. We found that the network is able to give a reasonable $PhaseGain$ in tasks including dereverberation (nr2d, nr2r). However, $PhaseGain$ was marginal for only-denoising-tasks (nr2r, n2d). We conjecture that this is because the network is not able to estimate a precise magnitude value for noisy mixture and this issue is left for futurework. A visualization of enhanced phase group delay tested on a reverberant source is shown in Fig. 3. Fig. 3 (b) shows the enhanced harmonic structure of phase group delay.

Following the constraint for real-time model suggested by the DNS challenge, we measured the elapsed time to compute a single frame. model20 that took 40 ms to compute a frame will be denoted as non-real-time (NRT) model, and model10 that took 4.32 ms to compute a frame will be denoted as real-time (RT) model. To compare the two models and how causal inference affects the model performance, we compare four combinations in nr2d task. Table 3 shows that both non-causal inference and model size are significant factors for performance.

Finally, we report the Mean Opinion Score (MOS) results from the DNS challenge based on the online subjective evaluation framework ITU-T P.808 [30]. For better perceptual quality, we linearly added the estimated direct source and reverberant source with a 15 dB ratio, and implemented a simple and zero-delay dynamic range compression to apply on it. Our causal-NRT and causal-RT model achieved a mean opinion score of 3.36 and 3.24, respectively.