Local Spectral Attention for Full-band Speech Enhancement

Contents

• Repository description
• Rquirements
• Network training
  • Data preparation
  • Start training
  • Inference
• Ablation study and experiment results
  • LSA on MTFAA and DPARN

Repository description

This repository conduct ablation studies on local attention (a.k.a band attention) applied in full-band spectrum, namely local spectral attention (LSA). Two full-band speech enhancement (SE) models with spectral attention replace the conventional attention (a global manner) with LSA that only looks at adjacent bands at a certain frequency (a local manner). One model is DPARN, whose source code can be found in https://github.com/Qinwen-Hu/dparn.
The other model is the Multi-Scale Temporal Frequency with Axial Attention (MTFAA) network, which ranked 1st in the DNS-4 challenge for full-band SE, and its detailed description can be found in paper https://ieeexplore.ieee.org/document/9746610. Here we release an unofficial pytorch implementation of MTFAA as well as its modification. This work have been submitted to Interspeech2023.

Rquirements

soundfile: 0.10.3
librosa: 0.8.1
torch: 3.7.10
numpy: 1.20.3
scipy: 1.7.2
pandas: 1.3.4
tqdm: 4.62.3

Network training

Data preparation

Split your speech and noise audios into 10 seconds segments and generate the .csv files to manage your data. Prepare your RIR audios of .wav format in one folder. Edit the .csv path in Dataloader.py:

   TRAIN_NOISE_CSV = './train_noise_data.csv'  
   VALID_CLEAN_CSV = './valid_clean_data.csv'  
   VALID_NOISE_CSV = './valid_noise_data.csv'  
   RIR_DIR = 'direction to RIR .wav audios'

where the .csv files for clean speech are organized as

file_dir	snr
./clean_0001.wav	4
./clean_0002.wav	-1
./clean_0003.wav	0
...	...

and the .csv files for noise are organized as

file_dir
./noise_0001.wav
./noise_0002.wav
./noise_0003.wav
...

the 'file_dir' and 'snr' denote the absolute direction to audios and signal-to-noise ratio(SNR) respectively.

Start training

After environment and data preparation, start to train the model by command:

python Network_Training_MTFAA_full.py -m model_to_train(including MTFAA, MTFAA_LSA or MTFAA_ASqBi) -c Dir_to_save_the_checkpoint_files -e Epochs_for_training(default is 300) -d Device_used_for_training(cuda:0)

Inference

Enhance noisy audios by command:

python Infer.py -m model_to_train(including MTFAA, MTFAA_LSA or MTFAA_ASqBi) -c path_to_load_the_checkpoint_files -t path_to_folder_containing_noisy_audios -s path_to_folder_saving_the_enhanced_clips -d Device_used_for_training(cuda:0)

Ablation study and experiment results

We demonstrate the effectiveness of our proposed method on the full-band dataset of the 4th DNS challenge. The total training dataset contains around 1000 hours of speech and 220 hours of noise. Room impulse responses are convolved with clean speech to generate simulated reverberant speech, which is preserved as training target. In the training stage, reverberant utterances are mixed with noise recordings with SNR ranging from -5 dB to 5 dB at1 dB intervals. For the test set, 800 clips of reverberant utterances are mixed with unseen noise types with SNR ranging from -5 dB to 15 dB. Each test clip is 5 seconds long. All utterances are sampled at 48 kHz in our experiments. We also conduct experiments on well-known VCTK-DEMAND dataset for comprehensive validation.

LSA on MTFAA and DPARN

The visualization of LSA mechanism can be seen in the figure below:

The unofficial Pytorch implementation of MTFAA and its LSA-based model can be seen in MTFAA_Net_full.py and MTFAA_Net_full_local_atten.py respectively. As for DPARN, readers may attend to https://github.com/Qinwen-Hu/dparn.
Firstly, we conduct experiments on different setting of N_l based on the VCTK-DEMAND dataset and the results can be seen in table below:

