LIMUSE: LIGHTWEIGHT MULTI-MODAL SPEAKER EXTRACTION

The audio encoder encodes the multi-channel audio mixture into spectrum-like representation in the time domain, and the audio decoder is used to transform the masked encoder representation to the target waveform. The encoder and the decoder follow the same structure in [27], which uses a 1-dimensional convolutional layer and a 1-dimensional transpose convolutional layer, respectively.

Formulating the audio signals in the time domain has some advantages over the traditional time-frequency representations obtained by Short-Time Fourier Transform (STFT), including trainable weights, finer coding granularity and avoidance of phase reconstruction [7].

The voiceprint models the characteristics of the target speaker and serves as prior knowledge to help the model filter out the target speaker’s speech from the mixture. The voiceprint encoder takes the speaker embeddings of the target speaker as input. In this paper, we use an off-the-shelf speaker diarization toolkit pyannote [28] to extract the speaker embeddings $\mathbf{s_{0}}\in\mathbb{R}^{U}$, where $U$ denotes the length of the speaker embedding feature dimension.

In order to align with the time dimension of the mixed audio stream, we stack the speaker embeddings along the time dimension to form $\tilde{S}=[\dots,s_{t},\dots]\in\mathbb{R}^{U\times T}$. The voiceprint encoder uses a fully connected (FC) layer to map the feature dimension of the voiceprint to obtain $\mathbf{S}\in\mathbb{R}^{N\times T}$, which is the same shape as the representation of the audio mixture stream.

We obtain face embeddings of the target speaker from videos using the network in [29], which extracts speech-related visual features from visual inputs explicitly by the adversarially disentangled method. The face embedding extraction network first learns joint visual and auditory representation from clean audio and video pairs to establish the correlation between the audio features and the visual features, and then disentangles the speech-related visual features from the joint audio-visual representation with the adversarial training method.

The pretrained face embedding extraction network is trained on LRW dataset [30] and MS-Celeb-1M dataset [31] to obtain $\mathbf{\tilde{V}}\in\mathbb{R}^{D\times E}$, where $D$ denotes the feature dimension of the face embedding and $E$ denotes the number of frames in each video stream. The visual encoder then maps the extracted face embeddings to match the feature dimension of the encoded audio features and obtain $\mathbf{v_{0}}\in\mathbb{R}^{N\times E}$. Moreover, since the time resolution of the video stream and the audio stream is different, we upsample the face embeddings along the temporal dimension to synchronize the audio stream and the video stream by nearest neighbor interpolation to obtain $\mathbf{V}\in\mathbb{R}^{N\times T}$. The face embeddings are then used to model the temporal synchronization and interaction between visemes and speech [32].

In this paper, we simply use a fully connected layer to map the face embedding rather than depth-wise separable convolution adopted by AV-ConvTasnet [8] or LSTM [33] adopted by AVMS [11].

In order to reduce the number of parameters, we apply Group Communication (GC) [18] to the separation modules. Given a feature $\mathbf{H}\in\mathbb{R}^{C}$, it can be decomposed into $K$ groups of low-dimensional feature vectors $\{\mathbf{g}^{i}\}_{i=1}^{K}$. A GC module is applied to capture the inter-group dependencies.

We use the Transform-Average-Concatenate (TAC) module [34] for the GC module. It is similar to the squeeze-and-excitation design paradigm [35], where there is a residual connection between a pool of features and a global-averaged feature which is generated to summarize the input feature pool. For ${\{\mathbf{g}^{i}\}}_{i=1}^{K}$, a fully-connected layer followed by PReLU activation [36] is used for the transformation step to generate $\mathbf{f}^{i}$. Then all $\mathbf{f}^{i}$s are averaged and passed through another fully connected layer with PReLU for the average step to get $\hat{\mathbf{f}}^{i}$. Finally, $\hat{\mathbf{f}}^{i}$ is concatenated with $\mathbf{f}^{i}$ and passed to the third fully connected layer with PReLU to generate $\hat{\mathbf{g}}^{i}$. A residual connection is added between $\mathbf{g}^{i}$ and $\hat{\mathbf{g}}^{i}$. It is demonstrated in [18] that using TAC for the Group Communication module achieves better performance and lower computational costs than other network structures, such as BLSTM and self-attention.

In this paper, Group Communication is applied before each Temporal Convolutional Network (TCN) block [27], which forms a GC-equipped TCN block, as is shown in Figure 1 (B). Each TCN block consists of a depth-wise separable convolution, PReLU activation [36] and a layer normalization operation. By introducing Group Communication to the TCN block, we can reduce the model width for the TCN block by $K$ (number of groups) times when there is no overlapping within the groups.

