PaperlayerView exact revision on arXiv

Source-linked AI summary

Mamba in Speech: Towards an Alternative to Self-Attention

Xiangyu Zhang, Qiquan Zhang, Hexin Liu, Tianyi Xiao, Xinyuan Qian, Beena Ahmed, Eliathamby Ambikairajah, Haizhou Li, Julien Epps

arXiv:2405.12609v6eess.AScs.SD

TL;DR

Mamba has been investigated less in speech than in language and vision, and it remains unclear how to use it for speech tasks requiring high-level semantics. This paper tests bidirectional Mamba designs and self-attention replacements across recognition and enhancement, finding that outcomes depend on task and architecture.

  • Problem

    Speech Mamba research has focused largely on low-level enhancement and separation, leaving its use for semantic tasks such as speech recognition unclear; vanilla Mamba is also unidirectional despite non-autoregressive speech's need for non-causal computation.

    Source: A. State Space Models and Mamba · p. 3, B. Task-aware model designs · p. 4

  • Method

    The paper compares two bidirectional Mamba designs and evaluates Mamba independently and as an MHSA replacement in Transformer and Conformer models.

    Source: §I · pp. 1–2, B. Task-aware model designs · p. 4

  • Results

    Independent BiMamba performs well in speech enhancement but not speech recognition, while ConExtBiMamba matched or exceeded SoTA Branchformer across multiple datasets.

    Source: §VI · p. 12

  • Takeaways & Limitations

    For high-level semantic speech tasks, the paper identifies BiMamba as more appropriate when replacing MHSA than when used independently, with additional nonlinearity needed.

    Source: §I · p. 2, §VI · p. 12

    The discussion assumes BiMamba is weakly nonlinear like MHSA and uses this assumption to explain why ASR may require greater nonlinearity.

    Source: D. Discussion · pp. 11–12

Abstract

from arXiv · show

Transformer and its derivatives have achieved success in diverse tasks across computer vision, natural language processing, and speech processing. To reduce the complexity of computations within the multi-head self-attention mechanism in Transformer, Selective State Space Models (i.e., Mamba) were proposed as an alternative. Mamba exhibited its effectiveness in natural language processing and computer vision tasks, but its superiority has rarely been investigated in speech signal processing. This paper explores solutions for applying Mamba to speech processing by discussing two typical speech processing tasks: speech recognition, which requires semantic and sequential information, and speech enhancement, which focuses primarily on sequential patterns. The experimental results confirm that bidirectional Mamba (BiMamba) consistently outperforms vanilla Mamba, highlighting the advantages of a bidirectional design for speech processing. Moreover, experiments demonstrate the effectiveness of BiMamba as an alternative to the self-attention module in the Transformer model and its derivates, particularly for the semantic-aware task. The crucial technologies for transferring Mamba to speech are then summarized in ablation studies and the discussion section, offering insights for extending this research to a broader scope of tasks.

I. INTRODUCTION

The paper investigates Mamba for speech tasks with different information needs, focusing on recognition and enhancement. It compares bidirectional designs and explores Mamba both independently and as a replacement for self-attention.

I. INTRODUCTION

  • Speech recognition and spoken language understanding remain less explored because they require high-level semantics, unlike enhancement and separation's focus on low-level speech information.Speech attention can also be computationally demanding because its context window typically spans the full signal and complexity grows quadratically with window size.Source: §I · p. 1
  • The study compares two bidirectional Mamba designs and tests them across speech tasks with different information requirements.Its examples are speech recognition and speech enhancement.Source: §I · p. 1
  • BiMamba needs additional nonlinearity to learn high-level semantic information, and replacing MHSA is identified as an effective approach for this setting.The paper evaluates BiMamba independently and as an MHSA replacement in Transformer and Conformer models.Source: §I · p. 2

A. State Space Models and Mamba

  • Mamba's input-dependent parameters support context-aware representations, but make its convolution kernels input-dependent and prevent parallel convolution operations.The selective mechanism computes Δ, B, and C as functions of the input.Source: A. State Space Models and Mamba · p. 2
  • Mamba remains unidirectional, whereas non-autoregressive speech models require non-causal computation similar to attention.The paper considers bidirectional designs to address this need.Source: §I · p. 1, A. State Space Models and Mamba · p. 3

B. Mamba for Audio Processing

Prior work reports Mamba's use in speech and audio tasks, including enhancement, separation, ASR, and TTS. This paper focuses on abstraction levels and Mamba as an alternative to self-attention.

