Towards Error-Resilient Neural Speech Coding

We take the low-latency neural speech coding network TFNet in our previous work [15] as the backbone. It takes the time-frequency spectrum with a 20ms window and a 5ms hop length as input with power-law compressed normalization on the magnitude. The encoder and decoder are composed of causal convolutions and deconvolutions for capturing frequency dependencies and two kinds of causal temporal filtering modules in-between them, i.e. dilated temporal convolution module (TCM) and group-wise gated recurrent unit (G-GRU) for capturing temporal dependencies. The two temporal filtering modules are organized in an interleaved way to efficiently extract both local and long-term temporal dependencies. For vector quantization, the latent embeddings $X^{S}\in\mathbb{R}^{T\times 1\times C}$ are split into $N$ groups. For each group, it uses an independent codebook containing $S$ codewords. At the target bitrate of 6kbps, we combine 4 overlapped frames (corresponding to 20ms new data) into a vector for quantization.

The FD-PLC module is composed of two kinds of causal modules stacked together, i.e, a group-wise temporal self-attention block (G-TSA) and a TCM module. The G-TSA block is similar to the multi-head self-attention block (MHSA) in transformer [22] and we turn it into a causal operation along the temporal dimension using a window of $N$ frames, i.e. each frame only has access to $N$ past frames without any look-ahead. It captures different attentiveness to different frames thus provides temporal adaptation to the content and the loss/non-loss properties. As shown in Figure 2, we use a pre-norm residual unit similar to that in [23] for the G-TSA block, with two convolutional layers with a kernel size of $1\times 1$ to reduce and increase the dimensions. The TCM block is similar to that used in the backbone network but with layer normalization preceding the convolutions and GELU as the activation function. Several TCM blocks with increasing dilation rates are stacked together to form a large TCM module with a large receptive field. We use G-TSA to extract the local correlations which is more important to the current frame at a fine granularity and the TCM is used to aggregate the G-TSA output features to catch long-term dependencies at a coarse granularity.

Adversarial training is widely used for restoration tasks to achieve good reconstruction audio quality. We employ two frequency-domain discriminators with discrimination capability at different granularity. They both take the magnitude spectrum, with a window length of $W=20ms$ and a hop length of $H=10ms$ as the input. The first is a segment-level discriminator used to judge the overall quality of an audio clip. It consists of four convolutional layers with a kernel size of (3, 3) and a stride of (2, 2), followed by normalization and Leaky Relu activation. The number of channels is progressively increased to 64 with the depth of network. Finally, a fully connected layer is used to aggregate all channels into one and we do an average pooling on both the (down-sampled) time and frequency dimensions to get just one logit at the output. The second discriminator targets at frame-level discrimination. For this purpose, we use a kernel size of (2, 5) with a stride of (1, 2) in convolution blocks so as to down-sample frequency bins while keeping the temporal resolution. Only frequency-dimension average pooling is used to obtain a 1-dimensional logits as the output. In both discriminators, we use spectral normalization [24] for the first convolution block and instance normalization [25] for others for stable training. Sigmoid activation is used on the output.

We take the Binary Cross Entropy (BCE) loss generally used for GAN [26], i.e.

$$\max\limits_{D}L_{adv}\left(D\right)=\mathbb{E}_{x}\left[log\left(D\left(x\right)\right)\right]+\mathbb{E}_{x}\left[log\left(1-D\left(G\left(x\right)\right)\right)\right],$$

(1)

$$\min\limits_{G}L_{adv}\left(G\right)=\mathbb{E}_{x}\left[log\left(1-D\left(G\left(x\right)\right)\right)\right].$$

(2)

We also use a feature matching loss [27] as an additional constraint for the generator, which is given by

$$L_{fm}(G)=\mathbb{E}_{x}\left[\sum\limits_{i=1}^{T}\dfrac{1}{N_{i}}\left\|D^{i}(x)-D^{i}(G(x))\right\|_{1}\right],\vspace{-0.1cm}$$

(3)

