A Two-Stage Training Framework for Joint Speech Compression and Enhancement

Audio codec is a fundamental technique in audio communication, which aims to compress audio signal at a given bit-rate with distortion as less as possible. Classic audio codecs include USAC[1], EVS[3] and OPUS[2]. USAC (Unified Speech and Audio Coding) is a low-delay, high-quality audio coding standard developed by MPEG. OPUS is an open-source audio codec with sampling rates ranging from 6 kbps to 510 kbps. It supports a wide range of audio applications from video conferencing to streaming media services. The performance of OPUS is excellent in high bit-rate compression. EVS (Enhanced Voice Services) is the latest codec developed by 3GPP standardization organization for mobile networks. It supports target bit-rates ranging from 5.9 kbps to 128 kbps and performs better than OPUS at medium and low target bit-rate.

Generative model has become a popular tool for audio reconstruction tasks due to its ability to generate high-quality audios. Wavenet[22, 57] is an auto-regressive sequence generation model, which can directly generate high-quality speech waveforms from input signal. Based on it, a series of generative models such as WaveRNN[24], WaveGAN[49], and Parallel WaveGAN[50] have been developed as vocoder for text-to-speech, which further improve the speed of processing and quality of the reconstructed audio. Lyra[6] utilizes an auto-regressive WaveGRU model as the decoder of the codec. It can achieve excellent performance at a low target bit-rate of 3 kbps. Another type of MelGAN-based generative models mainly focuses on improved discriminators for audio signal, which can improve perceptual quality by a large margin. For example, MelGAN [19] proposes to use multi-scale discriminators to achieve better perception quality. Furthermore, HIFI-GAN [12] designs a multi-period discriminator and proposes a multi-receptive field fusion method. Based on HIFI-GAN, Fre-GAN [51] uses DWT to retain high-frequency information and uses RCG to capture information of different frequency bands. These methods have shown effectiveness in improving naturalness of generated speech.

Recently, end-to-end neural audio codecs have shown promising performance. The work [31] proposes a neural network architecture based on VQ-VAE and WaveNet decoder. It can generate high-quality audio at a very low bit-rate of 1.6 kbps. The work [32] introduces a speaker encoder and speaker VQ codebook into the VQ-VAE architecture to generalize to unseen speakers or content. SoundStream [8] proposes a bit-rate scalable codec with VQ-VAE-2 [9], which can support scalable bit-rates from 3 kbps to 18 kbps with a single model by simply controlling the number of cascaded VQs during inference.

Speech enhancement, also known as noise reduction, aims to remove the background noise of observed noisy speech. It is widely used as a front-end processing technology in speech recognition, hearing aids, and telephone communication. Generally speaking, speech enhancement methods can be divided into two categories, the traditional statistical model based methods [59, 60, 61], and the recent DNN-based methods [35, 37, 36, 38, 7, 44, 45, 46, 47, 48, 53, 54, 58, 39, 40, 43, 41, 42]. Traditional methods, such as the spectral subtraction method, the statistical model based method, and the subspace based method, perform well on the suppression of stationary noise. However, these methods often perform poorly on non-stationary noise.

Benefited from the powerful deep learning technique, DNN-based methods can achieve significantly better performance than the traditional statistical model based methods [37, 36, 38, 7, 44, 45]. However, the most popular speech codecs, such as EVS and OPUS, do not consider speech denoising or use cascaded denoising module. OPUS applies discontinuous transmission (DTX) to reduce the impact of noise by reducing the bit-rate of silent frames or noisy frames and applies noise shaping filters to reduce noise. EVS uses DTX and Comfort Noise Generation (CNG) to handle noisy speech. It employs Voice Activity Detection (VAD) to detect silent segments or background noise, encodes the noise with dedicated cores and finally compensates the distortion in speech with CNG. Another popular speech codec Lyra still applies an additional module TasNet [7] for noise suppression before the encoding the signal. To improve this cascaded structure, SoundStream proposes that compression and enhancement can be jointly performed by the single model without increasing the overall delay. It uses Feature-wise Linear Modulation (FiLM) [10] to achieve controllable denoising. The FiLM layer performs a simple feature-wise affine transformation on the features of the intermediate layer of the neural network and dynamically activates and deactivates the denoising with an additional conditioning signal.