B. Mamba for Audio Processing

  • Audio Mamba studies report competitive audio-task performance, including DASS processing sequences up to 2.5 hours and AuM using roughly one-third the parameters of AST.SSAM reports strong masked-spectrogram modeling results, while SEMamba improves speech-enhancement state of the art through time-frequency Mamba.Source: B. Mamba for Audio Processing · p. 3
  • Prior studies compare Mamba encoders and decoders across TTS, ASR, SLU, and summarization, and incorporate Mamba into existing speech systems [36].The related-work discussion also notes SEMamba's speech-enhancement results.Source: B. Mamba for Audio Processing · pp. 3–4
  • This paper examines Mamba across speech-information abstraction levels, its use as a self-attention alternative, and nonlinearity as a factor in independent Mamba's lower ASR performance.The paper distinguishes its focus from work integrating Mamba with other architectures, [36].Source: B. Mamba for Audio Processing · p. 4

A. Bidirectional processing

Because speech models receive the complete signal, the paper explores bidirectional Mamba designs and progressively adds nonlinearity for tasks requiring higher abstraction. It also distinguishes the two bidirectional designs by how they handle directional projections.

A. Bidirectional processing

  • Bidirectional processing addresses Mamba’s use of historical information alone when speech tasks provide the complete signal and require global dependencies.InnBiMamba shares projections across directions, whereas ExtBiMamba uses separate directional projections intended to learn direction-specific, complementary features.Source: A. Bidirectional processing · p. 4

B. Task-aware model designs

  • The task-aware designs progress from independent Mamba/BiMamba layers to Transformer and Conformer MHSA replacements that add nonlinearity.The Transformer FFN and layer normalization provide nonlinearity; the Conformer variant additionally uses convolution to capture local information.Source: B. Task-aware model designs · p. 4

C. Algorithms of Bidirectional Mamba modules

  • InnBiMamba shares its input and output projections across directions, while ExtBiMamba uses separate directional input and output projections.The algorithms describe tensors by batch size B, sequence length L, input hidden dimension D, expanded hidden dimension E, and SSM state dimension N.Source: C. Algorithms of Bidirectional Mamba… · p. 4

IV. EXPERIMENTAL SETUP

The speech-enhancement experiments use LibriSpeech clean speech mixed with noise, compare Transformer, Conformer, and Mamba-based configurations, and assess enhancement with five speech-quality metrics.

A. Speech Enhancement

  • Training uses LibriSpeech train-clean-100 and noise from several datasets; validation uses 1,000 mixtures, and evaluation uses 400 mixtures with four held-out real-world noises.The evaluation mixes 20 clean test clips per noise at five SNR levels; validation mixtures use randomly selected clean clips and noise at SNRs from −10 to 20 dB.Source: A. Speech Enhancement · pp. 4–5
  • Audio is sampled at 16 kHz and represented by 257-point single-sided STFT spectral magnitudes using 32 ms windows and 16 ms hops.The transform uses a 512-sample square-root-Hann window and a 256-sample hop.Source: A. Speech Enhancement · p. 5
  • The experiments use a shared backbone pattern with Transformer and Conformer baselines across causal and non-causal configurations, alongside Mamba and BiMamba layers.Transformer and Conformer models use four and six stacked layers, respectively; the reported Mamba settings include hidden dimension 256, SSM state dimension 16, convolution width 4, and expanded dimension 512.Source: A. Speech Enhancement · pp. 5–6
  • Models are trained for 150 epochs using MSE on power-law-compressed spectral magnitudes, dynamically mixed noisy speech, Adam, gradient clipping, and a 40,000-step warm-up.Training mixtures use random noise sections and SNRs sampled from −10 to 20 dB; gradient values are clipped to the range −1 to 1.Source: A. Speech Enhancement · p. 6
  • Enhancement is assessed with W-PESQ, N-PESQ, ESTOI, CBAK, CSIG, and COVL, with preferred score direction varying by metric.Higher values are preferable for each metric except where the paper indicates a lower-is-better direction.Source: A. Speech Enhancement · p. 6

B. Speech Recognition

The ASR experiments span monolingual and English-Mandarin code-switching datasets, including a small-data setting. They report WER for monolingual recognition and MER for code-switching recognition.

B. Speech Recognition

  • ASR evaluation uses LibriSpeech, AN4, SEAME, and ASRU; LibriSpeech100 supports ablations, AN4 tests a small-data setting, and SEAME and ASRU provide code-switching scenarios.The ASRU code-switching training set is used, while the dataset also contains a separate Mandarin training set.Source: B. Speech Recognition · p. 6
  • Mamba and BiMamba use default settings with SSM state dimension 16, local convolution width 4, and expansion factor 2.LibriSpeech960 experiments ran on an NVIDIA A100 80GB; the other experiments ran on an NVIDIA V100 32GB.Source: B. Speech Recognition · pp. 6–7
  • The study uses WER for monolingual ASR and MER for code-switching ASR, with lower scores preferred for both.MER combines English word error rate and Mandarin character error rate.Source: B. Speech Recognition · p. 7

A. Speech Enhancement

