Adaptive Unitary SSM (AUSSM)

• AUSSM is defined by an input-dependent recurrence using a skew-symmetric generator that ensures unitary, norm-preserving transitions for efficient long-sequence modeling.
• It employs a separable convolution strategy with custom CUDA implementation to reduce memory and compute costs, achieving near-linear scaling.
• Hybrid stacks combining AUSSM with Mamba layers extend model expressivity to recognize solvable regular languages, validated through parity, cycle, and long-sequence benchmarks.

Searching arXiv for the AUSSM paper and closely related SSM papers mentioned in the provided data. First, locating the primary AUSSM paper by title and arXiv id. Querying (Karuvally et al., 7 Jul 2025) Bridging Expressivity and Scalability with Adaptive Unitary SSMs Adaptive Unitary State Space Model (AUSSM) is a class of state space models introduced to reconcile two objectives that are often in tension in long-sequence modeling: the scalability of efficient SSMs and the expressivity required for input-dependent computation and formal-language recognition. In "Bridging Expressivity and Scalability with Adaptive Unitary SSMs" (Karuvally et al., 7 Jul 2025), AUSSM is defined by an input-dependent recurrence together with a skew-symmetric generator, so that the resulting discrete transition is unitary (orthogonal in the real parameterization) and norm-preserving. The construction is motivated by the observation that standard Linear Time-Invariant SSMs, despite fast convolutions and stable training, cannot implement counting or recognize formal languages such as parity or balanced parentheses with constant hidden width, because their recurrence is time-invariant and real-valued (Karuvally et al., 7 Jul 2025).

1. Position within the SSM landscape

Standard state-space models for long-sequence processing are written in discrete time as

$x_t = A x_{t-1} + B u_t,\qquad y_t = C x_t,$

with time-invariant, real-valued $A$ and fixed $B, C$ (Karuvally et al., 7 Jul 2025). Such Linear Time-Invariant systems support fast $O(L \log L)$ convolutions and stable optimization, but they do not track input-dependent dynamics and are therefore limited relative to Linear Time-Varying systems.

AUSSM modifies this setting in two coupled ways. First, it makes the recurrence input-dependent, with $A_t = f(u_t)$, so the hidden evolution adapts to each new symbol or input. Second, it constrains the recurrence through a skew-symmetric generator, $A_t^\top = -A_t$, so that after discretization the transition has a unitary spectrum and preserves norm (Karuvally et al., 7 Jul 2025). The intended effect is a middle ground between fully adaptive recurrence and trainability: gradients remain stable because the transition is norm-preserving, while the state dynamics can still encode symbol-conditioned rotations and counting behavior.

This design is presented as being inspired by adaptive and structured dynamics observed in biological neural systems, in particular cortical travelling waves with input-modulated frequencies (Karuvally et al., 7 Jul 2025). A plausible implication is that the model’s inductive bias is not merely spectral stability, but controlled modulation of oscillatory modes by the input stream.

2. Mathematical formulation

The continuous-time formulation is

$$\frac{dx(t)}{dt} = A(u(t)) x(t) + B u(t),\qquad y(t)=C x(t),\qquad A(u)^\top=-A(u),$$

with skew-symmetry imposed pointwise in the control-conditioned generator (Karuvally et al., 7 Jul 2025).

Under zero-order-hold discretization with step $\Delta_t \ge 0$, AUSSM takes the form

$$x_t = \Phi_t x_{t-1} + \Delta_t B u_t,\qquad y_t=C x_t,$$

where

$\Phi_t=\exp(\Delta_t A_t).$

Because $A$0 is real skew-symmetric, all eigenvalues of $A$1 lie exactly on the complex unit circle and $A$2 (Karuvally et al., 7 Jul 2025). This is the basis for the claim that AUSSM supports stable gradients without vanishing or exploding induced by the recurrence itself.

The paper further introduces a polar or diagonal parameterization under the assumption that all $A$3 are simultaneously diagonalizable:

$A$4