In speech codec, the objective is to achieve lower distortion and higher perceptual quality at a given bit-rate. Recent studies show that distortion and perception quality are at odd with each other. More specifically, minimizing distortion only would not necessarily lead to good perception quality. Imposing perfect perception quality constraint would lead to increase of the lowest achievable distortion [11, 23]. That is high perception quality can only be achieved at the cost of increased distortion. In the field of audio processing, the commonly used spectral reconstruction loss is also a kind of distortion loss. Therefore, only considering the spectral distortion of the reconstruction is not enough for achieving high perception quality. Mathematically, perceptual quality can be expressed as the deviation between the distribution of the decoded signal $\hat{X}$ and that of the source signal $X$, as [11]

where $d(\cdot,\cdot)$ is a divergence measures the deviation between two distributions, such as the KL divergence or Wasserstein distance, which satisfies $d(p,q)=0$ if and only if $p=q$. When $d(p_{X},p_{\hat{X}})=0$, $\hat{X}$ and $X$ have the same distribution as $p_{X}=p_{\hat{X}}$. In this case, the reconstruction has perfect perception quality.

As generative adversarial training is effective in aligning distributions, adversarial loss is widely used in learning deep neural networks based codecs to minimize the deviation between the distributions of the decoded output $\hat{X}$ and the source ${X}$, i.e., $d(p_{X},p_{\hat{X}})=0$. It has been shown that, such adversarial training can significantly improve the naturalness of the generated audio. However, using only adversarial loss can result in excessive increase of the reconstruction distortion, e.g, though the decoded output can have good perception quality but with large distortion from the input signal. Therefore, adversarial loss is usually used in combination with distortion loss such as spectral reconstruction loss, with a balance between these two losses to achieve a trade-off between distortion and perception [19, 12, 8]

Recently, the works [13, 14] have studied the training framework for perceptual lossy compression under perception constraint. A method for flexibly controllable perception-distortion trade-off has been proposed in [14]. For audio codec, there also exist a perception-distortion trade-off problem, as empirically shown in [8]. As the work [13] considers the lossy compression problem in the absence of noise, the results do not apply to audio codec in processing noisy speech signal. Speech noise is inevitable in many real-world applications. In this work, we extend the work [13] to consider the joint compression and enhancement of noisy speech signal.

We first review existing results on the compression problem in the absence of noise and, then, present extended results on the joint compression and enhancement problem in the presence of noise.

Suppose that $X$ is a deterministic discrete source to be compressed. In the speech compression task without considering noise, an encoder $E$ encodes the source signal $X$ into a sequential compressed representation $Z$ at a given bit-rate $R$, and a decoder $G$ decodes the compressed representation $Z$ to obtain $\hat{X}$, which can be written as

At a given bit-rate $R$, the rate-distortion-perception (R-D-P) function for perceptual lossy compression, which additionally takes perception quality constraint into consideration in the Shannon’s rate-distortion function, can be expressed as

where $D$ stands for the distortion constraint, $P$ stands for the perception quality constraint based on the divergence between $\hat{X}$ and $X$, and $\Delta$ is a distortion measure function. For the R-D-P problem (2), the following result has been derived in [13].

Theorem 1 [13]: Suppose that $X$ is a deterministic discrete source, and the distortion function $\Delta$ measures the squared error. Let $(E_{d},G_{d})$ be an optimal encoder-decoder pair to problem (2) in the case without any constraint on perception, i.e. $P=+\infty$ in (2), then $E_{d}$ is also an optimal encoder in the case perfect perception constraint i.e. $P=0$ in (2). Furthermore, denote $X_{mse}:=G_{d}(E_{d}(X))$, there holds the following equations:

