Selector-Enhancer: Learning Dynamic Selection of Local and Non-local Attention Operation for Speech Enhancement

Long-range dependencies are important prior to SE, and this prior can be effectively obtained by non-local attention operation, which calculates the mutual similarity between T-F units in each frame, which is helpful for capturing the global information in the frequency domain with a slightly increased computing complexity. In general, non-local attention is defined as follows (Wang et al. 2018b):

where $x$ and $y$, respectively, represent the input and output tensor of the operation with the same size, $f$ denotes the pairwise function to calculate the correlation between the locations of the feature map, $g$ signifies the unary input function for information transform, and $\mathcal{C}(x)$ is a normalization factor.

Recently, deep neural networks (DNNs) have been prevalent in the community of SE, and non-local attention-based SE models have been proposed for modeling long-term dependencies of speech signals. In (Li et al. 2019a), embedding function $g$ is a convolutional layer that can be viewed as a linear embedding function, and dot product, i.e., $\phi(\textbf{y}_{i},\textbf{y}_{j})=\textbf{y}_{i}^{\top}\textbf{y}_{j}$ is applied to compute the similarity between query and keys. However, non-local attention is unreliable when speech is heavily corrupted and colored noise is contained in the background mixture. In addition, a transformer with Gaussian-weighted non-local attention is proposed (Kim, El-Khamy, and Lee 2020), whose attention weights are attenuated according to the distance between correlated T-F units. The attenuation is determined by the Gaussian variance which can be learned during training. While all these methods demonstrate the benefits of non-local attention, their approach is unreliable when speech is heavily corrupted and has high computational costs.

Local attention focuses on the important speech components in an audio stream with “high attention” while perceiving the unimportant region (e.g., noise or interference) in “low attention”, according to the past and current time frame, thus adjusting the focal point over time. In (Hao et al. 2019), a lightweight casual transformer with local self-attention is proposed for real-time SE in computation resource-limited environments. Local attention based RNN-LSTM (Oostermeijer, Wang, and Du 2021) is proposed for superior SE performances by modeling important sequential information, in which the RNN model learns the weights of past input features implicitly when predicting an enhanced frame, the attention mechanism calculates correlations between past frames and the current frame to be enhanced and give weights to past frames explicitly. Compared with non-local attention, local attention is more suitable for real-time settings and more lightweight (Hao et al. 2019; Oostermeijer, Wang, and Du 2021), since they do not rely on future information. However, compared with non-local operations, the main drawback of the local attention mechanism during SE is the relatively small receptive field.

Dynamic networks are effective to achieve a trade-off between speech and performance in different tasks. Recently, several approaches (Wu et al. 2018; Shalit, Johansson, and Sontag 2017; Yang et al. 2022) are applied to dynamic networks to reduce the computational cost in ResNet by dynamically skipping some residual blocks. Although the aforementioned methods mainly focus on a single task, a routing network (Rosenbaum, Klinger, and Riemer 2018) is developed to facilitate multi-task learning, and their method seeks an appropriate network path for each specific task. Existing dynamic networks successfully explore the route policy that depends on either the task to address. However, we aim at selecting different regions of input speech spectrogram to send to specific attention operations, in which noise information (noise type and signal-to-noise ratio (SNR)) and speech information (speaker identity and speech content) have a weighty influence on the performance and they should be both considered in the selection operation.

We aim to design a framework that can offer different attention mechanism options for SE. To this end, as shown in Figure 2, Selector-Enhancer is composed of 3 convolutional blocks at the start- and end-point for extracting features and reconstructing the target speech, and $N$ dynamic blocks in the middle. Each convolutional block contains a stacked 2D convolutional layer which is followed by batch normalization, and exponential linear unit (ELU). Additionally, skip connections are utilized to concatenate the output of each convolutional block to the input of the corresponding deconvolutional block.

In the $i$-th dynamic block $f^{i}_{DB}$, there is a shared path $f_{S}^{i}$ containing two convolutional blocks for extracting deeper speech features, which every feature map should pass through. Paralleling the shared path, a feature-filter generates a probabilistic distribution of plausible paths for selection by each T-F unit of the input feature map. Following the shared path, there are two paths for local and non-local attention denoted by $f^{i}_{LA}$ and $f^{i}_{NA}$. According to the output of the feature-filter, the path with the highest probability is activated, where the path index is represented by action $\alpha_{i}$. Since the path selection is for every T-F unit rather than a feature map, the action $\alpha_{i}$ is a two-channel mask rather than a scalar. The first channel $m^{i}_{L}$ implies the regions that go through the path for local attention, and similarly, the second channel $m^{i}_{N}$ corresponds to the path for non-local attention. The whole process is formulated as:

