A deep representation learning speech enhancement method using $\beta$-VAE

In previous work, we proposed a variational autoencoder-based (VAE) Bayesian permutation training speech enhancement (SE) method (PVAE) which indicated that the SE performance of the traditional deep neural network-based (DNN) method could be improved by deep representation learning (DRL). Based on our previous work, we in this paper propose to use $\beta$-VAE to further improve PVAE's ability of representation learning. More specifically, our $\beta$-VAE can improve PVAE's capacity of disentangling different latent variables from the observed signal without the trade-off problem between disentanglement and signal reconstruction. This trade-off problem widely exists in previous $\beta$-VAE algorithms. Unlike the previous $\beta$-VAE algorithms, the proposed $\beta$-VAE strategy can also be used to optimize the DNN's structure. This means that the proposed method can not only improve PVAE's SE performance but also reduce the number of PVAE training parameters. The experimental results show that the proposed method can acquire better speech and noise latent representation than PVAE. Meanwhile, it also obtains a higher scale-invariant signal-to-distortion ratio, speech quality, and speech intelligibility.


I. INTRODUCTION
The aim of speech enhancement (SE) is to remove background noise from the observed speech signal. In general, SE is mainly used to reduce the word error rate of the automatic speech recognition system [1] or improve speech quality and intelligibility for human listening [2]. Recently, with the wide application of online meeting systems, SE is required to reduce the WER for accurate live caption when providing high-quality speech audio under various complex noise conditions [3]. Thus, SE research is becoming more and more challenging.
During the past decades, many single-channel SE algorithms have been developed, including signal subspace methods [4], non-negative matrix factorization methods [5], [6], and codebook-based methods [7]. In recent years, deep neural networks (DNN) have shown great potential for SE [2], [8]- [14] because DNNs can use a non-linear process to model complex high-dimensional signals, which is more reasonable in practical applications [15]. Thus, DNN-based methods usually have a better SE performance than these previous linear models [4]- [7].
However, most of the regression-based SE algorithms [2], [8]- [10] do not consider applying DNNs to obtain better speech representations when conducting SE. Instead, they This work was partly supported by Innovation Fund Denmark (Grant No.9065-00046).
usually use DNN to directly predict pre-defined targets for SE [2]. Although this approach can avoid inaccurate assumptions [8], it cannot ensure that these methods always work in environments with complex noise [2]. In general, deep representation learning (DRL) is important for DNN because DRL can obtain good signal representations in an unsupervised way and can, potentially, improve DNN's ability to extract useful information in complex environments [15], [16]. Additionally, a better signal representation usually leads to better predictions for DNNs [15]. Thus, DRL has a huge potential for DNNbased SE algorithms and makes them more robust. Moreover, the lack of a good DRL strategy may cause poor generalization of DNN-based SE algorithms [2], [15]. A good DRL algorithm can also disentangle various latent representations [15] of speech signals (e.g., speaker and phoneme information), which can also help DNN achieve a better SE performance.
Recently, to improve traditional DNN's generalization ability, DRL-based SE algorithms are proposed [17]- [22]. The basic idea of these methods is that they use a variational autoencoder (VAE) [23] to learn speech representations when modeling speech, and apply a non-negative matrix factorization (NMF) to model noise. VAE is a DRL model and can perform efficient approximate posterior inference. Additionally, VAE can also learn the probability distribution of complex data. Thus, VAE is suitable for various speech generative tasks [23]- [25]. These VAE-based algorithms can effectively improve DNN's generalization ability, but they have difficulty obtaining good speech representations from the observed signal because they cannot disentangle speech representations from other latent representations [15], [17]- [22]. This causes the need to use a linear NMF to model noise, so their noise modeling ability is limited compared with these non-linear DNN-based methods [23]. And their SE performance is not always satisfactory in a complex noisy environment [18].
To obtain a better speech representation from the observed signal, a novel VAE-based SE method (named PVAE) is proposed [26]. This method applies an unsupervised method to learn signal representations and derives a novel VAE lower bound, which ensures that VAE can disentangle different latent variables from the observed signal. Compared to the previous VAE-based SE algorithms, PVAE can use non-linear DNNs to model noise, which improves the noise modeling ability. Additionally, this method can adopt various DNN structures [2], so the DNN-based SE algorithms [2] can be directly optimized by PVAE. This is not achieved by VAE-NMFbased algorithms [17]- [22]. The experimental results [26] indicate that the SE performance of the traditional DNN-based methods can be improved by introducing this PVAE-based DRL algorithm.
Inspired by previous works, in this paper we propose a novel β-VAE strategy to improve PVAE's representation learning and disentangling performance [15] with fewer DNN parameters. β-VAE [27], [28] is originally designed to push VAE to learn a more efficient latent representation of the data, which is disentangled if the data contains at least some underlying factors of variation [27]. However, in general, β-VAE has a trade-off problem [28]. A better disentanglement within the latent representations usually causes worse signal reconstruction. In this work, based on the VAE's application in SE [26], we propose a strategy to address this trade-off problem to obtain better speech and noise representation. As a result, our β-VAE can improve disentangling and representation performance without signal reconstruction loss. Moreover, the proposed β-VAE can also optimize the neural network structure of the original PVAE. This means that the proposed β-VAE (named β-PVAE) can possibly achieve a better SE performance with fewer training parameters compared to PVAE.