This result implies that any optimal encoder achieving the minimum mean squared error (MSE) in the case without perception constraint is also an optimal encoder in the case with perfect perception constraint. That is any optimal encoder to $R^{(I)}(D,+\infty)$ is also an optimal encoder to $R^{(I)}(D,0)$. Next we consider the condition in the presence of noise and provide extended result.

Theorem 2: Let $X$ be a deterministic discrete source, and $X^{\prime}$ be a noisy observation from $X$ as $X^{\prime}=X+N$, where $N$ is noise. Suppose that $N$ is independent with $X$ and the distortion function $\Delta$ measures the squared error, then, any optimal encoder under the condition without any perception constraint, i.e., $P=+\infty$ in (2), is also an optimal encoder under the condition with perfect perception constraint, i.e., $P=0$ in (2).

Proof: The minimum MSE codec process in the presence of noise can be expressed as

Denote the output of the optimal noise reduction by

Under the condition that $\Delta$ measures the squared error, the optimal process for joint compression and noise reduction can be written as

where $X_{mse}$ is the minimum MSE codec output from $X_{deN}$. Given $X_{deN}$, $X$ and $X_{mse}$ are independent of each other since the data process $X\longrightarrow X_{deN}\longrightarrow X_{mse}$ is a Markov chain. Then for the distortion term, we have

Under the condition that the noise $N$ is independent of the clean source $X$, we have $\mathbb{E}[X-X_{deN}]=0$ and then it follows from (4) that

It implies that the an end-to-end optimization from $X^{\prime}$ to $\hat{X^{\prime}}$ is equivalent to the two-stage optimization process which firstly conduct noise reduction to obtain $X_{deN}$ and then compresses $X_{deN}$ by the codec. Thus, we can conclude that under the assumption of independent noise, Theorem 2 holds, and the denoised minimum MSE encoder is still an optimal encoder under the condition with perfect perception constraint.

In speech codec model training, spectral reconstruction loss is usually used rather than the MSE loss for distortion measurement. Particularly, the multi-scale spectral reconstruction loss has been widely used in high-fidelity audio synthesis and has show superior performance. The multi-scale spectral reconstruction loss is typically defined as [25]

where $X_{t}^{s}$ denotes the $t$-th frame of the time-frequency spectrum (or Mel-spectrogram) of the input signal $X$ with a window length equals to $s$. $\hat{X}^{s}$ denotes the spectrogram of the reconstructed signal $\hat{X}$. $\alpha_{s}$ is a balance parameter which is typically chosen to be $\alpha_{s}=\sqrt{s/2}$ [8]. The loss (6) computes the $\ell_{1}$ distortion of the spectrum and the $\ell_{2}$ distortion of the log-scale spectrum for multiple scales in the spectral domain. Though Theorem 2 is only derived for the case of the MSE loss, it still sheds some light on the cases with other distortion losses such as (6).

The above analysis considers the MSE loss. For speech signal processing, the multi-scale spectral reconstruction loss (6), which consists of $\ell_{1}$ distance on the spectrogram slice and $\ell_{2}$ distance on the log of the spectrogram slice, has shown superiority and hence be widely used in practice. Next, we provide further analysis for the multi-scale spectral reconstruction loss to show that Theorem 1 in the MSE distortion case can be extended to $\ell_{1}$ norm case and logarithmic MSE case.

Let $\hat{X}$ be the decoding output which minimizes a generalized distortion such as the $\ell_{1}$ distance with respect to the spectrogram slice or the $\ell_{2}$ distance with respect to the log of the spectrogram slice. Meanwhile, let $\hat{X}_{p}$ denote the output with perfect perception constraint. According to the principle of Markov chain, $X$ and $\hat{X}_{p}$ are independent of each other.

For the case of the $\ell_{1}$ loss, the distortion is $\ell_{1}$ distance with respect to the spectrogram slice as

For the case of the log MSE loss, the distortion is $\ell_{2}$ distance with respect to the log of the spectrogram slice as