In speech enhancement, ExtBiMamba generally improves on InnBiMamba and vanilla Mamba, and Mamba-based modules improve results when substituted for MHSA. ExtBiMamba also reduces computational costs for longer inputs, though its comparison with Conformer has a short-input RTF exception.

InnBiMamba vs. ExtBiMamba

  • ExtBiMamba consistently outperforms InnBiMamba across model sizes and is faster to train and infer, with lower computational complexity at similar sizes.The comparisons cover six quality metrics, training time per step, inference speed, and MACs.Source: A. Speech Enhancement · p. 7, Table IV · p. 7, Table V · p. 7

Mamba vs. Transformer

  • Mamba models outperform similarly sized Transformer models across causal and non-causal configurations, and ExtBiMamba improves on unidirectional Mamba across all six metrics.ExtBiMamba-5 (4.51M) improves over Mamba-10 (4.51M) by 0.18 N-PESQ, 0.21 W-PESQ, 2.92% ESTOI, 0.16 CSIG, 0.13 CBAK, and 0.19 COVL.Source: A. Speech Enhancement · p. 7, Table VI · p. 7

Mamba vs. Conformer

  • Causal Mamba-20 (8.89M) outperforms causal Conformer-6 (9.26M) across all six metrics, while ExtBiMamba is slightly better or comparable to non-causal Conformer.ExtBiMamba also shows substantial gains over Mamba in the reported matched-size comparisons.Source: A. Speech Enhancement · pp. 7–8, Table VI · p. 7

Efficiency comparison

  • For 10s and 40s inference, ExtBiMamba-4 (3.64M) has RTF values of 0.93ˆ and 0.23ˆ and MACs of 0.85ˆ and 0.46ˆ versus Transformer-4 (3.29M).Its efficiency advantage grows with longer inputs; compared with Conformer, RTF is marginally higher at 10s but significantly lower at 20s and 40s, with lower MACs at those longer lengths.Source: A. Speech Enhancement · p. 8, Table VII · p. 8

Mamba vs. MHSA

  • Replacing Transformer MHSA with Mamba or BiMamba improves speech-enhancement scores, with TransExtBiMamba generally slightly ahead of TransInnBiMamba.TransExtBiMamba-4 (5.74M) scores slightly higher in W-PESQ, CSIG, CBAK, and COVL, but slightly lower in N-PESQ and ESTOI, than TransInnBiMamba-6 (6.19M).Source: A. Speech Enhancement · p. 8, Table VIII · p. 8
  • ConExtBiMamba-4 and -6 outperform non-causal Conformer-4 and -6 across six speech-enhancement metrics, while MHSA replacement also lowers RTF, particularly for longer inputs.ConExtBiMamba-4 (8.67M) also slightly outperforms non-causal Conformer-6 (9.26M).Source: A. Speech Enhancement · p. 8, Table VIII · p. 8, Table IX · p. 8

VoiceBank+DEMAND benchmark

  • Replacing Conformer with ExtBiMamba in MP-SENet yields consistent performance improvements, although MP-SENet-ExtBiMamba does not achieve the highest performance overall.The paper cautions that the independently evaluated ExtBiMamba is not directly comparable with MP-SENet-Conformer because of the latter's more complex design.Source: A. Speech Enhancement · p. 9

B. Speech Recognition

Independent Mamba and BiMamba perform poorly for ASR relative to Transformer and Conformer, but replacing MHSA with ExtBiMamba in Conformer yields stronger results. Bidirectional modeling helps, while matching parameter counts by increasing depth does not resolve independent Mamba's performance gap.

Independent vs. Substitute for MHSA

  • Independent Mamba and BiMamba underperform Transformer and Conformer at the same layer count, and increasing Mamba depth to match Conformer parameters leaves performance undesirable.The experiments report that increasing layer count to match parameter size does not proportionally improve performance.Source: B. Speech Recognition · p. 9
  • Replacing Conformer MHSA with ExtBiMamba produces higher ASR performance than Conformer and Branchformer across multiple datasets under the same training setups.ConExtBiMamba also trains and infers faster than Conformer; increasing Conformer size to match its parameter count still leaves it behind ConExtBiMamba and ConInnBiMamba.Source: B. Speech Recognition · p. 9, Table XI · p. 10, Table XII · p. 10

Discussion

  • Mamba-based models need stronger nonlinear capabilities for high-level ASR abstractions; Conformer's convolutional modules and nonlinear activations can compensate for Mamba's limited inherent nonlinearity.This architectural difference is offered to explain why BiMamba integration benefits Conformer more than Transformer for ASR.Source: B. Speech Recognition · p. 9

Independent vs. Substitute for MHSA

  • Independent Mamba's weaker ASR performance is accompanied by training-stability issues, which the paper relates to difficulty capturing high-level abstractions.The discussion contrasts this with using Mamba as an MHSA replacement, which yields significantly improved ASR performance.Source: B. Speech Recognition · pp. 9–10

