Source-Aware Neural Speech Coding for Noisy Speech Compression

As a codec system, the model consists of an encoder that converts time-domain mixture frame $\bm{x}=\in\mathbb{R}^{N}$ into code vector $\bm{z}\in\mathbb{R}^{D}$: $\bm{z}\leftarrow{\mathcal{F}}_{\text{enc}}(\bm{x})$. To marry the source separation concept, we assume $K$ mask vectors $\bm{m}^{(k)}\in\mathbb{R}^{D}$ that can decompose the code vector into $K$ components: $\sum_{k=1}^{K}m^{(k)}_{d}=1$. Note that $m^{(k)}_{d}$ is the probability of $d$-th code value belonging to $k$-th source. In addition, we further assume that this probabilistic code assignments to the sources are actually determined by one-hot vectors, so that the masking process assigns each code value to only one source, i.e.,

$$m_{d}^{(k)}=\left\{\begin{array}[]{rl}1&\text{ if }\operatornamewithlimits{\arg\max}_{j}m_{d}^{(j)}=k\\ 0&\text{ otherwise.}\end{array}\right.$$

(1)

The TasNet models estimate similar masking vectors via a separate neural network module, which led to the state-of-the-art separation performance. In there, the estimated mask values are indeed somewhat drastically distributed to either near zero or one, making the masked code vectors nearly orthogonal from each other. However, the sigmoid-activated masks in the TasNet architecture and our probabilistic masks do not specifically assume a hard assignment.

From now on, we assume orthogonal code vectors per source as a result of hard masking, i.e., $\bm{z}^{(1)}\!\perp\!\bm{z}^{(2)}\!\perp\!\cdots\!\perp\!\bm{z}^{(K)}$, where $\bm{z}^{(k)}=\bm{m}^{(k)}\odot\bm{z}$, $k$-th source’s code vector defined by the Hadamard product $\odot$ between the mask and the mixture code.

The proposed orthogonality leads us to a meaningful structural innovation. Instead of estimating the mask vector for every input frame, we can use structured masking vectors that force the code values to be grouped into $K$ exclusive and consecutive subsets. For example, for a two-source case with $D=8$, $\bm{m}^{(1)}=[1,1,1,1,0,0,0,0]^{\top}$. Hence, we can safely discard the masked-out elements (the latter four elements), by defining its truncated version as $\mathbf{z}^{(k)}\in\mathbb{R}^{D/K}$. The concatenation of the truncated code vectors determines the final code vector: $\bm{z}=[{\mathbf{z}^{(1)}}^{\top},{\mathbf{z}^{(2)}}^{\top},\cdots,{\mathbf{z}^{(K)}}^{\top}]^{\top}$.

In practice, we implement the encoder as a 1-d convolutional neural network (CNN) whose output is a $2L\times P$ matrix, where $2L$ is the number of output channels (see Figure 1, where $L=6$, $P=256$). We collect the first $L$ channels of this feature map (dark red bars) as our codes for speech, i.e., $\mathbf{z}^{(1)}$ corresponds to the vectorized version of the upper half of the feature map of size $L\times P$, or $LP=D/K$, where $D$ should be an integer multiple of $K$. The other half is for the noise source. Since the decoders are learned to predict individual sources, this implicit masking process can still work for source separation.

Quantization is a mapping process that replaces a continuous variable into its closest discrete representative. Since it is not a differentiable process, incorporating it into a neural network requires careful consideration. Soft-to-hard quantization showed successful performance both in image and speech compression models [2, 6, 7]. The idea is to formulate this cluster assignment process as a softmax classification during the feedforward process, which finds the nearest one among $M$ total representatives $\bm{\mu}_{m}$ for the given code vector $\bm{y}$ as follows:

	$\displaystyle d_{m}$	$\displaystyle={\mathcal{E}}(\bm{y}||\bm{\mu}_{m}),\quad\bm{p}=\text{Softmax}(-\alpha\bm{d}),$		(2)
	$\displaystyle\text{Testing: }\bar{\bm{y}}$	$\displaystyle=\bm{\mu}_{\operatornamewithlimits{\arg\max}_{m}p_{m}},\quad\text{Training: }\bar{\bm{y}}=\sum_{m=1}^{M}p_{m}\bm{\mu}_{m},$

where the algorithm first computes the Euclidean distance vector $\bm{d}$ against all the representatives (i.e., the cluster means), whose negative value works like a similarity score for the softmax function. Using the softmax result, the probability vector $\bm{p}$ of the cluster membership, we can construct the quantized code vector $\bar{\bm{y}}$: during test time, simply choosing the closest one will do a proper quantization. However, since the $\operatornamewithlimits{\arg\max}$ operation is not differentiable, for training, we do a convex combination of the cluster centroids to represent the quantized code, as a differentiable surrogate of the hard assignment process. The discrepancy between training and testing is reduced by controlling the scaling hyperparameter $\alpha$, which makes the softmax probabilities more drastic once it is large enough (e.g., a one-hot vector in the extreme case). Note that it also learns the cluster centroids $\bm{\mu}$ as a part of the learnable network parameters rather than employing a separate clustering process to define them.

In previous work the quantization has been on scalar variables, i.e., $\bm{y}\in\mathbb{R}^{1}$ [6, 7, 8]. In this work, the soft-to-hard quantization performs vector quantization (VQ). We denote the CNN encoder output by $\bm{Z}\in\mathbb{R}^{2L\times P}$, which consists of $K=2$ code blocks: $\bm{Z}=[\bm{Z}^{(1)};\bm{Z}^{(2)}]$. Then, each code vector for quantization is defined by the $p$-th feature spanning over $L$ channels: $\bm{y}=\bm{Z}_{1:L,p}^{(k)}$, having $L=6$ as the VQ dimension in our case.

The theoretical lower bound of the bitrate, as a result of Huffman coding, can be defined by the entropy of the quantized codes. The frequency of the cluster means defines the entropy of the source-specific codes: ${\mathcal{H}}\big{(}\bm{\mu}^{(k)}\big{)}=-\sum_{m=1}^{M}q_{m}^{(k)}\log q_{m}^{(k)}$, where $q_{m}^{(k)}$ denotes the frequency of $m$-th mean for the $k$-th source. Meanwhile, the entropy for the code of the mixture signal is smaller than or equal to the sum of the entropy of all sources: ${\mathcal{H}}(\ddot{\bm{\mu}})\leq\sum_{k=1}^{K}{\mathcal{H}}\big{(}\bm{\mu}^{(k)}\big{)}$, where $\ddot{\bm{\mu}}$ is the set of quantization vector centroids learned in a source-agnostic way, i.e., directly from the mixture signals. Therefore, in theory, SANAC cannot achieve a better coding gain than a codec that works directly on the mixture.

However, SANAC can still benefit from the source-wise coding, especially by exploiting the perceptual factors. Our main assumption in this work is that the perceptual importance differs from source-by-source, leading to a coding system that can assign different bitrates to different sources. For noisy speech, for example, we will try to assign more bits to the speech source. Consequently, although the user eventually listens to the recovered mixture of speech and noise (a) the perceptual quality of the speech component is relatively higher (b) the codec can achieve a better coding gain if the noise source’ statistical characteristics is robust to low bitrates.

Our argument is based on the codec’s ability to control the entropy of the source-specific codes. In SANAC, we adopt the entropy control mechanism proposed in [2], but by setting up a per-source loss between the target $\xi^{(k)}$ and the actual entropy values: $\big{(}\xi^{(k)}-{\mathcal{H}}\big{(}\bm{\mu}^{(k)}\big{)}\big{)}^{2}$. While this loss does not guarantee the exact bitrate during the test time, in practice, we observe that the actual bitrate is not significantly different from the target.

The source-specific truncated codes, after the quantization, $\bar{\bm{Z}}^{(k)}$, are fed to the decoder part of the network. The decoder function works similarly to Conv-TasNet [18] in that the decoder runs $K$ times to predict $K$ individual source reconstructions from $K$ sourse-specific feature maps as the input. However, SANAC’s decoding is different from Conv-TasNet’s as the decoder input is the quantized codes. In addition, our model cares about the quality of the recovered mixture, not only the separation quality.

Our training loss considers all these goals, consisting of the main mean squared error (MSE)-based reconstruction term and the entropy control terms. More specifically, for the noisy speech case $\bm{x}=\bm{s}+\bm{n}$ ($k=1$ for speech and $k=2$ for noise), the MSE loss is for the speech source reconstruction $\hat{\bm{s}}$ and the mixture reconstruction $\hat{\bm{x}}$, while the noise source reconstruction $\hat{\bm{n}}$ is implied in there. We regularize the total entropy as well as the ratio between the two source-wise entropy values:

$$\begin{split}{\mathcal{L}}&=\lambda_{\text{MSE}}\big{(}{\mathcal{E}}_{\text{MSE}}(\bm{s}||\hat{\bm{s}})+{\mathcal{E}}_{\text{MSE}}(\bm{x}||\hat{\bm{x}})\big{)}\\ &+\lambda_{\text{EntTot}}\Big{(}\xi-{\mathcal{H}}\big{(}\bm{\mu}^{(1)}\big{)}-{\mathcal{H}}\big{(}\bm{\mu}^{(2)}\big{)}\Big{)}^{2}\\ &+\lambda_{\text{Ratio}}\left(\psi-\frac{{\mathcal{H}}\big{(}\bm{\mu}^{(1)}\big{)}}{{\mathcal{H}}\big{(}\bm{\mu}^{(2)}\big{)}}\right)^{2},\end{split}$$

(3)

where $\xi$ and $\psi$ are the target total entropy and the target ratio, respectively.

500 and 50 utterances are randomly selected from training and test set of TIMIT corpus [19]. We generate 10 contaminated samples out of each clean utterance by adding 10 different non-stationary background sources, {bird singing, casino, cicadas, typing, chip eating, frogs, jungle, machine gun, motorcycle, ocean} used in [20]. Every contaminated speech waveform was segmented into frames of 512 samples (32ms), with overlap of 64 samples. We apply a Hann window of size 128 samples only to the overlapping regions. Since there are $16000/448$ frames per second and each frame produces $P$ code vectors for VQ, for the entropy of a source-specific codebook $\xi$, the bitrate is $16000P\xi/488$, e.g., 9.14kbps when $P=256$ and $\xi=1$.

Adam optimizer with an initial learning rate 0.0001 trains the models [21]. Both SANAC and the baseline are trained in three stages. Every jump to next stage is triggered when the validation loss stops improving in 3 consecutive epochs. We stop the training updates after validation loss does not improve for 20 epochs.

• •

  Stage 1: For the first three epochs, the model trains the encoder to separate the input into the speech and background sources, that are represented by the two orthogonal code vectors. No quantization is involved in yet, but this stage better initializes the parameters for the quantization process. The encoder consists of a few bottleneck blocks that are commonly used in ResNet [22]. With the bottleneck structure, the input and output feature maps can be connected via an identity shortcut with less filters to learn in between. The encoder module also employs a 1-d convolution layer to downsample the feature map from 512 to 256, followed by another bottleneck block and a channel changing layer to yield two sets of code maps of $6\times 256$ each. The learned source-specific feature maps are fed directly to the channel changer with no quantization, followed by an upsampler and the final decoder. In terms of upsampling, we interlace two adjacent channels into one, doubling the number of features up to 512 while halving the number of channels as introduced in [23], and then adopted for neural speech coding [6, 7].

• •

  Stage 2: In this stage, the model starts to quantize the encoder output using the soft-to-hard VQ mechanism. The VQ is done with $M=128$ centroids. We set the scale of the softmax function $\alpha=10$, and increase it exponentially until it reaches 500 to gradually introduce more hardness to the softmax function. Meanwhile, the other modules in the network are also updated accordingly to absorb the quantization error. As the model stabilizes, we introduce entropy control terms into the loss function by setting up the regularization weights $\lambda_{\text{EntTot}}=1/5$ and $\lambda_{\text{Ratio}}=1/60$. We set $\xi$ to be 1, 2, and 3, which correspond to three bitrates 9.14, 18.29, and 27.43kbps. The target ratio between speech and noise bitrates is trainable, which we set to be $\psi=3$ at the beginning.

• •

  Stage 3: Finally, the feature maps go through a few more ResNet blocks that strengthen the separation (the “Separator” module in Figure 1), followed by the decoder that runs twice for both sources. The “Separator” and “Decoder” consist of two bottleneck blocks with different input channels, 60 and 30, respectively. “Decoder” needs an additional channel changer.

• •

  The Baseline: Our baseline system shown in Fig. 2 is similar to the proposed architecture, except that it discards the first separation before coding, meaning that the codec system is not aware of the different sources. Hence, the baseline only generates one type of codes representing the mixture signal without any control of the entropy ratio in the loss function. We still use the two-headed decoder architecture to benefit from the two reconstruction loss terms.

We first show the average MOS-LQO (PESQ) score for SANAC, baseline and two standard codecs (AMR-WB [9] and Opus [24]), at 3 different target bitrates in Table 1. We evaluate the reconstruction of the input noisy speech. For each decoded sample, both the uncompressed mixture input and its corresponding clean speech source are considered as the reference, yielding the left and right hand scores, respectively. Hence, the right hand score measures the speech enhancement performance, while the left hand score indicates the level of overall reconstruction quality. Since our neural codecs are not designed to degrade noise source during compression, the right hand scores are lower than those from standard codecs where the added noise is somewhat suppressed. In terms of the left hand score, neural codecs outperform those counterparts in most cases. It is known, however, that PESQ is not well defined for the perceptual evaluation between noisy speech signals, which to a certain degree explains the very similar left hand scores from the baseline and SANAC.

Additionally, the baseline and proposed SANAC are evaluated based on the scale-invariant signal-to-distortion ratio (SiSDR) [25] and short-time objective intelligibility (STOI) [26]. The clean speech is used as the reference when computing STOI. Similar comparison results are observed for both STOI (Fig.3(a)) and SiSDR (Fig.3(b)): overall, SANAC achieves higher scores in most cases. For STOI, we find that the performance gap becomes more prominent in the lower bitrates, thanks to the more speech-oriented setup we used for SANAC. With respect to SiSDR, SANAC shows better mixture reconstruction performance, too. The margin against the baseline is more noticeable when the SNR level of the input mixture is higher.

Finally, we conduct a subjective listening test with eight audio experts participated. The A/B tests are designed to compare our own baseline and SANAC; as a controlled experiment, where all other system aspects are fixed, the test focuses on the impact of the source-aware coding scheme. The test consists of three sessions with different bitrates ($\sim 9$, $\sim 18$, and $\sim 27$kbps). With the uncompressed mixture signal as the reference, listeners are asked to designate the sample from 2 competing systems that sounds more similar to the reference. Results are presented in Table 2, where the preference rate denotes how likely an average listener prefers SANAC. We also report the standard deviation to measure the variation of the preference rate. In all cases, SANAC is favored by most listeners. As the bitrate gets lower, the listener tends to have a higher and more determined preference rate for SANAC over the baseline.