Under the perfect perceptual constraint, i.e. $d(p_{X},p_{\hat{X}_{p}})\leq 0$, we can derive perfect perceptual quality by $p_{\hat{X}_{p}|\hat{X}}=p_{X|\hat{X}}$. In the final expansion of (8), the angle between the two inner products of the third term is 180°, therefore the third term is negative. Thus it follows from (8) that

In (7) and (9), the first term is the distortion between the source and the minimum distortion decoding output, and the second term is the distortion caused by mapping $\hat{X}$ to $\hat{X}_{p}$ with the same distribution as the source. According to Theorem 1, optimizing the encoder $E_{d}$ in the case of $P=+\infty$ results in minimizing the upper bound of the minimum distortion in the case of $P=0$. Therefore, for the multi-scale spectral reconstruction loss $L_{dis}$ in (6), which is composed of $\ell_{1}$ loss and log MSE, the minimum distortion encoder $E_{d}$ in the case without any perception constraint (i.e. $P=+\infty$), can also be used as an approximately optimal encoder in the case with perfect perception constraint (i.e. $P=0$).

From the above analysis, in this subsection we propose a two-stage training framework for joint speech compression and enhancement based on generative adversarial training, in which a generator consists of an encoder-decoder pair is the codec model to be learned, as shown in Fig. 1.

The generator includes an encoder, a residual vector quantizer and a decoder. It firstly compresses the source signal into compressed representation, then encodes it with a series of cascaded vector quantizers, and finally reconstructs the source signal with the decoder. The encoder consists of a one-dimensional convolution layer and four encoding modules with stride $(2,4,5,8)$. Each encoding module consists of three residual units and a down-sampling layer. Finally, a one-dimensional convolution layer is used to obtain features with a dimension of 256. In order to enable real-time inference, all convolutions are causal. The decoder adopts a similar up-sampling structure. The residual vector quantizer consists of $N_{q}$ cascaded vector quantizers as in [8]. The target bit-rate $R$ of the codec can be calculated by $R=N_{q}\cdot S\cdot\log_{2}(N)$, where $N$ is the size of codebook and $S$ is the number of frames per second. We also use quantizer dropout in training to enable bit-rate scalability as in SoundStream [8].

The discriminator is employed for adversarial training to improve the perception quality of reconstructed audio. It consists of a waveform-based discriminator and an STFT-based discriminator. The waveform-based discriminator inputs the original audio waveform, two times down-sampled waveform and four times down-sampled waveform, respectively, extracting waveform features from different scales. The STFT-based discriminator first performs STFT transform on the waveform, and then performs 2-D convolution on the time-frequency spectrum to extract features.

Based on the Theorem 2, we propose a two-stage training approach for end-to-end speech codec learning to achieve joint compression and enhancement, as shown in Fig. 1. The generator is an encoder-decoder pair composed of an encoder $E$, a quantizer and a decoder $G$. From Theorem 2, an optimal training procedure contains two steps:

i) In the first step, the encoder-decoder pair $(E_{d},G_{d})$ is trained using distortion loss only, e.g., with the multi-scale spectral reconstruction loss $L_{dis}(\hat{X},X)$ as (6), without using any discriminator.

ii) In the second step, the encoder $E_{d}$ and the vector quantizer learned in the first stage are frozen, and a new decoder $G_{p}$ is trained by a combination of distortion loss and adversarial loss.

Though in theory the perceptual decoder can be trained only by adversarial loss, intensive experiments show that a combination of distortion loss and adversarial loss can yield better performance. This is because it is difficult to train a decoder to reconstruct the signal from compressed representation only using adversarial loss. In the second stage, the overall loss for the generator consists of three components

where $L_{dis}$ is the multi-scale spectral reconstruction loss (6). Similar to [8], we consider 4 individual discriminators, $\mathcal{D}_{k}$ for $k\in\{0,\cdots,3\}$, with the index $k=0$ denoting the STFT-based discriminator and $k\in\{1,\cdots,3\}$ denoting waveform-based discriminator for different resolutions. With these discriminators, $L_{adv}$ is an adversarial loss for the generator as

