Research map / Architectures and efficiency / RNNs and normalization
Bridging Expressivity and Scalability with Adaptive Unitary SSMs
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes AUSSM, a fully adaptive diagonal SSM with input-dependent skew-symmetric (unitary) recurrence, plus a separable convolution formulation and CUDA kernel, and shows that hybridizing it with Mamba recognizes solvable regular languages and improves on algorithmic, UEA time-series, and weather forecasting benchmarks.
Problem
State space models like Mamba have limited expressivity for formal languages (e.g., parity, modular counting) because their recurrences are time-invariant or only partially adaptive and have positive real eigenvalues, while fully adaptive general LTV models are computationally unscalable.
Approach
Make the SSM state matrix input-dependent and skew-symmetric so its exponential is unitary with input-controlled rotation frequencies; show the resulting kernel is separable, enabling linear-time parallel computation, and combine it with Mamba blocks to recognize all solvable regular languages.
Key results
- On the weather forecasting benchmark, AUSSM Hybrid reaches MAE 0.342 versus 0.464 for Mamba and 0.508 for LinOSS.
- AUSSM achieves perfect scaled accuracy (1.0) on parity and cycle navigation, where Mamba reaches 0.13 and 0.86 respectively.
- On UEA classification, the hybrid model averages 35.2 versus 28.8 for Mamba, but linOSS achieves the best average (44.2).
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use AUSSM Hybrid instead of Mamba | algorithmic task models · Repetition | scaled test accuracy: 0.15 -> 0.45 | Helped |
| Use AUSSM instead of Mamba | algorithmic task models · Bucket sort | scaled test accuracy: 0.69 -> 0.92 | Helped |
| Use AUSSM Hybrid instead of Mamba | algorithmic task models · Majority count | scaled test accuracy: 0.45 -> 0.37 | Hurt |
| Use AUSSM Hybrid instead of Mamba | algorithmic task models · Majority | scaled test accuracy: 0.69 -> 0.64 | Hurt |
| Use AUSSM Hybrid instead of Mamba | algorithmic task models · Solve equation | scaled test accuracy: 0.05 -> 0.07 | Helped |
| Use AUSSM Hybrid instead of Mamba | algorithmic task models · Mod arithmetic | scaled test accuracy: 0.04 -> 0.23 | Helped |
| Use AUSSM Hybrid instead of Mamba | algorithmic task models · Mod arithmetic without brackets | scaled test accuracy: 0.13 -> 0.53 | Helped |
| Use AUSSM instead of Mamba | algorithmic task models · Cycle navigation | scaled test accuracy: 0.86 -> 1.0 | Helped |
| Use AUSSM instead of Mamba | algorithmic task models · Parity | scaled test accuracy: 0.13 -> 1.0 | Helped |
| Replace selected Mamba blocks with AUSSM blocks | Mamba backbone · UEA time-series classification (Avg) | scaled test accuracy: 28.8 -> 35.2 | Helped |
| Replace selected Mamba blocks with AUSSM blocks | Mamba backbone · Weather (720-step forecasting) | Mean Absolute Error: 0.464 -> 0.342 | Helped |
| Use separable convolution instead of recurrent PyTorch implementation | AUSSM · Scalability evaluation | runtime efficiency | Helped |
Surprising findings
- AUSSM alone performs poorly on majority count (0.096) and equation solving, where dissipative dynamics seem needed, and linOSS outperforms the hybrid on UEA average.
- Mamba with complex eigenvalues reportedly fails to learn simple formal tasks that the hybrid performs perfectly on.
Limitations
- Assumes simultaneously diagonalizable recurrent matrices, so only solvable regular languages are covered, not full regular languages.
- Separable kernel retains a constant-factor overhead over LTI models; the best way to combine AUSSM and Mamba blocks is unclear.
- Experiments are modest-scale and not validated at foundation-model scale.
Experimental setup
Datasets: algorithmic reasoning tasks (parity, modular arithmetic, repetition, bucket sort, majority, cycle navigation, equation s UEA long time-series classification benchmarks (Heartbeat, SCP1, SCP2, Ethanol, Motor, Worms) Weather forecasting dataset
Benchmarks: algorithmic reasoning tasks (Chomsky hierarchy suite) UEA time-series classification (six datasets) Weather forecasting (720-step)
Metrics: scaled accuracy mean absolute error runtime peak memory
Baselines: Mamba xLSTM S5 S6 linOSS Informer LogTrans LSTnet S4 LSTMa
Builds on: Mamba S4 jPCA Krohn-Rhodes theory linOSS
Closest papers by meaning
- On the Expressiveness and Length Generalization of Selective State Space Models on Regular Languages (AAAI 2025 · similarity 0.88)
- Longhorn: State Space Models are Amortized Online Learners (ICLR 2025 · similarity 0.87)
- MaRK: Markov-adapted Recurrent Kernels for Dynamic Operator Conditioning in State Space Models (NeurIPS 2026 · similarity 0.87)
- Universality of Linear Recurrences Followed by Non-linear Projections: Finite-Width Guarantees and Benefits of Complex Eigenvalues (ICML 2024 · similarity 0.87)
- The Expressive Capacity of State Space Models: A Formal Language Perspective (NeurIPS 2024 · similarity 0.87)
- Theoretical Foundations of Deep Selective State-Space Models (NeurIPS 2024 · similarity 0.86)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.