On one hand, speech has a rich temporal structure over multiple time scales [37]. On the other hand, shorter kernel length of the convolution in the audio encoder leads to better separation performance and higher model complexity at the same time. Motivated by the above two characteristics, Context Codec [18] is proposed to decrease sequence length, which is similar to a nonlinear downsampling step by compressing the context of feature vectors into a summary one and then decompress it back to the context. A Context Codec module typically contains a context encoder and a context decoder. In this paper, instead of using RNN [18], we propose to use the GC-equipped TCN network for the Context Codec module.

To be specific, given a sequence of feature $\mathbf{H}\in\mathbb{R}^{C\times L}$, where $C$ denotes the number of feature channels, the context encoder first splits it along the temporal dimension into blocks ${\{\mathbf{D}^{i}\}}_{i=1}^{T}\in\mathbb{R}^{C\times S}$ with an overlap ratio of 50%, where $T$ denotes the number of context blocks and $S$ denotes the context size as illustrated in Figure 1 (C). Then the context encoder applies a GC-equipped TCN network (Section 2.2.1) on each $\mathbf{D}^{i}$ to establish inter-group correlation and obtains ${\hat{\mathbf{D}}}^{i}$. At last, a mean operation is applied on the context dimension of $\hat{\mathbf{D}}$ to produce a sequence of summarization vectors ${\{\mathbf{p}^{i}\}}_{i=1}^{T}$ which represent the whole feature sequence. The compressed sequence $P\triangleq{\{\mathbf{p}^{i}\}}_{i=1}^{T}\in\mathbb{R}^{C\times T}$ can then be sent into feature processing module with $T\ll L$. The context decoder adds the output of feature processing module to each time step in $\mathbf{D}^{i}$ and applies another GC-equipped TCN network for the nonlinear transformation. Overlap-and-add is then applied on ${\{\hat{\mathbf{D}}^{i}\}}_{i=1}^{T}$ to form the sequence of feature of the original length ${\hat{\mathbf{H}}}\in\mathbb{R}^{C\times L}$. The implementation details can be found in our released source code.

In order to capture the long-term dependencies in the encoded audio features, the audio block is composed of several GC-equipped TCN blocks. In the audio block, we use two repeats of the TCN networks, and each repeat of the TCN network is composed of several TCN blocks. The dilation factor of depth-wise convolution increases exponentially by two in each repeat, which enlarges the receptive field of the module to capture the correlation in a larger time scale.

The output of the audio block is gathered with the encoded voiceprint and visual features through concatenation. Then we use a repeat of GC-equipped TCN blocks with an exponential increasing dilation factor of depth-wise convolution in the fusion block to process the fused features. The voiceprint and visual cues serve as auxiliary information to help the model extract the target speech from the mixture. A convolutional layer is followed by the fusion block to generate the mask of the target speaker.

For weight quantization, we propose to use the quantization function [38] to quantize weights to ultra-low bits for multi-modal speaker extraction tasks. The quantization function combines a linear combination of several simple functions to reformulate the ideal quantization operation and gradually approximates it by gradually increasing the temperature parameter during training.

Suppose that $x$ is the full-precision value to be quantized, $y$ is the quantized integer constrained to a predefined set $\gamma=\{\gamma_{1},\gamma_{2},\dots,\gamma_{n}\}$. For example, if we perform 3-bit weight quantization, then $y$ can be defined as $\gamma=\{-3,-2,-1,0,1,2,3\}$. The ideal quantization operation can be reformulated as a combination of unit step functions as follows:

$$y=\sum_{i=1}^{n}{s_{i}A(\beta x-b_{i})-o},$$

(1)