where $z_{i}$ and $z_{i+1}$ represent the input and output of the $i$-th dynamic block, respectively, and $f^{i}_{R}$ denotes the decoder of dynamic block. Note that “$\cdot$” denotes the element-wise multiplication, and each channel of feature $f_{S}^{i}(z_{i})$ performs the same element-wise multiplication with mask $m^{i}$. Specially, the mask $m^{i}$ utilizes two branches to generate a 1-D frequency-dimension mask $\textbf{F}_{M}\in\mathbf{R}^{1\times F}$ and a 1-D time-frame mask $\textbf{T}_{M}\in\mathbf{R}^{T\times 1}$ in parallel, then combines them with a matrix multiplication to obtain the final 2-D mask.

The local and non-local attention operations are demonstrated as follows:

The architecture of the feature filter is illustrated in Figure 2, which is the key component to achieve the selection operation. The feature filter contains three convolutional blocks, two LSTMs for capturing the correlations of the feature filter in different dynamic blocks, and a convolutional block followed by a sigmoid activation function. In addition, the selection operation is non-differentiable, thus we adopt Markov Decision Process (MDP) and reinforcement learning (RL) to train the feature filter. The state and action of this MDP and the developed difficulty-regulated reward for the feature filter training are clarified as follows:

The Selector-Enhancer is composed of dual-attention based CNN and feature filter, and the training process of Selector-Enhancer consists of two stages. In the first stage, we train a dual-attention based CNN pre-trained model with a random policy mask. In the second stage, we train the feature filter and dual-attention based CNN together.

To evaluate the proposed Selector-Enhancer, we use two types of datasets: VoiceBank + DEMAND and AVSpeech + AudioSet.

For the test dataset, we randomly select 840 utterances from 84 speakers in WSJ0 SI-84 dataset (Paul and Baker 1992) and mingle with unknown types of non-stationary noises from NOISEX92 (Varga and Steeneken 1993) under three SNR levels $\{-5,0,5\}$ dB. The choice of the training set and test set is to let the Selector-Enhancer adapt to different speakers and different acoustic environments.

As for the number of dynamic blocks, we use $N=4$ for a fair comparison with the baseline methods. The threshold of difficulty $L_{t}$ and the reward penalty $\gamma$ are, respectively, set to $6\times 10^{-2}$ and $8\times 10^{-2}$ based on the specific task. Specifically, when the noise corruption is severe, the threshold and reward penalty is required to be larger. In the first training stage, the coefficient of intermediate loss $\beta$ is set to 0.5. All the utterances are sampled at 16 kHz and features are extracted by using frames of length 512 with frame-shift of 256, and Hann windowing followed by STFT of size K = 512 with respective zero-padding. We train the model with a minibatch size of 16. We use Adam as an optimizer and implement our method on Pytorch.

In our experiments, the performance is assessed by using standard speech enhancement metrics short-term objective intelligibility (STOI), and perceptual evaluation of speech quality (PESQ). STOI typically has a value range from 0 to 1 and can be interpreted as percent correct, and PESQ has a value range from -0.5 to 4.5.

We show the ablation study in Table 1 to investigate the effect of different components in Selector-Enhancer. Note that Case 0 denotes the proposed basic Selector-Enhancer with the default setting. In the ablation study, we compare our method with several models: 1) Case 1: we remove all non-local attention mechanisms, 2) Case 2: we remove all local attention mechanisms, 3) Case 3: We remove the feature filters and adopt concatenation operation to fuse results of local and non-local operations, 4) Case 4: we remove feature filters and adopt selection network (Li et al. 2019b) to fuse results of local and non-local operations, 5) Case 5: we remove all components of Selector-Enhancer. We show the detailed analysis of ablation below.

In order for a better understanding of the feature-filter, we visualize the masked feature maps processed by the feature filter for local and non-local attention mechanisms and processed feature maps produced by local and non-local attention operation in Figure 5. We show that the corrupted speech components are scaled and masked by feature-filter to separately feed into local and non-local attention operations.

We carry out the system comparison on both datasets mentioned in section 4.1.

Realizing that the number of trainable parameters can not completely reflect the complexity of a model, especially for deep non-local methods. Thus we employ the floating-point operations (FLOPs), the number of trainable parameters, and real-time factor (RTF) to evaluate the complexity of different methods. The complexity analysis of several representative methods is reported in Table 2. The proposed Selector-Enhancer has an attractive complexity and low computational costs and achieves quite promising performance.

To demonstrate the effectiveness and flexibility of the Selector-Enhancer, we further explore and apply the Selector-Enhancer to circumstances when multiple noise types are included in a noisy background. We randomly select 150 clean utterances from the test set and 3 challenging noises, babble (B), factory (F), and cafeteria (C) from NOISEX92, which are explained in Section 4.1, to produce the multiple types of noises with clean utterances at -5 dB, 0 dB, and 5 dB SNR levels. Figure 6 shows the STOI and PESQ results over three mentioned SNR levels on multiple noise types, where B, F, and C represent babble, factory, and cafeteria noise types, respectively. According to the Table, we observe that the results using 1, 2, and 3 noise types in the noisy background are comparable, although the STOI and PESQ scores are decreasing when the number of noise types is increase.