where $T_{k}$ is the number of logits at the output of the $k$-th discriminator along the time dimension.

$L_{feat}$ is a feature loss computed based on the discrepancy between the internal layer outputs of the discriminators, i.e., the difference on the features of the discriminators, as

where $K$ is the number of discriminators, $L$ is the number of internal layers of each discriminator. $\mathcal{D}_{k,l}$ is the output of $l$-th layer of the $k$-th discriminator. The discriminator is trained by a loss as

In the first stage of training, we aim to minimize the distortion by the multi-scale spectral reconstruction loss and achieve maximal noise suppression in the encoder $E_{d}$. In the second stage of training, we use generative adversarial training so that the decoder can generate high perception quality audio. In order to avoid excessive distortion elevation in adversarial training of the perceptual decoder, the distortion loss is also used in combination with the adversarial loss to achieve a satisfactory performance in practice.

We train our model on various noise conditions using clean and noisy speech datasets. For clean speech, we use the train-clean-100 and train-clean-360 of LibriTTS [15] as clean speech dataset, which has a sampling rate of 24 kHz. LibriTTS corpus is constructed from audiobooks and includes the speech data from 2456 different speakers. In data pre-processing, we filter out the samples which are resampled from 16 kHz to 24 kHz in the LibriTTS, resulting in 147k clean samples. We generate noisy speech by mixing natural noise with clean speech at a sampling rate of 24 kHz. We use Freesound [16] as natural environmental noise. Freesound contains music, mechanical noise, natural atmosphere and many other kinds of audio. We choose unlicensed audio without human voice as noise data. Different noise segments are inserted into clean speech every three seconds. The signal-to-noise ratio (SNR) is uniformly distributed between 0 dB and 15 dB. Finally, we train our model on 147k clean speech and 33k noisy speech.

To enable bit-rate scalability without retraining the model, Residual Vector Quantizer[20] and Quantizer Dropout[8] are used in the training of all our model. We cascade 24 vector quantizers, each with a codebook size of 1024, and then quantize the residuals iteratively. Under this setting, our model can support up to 18 kbps target rate. By dropping out different numbers of cascaded VQs in training, our model can support a target rate between 3 kbps and 18 kbps.

We train our model SEStream with Adam for 500 epochs, with a learning rate of $10^{-4}$ and batch size of 64. For fair comparison, we train SoundStream and SEStream with the same dataset, the same network structure, and the same hyper parameters in the loss function. The denoising flag of SoundStream is turned on 50% of the time during training. The input of the model is a randomly sampled 360 millisecond audio waveform. The ground truth is the corresponding clean speech. The peak of each fragment is normalized to 0.95 and multiplied by a gain between 0.3 and 1.0.

In this work, we use Virtual Speech Quality Objective Listener (ViSQOL) [18] as the evaluation metric for objective quality. It uses the similarity of the time-frequency spectrum between the reference audio and the test audio to model the perceptual quality of human speech, which is highly correlated with distortion loss. In addition, we use other objective evaluation metrics, including Short-Time Objective Intelligibility (STOI) [21], CSIG, CBAK and COVL, to evaluate the quality of reconstructed audio. STOI is used to evaluate the intelligibility of noisy speech, which is masked in time domain or weighted in frequency domain after short time Fourier transform. CSIG predicts the mean opinion score (MOS) of speech distortion, CBAK predicts the MOS of background noise, whilst COVL predicts the MOS of overall processed speech quality. STOI ranges from 0.0 to 1.0 and other metrics range from 0.0 to 5.0. For all these evaluation metrics, higher value represents better performance.

For subjective quality, we use Multi-Stimulus Test with Hidden Reference and Anchor (MUSHRA) [17] to compare our model and other representative audio codecs. MUSHRA is a double-blind listening test, in which test listeners rate the relative quality of the output of each codec against a hidden reference. The reference audio is clean speech, together with decoded results by different codecs in the presence of various background noise, are rated by the listeners. The scores are ranged from 0 to 100 and a higher score reflects a higher subjective quality.