It can be seen that the setting of N_l affects different models differently and we choose the setting achieving the best performance for each model, i.e. sqrt(F') for MTFAA and F'/2 for DPARN. Next, we train the models with the larger DNS4 dataset and the training process can be seen in figures below, where both LSA-based models achieve better convergence compared with the original models.

The objective test results can be seen in table below

The proposed LSA improves the enhancement performance of both the casual DPARN and MTFAA models in terms of all objective metrics. To reveal the benefit of LSA mechanism, we visualize the normalized average spectral attention plots, generated from audios in the test set, of attention blocks in both original MTFAA and LSA-based MTFAA, as shown in figures below

It can be seen from the fifth layer of attention that the LSA-based model more effectively emphasizes the structural features of harmonics in low bands (marked with red boxes) and the almost randomly distributed components in high bands (marked with black boxes). Furthermore, it can be seen from the blue boxes that LSA can also effectively alleviate the modeling of the invalid correlation between the low bands and the high bands. Hence, the speech pattern in spectrum can be better modeled by LSA. Further investigation of the enhanced signals reveals that the global attention in frequency domain is more likely to inflict distortion to speech components or produce excessive residual noise in non-speech segments, while this problem can be effectively alleviated by the proposed LSA. Two typical examples are shown in Figure 3, where the benefit of LSA can be clearly seen. A possible explanation is that the better exploitation to speech pattern helps LSA-based model more effectively discriminate speech and noise components especially in low-SNR environments.

To further demonstrate the importance of modeling local correlation in spectrum for full-band SE tasks, we also compare local attention with a recently proposed biased attention method, namely Attention with Linear Biases (ALiBi), which negatively biases attention scores with a linearly decreasing penalty proportional to the distance between the relevant key and query for efficient extrapolation. Its application on spectral attention can be seen in the figure below

We modify the penalty bias to decrease in a square manner for better performance and name the method as ASqBi, indicated in the figures below.

The modified method is combined with MTFAA MTFAA_Net_full_F_ASqbi.py and the ablation test results can be found in the last row of table below. It can be seen that the overall performance degrades compared with LSA. It may be explained that the negative bias added to local attention region weakens the model capability to extract local intercorrelation.

We also conduct subjective listening preference test on MTFAA model to validate the benifit of LSA mechanism. 50 enhanced samples as well as their reference target speech are randomly selected from the test set. 15 listeners with normal hearing compare the enhanced results based on the reference speech and choose the preferred result. Each sample is evaluated by at least 3 listeners. The subjective listening preference test results can be seen in table below. Over 60% samples enhanced by LSA-based MTFAA are considered to have better perceptual quality and lower noise levels, which demonstrates the efficiency of LSA in full-band SE tasks.

The proposed method also reduces computational complexity in spectral attention and the statistics are given in Table below

We compare the modified MTFAA model with previous full-band SOTA methods on VCTK-DEMAND dataset and the results are listed in table below

To further investigate the patterns of spectral attention, we firstly plot the attention figures generated from clean speech of male and female,

it can be seen that the patterns are related to speech characteristics that harmonics are basically distributed in low bands (top-left corner of attention plot) and consonants (almost randomly distributed) are in high bands (bottom-right corner of attention plot). It can also be seen that female's pitches are higher than male's with larger intervals. To be clearly, the harmonics-related lines are not parallel to each other for the frequencies are in the ERB scale.

Then we plot the attention figures generated from noisy speech and the attention plots in decoder are given below, where the harmonic-related features are highlighted,

Ignited by ALiBi, we also conduct experiment on the multi-scale local spectral attention (MSLSA) as shown in the figure below,

the performance of MSLSA can be seen in the table below,

The best performance of MSLSA is slightly worse than that of conventional LSA, while it can be seen from the training process that MSLSA may achieve a more stable result with a higher mean score and lower variance of PESQ in the validation set, which can be seen in the figure below (statistics based on the last 100 epochs)

Average attention plots of different heads in MSLSA are also given and it can be seen that the global attention head (N_l=F') cannot exploit clear spectral features while the local one (N_l=F'/4) makes it.