TLDR: This paper introduces RWSA-MambaUNet, a novel hybrid model combining Mamba and multi-head attention within a U-Net structure for speech enhancement. It leverages “resolution-wise shared attention” to significantly improve cross-corpus generalization performance, meaning it works well in new, unseen acoustic conditions. The model achieves state-of-the-art results on out-of-domain test sets, even outperforming existing models with fewer parameters and lower computational complexity, demonstrating its efficiency and effectiveness in noise reduction.
Speech enhancement, the process of removing background noise from speech signals to improve clarity and quality, is crucial for many applications, from hearing aids to mobile communication. Recent advancements in deep learning have significantly pushed the boundaries of this field, with models combining different neural architectures showing promising results.
A new research paper, “EXPLORING RESOLUTION-WISE SHARED ATTENTION IN HYBRID MAMBA-U-NETS FOR IMPROVED CROSS-CORPUS SPEECH ENHANCEMENT,” introduces a novel and efficient hybrid model called RWSA-MambaUNet. Authored by Nikolai Lund K¨uhne, Jesper Jensen, Jan Østergaard, and Zheng-Hua Tan, this work addresses a key challenge in speech enhancement: ensuring models perform well even when faced with acoustic conditions, noise types, and speakers not encountered during training. This is known as cross-corpus generalization. You can read the full paper here: https://arxiv.org/pdf/2510.01958.
The Challenge of Generalization
While many deep learning models, especially those based on attention mechanisms like Transformers, have achieved state-of-the-art performance in speech enhancement, they often struggle with computational complexity, scaling quadratically with input size. Recurrent models like Mamba and xLSTM offer linear scaling, but have shown weaker generalization performance across different datasets compared to purely attention-based models. The goal is to combine the strengths of both: the efficiency of recurrent models and the strong generalization of attention mechanisms.
Introducing RWSA-MambaUNet
Inspired by the success of models that integrate Mamba into a U-Net structure and the superior generalization of hybrid Mamba-Attention models, the researchers propose RWSA-MambaUNet. This model combines Mamba and multi-head attention within a U-Net architecture. A core innovation is “Resolution-Wise Shared Attention” (RWSA). RWSA involves sharing attention mechanisms across corresponding time and frequency resolutions within the U-Net’s downsampling and upsampling paths. This explicit alignment of global time and frequency relationships is hypothesized to be critical for robust cross-corpus speech enhancement.
How it Works
The RWSA-MambaUNet processes noisy speech by first converting it into a complex spectrogram. A feature encoder then prepares these features for the U-Net. The U-Net architecture employs multiple MambAttention blocks, which are themselves hybrid blocks combining bidirectional Mamba blocks (for time and frequency processing) with shared time- and frequency-based multi-head attention (MHA) modules. The RWSA mechanism ensures that these MHA modules are not only shared within each MambAttention layer but also across layers at matching resolution levels in both the downsampling and upsampling parts of the U-Net. This joint alignment of temporal and spectral dependencies across multiple scales helps the model learn more robust features for generalization.
After processing through the U-Net, the model uses magnitude and phase refinement layers, followed by decoders, to estimate the clean magnitude and phase spectra. Finally, an inverse short-time Fourier transform (iSTFT) reconstructs the enhanced speech waveform.
Impressive Results
The RWSA-MambaUNet models were rigorously tested on both in-domain (VB-DemandEx) and out-of-domain datasets (DNS 2020 and EARS-WHAM v2), which feature different speakers, noise types, and recording conditions. The results are highly encouraging:
- The smallest model, RWSA-MambaUNet-XS, significantly outperformed several baselines (LSTM-SENet, xLSTM-SENet, SEMamba, MP-SENet, and MambAttention) on the out-of-domain DNS 2020 test set across PESQ, SSNR, and ESTOI metrics. It also surpassed baselines on the EARS-WHAM v2 test set for SSNR, ESTOI, and SI-SDR. Remarkably, it achieved this with less than half the model parameters and a fraction of the computational operations (FLOPs).
- Larger versions of the model, RWSA-MambaUNet-S and RWSA-MambaUNet-M, further improved performance, with the medium-sized model (RWSA-MambaUNet-M) outperforming all baselines across all evaluation metrics on both out-of-domain test sets, while still requiring significantly fewer FLOPs.
- An ablation study confirmed the importance of RWSA and the MHA modules for the model’s superior generalization performance. Removing either component led to a decrease in performance.
Also Read:
- Enhancing Audio Clarity: UniverSR’s Vocoder-Free Method for Superior Sound
- VoiceBridge: A Unified System for High-Fidelity Speech Restoration Across Diverse Distortions
Conclusion
The RWSA-MambaUNet represents a significant step forward in speech enhancement, particularly for cross-corpus generalization. By cleverly combining Mamba and multi-head attention within a U-Net structure and introducing resolution-wise shared attention, the researchers have developed a model that is not only highly effective at removing noise but also remarkably efficient in terms of model size and computational cost. This makes it a promising solution for real-world applications where speech enhancement systems need to operate reliably across diverse acoustic environments.
