Results June 26
This projects is about testing LinOSS (Linear Oscillatory State Spaces) for speech enhancement. LinOSS is a State Space Model based on a second-order ODE, meaning its hidden state evolves with the dynamics of forced harmonic oscillators. Practically, this means LinOSS carries a bank of oscillators as a hidden state, where the input drives the oscillations. This is inspired by the cortical oscillations found in human brains. LinOSS particularly well suited for long-sequence predictions that are of periodic nature. For discretization, LinOSS offers two options: A decaying implicit discretization (IM) and a symplectic implicit-explicit discretization (IMEX). Dampened LinOSS (d-LinOSS) introduces dampening into the IMEX-discretized state dynamics, so the model can learn symplectic and decaying oscillations, combining the IMEX and IM discretization. In the following work, LinOSS will always refer to d-LinOSS.
Speech is similarly composed of glottal harmonic oscillations, filtered by the shape of the vocal tract. Furthermore, we observe that speech modulates brain waves to concentrate on the theta band, which locks to the slow moving magnitude envelope of the time-domain waveform of speech, corresponding to syllables in speech. Thus, the LinOSS oscillator bank might lock on to these oscillatory dynamics in speech, capturing the underlying system better than a first-order model.
Experiments¶
The first experimental setup integrates LinOSS in SEMamba, replacing the selective SSM backbone with an oscillatory SSM. SEMamba acts in the spetro-temporal STFT features, which contain a the periodic structure of voice harmonics along the frequency axis and a quasi-periodic structure over time.
Time-Frequency Dynamics¶
SEMamba uses a DenseNet to encode the STFT magnitude and phase into a 3-dimensional latent speech representation, with the shape (T, F // 2, C). Each time-frequency bin is thus embedded as a vector of size C. The SSM is then run bi-directionally across the time-axis (observing frames of shape (F//2, C)) and across the frequency axis (observing frames of shape (T, C)).
Replacing Mamba with LinOSS has slightly different effects for the time and frequency axis: - On the time-axis, LinOSS almost matches Mambas performance. PESQ slightly reduces, but MOS scores stay similar. - On the frequency-axis, LinOSS has more detrimental effects, worsening PESQ even more, and also reducing MOS scores and Phase losses. - Adding LinOSS on both axes accumulates the negative impact.
Comparison to other SSMs¶
In a second experiment, other SSMs are implemented and matched for parameters with Mamba/LinOSS.
First-Order & Non-Selective: S4D (SISO)¶
The SSM S4D is very similar to LinOSS except that its hidden state evolves with first-order dynamics. In comparison to S4D, LinOSS slightly outperforms S4D in terms of perceptual scores and SI-SDR. LinOSS significantly outperforms S4D in terms of matching the phase of the clean waveform. This suggests an advantage of oscillatory dynamics for modeling phase information in speech.
First-order Complex & Selective: Selective LRU (similar to GateLoop & Mamba 3)¶
Important Note: These results were gathered from MIMO selective LRU
The SSM selective Linear Recurrent Unit is a first-order complex selective SSM. It uses scaled rotations for the transition matrix, where each scale and rotation a independently selective to the input (\(h_{k}=\nu_{k}e^{i\theta_{k}}h_{k-1}+b_{k}u_{k}\)). The complex nature allows the model to learn rotational dynamics, and thus better conforms to oscillatory signals as well. However, in comparison to LinOSS, it lacks the inductive bias of physical harmonic oscillators.
This model strongly outperforms LinOSS and performs within ~0.1 PESQ below Mamba and maintains phase slightly better. Along other metrics, both models perform about equal. Again, this strengthens the hypothesis that oscillatory dynamics help modeling phase information in speech.
Mamba 3 actually incorporates rotational dynamics as well, and GateLoop basically the same SSM as this model, so the theoretical novelty is rather limited. Executing this model in parallel SISO style, like Mamba, the performance is slightly worse, which matches the findings in Mamba 3.
Summary¶
This experiment suggests that complex state dynamics in SSMs help in phase modeling, which is a logical connection. While the complex selective SSMs can model oscillations as well, they lack the inductive bias of an actual bank of harmonic oscillators, which is the bio-physical foundation of LinOSS. Unfortunately, an input-selective LinOSS-style oscillatory bank is difficult to construct due to system instability: The IMEX discretization in LinOSS has a closed-form solution for the dampened second-order LTI dynamics, with eigenvalues within the unit-circle. But once time-varying dynamics are introduced, a Lyapunov stability proof is required, which I didn't solve yet.
Conclusion and Future Ideas¶
LinOSS has not proven superior to other SSMs for speech processing. It shows improvements in phase-modeling, but this is likely reproduced by a complex state-space as well. The inductive bias of harmonic oscillators in particular have not shown any advantage.
Potential explanations - For frequency modeling, the SSM must be able to adapt to a time-varying pitch, which LinOSS struggles most with. Also formants and plosives contain little oscillatory structure which LinOOS can adapt to. - For time modeling, LinOSS is better suited due to its long-range modeling capabilities. However, its still outperformed by Mamba, and a selective SSM with similar long-range modeling capabilities (e.g. Mamba 3 with RoPE) likely outperforms both. - For raw waveform modeling, LinOSS has trouble adapting its oscillator banks to the complex waveform dynamics. This might be caused by lack of hyperparameter tuning, or might be a general shortcoming of the model.
Future ideas: - Incorporating LinOSS for specialized tasks: - Use in combination with Mamba blocks in SEMamba to improve tracking of periodic structure (along frequency-axis) and quasi-periodic structures (along time-axis) - Explore parallel LinOSS-Mamba blocks, where input is passed through both SSMs and their output is combined.