Cross-attention Inspired Selective State Space Models for Target Sound Extraction

2024-09-07

IEEE International Conference on Acoustics, Speech, and Signal Processing

AI Summary

Transformer-based cross-attention incurs high computational overhead in target speech extraction, while state-space models (e.g., Mamba) lack inherent cross-sequence modeling capability. Method: We propose Mamba-Cross—the first framework to reformulate Mamba’s implicit state update as a cross-attention–like mechanism: cue signals generate Queries, and mixed audio generates Keys/Values, enabling directed cross-sequence dependency modeling. This design synergistically integrates the linear-time complexity of state-space models with the semantic alignment prior of cross-attention. Contribution/Results: On two mainstream benchmarks, Mamba-Cross matches Transformer performance, substantially outperforms vanilla Mamba, and achieves >40% inference speedup—demonstrating superior trade-offs between representational capacity and computational efficiency.

Technology Category

• Machine Learning: Mixture of Experts (MoE)
• Computer Vision: Diffusion Models for Vision
• Cognitive Modeling & Cognitive Systems: Neural Spike Coding

Application Category

• Search and Retrieval-Augmented AI: Multilingual and cross-lingual Web search
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs

Abstract

The Transformer model, particularly its cross-attention module, is widely used for feature fusion in target sound extraction which extracts the signal of interest based on given clues. Despite its effectiveness, this approach suffers from low computational efficiency. Recent advancements in state space models, notably the latest work Mamba, have shown comparable performance to Transformer-based methods while significantly reducing computational complexity in various tasks. However, Mamba's applicability in target sound extraction is limited due to its inability to capture dependencies between different sequences as the cross-attention does. In this paper, we propose CrossMamba for target sound extraction, which leverages the hidden attention mechanism of Mamba to compute dependencies between the given clues and the audio mixture. The calculation of Mamba can be divided to the query, key and value. We utilize the clue to generate the query and the audio mixture to derive the key and value, adhering to the principle of the cross-attention mechanism in Transformers. Experimental results from two representative target sound extraction methods validate the efficacy of the proposed CrossMamba.

Problem

• Improve computational efficiency in target sound extraction
• Enhance sequence dependency capture in state space models
• Integrate Mamba with cross-attention for sound extraction

Innovation

• Leverages hidden attention in Mamba
• Generates query from given clues
• Derives key and value from audio