where $T$ denotes the number of layers in the discriminator that are used for the feature loss. $D^{i}$ and $N^{i}$ denote the features and feature sizes in the $i$-th layer of the discriminator, respectively.

The neural coding network works as the generator in adversarial training. We use a combination of several loss terms in guiding it towards a good decoded audio quality as follows

$$\begin{split}L_{G}=\lambda_{plc}L_{plc}(G)+\lambda_{bin}L_{bin}(G)+\lambda_{mel}L_{mel}(G)\\ +\lambda_{adv}L_{adv}(G)+\lambda_{fm}L_{fm}(G),\end{split}$$

(4)

where the scalars $\lambda_{plc},\lambda_{bin},\lambda_{mel},\lambda_{adv},\lambda_{fm}$ are weights to balance different terms and set by $\lambda_{plc}=1,\lambda_{bin}=1,\lambda_{mel}=0.25,\lambda_{adv}=1e^{-3},\lambda_{fm}=2e^{-5}$ in our implementation. $L_{adv}(G)$ and $L_{fm}(G)$ are adversarial and feature loss terms in Eq. 2 and 3.

The first term $L_{plc}(G)$ adds the supervision on guiding the FD-PLC module to recover the lost frame from $X^{Q}$. During training, we add proportional lost frames and non-lost frames to make the data balance. Let $S$ denote this set and the loss term is given by

$$\vspace{-0.1cm}L_{plc}\left(G\right)=\mathbb{E}_{x}\left[\dfrac{1}{\left|S\right|}\sum\limits_{t\in{S}}\left\|X^{Q}(t)-\hat{X}^{Q}(t)\right\|_{1}\right],\vspace{-0.1cm}$$

(5)

where $X^{Q}(t)$ and $\hat{X}^{Q}(t)$ are $t$-th frame of $X^{Q}$ and $\hat{X}^{Q}$, respectively. $\left|S\right|$ denotes the number of frames in $S$. The L1 loss is used to measure the distance between $\hat{X}^{Q}$ and $X^{Q}$. In our implementation, we take the error-free baseline as the pretrained model for encoder, the codebook and $X^{Q}$ and train the FD-PLC and the decoder jointly. More ablation studies on the training algorithm could be found in section 4.2.

The second $L_{bin}(G)$ and the third terms $L_{mel}(G)$ are quality terms at each frequency-bin and mel-band. As shown in [28] that time-frequency distribution can be effectively captured by jointly optimizing multi-resolution spectrum and adversarial loss functions, we use $L_{mel}(G)$ at multiple resolutions to achieve a high perceptual quality, which is given by

$$L_{mel}\left(G\right)=\mathbb{E}_{x,\hat{x}}\left[\dfrac{1}{R}\sum\limits_{r=1}^{R}\left\|\phi^{r}(x)-\phi^{r}(\hat{x})\right\|_{1}\right].$$

(6)

$\phi^{r}(\cdot)$ is the $r$-th mel-scale band. To achieve high-fidelity, we also use the frequency-bin wise quality term $L_{bin}(G)$ by

$$L_{bin}\left(G\right)=\mathbb{E}_{x,\hat{x}}\left[\left\|\psi(x)-\psi(\hat{x})\right\|^{2}_{2}\right].\vspace{-0.1cm}$$

(7)

$\psi(x)$ is the power-law compressed magnitude of $x$. We use L2 distance metric here.

We use the 16khz raw clean speech data from Deep Noise Suppression Challenge at ICASSP 2021[29]. It includes multilingual speech and emotional clips. For packet loss, we simulate with a random loss rate from $\{10\%,20\%,30\%,40\%,50\%\}$, whose maximum burst loss length is $\{60,80,120,160,220\}$ milliseconds, respectively. Besides, we simulate WLAN packet loss pattern with three-state Markov models[30]. For training, we synthesized 600 hours of data, 100 hours for each loss rate category. For testing, we use the blind test set from Audio Deep Packet Loss Concealment Challenge at INTERSPEECH 2022[31]. This test set uses real packet loss traces captured in real Microsoft Teams calls. Its maximum burst loss length is up to 1000 milliseconds. We also use another test sets with synthetic traces with a random loss rate from $10\%$ to $50\%$ for a deep investigation.

