Omni-directional attention mechanism based on Mamba for speech separation
Ke Xue, Chang Sun, Rongfei Fan, Jing Wang, Han Hu
TL;DR
This work introduces an Omni-directional Attention (OA) mechanism built on the selective state-space model Mamba to address the global 2D dependencies in time-frequency speech spectrograms. OA projects a 2D spectrogram into ten directional streams across temporal, frequency, and channel axes and processes them with Mamba to maintain linear complexity, enabling global modeling without the quadratic costs of traditional self-attention. By replacing self-attention with OA in TF-GridNet and SPMamba, the approach achieves state-of-the-art results on WSJ0-2Mix, WHAM!, and Libri2Mix with notable SI-SDRi gains while preserving linear scaling. The findings highlight the practical potential of linear-complexity omni-directional attention for efficient long-range modeling in speech separation, with strong implications for scalable real-time systems.
Abstract
Mamba, a selective state-space model (SSM), has emerged as an efficient alternative to Transformers for speech modeling, enabling long-sequence processing with linear complexity. While effective in speech separation, existing approaches, whether in the time or time-frequency domain, typically decompose the input along a single dimension into short one-dimensional sequences before processing them with Mamba, which restricts it to local 1D modeling and limits its ability to capture global dependencies across the 2D spectrogram. In this work, we propose an efficient omni-directional attention (OA) mechanism built upon unidirectional Mamba, which models global dependencies from ten different directions on the spectrogram. We expand the proposed mechanism into two baseline separation models and evaluate on three public datasets. Experimental results show that our approach consistently achieves significant performance gains over the baselines while preserving linear complexity, outperforming existing state-of-the-art (SOTA) systems.