$$A(x)=\left\{\begin{array}[]{rcl}1&{x\geq 0,}\\ 0&{x<0,}\end{array}\right.$$

(2)

where $n=|\gamma|-1$, $\beta$ is the scale factor of the input. Eq 2 defines the unit step function, where $s_{i}={\gamma}_{i+1}-{\gamma}_{i}$ and $b_{i}$ is the bias for the unit step function. The global offset is defined as $o=\frac{1}{2}\sum_{i}^{n}s_{i}$.

However, the unit step function in Eq 2 is non-differentiable and thus can not be trained from end to end. The quantization function replaces the unit step function with differentiable sigmoid function in the training procedure as follows:

$$y=Q(x)=\alpha\left(\sum_{i=1}^{n}{s_{i}\sigma(T(\beta x-b_{i}))-o}\right),$$

(3)

$$\sigma(Tx)=\frac{1}{1+\exp(-Tx)},$$

(4)

where $T$ denotes the temperature parameter, $\alpha$ and $\beta$ are learnable scale factors of the output and input, respectively. While in the inference phase, the ideal quantization function in Eq 1 is used. Performance may suffer if different quantization functions are used throughout the training and inference phases. To narrow the gap, temperature $T$ is introduced to the sigmoid function. When $T$ is large, the gap between the sigmoid function and the unit step function is minimal, and vice versa. Starting from a small temperature and gradually increasing it, the quantized neural networks can be well learned.

In this paper, we perform layer-wise non-uniform quantization, in which weights from different layers use different quantization functions. We quantize all neural network modules except the audio decoder. We also do not quantize the PReLU activation and the layer normalization function because they may be sensitive to quantization and have a small number of parameters. To perform non-uniform quantization, we use K-means clustering on full-precision weights to obtain $n+1$ ranked cluster centers $c=\{c_{1},c_{2},\dots,c_{n+1}\}$ in the ascending order. The bias $b_{i}$ is initialized as $b_{i}=\frac{c_{i}+c_{i+1}}{2}$ and gets fixed during training.

We also quantize the activations of the corresponding quantized layers to make matrix multiplication run faster. For activation quantization, we use the most commonly used min-max linear quantization [39], which is defined as:

$$Q(x)=\operatorname{round}\left(\frac{x-\mathrm{X}_{\mbox{\tiny{min}}}}{s}\right),$$

(5)

$$s=\frac{\mathrm{X}_{\mbox{\tiny{max}}}-\mathrm{X}_{\mbox{\tiny{min}}}}{2^{p}-1},$$

(6)

where $\mathrm{X}$ is the tensor to be quantized, $x$ is the element in $\mathrm{X}$, $\mathrm{X_{min}}$ and $\mathrm{X_{max}}$ are the minimum and maximum elements, respectively. In this paper, the activation quantization bits are $8$, then $p=8$ in equation 6. Straight-Through Estimator (STE) [40] is used to backpropagate the gradients through quantized activations.

We investigate the effects of the Group Communication, Context Codec and the ultra-low bit quantization. The results are listed in Table 6.

The incorporation of Group Communication (GC) and Context Codec (CC) allows LiMuSE to maintain on par or even better performance under a suitable number of groups with a much smaller model size and less model complexity. We can see from the table that LiMuSE (GC) slightly performs better than LiMuSE (vanilla) when $K=16$ with only 0.37M parameters. LiMuSE (GC+CC) with $K=16$ further improves the performance with 0.04M extra parameters but reduces the MACs from 11.77G to 7.52G compared to LiMuSE (GC). It suggests that an appropriate number of groups may promote the inter-modal information exchange and achieve considerable improvements. When $K=32$, the number of parameters and the MACs of LiMuSE (GC) and LiMuSE (GC+CC) are further reduced, and the performance degradation of LiMuSE (GC+CC) is very limited compared to LiMuSE (vanilla). Intuitively, we can empirically conclude that the larger the value of K, the fewer the model parameters with some performance degradation in most cases. A trade-off can be considered between the model complexity and performance according to the needs of production.

Furthermore, we apply ultra-low bit quantization to compress the model size and achieves as much as 180$\times$ of compression ratio when $K=32$. In an extreme case, the model size of LiMuSE is compressed to 0.19MB. Combing the results of Table 2, the quantized LiMuSE still achieves competitive results over the audio-only baseline model TasNet [27] and audio-visual baseline models AVMS [11] and AVDC [9]. Though in [26], the authors have shown that the ultra-low bit quantization technology can achieve on-par performance on audio-only TCN-based speaker extraction models, applying the ultra-low quantization to our audio-visual TCN-based LiMuSE model leads to some performance degradation when the weights are quantized to 3 bits. While deploying the model on resource-constrained devices, it is recommended that an appropriate quantization bit be chosen according to the trade-off between the performance and model size as is shown in Table 7.

Based on the experimental results in Section 4.4.1, we investigate the effects of the visual cue and the voiceprint cue on the unquantized LiMuSE (GC+CC). The results are listed in Table 8. We can see from the table that the performance decreases after discarding the voiceprint cue or the visual cue, especially the visual cue. LiMuSE (GC+CC-Vis) suffers from a larger performance gap compared to LiMuSE (GC+CC-Vp). There is about a 10dB performance gap when $K=32$ and a 7dB performance gap when $K=16$ between LiMuSE and LiMuSE (GC+CC-Vis). It suggests that the visual cue is more informative than the voiceprint cue. It can be further concluded that the voiceprint cue and the visual cue serve as complementary information for each other and introducing both of them into LiMuSE leads to the best performance compared with LiMuSE (GC+CC-Vp) and LiMuSE (GC+CC-Vis).