In that basis, each channel behaves as a pure rotation by angle $A$5 (Karuvally et al., 7 Jul 2025). This assumption is structururally important: it underlies both the efficient implementation and the automata-theoretic characterization, but it also restricts the class of formal languages reachable by the model.

3. Separable convolution and computational scaling

A central obstacle for adaptive recurrence is that naive time-varying kernels are expensive to materialize. AUSSM addresses this by exploiting a separable factorization of the kernel. For each channel $A$6,

$A$7

where

$A$8

The output can then be computed through two parallel scans (prefix sums) plus an outer-product:

$A$9

This reduces the storage requirement to $B, C$0 memory for the relevant sequences and yields $B, C$1 or even $B, C$2 time on modern GPUs (Karuvally et al., 7 Jul 2025).

The implementation section emphasizes that the asymptotic improvement is not by itself sufficient, because the linear-memory representation still carries large constants through tensors such as $B, C$3. The paper therefore describes a custom CUDA kernel that fuses the two prefix scans and the outer-product. In the reported implementation, one CUDA thread block processes a fixed minibatch chunk, for example $B, C$4, over hidden channels $B, C$5 in registers. The forward pass streams through $B, C$6, maintaining running partial sums of $B, C$7 and prefix sums of $B, C$8, while the backward pass recomputes forward intermediates on the fly to avoid extra memory (Karuvally et al., 7 Jul 2025).

On a 2080 Ti GPU, the paper reports three regimes. A recurrent PyTorch implementation has $B, C$9 GPU time and $O(L \log L)$0 memory and becomes unusable beyond approximately $O(L \log L)$1. A PyTorch implementation using separable scans has $O(L \log L)$2 time and $O(L \log L)$3 memory and reaches modest lengths around $O(L \log L)$4. The custom CUDA implementation has $O(L \log L)$5 time times a constant factor and $O(L \log L)$6 memory, with peak memory approximately Mamba’s footprint, and runs to $O(L \log L)$7 (Karuvally et al., 7 Jul 2025).

4. Expressivity and automata-theoretic characterization

The theoretical contribution is framed through algebraic automata theory, specifically Krohn–Rhodes decomposition (Karuvally et al., 7 Jul 2025). The first basic result is a modulo-counting construction. To count the number of ones modulo $O(L \log L)$8, the paper uses a one-dimensional AUSSM with

$O(L \log L)$9

When $A_t = f(u_t)$0, the hidden state rotates by $A_t = f(u_t)$1; when $A_t = f(u_t)$2, it remains fixed. Consequently, $A_t = f(u_t)$3 encodes the count modulo $A_t = f(u_t)$4 (Karuvally et al., 7 Jul 2025). This establishes a counting mechanism absent from standard real-valued diagonal SSMs.

The second step concerns cascade products. The paper states that Mamba-type layers with positive real $A_t = f(u_t)$5 simulate set-reset automata, whereas AUSSM layers simulate cyclic permutation automata. By alternating AUSSM and Mamba layers, one can implement the Krohn–Rhodes cascade of set-reset and cyclic automata (Karuvally et al., 7 Jul 2025). The corollary is that a hybrid stack of Mamba and AUSSM layers recognizes exactly the class of solvable regular languages, defined as those whose syntactic monoid has no non-solvable subgroup. The paper further states that this strictly exceeds the star-free languages accessible to any diagonal SSM with real, positive eigenvalues (Karuvally et al., 7 Jul 2025).

The finite-precision discussion is unusually explicit. Repeated complex multiplications accumulate error on the order of $A_t = f(u_t)$6, but adjacent $A_t = f(u_t)$7-th roots of unity are approximately $A_t = f(u_t)$8 apart, so reliable counting is possible up to $A_t = f(u_t)$9 in double precision (Karuvally et al., 7 Jul 2025). This suggests that the theoretical counter construction remains numerically meaningful far beyond practical sequence lengths.

5. Empirical behavior on algorithmic and long-sequence tasks

The algorithmic evaluation uses length generalization to $A_t^\top = -A_t$0 on tasks spanning multiple Chomsky-hierarchy levels: parity, cycle navigation, modular arithmetic, majority, repetition, and equation solving (Karuvally et al., 7 Jul 2025). Performance is reported in scaled accuracy, defined as

