[論文レビュー] 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はTF-GridNetのTransformerを bidirectional Mambaモジュールに置き換え、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.

研究の動機と目的

• 状態空間モデル（SSM）を用いて、CNN-および Transformerベースの手法における長距離シーケンス音声分離の制約を克服する動機づけ。
• TF-GridNetのTransformerコンポーネントをbidirectional Mambaモジュールで置換することによりSPMambaを提案。
• ノイズと残響を含む Librispeechベースのデータセットで、分離性能と効率性の改善を実証。

提案手法

• TF-GridNetをベースフレームワークとして採用し、BLSTM/TransformerコンポーネントをBMambaに置換して双方向の文脈を提供。
• 前方および後方のシーケンスを処理するBMambaを導入し、非因果でBLSTMのような情報統合を実現。
• TF-GridNet設計に従い、BMamba層を用いて時分割モジュール、周波数ドメインモジュール、時-周波数アテンションモジュールでSPMambaを構築。
• PIT（Permutation Invariant Training）とSNR損失を用いてソース分離品質を最適化して訓練。
• SI-SNRiとSDRiで評価し、パラメータ数とMACsを最先端モデルと比較。

実験結果

リサーチクエスチョン

• RQ1 難易度の高いノイズ/残響データセットにおいて、SPMambaはTF-GridNetおよび他のベースラインをSDRiとSI-SNRiで上回るか。
• RQ2 双方向のMambaはTransformerコンポーネントを効果的に置換して、より少ないパラメータと低い計算量で性能を維持または向上できるか。
• RQ3 SPMambaのパラメータ数とMACsの相対的な効率性は、TF-GridNetおよび他の先行モデルと比較してどうか。
• RQ4 BMambaはTF-GridNetフレームワーク内で、時刻・周波数領域の長距離依存性のモデリングにどのように寄与するか。

主な発見

モデル	SDR	SDRi	SI-SNR	SI-SNRi	Params(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はテスト済みモデルの中で最先端の性能を示す。
• Transformerをbidirectional Mambaに置換しても高い性能を維持し、計算負荷を低減。
• 長距離シーケンスの音声分離処理におけるMambaベースアーキテクチャの重要性を示す。