SSM Domain Evaluation
Project goal: Compare different SSM backbones in different domains for speech enhancement.
Motivation: LinOSS has emerged as a competitive alternative Mamba for other tasks, while beeing LTI and well-interpretable. Also, speech enhancement has adopted Mamba without deeper consideration for what attributes make SSMs suited for speech enhancement. This work can elucidate the differences in SSM architecture across different speech enhancement domains.
Architecture¶
Keep it simple. The model is just an encoder -> N x SSM -> decoder structure. The encoder projects the raw tme-domain waveform in a given domain, which is then passed through \(N\) SSM blocks (Mamba-style block with different SSM cores), then projected back into time-domain. Any variation in performance can then be attributed to the SSM core in a given domain, without confounding architectural decisions.
Domains¶
Potential domains for speech enhancement include time-domain, spatio-temporal STFT and filterbank domain, and the token-domain of Neural Audio Codecs like Encodec. The latter is of particular interest due to an increasing focus on neural audio token inference, which has potential to leverage the progress of the LLM field.
| Domain | Reference |
|---|---|
| Time | Attenuate |
| STFT (real/imag) | TF-GridNet |
| STFT (mag/phase) | SEMamba |
| Filterbank (Gammtone) | Conv-TasNet |
| NAC (continuous) | C-AR |
| NAC (discrete) | D-AR, SELM |
One issue is that SSMs require dense representations of the input. - Time-domain models like Attenuate achieve this by a U-Net style bottleneck, where temporal resolution is traded for increasing depth in the bottleneck layers. - STFT models like SEMamba use a DenseNet encoder to produce deep latent representations. Inital experiments show that replacing the DenseNet with a simple 1D convolution drops PESQ by ~0.6.
Possible solutions: - Use U-Net model, which has been shown to work for time-domain and STFT domain. However, this kind of model is unsuited for NAC domain next token prediction. - Use dense encoder for STFT domain, like SEMamba and most its follow-up papers. - Use simple convolutional encoder (like Conv-TasNet (time) and TF-GridNet (STFT)) and accept worse performance. - Do not use any encoder. Just process 1D waveform or 2D STFT as they are. While performance would be much worse, there is absolutely no learnable encoder parameters, and thus no confounding encoder/decoder.
SSMs¶
Several classes of SSMs are compared: - Non-selective LTI models (S4D, LinOSS) vs. selective LTV models (Mamba, Selective LinOSS) - Decaying (S4D-Real, Mamba) vs. oscillatory models (S4D-Lin, complex-valued Mamba) vs. harmonic oscillatory models (LinOSS, Selective LinOSS)
| 1st-Order | 1st-Order | 2nd-Order | |
|---|---|---|---|
| \(\mathbb{R}\) | \(\mathbb{C}\) | ||
| LTI | S4D-Real | S4D-Lin | d-LinOSS |
| LTV | Mamba | C-Mamba | Selective LinOSS |
This allows for assessing the roles of selectivity, simple oscillations, and physically grounded harmonic oscillations for speech signal processing.
Note that C-Mamba and Selective LinOSS do not exist. For C-Mamba, the implementation and stability-proof is straight-forward and follows from the stability-proof of real Mamba. The resulting model is similar to GateLoop without K and Q parameters. The Selective LinOSS model can be shown to be stable for the IM-discretization. Implementation and proof follow later.
MIMO-SISO Confound¶
One confound is the MIMO (dense B,C) vs per-channel SISO (small per-channel B,C) structure. The S4D and Mamba models are SISO by design and applied per-channel, while LinOSS is a MIMO model, explicitly representing a bank of harmonic oscillators. This confound could be accepted, or a SISO variant of LinOSS be designed. This would imply that, instead of a multi-dimensional input forcing a large bank of oscillators, each input channel would force a much smaller bank of oscillators. Like in Mamba, this leads to a massive increase in overall state size but removes any channel mixing.
Making LinOSS SISO trades neuro-mechanical explainability (neural population coding) for computational efficiency. A single oscillator bank is well interpretable in terms of poles and resonances, and fits well with cortical models of auditory processing. This project is directly associated with Panzeri, so this is a relevant consideration.
The best option is likely to just make SISO vs MIMO an ablation: Create both variants so we can compare SISO SSMs, and ablate the MIMO version of LinOSS to compare performance.