II. RELATED WORK
Signal Model: in an additive noisy environment, using the short-time Fourier transform, the observed signal y f,n ∈ C, speech signal x f,n ∈ C, and noise d f,n ∈ C can be written as where frequency bin f ∈ [1, F ] and time frame index n ∈ [1, N ]. N and F denote the number of time frames and frequency bins, respectively. Their log-power spectrum (LPS) vector [8] at each frame can be represented as y, x, and d, respectively, where we omit the frequency and time frame index for simplicity. In [26], we assume that y is generated from a random process involving the speech latent variables z x ∈ R L and the noise latent variables z d ∈ R L . L is the dimension of latent variables. The latent variables z x and z d are independent. Similarly, x and d are independently generated by z x and z d , respectively. Fig. 1(a) shows the generative process. In [26], it is assumed that z x and z d can be estimated from speech and noise posterior distributions p(z x |x) and p(z d |d), respectively, and that they can also be estimated from the noisy speech posterior distributions p(z x |y) and p(z d |y). To disentangle latent variables, we assume that p(z x , z d |y) = p(z x |y)p(z d |y). Although this assumption is not always accurate in practical environments, it simplifies derivations, and helps us obtain a better signal model. Additionally, its effect towards signal estimation is not significant [26] (related analysis will be also given in Section IV). Fig. 1(b) shows the recognition process. VAE and β-VAE: the original VAE [23] defines a probabilistic generative process between the observed signal and its latent variables, and provides a principled method to jointly learn latent variables, generative and recognition models. The where D KL (||) denotes the Kullback-Leibler (KL) divergence. z y ∈ R L is the noisy latent variable. Maximizing this lower bound is equivalent to minimizing L n . β-VAE [27] is a modification of the original VAE framework, which introduces an adjustable hyperparameter β in the KL divergence term: In general, β > 1 results in more disentangled latent representations [27]. Higher values of β can encourage learning a more disentangled representation. However, β-VAE usually has a trade-off problem between the latent representation disentanglement and signal reconstruction. Bayesian permutation training VAE (PVAE) for SE: Although the VAE-based algorithms [23], [27] can learn signal representations and disentangle latent representations in a self-supervised way, their performance is limited when disentangling desired latent representations for SE application. Therefore, a Bayesian permutation training VAE (PVAE) [26] is proposed for SE. PVAE is a semi-supervised DRL method, which introduces multiple latent variables in VAE and disentangles them in a semi-supervised way. Fig. 2 shows the PVAE framework. It can be seen that PVAE includes three VAE structures: clean speech VAE (C-VAE), noise VAE (N-VAE), and noisy VAE (NS-VAE). C-VAE and N-VAE are trained without supervision to obtain speech and noise latent representations and their posterior estimates p(z x |x), p(z d |d), respectively. This is achieved by minimizing the following VAE loss function: where θ x , ϕ x , θ d , ϕ d are the DNN parameters for the related probability estimation [26]. Additionally, NS-VAE is trained under the supervision of C-VAE and N-VAE's encoders. Based on the derivation in [26], the NS-VAE's training loss function is where θ y , ϕ y are the NS-VAE's network parameters.
In the online SE stage, we assume that the z x , z d sampled from p(z x |x) and p(z d |d) are approximately equal to the sample z x , z d sampled from p(z x |y), p(z d |y), respectively. So, we separately use the NS-VAE encoder's two outputs as input of C-VAE and N-VAE to estimate related signals for SE.

