Bridging Expressivity and Scalability with Adaptive Unitary SSMs

TL;DR

This paper introduces the Adaptive Unitary State Space Model (AUSSM), which enhances the expressivity and scalability of SSMs by incorporating adaptive, input-dependent recurrence inspired by biological neural systems, enabling better modeling of formal languages and long sequences.

Contribution

The paper proposes AUSSM, a novel SSM with adaptive, input-dependent recurrence that can perform modulo counting and simulate automata, improving expressivity over traditional SSMs.

Findings

AUSSM can perform modulo counting and simulate automata.

AUSSM outperforms prior SSMs on formal algorithmic tasks.

AUSSM achieves good performance on long time-series benchmarks.

Abstract

Recent work has revealed that state space models (SSMs), while efficient for long-sequence processing, are fundamentally limited in their ability to represent formal languages-particularly due to time-invariant and real-valued recurrence structures. In this work, we draw inspiration from adaptive and structured dynamics observed in biological neural systems and introduce the Adaptive Unitary State Space Model (AUSSM): a novel class of SSMs that leverages skew-symmetric, input-dependent recurrence to achieve unitary evolution and high expressive power. Using algebraic automata theory, we prove that AUSSM can perform modulo counting and simulate solvable group automata at precision logarithmically bounded in the input length, enabling SSMs to model a broad class of regular languages out of reach for other SSM architectures. To overcome the practical inefficiencies of adaptive recurrence, we develop a separable convolution formulation and a CUDA implementation that enables scalable parallel training. Empirically, we show that AUSSM and its hybrid variant-interleaved with Mamba-outperform prior SSMs on formal algorithmic tasks such as parity and modular arithmetic, and achieve competent performance on real-world long time-series classification benchmarks. Our results demonstrate that adaptive unitary recurrence provides a powerful and efficient inductive bias for both symbolic and continuous sequence modeling. The code is available at https://github.com/arjunkaruvally/AUSSM