Unidirectional vs. Bidirectional

  • InnBiMamba and ExtBiMamba improve over unidirectional Mamba when replacing MHSA, and ExtBiMamba remains stronger than InnBiMamba even with parameter counts matched by added InnBiMamba layers.Comparisons with larger Conformer and ConInnBiMamba models also indicate that the ASR gains are not solely due to more parameters.Source: B. Speech Recognition · p. 10

C. Ablation Study on ConExtBiMamba

The ablations examine ConExtBiMamba's hyperparameters and small-dataset behavior. Gaussian-noise initialization of the diagonal A matrix performs best, Swish and Macaron-style feed-forward layers improve performance, and ConExtBiMamba outperforms Conformer on AN4.

Ablation setup

Hyper-Parameter for BiMamba in ConExtBiMamba

Hyper-Parameters for ConExtBiMamba

  • Swish activation and Macaron-style feed-forward layers improve performance, while positional encoding and dropout have no significant impact.The paper suggests Mamba's recurrent-like dependencies already embed positional information, reducing the effect of positional encoding.Source: C. Ablation Study on ConExtBiMamba · p. 11, Table XIV · p. 11

Performance on Extremely Small Datasets

  • WER averages 10.02 for Conformer and 4.54 for ConExtBiMamba on AN4; WER variance is 77.71 versus 1.11, respectively.Across five random seeds, both models converge on the development set, with more stable convergence for ConExtBiMamba; Fig. 3 shows train and dev set loss.Source: C. Ablation Study on ConExtBiMamba · p. 11, Table XV · p. 11, Figure 3 · p. 11
  • ConExtBiMamba appears less prone to overfitting on small datasets than Conformer, and the authors suggest recursive temporal processing helps in limited-data settings.The discussion also states that this framework avoids the gradient vanishing and explosion issues typically associated with RNNs.Source: C. Ablation Study on ConExtBiMamba · p. 11

D. Discussion

The discussion argues that BiMamba works better as an MHSA replacement than as an independent model for ASR, where greater nonlinearity is needed. It relates this finding to ablations, a synthetic decision-boundary experiment, training stability, and task abstraction levels.

D. Discussion

  • Replacing MHSA with BiMamba layers, alongside an FFN and residual connection, outperforms vanilla Transformer and Conformer models on ASR.The authors examine the factors behind this improvement through ablation studies.Source: D. Discussion · p. 11
  • The authors treat BiMamba as weakly nonlinear and argue ASR needs more nonlinearity than enhancement because ASR also involves semantic information and token alignment.They suggest Conformer convolution modules provide additional nonlinearities compared with the Transformer.Source: D. Discussion · pp. 11–12
  • On synthetic complex interleaved spirals, BiMamba struggles to learn the decision boundary without an FFN, unlike on simpler Gaussian clusters.The authors present this result as support for their nonlinearity assumption.Source: D. Discussion · p. 12
  • Removing the FFN and residual connection from Transformer ASR causes gradient vanishing, reduces nonlinearity, and substantially degrades performance.The authors cite prior work on normalization for Mamba training stability and on adding components such as decoders or feed-forward layers for speech recognition.Source: D. Discussion · p. 12
  • Prior independent-Mamba ASR work reported occasional training instability; the authors also observed instability and could not exactly reproduce its results.They transferred optimal attention-model parameters for fair comparisons and found their baseline improved by approximately 10% with AdamW.Source: D. Discussion · p. 12
  • The discussion groups phoneme recognition, separation, and voice activity detection with low-abstraction tasks, and translation, text-to-speech, and emotion recognition with high-abstraction tasks.These groupings extend the authors’ discussion beyond speech enhancement and recognition.Source: D. Discussion · p. 12

VI. CONCLUSION

The conclusion contrasts independent BiMamba’s strength on speech enhancement with its difficulty on semantic ASR. It reports that replacing Conformer’s MHSA with BiMamba matches or exceeds Branchformer across multiple datasets, with ablations supporting this use for high-abstraction tasks.

VI. CONCLUSION

  • Independent BiMamba performs well on speech enhancement but does not achieve speech recognition, which requires semantic information.The conclusion characterizes enhancement as involving low-abstraction spectral information.Source: §VI · p. 12
  • ConExtBiMamba, which replaces Conformer MHSA, matched or exceeded SoTA Branchformer performance across multiple datasets.Source: §VI · p. 12
  • Ablations suggest BiMamba is more appropriate as an MHSA replacement than as an independent model for tasks requiring high-abstraction information.The authors attribute this suitability to greater nonlinearity compared with independent BiMamba.Source: §VI · p. 12
Loading 2405.12609v6…