III. β-VAE-BASED SPEECH ENHANCEMENT
Inspired by β-VAE, we propose a novel β-VAE strategy (named β-PVAE) to further improve PVAE's SE performance. More specifically, β-VAE is used to improve PVAE's representation learning ability that can better disentangle speech and noise latent variables from the observed signal, which can help PVAE obtain better SE performance. In PVAE, all the PVAE's decoders are trained in an unsupervised way [26]. The accuracy of the restored signal depends on the quality of latent representations. This means that the SE performance in PVAE is determined by the quality of speech and noise latent variables.
In [26], we derived a novel evidence lower bound (ELBO) (E y∼p(y) [log q(y)] ≥ −L p ). Additionally, β-VAE [27] applies an adjustable hyperparameter β in original VAE's [23] KL divergence term. Following β-VAE's property and PVAE's derivation [26], we apply this hyperparameter in the derived ELBO [26], the (6) can be written as In (7), we introduce a hyperparameter α in the restoration term. The purpose is to better analyze β-VAE [27] in PVAE. Note, α will not generate any effects for the original β-VAE's property because what is important in (7) is the weight ratio β : α. This weight ratio can also be written as: γ = β : α = (β/α) : 1, which is equal to the original β-VAE's loss function in (3). β-VAE [27] indicates that a higher value of β encourages VAE learning a more disentangled representation. Thus, we hypothesize that a higher value of β : α in (7) can cause a better disentangling performance for speech and noise latent variables. This point will be verified by later experiments. β-VAE usually has a trade-off problem between the disentanglement and signal reconstruction [27], which means that a good disentangled representation usually leads to poor signal reconstruction performance. In NS-VAE (as shown in Fig. 2), this trade-off is between the quality of observed signal reconstruction and the disentanglement of speech and noise latent variables. In SE application, we only need NS-VAE's disentanglement function, observed signal reconstruction is not useful (dashed part in Fig. 2). This means that we should set a very high weight ratio γ to obtain a better disentanglement performance [27]. Ideally, γ → +∞. One strategy to achieve this purpose is to set α = 0, so the loss function (7) can be rewritten as In (8), it can be found that there is no reconstruction term. This means that we do not need to train the NS-VAE's decoder, which reduces the PVAE's training parameters. The dashed part in Fig. 2 is removed in the proposed β-PVAE framework.
Comparing the PVAE and proposed β-PVAE, we can find that the β-VAE can be used to optimize the PVAE's network structure and β-PVAE also addresses the β-VAE's trade-off problem for SE application. All in all, the combination of β-VAE and PVAE can not only improve PVAE's disentanglement performance, but also simplify its framework. To summarize, the proposed β-PVAE includes a training and an enhancement stage for the SE application, which is similar to PVAE [26]. In the training stage, C-VAE and N-VAE are separately pre-trained by self-supervision using (4) and (5). After that, we apply (8) to train NS-VAE. In the enhancement stage, we can separately use the NS-VAE encoder's two outputs as input of C-VAE and N-VAE to obtain the prior distributions q(x|z x ) and q(d|z d ) for SE. Moreover, to calculate (8), related prior and posterior distributions need to be determined. Here, all the estimations of these distributions are the same as PVAE. More details can be found in [26].

