SPMamba: State-space model is all you need in speech separation

Kai Li,Guo Chen,Runxuan Yang,Xiaolin Hu

2024-09-10

Abstract:Existing CNN-based speech separation models face local receptive field limitations and cannot effectively capture long time dependencies. Although LSTM and Transformer-based speech separation models can avoid this problem, their high complexity makes them face the challenge of computational resources and inference efficiency when dealing with long audio. To address this challenge, we introduce an innovative speech separation method called SPMamba. This model builds upon the robust TF-GridNet architecture, replacing its traditional BLSTM modules with bidirectional Mamba modules. These modules effectively model the spatiotemporal relationships between the time and frequency dimensions, allowing SPMamba to capture long-range dependencies with linear computational complexity. Specifically, the bidirectional processing within the Mamba modules enables the model to utilize both past and future contextual information, thereby enhancing separation performance. Extensive experiments conducted on public datasets, including WSJ0-2Mix, WHAM!, and Libri2Mix, as well as the newly constructed Echo2Mix dataset, demonstrated that SPMamba significantly outperformed existing state-of-the-art models, achieving superior results while also reducing computational complexity. These findings highlighted the effectiveness of SPMamba in tackling the intricate challenges of speech separation in complex environments.

Sound,Artificial Intelligence,Audio and Speech Processing

What problem does this paper attempt to address?

The paper aims to address the issue of capturing long-term dependencies in the task of speech separation. Specifically, existing speech separation models based on CNN, RNN (such as LSTM), and Transformer have limitations when processing long audio sequences. For example, CNN is limited by its local receptive field and cannot effectively capture long-term dependencies; LSTM can handle long-term dependencies but still faces challenges in capturing distant dependencies; Transformer can model long-range dependencies, but its self-attention mechanism results in quadratic complexity with respect to sequence length, making it computationally expensive for real-time applications.

To address these issues, the paper proposes a new speech separation method called SPMamba. This method is based on the TF-GridNet architecture and replaces the traditional BLSTM modules with bidirectional Mamba modules. These modules can effectively model the spatiotemporal relationships between the time and frequency dimensions, allowing SPMamba to capture long-range dependencies while maintaining linear computational complexity. Experimental results show that SPMamba significantly outperforms existing state-of-the-art models on multiple public datasets (including WSJ0-2Mix, WHAM!, Libri2Mix, and the newly constructed Echo2Mix) and also has advantages in terms of computational complexity and the number of parameters.