$A_t^\top = -A_t$1

Task	Reported outcome
parity	xLSTM: 0.00, Mamba: 0.13, AUSSM: 1.00, Hybrid: 1.00
cycle-nav	xLSTM: 0.80, Mamba: 0.86, AUSSM: 1.00, Hybrid: 1.00
mod-arith	xLSTM: 1.00, Mamba: 0.13, AUSSM: 0.48, Hybrid: 0.53
majority	xLSTM: 0.64, Mamba: 0.69, AUSSM: 0.57, Hybrid: 0.64
solve-eq	xLSTM: 0.24, Mamba: 0.05, AUSSM: 0.07, Hybrid: 0.07

The reported interpretation is specific. AUSSM perfectly solves parity and cycle counting, which the paper presents as validation of the modulo-counter theory. At the same time, AUSSM struggles on majority and equation tasks, which are described as requiring dissipative forgetting rather than purely norm-preserving recurrence (Karuvally et al., 7 Jul 2025). The hybrid AUSSM+Mamba model consistently achieves the best or near-best performance across tasks, supporting the claim that alternating adaptive unitary and real dissipative layers expands the effective expressivity of diagonal SSM stacks.

On long-sequence benchmarks, the paper reports results on UEA-long classification and weather forecasting. For UEA-long classification, covering six datasets from Heartbeat with length approximately $A_t^\top = -A_t$2 to Worms with length approximately $A_t^\top = -A_t$3, the hybrid AUSSM+Mamba model raises mean scaled accuracy from Mamba’s $A_t^\top = -A_t$4 to $A_t^\top = -A_t$5, compared with linOSS at $A_t^\top = -A_t$6 and S6 in Mamba at $A_t^\top = -A_t$7 (Karuvally et al., 7 Jul 2025). For weather forecasting with prediction of the next $A_t^\top = -A_t$8 steps, the hybrid model achieves $A_t^\top = -A_t$9 versus Mamba at $$\frac{dx(t)}{dt} = A(u(t)) x(t) + B u(t),\qquad y(t)=C x(t),\qquad A(u)^\top=-A(u),$$0 and S4 at $$\frac{dx(t)}{dt} = A(u(t)) x(t) + B u(t),\qquad y(t)=C x(t),\qquad A(u)^\top=-A(u),$$1 (Karuvally et al., 7 Jul 2025). All runs were conducted on a single 2080 Ti.

6. Inductive biases, limitations, and open problems

The inductive bias of AUSSM is summarized in the paper as adaptive unitary recurrence that preserves norm while enabling input-driven gating of frequency modes and rich, multiscale memory (Karuvally et al., 7 Jul 2025). In practical terms, this means the model is not merely a stable oscillator: the input sequence modulates which rotational modes are active, which is why the architecture is effective on parity and cycle-counting tasks.

The limitations are equally explicit. AUSSM requires the matrices $$\frac{dx(t)}{dt} = A(u(t)) x(t) + B u(t),\qquad y(t)=C x(t),\qquad A(u)^\top=-A(u),$$2 to be simultaneously diagonalizable, and the paper identifies this as a restriction to solvable languages. Its per-step compute is more costly than that of purely LTI SSMs. The best block ordering in deep hybrid stacks is not fully understood. Large-scale language modeling, for example at billions-token scale, remains to be demonstrated (Karuvally et al., 7 Jul 2025).

Several common misconceptions are directly addressed by these results. Unitary recurrence does not imply universal superiority on sequence tasks: the majority and equation-solving outcomes indicate that dissipative mechanisms remain important. Conversely, adaptive recurrence does not necessarily preclude efficient training: the separable convolution and CUDA implementation show that fully input-dependent recurrence can be made scalable. A further misconception would be to read the automata-theoretic results as establishing coverage of all regular languages. The paper does not claim that; rather, it characterizes a hybrid Mamba+AUSSM stack as recognizing the class of solvable regular languages under its stated assumptions (Karuvally et al., 7 Jul 2025).