IV. EXPERIMENTS
In this section, we report two experiments. First, we will investigate the disentanglement ability of the latent variables in the proposed algorithm. In addition, β-PVAE's SE performance will be indicated. Datasets: In this work, we use the DNS challenge 2021 corpus [29] to evaluate the performance of the proposed algorithm. We select English speakers and randomly split 70% of speakers for training, 20% for validation, and 10% for evaluation. Then, all the noise from the DNS noise corpus are randomly divided into training, validation, and test noise in a proportion similar to that used for speech utterances. Next, the corresponding training, validation, and test corpus for speech and noise are randomly mixed using DNS script [29] with random signal-to-noise ratio (SNR) levels (SNR range is from -10dB to 15dB). Other parameters of signal mixing are the default values in the DNS script [29]. Finally, we randomly choose 20 hours mixed training utterances, 5 hours mixed validation utterances, and 1 hour mixed test utterances to build experimental dataset. All signals are down-sampled to 16 kHz.
Experimental settings: In the experiments, the neural structures for C-VAE and N-VAE are the same. Their encoders include four hidden 1D convolutional layers [11]. The number of channels in each layer is 32, 64, 128, and 256. The size of each convolving kernel is 3. The two output layers of the encoders are fully connected layers with 128 nodes. Their decoders consist of four hidden 1D convolutional layers (the channel number of each layer is 256, 128, 64, and 32 with 3 kernel) and two fully connected output layers with 257 nodes. For NS-VAE, its encoder's hidden layer setting is the same as C-VAE. NS-VAE's encoder has four output layers with 128 nodes. For C-VAE, N-VAE, and NS-VAE, their activation functions in the hidden and output layer are ReLU and linear activation function, respectively. All networks are trained by the Adam algorithm with a 128 mini-batch size.
First, we will investigate β-PVAE's performance in disentangling speech and noise latent variables. Based on our previous derivation and analysis [26], β-PVAE's SE performance is determined by disentanglement performance. Table. I  shows the SE performance with a 95% confidence interval if latent variables are completely disentangled. Here, the signals are reconstructed by mask estimation [9]. The complete disentanglement means that they have the same posterior forms: p(z x |x) = p(z x |y) and p(z d |d) = p(z d |y). This is because p(z x |x) and p(z d |d) are learned in an unsupervised way with speech or noise only, which ensures that their latent representation only contains speech or noise representation. 'Oracle' results indicate that β-PVAE achieves a very satisfactory SE performance in SI-SDR, STOI, and PESQ, which shows the importance of disentangling latent variable for achieving excellent SE performance. The NS-VAE's purpose is to disentangle different representations from the observed signal and obtain the closest possible speech and noise posterior. Next, we use KL divergence to evaluate the practical disentanglement performance in latent space. A better disentanglement can lead to a lower KL divergence (D KL (p(z d |y)||p(z d |d)) and D KL (p(z x |y)||p(z x |x))). Fig. 3 shows the average KL divergence comparison of validation samples for using different ratios γ = β : α in loss function (7) to train NS-VAE. In (7), we keep α = 1 and change different β to determine ratio γ, and γ = +∞ means that α = 0, β = 1, which is equal to (8). In Fig. 3, we see that the KL divergence decreases with the increase of γ for both speech and noise latent variables, which means that the disentangled posteriors get closer to the true posteriors and the NS-VAE achieves a better disentanglement performance. When NS-VAE's decoder is removed (γ = +∞, α = 0, β = 1), NS-VAE can acquire the best posterior estimation. This verifies our hypothesis and deduction in Section 3. Additionally, although we have an inaccurate posterior conditional assumption p(z x , z d |y) = p(z x |y)p(z d |y), Fig. 3 shows that NS-VAE can still estimate a satisfactory posterior with a low KL divergence. However, this inaccurate assumption may hinder NS-VAE from obtaining a lower KL divergence when γ = +∞.
Next, we will evaluate the SE performance of the proposed β-PVAE. We use basic PVAE [26] as the reference method, which can be more direct to find the effects of β-VAE for the previous PVAE. The enhanced speech is obtained by mask estimation [9]. Table. I shows the experimental results. We find that β-PVAE achieves a very significant STOI, PESQ, and SI-SDR improvement over PVAE (from β = 1 to β = 2 ). This indicates that good disentanglement performance in latent space can directly lead to an improvement in speech quality and intelligibility. In addition, β-PVAE achieves the best SE performance when β = +∞. This illustrates that the proposed β-PVAE can effectively improve PVAE's SE performance with a simpler network structure.
V. CONCLUSIONS In this paper, a β-PVAE-based SE method is proposed to improve previous PVAE's SE performance. More specifically, β-PVAE can improve PVAE's ability to disentangle speech and noise latent variables from the observed signal. In addition, based on VAE's application in SE, the proposed β-PVAE addresses the trade-off problem between disentanglement and signal reconstruction, which widely exists in β-VAE. Compared with the previous PVAE algorithm, β-PVAE also simplifies its neural network and reduces the number of training parameters when improving the SE performance. Experimental results indicate that a good signal representation can achieve a very satisfactory SE performance. Moreover, β-PVAE obtains a better disentanglement performance and achieves higher SI-SDR, PESQ, and STOI scores than PVAE. In future work, we believe that β-PVAE can achieve better SE performance by improving the latent space disentanglement performance or the decoder's signal reconstruction ability.