For training, we use the Adam optimizer[32] for both generator and discriminator. A learning rate of $4e^{-4}$ is employed for generator while for discriminator the learning rate is decayed by $0.999$ for every epoch with an initial rate of $4e^{-4}$. The temporal window size $N$ of the FD-PLC module is set to 32 frames. The network is trained for 60 epochs with a batch size of 200.

We measure signal quality at various packet loss rates while the bitrate is fixed to 6kbps. For a generative task, speech distortion, perceptual quality and naturalness are important factor to measure the quality. To this end, we employ mel cepstral distortion (MCD; in dBs) [33] to measure speech disortion which focuses on perceptually relevant speech characteristics of the short-term speech spectrum. To measure the overall quality, we use the PLCMOS, the evaluation tool for the PLC Challenge at INTERSPEECH 2022 [31] and NISQA[34]. NISQA is a deep learning framework for speech quality evaluation covering several sources of degradation, including packet losses and audio compression. We use NISQA-MOS to evaluate the overall quality, NISQA-Discontinuity for the audio continuity, and NISQA-TTS for the naturalness of synthesized speech. Our experiments find that they reasonably match our perception when using the same neural audio codec.

We compare the proposed FD-PLC scheme with several baselines to verify its effectiveness, the error-free baseline, the error-resilient baseline, Baseline-GAN, Atten-GAN, Post-PLC and the FD-PLC as shown in Fig. 3 and Table. 1. The error-free and error-resilience baselines are trained without and with consideration on packet losses. The Baseline-GAN differs from the Error-resilient baseline by adding adversarial training. The Atten-GAN further introduces the FD-PLC module into Baseline-GAN so that it uses the same network as the proposed FD-PLC. It differs from the proposed FD-PLC only in that no $L_{plc}(G)$ loss is used. What’s more, the Post-PLC moves the PLC module to waveform domain, acting as a post-processing of Baseline-GAN. The Post-PLC module takes a U-Net structure with causal convolutional encoder and decoder and skip connections between them. TCM and TSA blocks similar to that in FD-PLC are used in-between the encoder and decoder. It is designed with the similar model size as the FD-PLC module but with much larger receptive field.

Fig.3 and 3 show the MCD and NISQA-TTS comparison on synthetic traces at different packet loss rates. They measure the signal distortion and the naturalness of the concealed speech, respectively. It can be seen that the Error-free baseline works well when there is no packet losses but the quality drops sharply when there is packet loss, showing its sensitiveness to channel noises. Other error-resilient schemes surpass the error-free baseline not only in loss scenario but also in the non-loss scenario, showing their stronger robustness and restoration capability against the error-free counterpart. Among these error-resilient schemes, except for Error-resilient baseline and Post-PLC, others are comparable in terms of MCD. The proposed FD-PLC scheme achieves relatively lower MCD and highest NISQA-TTS MOS scores consistently over all loss rates. The Post-PLC has much larger MCD, indicating large difference on the signal level. The Atten-GAN also shows promising results, especially at high packet loss rate. Compared with the model without attention module, we observe that more frequency bins are restored since more attention has been put on the lost part. But the content it generates is not as coherent as the proposed scheme. Similar results can be found in Table 1 evaluated on the blind test set with real traces. The proposed FD-PLC scheme achieves best NISQA-TTS MOS and comparable results on other metrics with the top one. It also outperforms the Post-PLC on the signal fidelity and perceptual quality.

Here we investigate several training schemes, the end-to-end training, multi-stage training and the proposed. In multi-stage training, the encoder and codebook are pretrained as that in proposed but it trains the FD-PLC module first in decoding and finetunes the decoder after that. As Table. 2 shows, end-to-end training is the worst. This is because the model will be confused by the PLC task when the target quantized features are still at a preliminary stage. The proposed joint training of the FD-PLC and the decoder provides more room for the trade-off between packet loss recovery and quantization loss recovery.