[论文解读] SPMamba: State-space model is all you need in speech separation

Kai Li, Chen Guo+2 more|arXiv (Cornell University)|Apr 2, 2024

Speech Recognition and Synthesis被引用 8

一句话总结

SPMamba 用双向 Mamba 模块替代 TF-GridNet 的 Transformer，在基于 Librispeech 的含噪与混响数据上，以更少的参数和更低的计算成本实现了最先进的语音分离。

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.

研究动机与目标

• 激发使用状态空间模型（SSMs）来解决基于 CNN 和 Transformer 的方法在长序列语音分离方面的局限性。
• 通过在 TF-GridNet 中用双向 Mamba 模块替换 Transformer 组件来提出 SPMamba。
• 在包含噪声和混响的 Librispeech 基础数据集上演示改进的分离性能和效率。

提出的方法

• 以 TF-GridNet 作为基框架，并用 BMamba 替换 BLSTM/Transformer 组件以实现双向上下文。
• 引入 BMamba 以处理前向和后向序列，实现非因果、类似 BLSTM 的信息聚合。
• 按照 TF-GridNet 的设计，在 SPMamba 中包含时域模块、频域模块和时频注意力模块，并使用 BMamba 层。
• 使用置换不变训练（PIT）并结合信噪比（SNR）损失来优化源分离质量。
• 使用 SI-SNRi 和 SDRi 进行评估，并将参数量和 MACs 与最先进模型进行比较。

实验结果

研究问题

• RQ1 在一个具有挑战性的嘈杂/混响数据集上，SPMamba 是否在 SDRi 和 SI-SNRi 上超越 TF-GridNet 及其他基线？
• RQ2 双向 Mamba 是否能够有效替代 Transformer 组件，在更少参数和更低计算量的前提下保持或提高性能？
• RQ3 相对于 TF-GridNet 和其他领先模型，SPMamba 的相对效率（参数量与 MACs）如何？
• RQ4 BMamba 如何在 TF-GridNet 框架内对时域和频域的长程依赖建模做出贡献？

主要发现

模型	SDR	SDRi	SI-SNR	SI-SNRi	参数量(M)	MACs (G/s)
Conv-TasNet	7.58	7.69	6.71	6.89	5.62	10.23
DualPathRNN	5.76	5.87	4.88	5.06	2.72	85.32
SudoRM-RF	7.59	7.70	6.66	6.84	2.72	4.60
A-FRCNN	9.53	9.64	8.58	8.76	6.13	81.20
TDANet	9.93	10.14	8.95	9.21	2.33	9.13
BSRNN	12.64	12.75	12.04	12.23	25.97	98.69
TF-GridNet	13.59	13.70	12.62	12.81	14.43	445.56
SPMamba (Ours)	16.01	16.14	15.20	15.33	6.14	78.69

• SPMamba 达到 SDR 16.01 dB 和 SI-SNRi 15.20 dB，分别比 TF-GridNet 高出 2.42 dB 和 2.58 dB。
• SPMamba 使用 6.14M 参数和 78.69 GMACs/s，显著少于 TF-GridNet 的 14.43M 参数、445.56 GMACs/s。
• 在包含噪声和混响的 Librispeech 基础数据集上，SPMamba 在所测试模型中取得最先进的性能。
• 用双向 Mamba 替换 Transformer 同时保持高性能并降低计算需求。
• 该模型展示了基于 Mamba 的体系结构在语音分离的长序列音频处理中的重要性。