In our experiment, we randomly select 200 clean samples from the test set of LibriTTS as clean test set, which is disjoint with the training set. To evaluate the denoising effect, natural noise with SNRs of 0 dB, 5 dB, 10 dB and 15 dB were added to the 200 samples respectively to generate the noisy test set. The synthesis criterion is the same as mentioned in Section IV-A. Each audio in the test set lasts 2 to 4 seconds.

To evaluate the objective quality, we compare the proposed SEStream with SoundStream in terms of ViSQOL test. Fig. 2 shows the ViSQOL scores evaluated on clean test set and several noise test sets with different SNRs. For SoundStream, the denoising flag is on for the conditions of noisy speech to realize denoising. By controlling the number of cascaded VQs during inference, we show the objective quality scores at target bit-rates from 3 kbps to 18 kbps. As shown in Figure 2, the proposed SEStream can achieve consistently better ViSQOL score than SoundStream at all the tested bit-rates, increasing the ViSQOL score by up to 0.089 (at 3 kbps bit-rate and clean set) and with an average improvement of about 0.057 on the test cases.

To further compare the denoising performance of the two methods, we select the target bit-rate of 6 kbps and use STOI, CSIG, CBAK, and COVL to evaluate the denoising quality under different noise levels. The result is shown in Table I. It can be seen that in the case of relatively large background noise, especially when the SNR is 0 dB, our proposed method has better performance than SoundStream in most cases in terms of the evaluation metrics. This indicates SEStream can generate speech with higher intelligibility and less background noise. In the cases of relatively low noise, e.g., with less background noise, the advantage of SEStream vanishes in terms of the STOI, CSIG, CBAK, and COVL scores, but SEStream still achieves better ViSQOL score.

Fig. 6 compares the time-frequency spectrum of two examples generated by SoundStream and SEStream under 0 dB SNR and 10 dB SNR. Both SoundStream and SEStream can reconstruct the speech from noisy speech that contains severe environmental noise. By comparing the spectrum of each sample longitudinally, it can be seen that SEStream can suppress the background noise better than SoundStream, as shown in the blue boxes. In the green boxes, SEStream retains more information of the clean spectrum. This phenomenon is relatively obvious in the case of severe noise.

We further compare our method with other representative audio codecs, including Lyra, EVS, and OPUS at various bit-rates. We compare the methods on 20 examples for each case of low bit-rate and medium bit-rate, which contain clean test set and noisy test set with SNRs ranging from 0 dB to 15 dB uniformly. The subjective quality score based on MUSHRA is shown in Fig. 3. Details are shown in Fig. 4. It can be seen that at low and medium bit-rates, our method achieves a performance comparable with SoundStream in terms of subjective quality score with a slight advantage. Using a low bit-rate, e.g., 3 kbps, SoundStream and SEStream can achieve scores close to that of EVS at a higher bit-rate, e.g., 9.6 kbps, while outperforming OPUS at a higher bit-rate of 12 kbps.

As can be seen from Fig. 4, the perceptual quality of our proposed method on clean speech is similar to that of SoundStream, and the advantage is in noisy speech at a low bit-rate. Compared with other codecs, it is obvious that the joint speech compression and enhancement has better perceptual quality on noisy speech.

Fig. 5 compares the subjective quality score in the cases of low bit-rate and high noise with 0 dB and 10 dB SNR. In these cases, the performance of traditional codecs deteriorates drastically. In comparison, the advantage of SoundStream and SEStream over other codecs is more conspicuous, reflecting the better performance of SoundStream and SEStream in the presence of noise.

Table II compares the codecs in terms of the objective metric ViSQOL. We select 200 examples for the clean test set and 200 examples for the noisy test set (with SNRs ranging from 0 dB to 15 dB). SEStream has better ViSQOL scores than SoundStream, indicating that the distortion of SEStream is lower. Although the denoised results of Lyra sound relatively natural, there exists more serious distortion in timbre, thus, with significantly worse ViSQOL scores.