Bridging Expressivity and Scalability with Adaptive Unitary SSMs
Authors: Arjun Karuvally
Franz Nowak
T. Anderson Keller
Carmen Amo Alonso
Terrence Sejnowski
Hava T Siegelmann
Venue: NeurIPS 2025 poster ()
Keywords: State Space Models
Unitary Recurrent Neural Networks
Adaptive Recurrent Neural Network
Primary area: deep_learning
TL;DR: Introduced a new SSM that is maximally expressive and scalable to long sequence modeling tasks

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

Bridging Expressivity and Scalability with Adaptive Unitary SSMs
Arjun Karuvally
Franz Nowak
Salk Institute for Biological Studies akaruvally@salk. edu
ETH Zürich franz nww.k@inf .ette.thcc. ch
T. Anderson Keller
Carmen Amo Alonso
The Kempner Institute for the Study of Natural and Artificial Intelligence at Harvard University
Computer Science Department
Stanford University
t. anderson keller@gmail. com
camoalon@stanfoom edu
Terrence J. Sejnowski Salk Institute for Biological Studies terry@salk.ed edu
Hava T. Siegelmann
University of Massachusetts Amherst hava@umass edu
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 gegguages-partiuulle
due to time-invariant and real-valued recurrence structures. 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 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 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-interleavee with Mamba-ouuperforn prior SSMs on formal algorithmic tasks such parity and modular arithmetic, and achieve competent performance on real-world long time-series classification benchmarks. Our results demonstrate that adaptive unitary recurrence provides powerful and efficient inductive bias for both symbolic and continuous sequence modeling. The code is available at https///ghhlbl.ooom///uuuuvvv//s
Introduction
Modeling long-range dependencies efficiently and expressively remains central challenge in quence modeling. While Transformer architectures have achieved state-of-the-ard results across domains such language modeling [ 2 3 forecasting 5) and protein design (6), their quadratic complexity with respect to sequence length limits scalability 77] In response, recent work has explored state space models [8 (SSMs) as scalable alternative, using linear-time convolutions and structured recurrence to enable efficient processing of long sequences 99 [0]. Despite the computational advantages SSMs offer, they fundamentally limited in their ability express general linear time-varying systems and formal languages efficiently. Even basic regular languages
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

that require counting, such as parity or balanced parentheses [1] are impossible for practically used SSM architectures like Mamba that have positive real eigenvalue spectra. Frontier SSMs are either Linear Time Invariant (LTTI) or partially Linear Time Varying (I.TV) resulting in weaker expressivity compared to more general LTV systems that are capable of approximating any non-linear dynamical systems 12] Two properties thus emerge necessary for increasing the expressivity of SSMs: general cigenvalue spectrum and adaptive recurrence. However, incorporating both these properties naively in SSMs introduces gradient instability through the exploding/vanise
gradient problem 13. and leads to quadratic computational complexity, severely limiting scalability.
In this work, we propose the Adaptive Unitary State Space Model (AUSSM) as a principled middle ground between scalability and expressivity. AUSSM is fully adaptive state space model with linear time-varying recurrence and a unitary eigenvalue spectrum, combining the theoretical benefits of vime-narying recurrence with the practical advantages of structured, conserved dynamics. We formally prove that AUSSM can perform modulo counting with constant- width hidden states, and that combining AUSSM with existing non-adaptive models like Mamba yields maximal expressivity within the class of diagonal SSMs. To make this architecture scalable, we introduce novel separable kernel formulation for adaptive SSMs that exposes efficient computational algorithms which reduce the quadratic cost of adaptive recurrence to linear time and space. Empirically, we validate the theoretical claims through suite of algorithmic tasks, demonstrating substantial performance gains over Mamba, and showing that AUSSM retains competitive efficiency through optimized CUDA implementation. Further, we evaluate the long-range modeling capabilities by testing on a suite of time series benchmarks.
Interestingly, structured unitary and adaptive dynamics are also found as emergent behavior in biological neural systems [14] and even trained non-linear recurrent neural networks [15, where they are believed to support flexible integration information over space and time [66. We take the computational role of these structured unitary dynamics motivation to derive AUSSM using skew-symmetric ODE used in identifying purely rotational features from data neuroscience [17.
AUSSM provides a new architectural foundation that bridges formal expressivity and practical scalability (Figure []] It expands the space of scalable SSMs by showing that adaptive (and time- varying) recurrences can be made computationally efficient, unlocking new capabilities for symbolic and longecontext sequence modeling that is grounded in biological principles and theory.
2 Background and Motivation
State Space Models (SSMs) have emerged efficient alternatives to Transformers for sequence modeling, particularly in long-context settings. SSMs compress arbitrarily long sequences into a fixed-dimensional hidden state using recurrent formulation and this can be computed in parallel using an efficient convolution formulation.
The most general SSMs are described by continuous-time Ordinary Differential Equation (ODE);
dx(() = Aat(t) + Bru(t), s(t) = Gar(t) dt
or its discrete counterpart:
(1) = A1z(11 Biu((,, s(1) = Cxx(1)
where x(t) IR" is the hidden state, u(f) is the input, and (f) is the output. The matrices A,, B,, Ca define the system dynamics and may vary over time. In the discrete system, these matrices have an equivalent discretized counterpart in A', B', C", respectively. The discrete recurrence can be reformulated as a parallel convolution:
B,,C",
v(t) (AL (..A)) Blu(k)
Ab+1)
However, this convolution kernel requires O(L?) memory for sequence length L, as each kernel entry must be stored for and l that index over the sequence length. To avoid this, most practical SSMs
Interestedd readers are referred to Appendix B partial LTV, and LTV.
Expressivity of Single Block SSMs for the definitions of LTI,

---

MAMBA
AUSSM
Regular
Non-diagon dingonal LRNN
Solvable
AUSSMM-MAMAA
Bybrid
A(4e)
Cyclic Group
Sear Free
AUSSM
MAMBA
B3
(a) Formal language classes recogniz- (b) SSM block structure of Mamba and AUSSM (input-dependent com- able by different architecturcs.
ponents are shaded in blue).
Figure 1: (a) Existing practical SSM blocks like Mamba use fast parallel algorithms for computing the output, resulting in tradeoff with expressivity. Non-diagonalizable Linear RNNs the most expressive (in formal language terms) but lack scalable computational algorithms and suffer from gradient issues. AUSSM balances the unitary recurrence. Fast SSMs with improved expressivity can be built by combining AUSSM with MAMBA blocks. (b) The AUSSM block uses the same block structure Mamba [TO], where the S6 SSM in Mamba is replaced with AUSSM. The main difference between AUSSM and S6 is the adaptive recurrence, where in the case of S6, C. and A are adaptive, whereas in AUSSM, and A are adaptive (see Section 3 for details). AUSSM blocks can be used as drop-in replacements for existing SSM backhones to provide higher expressivity (see Section 3.1 for theoretical and Section 5 for experimental validation).
tradeoff using a fully adaptive diagonal
allowing for a compressed storage of the kernel but significantly restricting expressivity. Recent SSMs like Mamba [TO] introduce partial adaptivity, where B, C, and step size A) are adaptive while keeping A fixed or diagonal. However, such models cannot simulate general Linear Time-Vare Varying (LTV) systems or perform counting- based tasks (e.g.. parity, modular arithmetic) with constant-width hidden states (see Appendix B These limitations prevent Mamba from modeling input-sensitive dynamics or general multiscale time-varying behavior (Appendix(B) B). There are other approaches that try to improve the expressivity of SSMs by generalizing real-valued recurrences. Notably, Linear Recurrent Units [8] generalize the real-valued eigenvalue spectra with initialization close to the unit circle on the imaginary plane. This formulation has been shown to be capable of universal approximation when interleaved with non-linear multi-layer perceptrons [1). However, this approximation relies on perfectly storing the dynamical system history without regard to resource constraints. General LTV systems are much more flexible they have the capability to gate information based on input, thereby retaining only selected information (compressed information) that necessary for processing, rather than lossless history. Another notable work is linear Oscillatory State Spaces (linOSS) (20), where simple harmonic ODEs are discretized derive novel oscillatory SSMs with conservation properties identical to AUSSM. The linOSS models more expressive than SSMs with purely real eigenvalues, but fall short of an LTV system. AUSSM balances the two the improved expressivity of diagonal LTV systems (using adaptive recurrence) and the scalability of separable convolutions.
assume time-invariant dynamics = BB
3 Adaptive Unitary State Space Model (AUSSM)
We tackle the problem of balancing expressivity with scalability in Adaptive State Space Models by introducing two features. Adaptive input-dependent recurrent matrix improving expressivity, and unitary dynamics addressing training scalability. this section, we derive AUSSM from the skew-symmetric ODE used to identify rotational features in the brain [17], then prove that the inputs control AUSSM rotational frequencies smoothly, enabling a stable and effective adaptive SSM. Next, prove that the AUSSM, combined with regular Mamba layers, is maximally expressive in the class of diagonal SSMs in terms of formal language recognition.
To derive the AUSSM model with purely rotational properties, we use the skew-symmetric Ordinary Differential Equation (ODE) used in the rotational Principal Component Analysis (jPCA) procedure- variant of Principal Component Analysis (PCA) used in neuroscience [17]. jPCA is used to identify
3

---

rotational features of dynamical system using observations from it. Since our requirement is to process an imput signal u(t) into the hidden state, we use version of the jPCA ODE with control input given by Equation [ with the additional constraint that the input matrix A, is skew-symmetric (with purely imaginary eigenvalues) and B, and C2 stay constant with time, i.e. R; = B and C₁ C.
We discretize the ODE following the procedure used in Mamba [T0] with a step size A, E IR to obtain a discrete dynamical system (See Appendix [ for details),
{ [x(() = exp(A,A.) A.) xxt- 1) (v()) = Cr(().
1) + A,Bu(t),
(4)
Note here that A, changes with time from adaptivity, and it is klew-symmetric matrix. We assume that A,, Vt belongs to a class of simultaneously diagonalizable matrices. Therefore cxp(A,A.) can be diagonalized to obtain ixp(AiiA,(t) the matrix A,. This implies that the final discrete dynamical system has purely unitary eigenvalues, i.e., eigenvalues exactly on the unit circle. The AUSSM ODE is a marginally stable, time-varying linear system where the input both drives and dynamically reshapes the system. The skew-symmetric nature of A, guarantees marginal stability by ensuring that all cigenvalues lic on the imaginary axis in continuous time, or on the unit circle after discretization (see Lem. in Appendix D). This structure enables long-term memory retention without gradient explosion or decay (see Lem. Appendixx
where (t) R and each (t) is the j eigenvalue of
P
The adaptivity of A, is enforced by making A, a function of input with A, f(u(t)) where f: IR + R is the function defining how the input influences the recurrent matrix. With adaptivity, the input acts as a control signal, shaping the rotational dynamics based on the instantaneous input, analogous to gain scheduling or bilinear control systems [2T 22]. This design allows the system to dynamically traverse spectrum of rotational behaviors in the state space, facilitating expressive temporal modeling driven by the input signal.
Theorem (iput-Moduulate
Rotation Frequencies via Exew-Symmetric Generator). Let A R - ROxR be smooth function such that A(u) is skew-symmetric for all u E RR. Then for each u E IR, all eigenvalues of A(u) lie on the imaginary axis, and the eigenvalues of the discrete-time transition matrix (u))= = exp( A ) lie on the complex unit circle. Furthermore, the eigenvalues of A(u) depend continuously on u, and thus the angular frequency of state-space rotation is smoothly and directly modulated by the input. See proof in Appendix D
Hence, by designing A(u) appropriately (e.g.. via a learnable function f(u),, the AUSSM can modulate the rotational speed and mode structure of the hidden state space based on the input signal in smooth and controlled manner. Further details on how the above SSM is practically implemented are in Section where we diagonalize the above ODE and introduce input adaptivity. The inputs also have a dimension of d, in which case the proposed SSM is applied to each input dimension following the approach used in Mamba.
3.1 Formal Expressivity
Given our construction of the AUSSM, how expressive is it for formal languages?
The goal of formal expressivity theory is determine, for a given architecture class, which functions or formal languages can be represented by some finite instantiation of that architecture. The quantification is over architectures -i..., over possible finite hyperparameter settings such as model dimension, input dimension, or transition rank-rather than over the parameters within fixed instantiation. Formal expressivity analysis goes beyond our earlier discussion of limitations of SSMs in expressing LTV dynamical systems (further detailed in Appendix B).
Representing simple formal languages has been found to be major weakness of SSMs. Recently, flurry of research has utilized algebraic automata theory. specifically Krohn- Rhodes theory [23] to analyze the types of formal languages expressible by different LLM architectures, notably transform- ers (24) and SSMs [1 25) The Krohn-Rhodes decomposition theorem states that any finite-state machine can be simulated by cascade of simpler automata drawn from two types: permutation au- tomata, which model reversible group-like behavior, and reset automata, which model state-resetting dynamics (the next state depends only on the input, not on the current state). This result implies
In practice. there will always be small deviations from the ideal theoretical behavior due to the limited precision of modern computers.

---

that complex regular languages can be recognized by composing SSMs that simulate these simple automata. There is a subset of finite-state automata whose decomposition contains only set-reset automata and cyclic permutation automata, which suffices to recognize a large subset of regular languages, the so-called solvable languages.
Most SSMs used in practice have diagonal or diagonalizable transition matrices A that can only have positive eigenvalues. As [I]] showed, this means they cannot perform modulo counting, restricting their expressivity to the subset of star-free regular languages
also outlines the necessary conditions for SSMs to overcome this limitation and represent a larger class of regular languages. This requires 1. the ability to implement modulo counters, and 2. the ability to implement Krohn- Rhodes cascade products. Here, we reiterate their relevant results and show that our implementation not only satisfies these conditions but can recognize any solvable regular language, a language class out of reach for most practical SSMs. For a unified overview situating our expressivity results within the SSM expressivity literature, see A
Fact 1 (II), Thm. 2). Diagonal (or diagonalizable) SSMs with only positive eigenvalues cannot per- form modulo counting at finite precision, which means they can only recognize star-free languages
Lemma 1. For any & E 2", one can construct a single-layer AUSSM that counts modulo k, which means AUSSMs can simulate arbitrary cyclic group automata.
Proof sketch. Assume we want to count the number of I's modulo 2 in a length-T input sequence
(u) E {0, 13" A single-layer AUSSM with by = 1, A(1) = -1, A(0) = 1, and ((0)) = B(1) = 0 will have hidden state of 2; -1 for odd counts and 2e 1 for even counts of 1 up to position f. Similarly, to count modulo 4, we can set A(1) to the fourth root of unity, i.e., either i or -i, and A(0) 1, (() (1) 0, before. This method can be extended to other mod & counters by setting A(1) to the kth root of unity, exp(2ri/K /k). An AUSSM can take on these parameters as it uses input-dependent A matrices whose eigenvalues lie the unit eircle of the complex plane This technique can be extended to perform modulo k addition, which allows the simulation of cyclic group automata (see
E
Lemma 2. An SSM consisting of interleaved Mamba and AUSSM blocks (hybrid Manba+AUSSM) can implement cascade products of automata simulated by Mamba SSMs and AUSSMs.
Proof sketch. [T] Lem. 19] showed that multilayer Mamba SSMs can implement cascade products of Mamba layers simulating set-reset automata, which, by Schützenberger'' theorem 266. means they can recognize any star-free language. This can easily be extended to show that any automaton simulated by Mamba or AUSSM layers can be joined into a cascade product within alternating Mamba and AUSSM blocks. This works because we can always add additional padding layers at any point in the hybrid SSM without changing the behavior of the remainder of the SSM.
Theorem 2. Hybrid MUSSM can recognize any solvable language, that is, any regular language whose syntactic monoid does not contain non-solvable subgroups.
Proof sketch. By Lem. 1 an AUSSM layer can simulate cyclic group automata, and 11 Lem. 19] showed that a Mamba layer can simulate set-reset automata. Now, the Krohn- Rhodes theorem states that every finite automaton divides a cascade of alternating aperiodic monoids (set-reset automata) and finite simple groups (permutation automata). A finite group is solvable iff its decomposition series contains only cyclic groups of prime order (cyclic group automata with prime-length cycles) [27] Ex. 3.4.8). By Lem. 2 hybrid Mamba+AUSSI can implement the Krohn- cascade product of set-reset automata (Mamba) and cyclic group automata (AUSSM). Therefore, it can recognize all solvable languages. (cf. Thm. 21]).
'Star-free languages are those languages that can be defined without the use of Kleene star, only using concatenation, union, and complement.
'Note that this theorem assumes finite precision; similar limitations expected to persist under logarithmi- cally bounded precision. though we do not attempt full extension here.
We assume an idealized setting in which the numerical precision is logarithmic in the sequence length. For most sequence lengths seen in practice, this reasonable assumption (see {E details).
5

---

Regular languages that require representing more complex non-solvable group transformations, such as the word problem in Ss or As, lie outside of this set, and according to widely held assumptions about computational expressivity theory, cannot be modeled by diagonal SSMs [28). This means combining Mamba with AUSSM maximizes the representational capacity of diagonal SSMs (short of lifting the diagonal transition restraint, which leads to poor scaling).
Separable Convolution Kernels for Scalable Adaptive SSMs
One of the main challenges in designing SSMs is the computational efficiency of the implementation. Simulating the discrete dynamical system in Equation 4 naively is not computationally efficient as it leads to quadratic memory scaling when it parallelized using the typical SSM convolution procedure (.p.endix(F.I).. In this section, we introduce a separable kernel formulation for the efficient computation of adaptive time-varying SSMs. Our formulation works directly in the convolution form of the SSM and instantly exposes the separability and is applicable to wider class of adaptive SSMs as shown below. We note here that the separable kernel formulation is not specialized for the AUSSM, but applies to any class of SSMs that are simultaneously diagonalizable (29]. We therefore formulate the theory in the general case and provide sufficient conditions to apply the theory in practice.
The general convolution formulation of general SSMs Equation is typically used to convert discrete dynamical system form of an SSM to an efficient parallel implementation. This form is abstracted as applying a convolution on the imput, following the equation
(t))
Σ k)t,k) k
The reason for the quadratic memory scaling of the convolution operation be observed in this abstracted form as storing the K(t, k)) convolution kernel requires, in general, ((1,2) memory, where I, is the sequence length.
Separable convolution kernels have the additional property that K(t,k) f(t) (k)) that enables writing the output
(t) = f(t) Σ (k))((k)
ASt
Storing the additional ftt and glk requires only additional ((L)) memory. This is comparable to the non-adaptive case, which has a scaling (T.), producing asymptotic memory efficiency matching that of the non-adaptive SSM with only constant factor increase in memory use. The above convolution formulation can be efficiently computed in (L)) time by using the parallel prefix sum algorithm (30). It instructive to apply this formulation to an existing SSM identify efficient computational structures we use the partially adaptive Selective State Space Model (S6) used in the popular Mamba model 1T)
Separable Convolution Formulation of Mamba Selective SSM (S6): In S6, the matrices C and B vary with input (making the SSM selective to input). in addition to the step size A varying with time. This generalization results in the output of the SSM written in the convolution form as:
Yr ΣΣ Cij exp((IA., Ag)))A Baj ukkk)
Here, the input u E Rd is a vector and the SSM is applied to each input dimension in parallel. The index is the imput dimension d, and indexes the hidden state dimension n. In the general convolution formulation we showed above, the S6 output is formulated as applying the convolution kernel (t,,k) Cij exp((tAn AA (A,) A A.: here that, unlike typical time-invariant SSMs, the S6 convolution kernel is unique to each y as the kernel is separable and we can use the procedure we introduced above to compute the S6 output in time and memory-efficient manner (see Appendix for an efficient PyTorch Implementation of the S6).
the inputs over time u, Note
A, B, C change with time. Since K(t,k) k) (C., exp(tA,;:A)) exp(-kAs,4))

---

Average Runtime
Peck GPU Memory Usage
AUSSM AUSSM-adh AUSSM--rpaccble
Mennar dal)
memmhs
A AUSSMereeurree
AUSSM-cuda
100
10"1
10-1
500
1000
1500
2000
2500
9000
3500
4000
0 500
1000
1800
2000
2000
3000
3500
0000
lengti
leagch
(a)
(b)
Figure 2: AUSSM with separable convolution achieves efficient runtime and memory scaling for fully adaptive SSMs. The runtime and peak memory usage of four implementations are compared: recurrent Py PyTorch AUSSM, separable PyTorch AUSSM, our optimized CUDA AUSSM kernel, and the Mamba CUDA kemel.. (a) The AUSSM CUDA implementation outperforms both PyTorch baselines in speed and memory efficiency, and approaches the memory efficiency of Mamba despite AUSSM's full adaptive recurrence. Notably, the PyTorch implementation of the separable convolution has better runtime efficiency compared to the recurrent implementation, albeit at a higher memory cost. (b) The AUSSM CUDA kernel has significantly lower memory footprint, identical to that of the partially adaptive and optimized Mamba CUDA kernel.
Separable Convolution Formulation of AUSSM:
In the case S6, the separable formulation was easily revealed directly from the convolution form. In the case of AUSSM, this separability is not possible for the most general case. However, when the set of recurrent matrices A, is simultaneously diagonalizable, the output of the AUSSM in Equation 4 be formulated as (See Appendix for details)
Vi ΣΣ E:D 3
C2 exp
i> I
A
(k()))
(8)
exp
E OAIEs)
hias is the angle argument of the unitary discretized A' matrix in the polar form. Note here that as the AUSSM has complex eigenvalues. the final output is also complex, and we use only the real part of the output with the function R[.). With this formulation of the AUSSM recurrence, a memory and time efficient computation of the adaptive SSM is obtained, however, implementing this convolution directly in PyTorch can still result high memory usage as the constant in the (T)) is bdn where b is the batch size, d is the input dimension and is the hidden dimension resulting in a large constant factor. We therefore create a CUDA kernel, where this additional complexity is hidden and the hidden state is only partially materialized in the CUDA kernel (Appendix G.2).
Here, A =E, Uk) + m.
Another approach to improving the performance of SSMs is tensor core optimization. In tensor core optimization, special hardware features in NVIDIA GPUs called tensor cores are used to speed up matrix computations inherent in SSM implementations. This approach is not an entirely new algorithm with improved scaling behavior, but an implementation approach that enables speed-up in the special case of GPU architectures where tensor cores are available which is most high-end GPUs available in the market. Experimental evaluations have also shown that tensor core optimization approaches can provide constant factor increase in performance on high-end GPU hardware, but retain the same scaling behavior the big-O scaling factor. Recent works have utilized this approach to improve the performance of time-varying SSMs, but side-step the fundamental algorithmic limitations of the problem. In contrast, our proposed algorithm for adaptive SSMs can applied in more general cases and still provide guaranteed algorithmic scaling behavior even in GPUs where tensor optimization not available for example, edge computing, other GPU makers.
Experiments
We empirically validate the theoretical claims of AUSSM by evaluating both its computational effi- ciency and expressivity. First, we benchmark runtime and memory usage across four implementations,
7

---

Mamba Complex
Mamba [-1,]]
AUSSM
Task
xLSTM
Mamba
AUSSM
Hybrid
repetition
0,09
0.10
0,09
0.15I
0.19932
0.45
bucket sort
0.21
0.912
0.7
0.69
0.92
0.833
C.S
majority count
0.19
0.31
0.5
0.45²
0.096
0.37
majority
0.13
0.63
0.64
0.691
0.57
0.64
solve equation mod arith
0.43
0.242
0.242
0.05
0.07
0.073
D.C.F
0.12
0.116
0.15F
0.04
0.133
0.23
mod arith wo bra
0.23
0.24
1.0'
0.13
0.48²
0.53²
Reg.
cycle nav
0.42
0.912
0.8
0.86
1.0*
1.0*
parity
0.27
1.0
1.0
0.132
1.0'
1.0'
Table 1: AUSSM and hybrid AUSSM+Mamba models outperform Mamba tasks requir- ing counting and structured memory. We evaluate xLSTM, Mamba, AUSSM, and a hybrid AUSSM+Mamba model suite of algorithmic reasoning tasks. The table shows the scaled test accuracies on each task. The tasks are grouped by their position in the Chomsky hierarchy (C.S: context-sensitive, D.C.F: Deterministic Context Free, Reg. Regular). AUSSM achieves perfect accuracy on tasks like parity and cycle navigation, which require modulo counting, validating its theoretical expressivity. While Mamba performs better tasks such majority count, the hy- brid model consistently achieves the best or near-best performance across most tasks, demonstrating that combining adaptive unitary dynamics with real-valued recurrence yields more expressive and general-purpose purpose architecture. The scaled accuracies for xLSTM and Mamba are obtained from (1]
including our CUDA-optimized AUSSM and Mamba. Second, we assess expressivity on suite of algorithmic tasks requiring formal language recognition, such as parity and modular arithmetic. Finally, we evaluate real-world applicability on long-sequence classification and regression tasks, demonstrating that the improved expressivity of AUSSM translates to practical performance gains.
5.1 Scalability Evaluation
We benchmark four implementations of AUSSM to assess efficiency: (1) naive Py PyTorc recurrent version, (2) a PyTorch version using separable convolutions with higher constant factor the linear scaling, (3) our custom CUDA kernel, and (4) the Mamba CUDA kernel bascline. Experiments were run on a single NVIDIA 2080 Ti GPU with 11 GB VRAM. As shown in Figure 2
the CUDA- based AUSSM achieves significantly lower memory usage and faster inference compared to the PyTorch variants, approaching the efficiency of Mamba despite full adaptivity. The separable PyTorch implementation improves runtime the recurrent baseline but incurs higher memory costs. Óverall, our separable formulation paired with a low-level CUDA kernel enables AUSSM to scale to long sequences efficiently, validating the theoretical benefits of scalability.
5.1.1 Expressivity Evaluation
To evaluate formal expressivity, we benchmark AUSSM, Mamba, XLSTM 31, and hybrid AUSSM+Mamba model on a suite of algorithmic tasks drawn from various levels of the Chomsky hierarchy. These include tasks requiring counting (e.g., parity, modular arithmetic), memory ma- nipulation (c.g., repetition), and symbolic reasoning (c.g., equation solving). Models are trained on sequences up to length 40 and tested on lengths up to 256 assess length generalization performance. evaluate all the models using scaled validation accuracies to account for the differing number of output classes in the algorithmic tasks.
The results are shown in Table The AUSSM achieves perfect accuracy on tasks that require modulo counting and cycle tracking. validating its theoretical ability to simulate cyclic group automata via unitary and adaptive dynamics. contrast, Mamba fails to generalize on these tasks, consistent with the limitations of partially adaptive and dissipative models. However, AUSSM performs poorly on tasks such as majority or equation solving, where dissipative dynamics may be required stability and information aggregation. Notably, the AUSSM hybrid model performs significantly better than

---

Dataset
Heartbeat
SCPI
SCP2
Ethanol 1751 4
Motor 3000 2
Worms 17984 5
Avg
Seq. len. # classes
405 2
896 2
1152
2
S5 S6
47.8 3.1
74.2 2.1
10.2 = 3.3
0.8 3.5
6.0 6.0000 3.9
79.9 4.1
36.4
53.0 = 8.3
65.6 2.7
-0.2 = 9.4
1,9 = 6.4
2.6 4.7
$1.36.22
34.0
linoss
3.7
75.6 = 2.6
17.8 8.1
6,5 + 0.6'
20.0 + 7.5'
93.8 4,4
44,2'
Mamba
52.4 3.8
61.4 1.4
-3.6 1 3.9
3.9 4.5
-4.6 4.5
63.6 1 15.8
28.8
Hybrid
53.0 3.8
64.2 4.9
4.2
6.8
4.7 1 4.1
2.6 5.5
82.6
1 3.4
35.2
Table 2: Hybrid AUSSM with Mamba achieves competent performance on long time-series classification benchmarks. We evaluate the hybrid model on six UEA datasets spanning wide range of sequence lengths and domains. The table shows the scaled test accuracies for the different models compared to the hybrid AUSSM. The hybrid AUSSM consistently outperforms the base Mamba and achieves competent accuracy across datasets. These results demonstrate that the increased expressivity AUSSM, when combined with Mamba's s stability, translates into strong real-world performance even on long and complex sequence data. Our model is evaluated a statistically rigorous test with 20 different seeds to obtain better estimate of test accuracy to reduce the reliance on arbitrary evaluation seeds used in prior works (20).
all existing RNNs, including the XLSTM, suggesting that AUSSMs and Mamba blocks are synergistic and exhibit performance benefits that neither individual model provides. These results empirically support our theoretical claim that hybrid models combining AUSSM and Mamba maximize the expressivity of diagonal SSMs under the Krohn-Rhodes framework.
5.2 Long Time-Series Classification and Regression Benchmark
To evaluate the practical benefits of our architecture, we test the hybrid AUSSM+Mamba model on a suite of UEA tinetiim-sseries classification benchmarks 32] and the challenging Weather regression benchmark. We take the AUSSM block as a drop-in replacement for an existing Mamba backbone. Specifically, we randomly selected fixed number of Mamba blocks in deep Mamba SSM model and replaced them with the AUSSM blocks. The UEA tasks feature much longer sequences than the algorithmic benchmark, with lengths ranging from 405 in the Heartbeat dataset over 17,000 in the Worms dataset. For regression, we use the challenging Weather dataset where climate variables are forecasted 720 steps into the future, given window of 720 timesteps. These benchmarks present a more realistic and diverse set of challenges, which includes physiological signals, chemical concentrations, motion data, and climate, each requiring the model to capture both local and global temporal dependencies.
For the UEA benchmarks and the weather dataset, used identical hyperparameter strategies to those used by the models we compare against. During testing, we found that previous works used five randomly chosen seeds evaluate the test performance. This is easily reproducible, as the particular choice of the seeds influences the specific test datasets that are chosen for evaluation and may produce biased results. We instead use a statistically rigorous technique where the best hyperparameter model is chosen based on the validation set performance on five random seeds, and the test accuracy evaluated on random train-test splits on the selected model with 20 different seeds to produce better test accuracy estimates. We scaled test accuracies with the baseline and report the results in Table The hybrid AUSSM+Mamba model achieves substantial improvements over the partially adaptive Mamba SSM on average, even under the modified testing protocol. The results demonstrate that the improved expressivity of AUSSM carries over to real-worlke tasks when appropriately combined with the stability and inductive biases of partially adaptive models. Notably, the hybrid model achieves these results while maintaining high efficiency: all experiments except EigenWorms were run on single NVIDIA 2080 Ti GPU with 11 GB of VRAM, in contrast to the large-scale hardware (e.g., 100 GB A00 GPUs) typically used for long- sequeeeee modeling. The Ei sigenworms dataset trained on an L4 GPU with 23 GB VRAM due to its larger size.

---

Model
Mean Absolute Error
Informer
0.731
LogTrane
0.773
LSTMa
1.109
LSTnet
0.757
S4
0.578
LinOSS
0,508
Mamha
0.464
AUSSM Hybrid
0.342
Table 3: Hybrid AUSSM+Mamba achieves state-of-the-art performance on long time-series weather forecasting benchmark. We evaluate AUSSM against 7 different models on the challenging weather forecasting benchmark, where climate variables are forecasted 720 timesteps into the future. AUSSM Hybrid achieves state-of-the-art performance on the task, improving the base Mamba model and all the other models.
Discussion
In this work, we address the expressivity- tradeoff in state space modeling. Existing SSMs like Mamba are scalable but limited in expressivity due to fixed or partially adaptive recurrence. the other hand, more general LTV SSMs are more expressive but not have an efficient and scalable parallel implementation. We introduce the Adaptive Unitary State Space Model (AUSSM), which uses input-dependent skew-symmetrre recurrence to achieve both unitary evolution and high expressivity. We showed that theoretically, AUSSM can implement modulo counters and simulate broad class of regular languages, maximizing expressivity among diagonal SSMs when combined with Mamba under the Krohn-Rhode framework. To ensure scalability, we develop separable convolution formulation and a custom CUDA kernel, enabling linear-time training despite full adaptivity. Experimental analysis standard benchmark tasks showed that AUSSM achieves strong performance on symbolic reasoning tasks and serves as an effective drop-in enhancement to Mamba for long-range sequence modeling. Together, these results suggest that adaptive unitary recurrence is a powerful inductive bias for both symbolic and continuous sequence tasks.
Limitations. One limitation related to expressivity is that AUSSM is capable of LTV recurrence only through its linearly adaptive (input-dependent) recurrent matrix. For the more general LTV recurrence, the recurrent matrix needs to have the capability for non-linear input-dependence that is additionally dependent on time.
Further, our separable kernel approach used for optimization relies on the assumption that recurrent matrices are simultaneously diagonalizable, limiting the ability to express languages beyond solvable regular languages. While the separable kernel has identical scaling behavior to efficient LTI models, still constant factor higher. Hybrid AUSSM+Mamba models show promise, but the best strategy for combining blocks is not yet well understood. Another limitation in the parallel scan- based algorithm we proposed is that, the case of high-end NVIDIA GPUs, alternate tensor core approaches may provide even better absolute speedup, although with the same scaling behavior, An alternate tensor core-based algorithm for AUSSM is an interesting avenue for future work in real- world applications. Finally, due resource limitations, our evaluations are limited to modest-scale tasks; further validation on foundation-mont scale benchmarks is needed.
7 Acknowledgements
We thank Reda Boumasmoud for his input and suggestions on earlier draft of this manuscript. We also thank Michael Hahn for useful conversation about SSM expressivity. T.A.K. acknowledges the Kempner Institute for the Study of Natural and Artificial Intelligence at Harvard University for funding during work on this article. T.J.S. acknowledges funding from ONR N00014-23-1-2066 A.K. and H.T.S. acknowledge NSF for EAGER: Neural Networks that Temporally Change (NOTCH).
10

---

References
[1] Abhimanyu Dubey, Abhinav Jauhri and Zhiwei Zhao. The llama herd of models. ArXiv, abs/2407.21783,, 2024.
[2] Gemma Team Morgane Riviere, Shreya Pathak... .... and Alek Andreev. Gemma 2: Improving open language models at a practical size. ArXiv, abs/2408.001 001 18, 2024.
[3] OpenAI, Josh Achiam, Steven Adler, Sandhini Agarwal and Barret Zoph. Gpt-4 technical report, 2024.
[4] Xihao Piao, Zheng Chen, Taichi Murayama. Yasuko Matsubara, and Yasushi Sakurai. Fred- former: Frequency debiased transformer for time series forecasting. Proceedings of the 30th ACM SIGKDD Conference on Knowledge Discovery and Data Mining, 2024.
[5] Haoyi Zhou, Shanghang Zhang, Jieqi Peng, Shuai Zhang, Jianxin Li, Hui Xiong, and Wancai Zhang. Informer: Beyond efficient transformer for long sequence time-series forecasting. In Proceedings of the AAAI conference on artificial intelligence, volume 35, pages 11106-11115,
2021.
[6] Noelia Ferruz, Steffen Schmidt, and Birte Höcker. ProtGPT2 is deep unsupervised language model for protein design. Nature Communications, 13(1):4348 July 2022.
[7] Feyza Duman Keles, Pruthuvi Mahesakya Wijewardena, and Chinmay Hegde, On the computa- tional complexity of self-attention. In International conference on algorithmic learning theory, pages 597-619. PMLR, 2023.
[8] Albert Gu, Isys Johnson, Karan Goel, Khaled Saab, Tri Dao, Atri Rudra, and Christopher Ré. Combining recurrent, convolutional, and continuous-time models with linear state space layers. Advances in neural information processing systems, 34:572-585, 2021.
[9] Albert Gu, Karan Goel, and Christopher Ré. Efficiently modeling long sequences with structured state spaces. arXiv preprint arXiv:2111.003990 2021.
[10] Albert Gu and Tri Dao. Mamba: Linear-time sequence modeling with selective state spaces. ArXiv, abs/2312.00752, 2023.
[11] Yash Sarrof, Yana Veitsman, and Michael Hahn. The expressive capacity of state space models:
A formal language perspective. In The ehirt Annual Conference Neural Information Processing Systems (NeurIPS 2024), June 2024.
[12] Dianlin Zhang. A Linear Time-varying Model for Nonlinear Systems. PhD thesis, Rice University, 1969.
[13] Sepp Hochreiter, Yoshua Bengio. Paolo Frasconi, Jürgen Schmidhuber, et al. Gradient flow in recurrent nets: the difficulty learning long-term dependencies, 2001.
[14] Zachary W Davis, Gabriel Benigno, Charlee Fletterman, Theo Desbordes, Christopher Steward, Terrence Sejnowski, John H. Reynolds, and Lyle Muller. Spontaneous traveling waves naturally emerge from horizontal fiber time delays and travel through locally asynchronous- irregular states. Nature Communications, 12(1):605) 2021.
(15) Arjun Karuvally, Terrence Sejnowski, and Hava T Siegelmann. Hidden traveling waves bind working memory variables in recurrent neural networks. In Ruslan Salakhutdinov, Zico Kolter, Katherine Heller, Adrian Weller, Nuria Oliver, Jonathan Scarlett, and Felix Berkenkamp, editors, Proceedings of the 41st International Conference on Machine Learning, volume 235 of Proceedings of Machine Learning Research, pages 23266-23299 PMLR, 21-27 Jul 2024.
[16] Lyle Muller, Frederic Chavane, John Reynolds, and Terrence J Sejnowski. Cortical travelling waves: mechanisms and computational principles. Nature Reviews Neuroscience, 19(5):255- 268, 2018.
[17] Mark M. Churchland, John P Cunningham, Matthew T. Kaufman, Justin D. Foster, Paul Nuyujukian, Stephen Ryu, and Krishna V. Shenoy. Neural population dynamics during reaching. Nature, 487:51 56, 2012.

---

[18] Antonio Orvieto, Samuel L Smith, Albert Gu, Anushan Fernando, Caglar Gulcehre, Razvan Pascanu, and Soham De. Resurrecting recurrent neural networks for long sequences. In International Conference on Machine Learning, pages 26670-26699 PMLR, 2023.
[19] Antonio Orvieto, Soham De, Caglar Gulcehre, Razvan Pascanu, and Samuel L Smith. Uni- versality of linear recurrences followed by non-linear projections: finite-width guarantees and benefits of complex eigenvalues. arXiv preprint arXiv:2307. 1888, 2023.
[20] Konstantin Rusch and Daniela Rus. Oscillatory state-space models. ArXiv, abs/2410.03943, 2024.
[21] Heithem Boufrioua and Boubekeur Boukhezzar. Gain scheduling: short review. In 2022 2nd International Conference on Advanced Electrical Engineering (ICAEE), pages 1-6, 2022.
[22] Luca Zancato, Arjun Seshadri, Yonatan Dukler, Aditya Sharad Golatkar, Yantao Shen, Benjamin Bowman, Matthew Trager, Alessandro Achille, and Stefano Soatto. B B'mojo: Hybrid state space realizations of foundation models with eidetic and fading memory. Advances in Neural Information Processing Systems, 37:1 177:33333-10062, 2024.
[23] Kenneth Krohn and John Rhodes. Algebraic theory of machines. i. prime decomposition theorem for finite semigroups and machines. Transactions of the American Mathematical Society, 1 116:450-464, 1965.
[24] Bingbin Liu, Jordan T. Ash, Surbhi Goel, Akshay Krishnamurthy, and Cyril Zhang. Trans- formers learn shortcuts to automata. In The Eleventh International Conference on Learning Representations, 2023.
[25] Riccardo Grazzi, Julien Siems, Jörg K.H. Franke, Arber Zela, Frank Hutter, and Massimiliano Pontil. Unlocking state- taate-tracking in linear RNNs through negative eigenvalues. In The Thirteenth International Conference on Learning Representations, 2025.
[26] N. Chomsky and M.P. Schützenberger. The algebraic theory of context-free languages*. In Braffort and D. Hirschberg, editors, Computer Programming and Formal Systems, volume 35 of Studies in Logic and the Foundations of Mathematics, pages 1 118-161. Elsevier, 1963,
[27] David S. Dummit and Richard M. Foote. Abstract algebra. Wiley, New York, 3rd ed edition, 2004.
[28] William Merrill, Jackson Petty, and Ashish Sabharwal. The illusion of state in state-space models. In Proceedings of the 41st International Conference on Machine Learning, ICML JogR. 2024.
[29] Howard Anton and Chris Rorres. Elementary linear algebra: applications version. John Wiley & Sons, 2013.
[30] Guy E Blelloch. Vector models for data-paralle computing, volume 2. MIT press Cambridge, 1990.
[31] Maximilian Beck, Korbinian Poppel, Markus Spanring, Andreas Auer, Oleksandra Prudnikova, Michael K Kopp, Günter Klambauer, Johannes Brandstetter, and Sepp Hochreiter. xlstm: Extended long short-term memory. ArXiv, abs/2405.04517, 2024.
[32] Benjamin Walker, Andrew D. McLeod, Tiexin Qin, Yichuan Cheng, Haoliang Li, and Terry Lyons. Log neural controlled differential equations: The lie brackets make a difference. ArXiv, abs/2402.18512. 2024.
[33] Martín Arjovsky, Amar Shah, and Yoshua Bengio. Unitary evolution recurrent neural networks. In International Conference Machine Learning, 2015.
[34] Scott Wisdom, Thomas Powers, John Hershey, Jonathan Le Roux, and Les Atlas. Full-capacity unitary recurrent neural networks. In D. Lee, M. Sugiyama, U. Luxburg, I. Guyon, and R. Garnett, editors, Advances in Neurul Information Processing Systems, volume 29. Curran Associates, Inc., 2016.
12

---

[35] Li Jing, Yichen Shen, Tena Dubcek, John Peurifoy, Scott Skirlo, Yann LeCun, Max Tegmark. and Marin Soljačíć. Tunable efficient unitary neural networks (EUNN) and their application to RNNs. In Doina Precup and Yee Whye Teh, editors, Proceedings of the 34th International Conference on Machine Learning, volume 70 of Proceedings of Machine Learning Research, pages 1733-1741. PMLR, 06-11 Aug 2017.
[36] Eugene Vorontsov, Chiheb Trabelsi, Samuel Kadoury, and Chris Pal. On orthogonality and learning recurrent networks with long term dependencies. In Doina Precup and Yee Whye Teh, editors, Proceedings of the 34th International Conference on Machine Learning, volume 70 of Proceedings of Machine Learning Research, pages 3570-3578. PMLR, 06-11 Aug 2017.
[37] Zakaria Mhammedi, Andrew Hellicar, Ashfaqur Rahman, and James Bailey. Efficient orthogonal parametrisation of recurrent neural networks using householder reflections. In Doina Precup and Yee Whye Teh, editors, Proceedings of the 34th International Conference on Machine Learning, volume 70 of Proceedings of Machine Learning Research, pages 2401-2409, PMLR, 06-11 Aug 2017.
[38] Mario Lezcano-Casado and David Martínez-Rubio. Cheap orthogonal constraints in neural networks: A simple parametrization of the orthogonal and unitary group. In Kamalika Chaud- huri and Ruslan Salakhutdinov, editors, Proceedings of the 36th International Conference on Machine Learning, volume 97 of Proceedings of Machine Learning Research, pages 3794-3803. PMLR, 09-15 Jun 2019.
[39] Nitzan Guberman. On complex valued convolutional neural networks. CoRR, abs/ 660.009046, 2016.
[40] Mikael Henaff, Arthur Szlam, and Yann LeCun. Recurrent orthogonal networks and long- memory tasks. In Maria Florina Balcan and Kilian Q. Weinberger, editors, Proceedings of The 33rd International Conference on Machine Learning, volume 48 of Proceedings of Machine Learning Research, pages 2034-2042, New York, New York, USA, 20-22 Jun 2016. PMLR.
[41] Bo Chang, Minmin Chen, Eldad Haber, and Ed H. Chi. Antisymmetricmm:: A dynamical system view on recurrent neural networks, 2019.
[42] Corentin Tallee and Yann Ollivier. Can recurrent neural networks warp time? CoRR, 2018.
[43] Mario Lezcano Casado. Trivializations for gradicnt- baddde optimization on manifolds. In Neural Information Processing Systems, 2019.
[44] Doug Rubino, Kay A. Robbins, and Nicholas G. Hatsopoulos. Propagating waves mediate information transfer in the motor cortex. Nature Neuroscience, 111159-1157,, 2006.
[45] Maria V. Sanchez-Vives and David A. McCormick. Cellular and network mechanisms of rhythmic recurrent activity in neocortex. Nature Neuroscience, 3: 1:044 2000.
[46] G. Bard Ermentrout and David Kleinfeld. Traveling electrical waves in cortex insights from phase dynamics and speculation computational role. Neuron, 29:33-44, 2001.
[47] Lyle E. Muller, Frédéric Chavane, John H. Reynolds, and Terrence J. Sejnowski. Cortical travelling waves: mechanisms and computational principles. Nature Reviews Neuroscience, 2222,5-268, 2018
[48] Juan A Gallego, Mathhew G Perich, Lee E Miller, and Sara A Solla. Neural manifolds for the control of movement. Neuron, (4(5):978-9844 2017.
[49] Valerio Mante, David Sussillo, Krishna V Shenoy, and William T Newsome, Context-dependent computation by recurrent dynamics in prefrontal cortex. Nature, 503((7744)778-888) 84, 2013.
[50] Mathew T Kaufman, Mark M Churchland, Stephen Ryu, and Krishna V Shenoy. Cortical activity the null space: permitting preparation without movement. Nature neuroscience, 17(3):440 44844448, 2014.
13

---

[51] Juan A Gallego, Matthew G Perich, Raeed H Chowdhury, Sara A Solla, and Lee E Miller. Long-term stability of cortical population dynamics underlying consistent behavior. Nature neuroscience, 23(2):260-213 2700 2020.
[52] Abigail A Russo, Sean R Bittner, Sean M Perkins, Jeffrey S Seely. Brian M London, Antonio H Lara, Andrew Miri, Naja J Marshall, Adam Kohn, Thomas M Jessell, et al. Motor cortex embeds muscle-like commands in untangled population response. Neuron, 97(4):953- 2018.
[53] Carl Van Vreeswijk and Haim Sompolinsky. Chaos in neuronal networks with balanced excitatory and inhibitory activity. Science, 274(5293) 1996.
[54] Guillaume Hennequin, Tim Vogels, and Wulfram Gerstner. Optimal control of transient dynam- ies in balanced networks supports generation of complex movements. Neuron, 12(6):134-1140 2014.
[55] Dean Buonomano and Wolfgang Maass. State-dependent computations: spatiotemporal processing in cortical networks. Nature Reviews Neuroscience, 10(2):133-125 2009.
[56) David Sussillo and Larry F Abbott. Generating coherent patterns of activity from chaotic neural networks. Neuron, 63(4):544-557) 2009.
[57] Carsen Stringer, Marius Pachitariu, Nicholas Steinmetz, Matteo Carandini, and Kenneth D Har- ris. High-dimensional geometry of population responses visual cortex. Nature, 571(7765):361- 365, 2019.
[58] Francesca Mastrogiuseppe and Srdjan Ostojic. Linking connectivity, dynamics, and computa- tions low-rank recurrent neural networks. Neuron, (9(3):609-633 2018.
[59] David Sussillo. Neural circuits as computational dynamical systems. neurobiology, 25:156-163, 2014.
Current opinion in
[60] Jeffrey L. Elman. Finding structure in time. Cognitive Science, 14(2):179-211, 1990.
[61] Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural Computation. 1735 1997.
[62] Kyunghyun Cho, Bart van Merriënboer, Dzmitry Bahdanau, and Yoshua Bengio. On the properties of neural machine translation: Encoder-decoder approaches. In Dekai Wu, Marine Carpuat, Xavier Carreras, and Eva Maria Vecchi, editors, Proceedings of SSST-8, Eighth Workshop on Syntax, Semantics and Structure in Statistical Translation, pages 103-111, Doha, Qatar, October 2014. Association for Computational Linguistics.
[63] William Merrill and Ashish Sabharwal. The parallelism tradeoff: Limitations of log-precision transformers. Transactions of the Association for Computational Linguistics, 11:531-545, 2023.
[64] Bo Peng, Ruichong Zhang, Daniel Goldstein, Eric Alcaide, Xingjian Du, Haowen Hou, Jiaju Lin, Jiaxing Liu, Janna Lu, William Merrill, et al. Rwkv-7" goose" with expressive dynamic state evolution. In Conference on Language Modeling, 2025.
[65] Julien Siems, Timur Carstensen, Arber Zela, Frank Hutter, Massimiliano Pontil, and Riccardo Grazzi. Deltaproduct: Improving state-trackin in linear rnns via householder products. In Neural Information Processing Systems, 2025.
[66] William Merrill, Gail Weiss, Yoav Goldberg, Roy Schwartz, Noah A. Smith, and Eran Yahav. A formal hierarchy of RNN architectures. In Dan Jurafsky, Joyce Chai, Natalie Schluter, and Joel Tetreault, editors, Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics, pages 443-459, Online, July 2020. Association for Computational Linguistics.
[67] Oded Maler. On the krohn-rhodes cascaded decomposition theorem. In Time for Verification: Essays in Memory of Amir Pnueli, pages 260-278. Springer, 2010.
[68] István Babesányi. Automata with finite congruence lattices. Acta Cybernetica, 18(1):155-165 2007.
14

---

[69] Richard Brent, Colin Percival, and Paul Zimmermann. Error bounds on complex floating-pont multiplication. Mathematics of Computation. 76(259):1469-1444 2007.
[70] Karl-Heinz Zimmermann. On krohn-rhodes theory for semiautomata. arXiv preprint 1.Xi::2010.162333 2020.
[71] John Hewitt, Michael Hahn, Surya Ganguli, Perey Liang, and Christopher D. Manning. RNNS can generate bounded hierarchical languages with optimal memory. Bonnie Webber, Trevor Cohn, Yulan He, and Yang Liu, editors, Proceedings of the 2020 Conference Empirical Methods in Natural Language Processing (EMNLP), pages 1978-2010, Online, November 2020. Association for Computational Linguistics.
[72] Thomas H. Cormen, Charles E. Leiserson, Ronald L. Rivest, and Clifford Stein, Introduction to Algorithms, Third Edition. The MIT Press, 3rd edition, 2009.
15

---

A Related Work
In this appendix section, we provide an extensive discussion of related work on unitary RNNs, conserved dynamics in the brain, and the computational complexity of RNNs and SSMs.
A.1 Unitary RNNs
Recurrent neural networks (RNNs) are powerful models for processing sequential data, but their training is often hindered by the vanishing and exploding gradient problem, which limits their ability to capture long-range dependencies [ A major source of this instability is the repeated multiplication of the hidden state by the recurrent weight matrix, which causes gradients to exponentially decay or grow depending on the spectral norm of the matrix. To address this, a prominent line of research constrains the recurrent dynamics to be unitary, ensuring that the hidden state evolution preserves its norm through time. Such porm-preser dynamics prevent gradient magnitude degradation and are reversible dynamical systems.
mathematically analogous to energy-conservee
The first major work in this direction was Unitary Evolution Recurrent Neural Network (uRNN) 33 that proposed parameterizing the recurrent weight matrix as product of structured unitary matrices (including diagonal phase matrices, Fourier transforms, and Givens rotations) ensure exact unitarity while retaining efficient computation and gradient flow. However, this parameterization limited expressivity due to its constrained structure. To overcome this, Full-Capacity Unitary RNN 34 optimized directly over the unitary group using manifold optimization techniques. By employing the Cayley transform and optimization on the Stiefel manifold, they achieved full representational power while preserving unitary constraints.
Subsequent research explored alternative approaches for maintaining orthogonality and improving trainability. Efficient Unitary Neural Networks (EUNN) [35] used parameter-efficien decompositions enabling flexible trade-offs between computational cost and expressivity. Vorontsov al. Orthogonal- ity regularization methods softly constrain the recurrent weight matrix toward being orthogonal rather than enforcing strict unitarity, allowing some deviation to improve learning flexibility [36]. Similarly, Mhammedi et al. [37] developed real-valued orthogonal RNN based on Householder reflections, which guarantees orthogonality through efficient matrix parameterizations. Later, Lezcano- Casado and Partinez-Rubio (38] introduced an elegant exponential parametrization of orthogonal and unitary matrices via skew-symmetric matrices, offering a smooth and numerically stable way to maintain orthogonality during training.
key challenge unitary and orthogonal RNNs the choice of suitable nonlincarities. Standard nonlincarities such as ReLU or tanh can break the norm-proe com--reserrvnng property of the recurrent map, leading to unstable or dissipative dynamics. To address this, (33) introduced modReLU, complex- valued nonlinearity that preserves the phase of hidden activations while applying a learned threshold on their magnitudes. Subsequent studies explored alternatives such as zReLU 33 scaled tanh, and
nonlinearities for complex- complex-valued valued networks. In 400, the authors analyzed the interplay between orthogonality, nonlinearity, and gradient flow, providing theoretical insights into how orthogonal constraints help maintain long-term dependencies even in nonlinear regimes. More recently, Chang et al. [41 proposed Antisymmetric RNNs, where the recurrent weight matrix constrained to be wat-s-kew--ymmmtttte thereby approximating a Hamiltonian or energy-conserving flow; this represents a bridge between strict unitarity and continuous-timg time neural dynamics.
The success of unitary and orthogonal RNNs has motivated several extensions. For example, Tallec and Ollivier 42 analyzed time- warping effects and timescale adaptation in RNNs, showing how orthogonality can help control effective memory timescales. Lezcano-Casado 45 further generalized the parameterization of orthogonal operators to improve optimization stability across different architectures. The insights from unitary evolution have also influenced more recent Structured State Space Models (SSMs) 9 which impose spectral stability and linear time-invariant dynamics to achieve long-range sequence modeling with recurrent efficiency. These connections underscore the broader relevance of porm-preserring and energy-conserving formulations for building stable, interpretable, and trainable dynamical models.
16

---

A.2 Unitary Dynamics in the Brain
growing body of experimental and theoretical work indicates that neural population activity often evolves according to dynamics that are, in important respects, conserved or weakly dissipative, with clear implications for how the brain stores and transforms information. Conserved dynamics here refers to trajectories or modes that preserve key quantities (e.g.. norms, phase relationships, low-dimensional energy-like functions) over behaviorally relevant timescales, producing smooth, reversible, or rotational population flows rather than rapidly diffusive purely dissipative responses. Empirically, such phenomena appear across modalities and brain areas: propagating and wave-like activity has been observed in sensory and motor cortices 44
45 46 477, low-dimensional structured trajectories that persist across trials and conditions have been reported in motor and prefrontal popu- lations 48 48 50 51 and cortical responses frequently exhibit coherent oscillatory/roaaiionll components that suggest near-une unitary evolution within a task-relevant subspace [49] 52]. These con- served modes are often embedded within a larger sigh-dimensional network state, but they dominate the behaviorally relevant dynamics and appear to support robust temporal computation, short-term memory, and smooth transformation between input and output representations.
Theoretical and modeling studies have proposed multiple mechanistic origins for conserved or weakly dissipative neural dynamics. Balanced excitatory-inhin inhibitory network regimes can produce rich transient dynamics with slow decay or quasi-conserved activity on short timescales; classic balanced- network analyses show how tightly coupled excitation and inhibition enable irregular activity while constraining macroscopic statistics, and follow-up work has shown how such balance supports structured transient trajectories [53 [54] Network architectures with antisymmetric
-kew--kew- symmetric connectivity produce rotational and energy-like flows that are approximately conservative; control- and dynamical-systemm -rreeet-ooieeeeddd analyses have demonstrated that small departures from strict antisymmetry (e.g., weak damping or inputs) permit flexible routing and readout while retaining the stability advantages of corm-preserving flows [55] 56). In parallel, reservoir- and recurrent-netwon modeling (including trained networks initialized in richly recurrent regimes) has shown that networks learn to generate low-dine wwwdimmnsional conserved trajectories that implement computations (e.g., context-dependent integration, short-term memory) with high robustness to noise and parameter
changes 157 58
Methodologically, the identification of conserved modes in neural data relies on dimensionality- reduction and dynamical- tools that explicitly search for rotational, dow-dimensional, or wave-like structure. Approaches range from linear subspace methods that highlight persistent modes to more specialized decompositions and dynamical fits that extract antisymmetric components, traveling-wave decompositions, or stable latent manifolds; these analyses have repeatedly revealed that relatively small number of conserved or ceas-coneerved modes often capture most of the task- relevant variance, even when single-neuron responses are heterogeneous [49 50 59). Importantly, conserved neural dynamics appear functionally beneficial: by concentrating computation in norm- preserving subspaces, the brain can transform and transmit signals with minimal degradation, enabling temporally-extender operations such sequence generation, motor command shaping, and transient working memory without continual external reinforcement.
Finally, the convergence of empirical findings and theoretical models has motivated viewing cortical and subcortical circuits through the lens of structured dynamical primitives-rotatiee and hearly-Hamiltoniee perspective helps explain phenomena such as reliable single-trial trajectories, robustness of latent dynamics across learning and perturbation, and the coexistence of fast irregular activity with slow conserved modes [8] 51 [4]. Recognizing conserved dynamics in the brain also provides a bridge to normative and engineering approaches (e.g., unitary or antisymmetrie RNNs, energy-basee network formulations) that aim to replicate the computational advantages of biological circuits while offering interpretable and stable mechanisms for tong-timescale processing.
waves, towss neither purely feedforward nor purely dissipative. This
17

---

Model family
Constant diagonal
Adaptive diagonal
Adaptive non-diagonal
(examples)
Mamba, AUSSM-hybrid (ours)
DeltaProduct, RWKV-7
Adaptive
yes
yes no complex permutation group langs.
Diagonal
yes
Eigenvalues
[0,1)
various
Languages
(subset of) star-free langs. star-free solvable langs.
Table 4: High-level comparison of three families of linear recurrent neural networks (LRNNs) by
adaptivity and diagonality, and their confirmed expressivity under standard assumptions.
Model
Mamba
Mamba-negative
AUSSM-hybrid (ours)
Mamba-complex
Adaptive
indirect
indirect
yes (AUSSM)
indirect
Diagonal
yes
yes
yes unitary (AUSSM)
yes
Eigenvalues
(0, 1)
(-1,
complex
Star-free
Thm.
Thm. 4]
11 Thm. 4]
Thm. 4]
Solvable
Thm. 4]
X 25
Thm.
/Thm. 2
Thm. 21]
Regular
28 Thm 4.4]
28
Thm 4.4]
x Thm 4.4]
Thm 4.4]
Table 5: More fine-grained comparison of design features and formal language expressivity of Mamba-like models.
A.3 Computational Expressivity Theory of Linear RNN Architectures
As linear RNNs (LRNNs), of which SSMs are special case, have gained in performance and become a viable alternative to transformers for sequence processing tasks, their formal expressive power has garnered interest. LRNN denotes the family of recurrent neural networks whose transition function is a linear or affine transformation of the hidden state (note, however, that the linear transition may be non-linear function of the input at each time step). In contrast, traditional RNNs such as the Elman RNN 660] 61 or GRU 62 update their hidden state non-linearly between time steps. State space models (SSMs) such S4 [ and Mamba [10] are subtype of LRNNS motivated by continuous linear dynamical systems that are discretized to work recurrently as LRNNs. terms of formal expressivity, (28) and [1] point out that most SSMs make architectural decisions that limit their ability to model formal languages. The main factors impacting expressivity are:
Adaptivity Whether the transition matrix is held constant over time or a (non-linear) function of the input at the current time step.
Diagonality Imposing that the transition recurrence is diagonal diagonalizable matrix (simultancously for all inputs).
Eigenvalue range Restricting the transition recurrence to matrices with eigenvalues in a specific range (e.g., non-negative, real between -1 and complex unitary, etc.).
As well as impacting their expressivity, these decisions also determine the scalability of the architec- ture, since certain restrictions allow for more efficient implementations, e.g., the product of diagonal matrices can be computed more efficiently than that of dense matrices. There appears to be distinct tradeoff between expressivity and scalability, informally dubbed the parallelism tradeoff 633. See Tab, for comparative high-level overview of different families of LRNNs and their expressivity and relative scalability.
Adaptivity Some of the earlier SSM variants, such as S4, were non non-adaptive (or invarianaai, making them very scalable through the use of convolution over the whole sequence using a pre- computable constant convolution kernel. As (28 points out, input-independdece ensures that the expressivity of SSMs is upper-bounded by the circuit complexity class TCO, while some regular languages require cireuit complexity NCI (it is widely assumed that TC0 NC1). More recent SSM architectures add adaptivity through various means, e.g., the Mamba recurrence is indirectly
18

---

input-dependent via its time discretization factor A. Similarly, LRNNs such as RWKV-7 64 or DeltaProduct [65], and non-linear RNNs such as xLSTM [I]]. are adaptive through mechanisms such as input-dependent gating factors. The transitions of our AUSSM component are directly input-dependent, making our architecture fully adaptive (see Tab. 5
Diagonality Another critical factor is whether the transition recurrence is diagonal or simultaneously diagonalizable for different inputs. [ Thm. 21J shows that such diagonal SSMs can recognize solvable languages. but [28] again shows that such SSMs are contained within circuit complexity class TCO, meaning they cannot recognize all regular languages, assuming TC0 != NC1. The upshot is that diagonal transition matrices mean quicker or less memory-intension computation for training and inference, making it a very attractive architectural decision as employed by S4, Mamba, and others. For this reason, we also choose to keep diagonality for our AUSSM and hybrid architecture, and mainly compare performance between diagonal SSMs.
In contrast, the design of DeltaProduct or RWKV-7 attempts to create non-diagonal LRNN architec- tures whose transition matrices are products of generalized householder matrices, which are diagonal matrices plus added rank-1 1 component. Such models can recognize all regular languages 25 64 albeit at the price of additional time and memory cost.
We also compare the performance of our architecture to that of XLSTM, which, as an extension of the traditional LSTM, an RNN but not an LRNN since the recurrence is non-linear in the hidden state. Since LSTMs can recognize more than just regular languages [66]. we use this as an upper-bound comparison stronger model. In fact, while the authors do not formally prove that xLSTMs can recognize all regular languages, the experimental results show strong performance on this class of languages.
Eigenvalue range Within the realm of adaptive diagonal (or diagonalizable) SSMs in particular, the eigenvalue range of the transition matrices plays a crucial role. This is because, as [1] Thm. 4] proves, non-negative eigenvalues restrict diagonal SSMs to the class of star-free regular languages, while negative eigenvalues allow for the recognition of non-star-free languages such as parity. [25] point out that Mamba can be trivially adapted to have negative eigenvalues without additional computational cost. However, negative eigenvalues alone are not enough to recognize all solvable languages; complex eigenvalues are required, e.g., for solvable languages like mod ni parity for n > 25 Thm. 2]. In non-diagonal LRNNs, multiplying generalized Householder matrices as in DeltaProduct or RWKV-7 can yield cigenvalues with non-zero imaginary components, raising their expressivity to include all solvable languages (and, indeed, all regular languages (655).
In order to recognize all solvable languages with diagonal SSMs, we need to extend the eigenvalue range to complex numbers. The simplest way to do this to just use Mamba with complex hidden states. This incurs additional overhead, however, because it increases the parameter count to include all possible complex values. As we show in S in order to accept all solvable languages, we only need to add SSM components with unitary complex eigenvalues, which is why we introduce AUSSM components to Mamba, allowing the hybrid architecture to model all solvable languages with minimal overhead. Additionally, while formally Mamba with complex values is just expressive architecture, our experiments showed that Mamba with complex eigenvalues fails to learn even simple formal tasks that hybrid architecture can perform with perfect accuracy, indicating that the additional restriction to unitary values is indeed helpful for learning formal languages. See Tab, for comparison of Mamba-like models with their relative advantages and disadvantages.
B Limitations of Non-AdaptiivePPrttalll Adaptive SSMs
The expressivity of different classes of SSMs is defined by the types of dynamical systems and formal languages they are able to simulate. Appendix D analyzes the formal language expressivity and limitations of different classes of SSMs. In this section, we analyze expressivity related to different kinds of dynamical systems. First, we show that in line with Figure I SSM expressivity can be arranged in the order LTI real spectra C LTI complex C LTV partial C LTV. The models that have higher expressivity can simulate the models lower in the expressivity scale. Since LTV w uni tary spectra cannot be arranged precisely in this scale, we show an example of class of multitimescale processes that a partial LTV model like Mamba cannot simulate in a fixed hidden state and layer limits.
19

---

Expressivity of Single Block SSMs The dynamical systems that can be simulated by single block SSMs without non non-lineann coneiincssssies can be arranged in the order LTI real spectra LTI complex C LTV partial C LTV.
Proof: For the proof, we start with the most general LTV SSM and show that the next lower class SSM is special case. We do the same for all the subsequent SSM classes in the expressivity chain.
The single block LTV SSM has the following discrete form:
da(() exp( (A, A.) xtt) + A,B,
dt y()) C2 (().
The next lower class of ssm: LTV partial has the following form
dm(t) exp(Art =tt) + A, y(t) C. x(t).
Note here that the A is a scalar that varies with time, but A is fixed matrix. This can be derived instance of the LTV with A, A where the equivalence between the holds only in the case where the dimensionality of the SSM is 1. Similarly, the next lower class LTI complex has the following form
d=(t)) = exp(AA) s(t) + AB u(t) Cr(().
This is an instance of the LTV partial with A = A, B and C2 = C, which means all the matrices time invariant. The final class LTI real spectra is instance of LTI Complex where the eigenvalues are further restricted to have angle in the imaginary plane.
LTV is the most general class, but it is computationally infeasible to simulate the most general case. The -ioooiagnnaliiabbiityy of general matrix classes requires performing a full O matrix computation at each time step. Hence LTV unitary spectra with simultaneously diagonalizable unitary matrices is chosen as a principled middle ground. is, however, not instantly apparent how LTV unitary spectra compares against LTV partial. To illustrate the difference, we introduce an example a multi-timescale process.
Multi- tilesclimeecale features: A time-series u(t) E R is said to have hidden state can be factorized into the following form
features if the
(t)
x11 + 1)
g(1)) x(1).
Where f(t) E C, gtt E C are general complex-valuee time-varying functions and f(t) * cg(t) for some constant That is, the timeseries exhibits at least two independent features denoting two different timescales.
partial LTV SSMs in multi-timescale timeseries: partial LTV SSMs not able represent multi-timescale features in data.
Proof: The proof is by contradiction. If partial LTV SSMs are able to solve multi-timescale
(1()
(t) 0 = A.A is true. Solving the system for A leads to a g(t)) ((()
timeseries, the following
constraint on ft) ct(t) where is some constant. This is true only when one of the functions is a constant multiple of the other, that is the two functions are dependent and have the same timescale (with a possible constant factor difference).
LTV w unitary spectra SSMs in calti-timmescaee timeseries: LTV unitary spectra SSMs can represent
features in data as long as the ft),, (t) cxp(i8) where 0 T, m]. f(t), (t)) can be independent.
Proof: We first substitute f(t) = exp(ie (t)) and gtt) = (t) The resulting dynamical system has ftt) and g(t) as eigenvalues, which have unit magnitude themselves. This is the definition LTV unitary spectra.
The results trivially extend systems with more than dimensions
20

---

To summarize, if f and g independent (e.g., f(t) - = g(t) = t). then the partial LTV system cannot represent the multi-timescale features in the hidden state. On the other hand, LTV w unitary spectra imposes weaker constraint where the only requirement is that ((t) and g(() are constrained to the unit circle the imaginary plane; the timescales of the two variables can be independent.
Note 1. In the main text, when say that AUSSM is diagonal LTV system, mean that AUSSM dependee))))) recurrent matrix. For the more
capuble of LTV recurrence through its adaptive general LTV recurrence, the recurrent matrix needs to have the capability for non-lineat dependence which AUSSM currently does not support.
In this section, we showed that single AUSSM block with fixed model dimension and hidden state size can represent functions that Mamba cannot represent with the same hyperparameters (it may need greater width more layers for the same function). At first glance, this seems to contradict the results shown in Tab. [] which posit that Mamba with complex entries is as expressive as our hybrid architecture. The explanation is that in the formal language expressivity analysis in $3.1 and $ is concerned with the expressivity of the whole architecture class over any finite parametrization, rather than specific model parametrization. The two analyses, therefore, keep different quantities constant: the expressivity analysis is about the existence of any finite instantiation of the model class that realizes a given language, while the fixed sxeerratperrrrreeeer single-layer measures relative capacity at constant size.
Note 2.
C AUSSM Derivation
We derive the AUSSM from controlled and adaptive version of the akew-symmetric ODE used in the jPCA procedure in computational neuroscience, given below. The skew-symmetric ODE first discretized using the Zero Order Hold procedure and then parameterized in polar coordinates. The steps to obtain the final AUSSM formulation are provided below.
dr((t) = A, x(t) + B u(t).
(9)
yt)) C x(t).
The above ODE is discretized following the Zero Order Hold procedure with step size of A, (note that the step size also time varying like the recurrent matrix)
xx) exp(A.AA) z(t-1) 1) + A,Bu(t), (t) C (tt).
(10)
The convolution form of the above system can be derived from this recurrence shown below (assuming (0) 0)
v1 CA,
(11)
(2 C cxp(A24))A A2CBu(())
(12)
(13)
v(() = c Σ k=
exp(A A Bu(k) + A,CBu(I) V=++1
II
(14)
Note that without additional assumptions cannot be simplified further, which can result in computationally inefficient approaches to compute the output. We draw motivation from the use of structured matrices in efficient SSM implementations and propose that A, belongs to class of matrices that are simultancously diagonalizable with the same basis. Let this diagonalizable basis be P.
A, the matrix exponential and the repeated products
t-1
v(t) = C
Σ 2=2
II exp(AAA(A)) P-IA,Bu(k)) + A,CBu(t), 1=++1
(15)
where A(At) is the diagonal matrix with the eigenvalues of A on the diagonal. Now, the repeated matrix product has a simplified form as shown below.
ACA P P-'AnBu(k) + A.CBu(tt)
(t))= = CP
Σ
exp Σ
(16)
2=2
1=k+1
21

---

For new set of B' and C' such that C" = CP and B' = p-1B. B, we get
:- Σ 1=++1
AAA )) A: + A,C'B'u())
y(1)= = Σ exp
(17)
The above equation undergoes one additional simplification, which reveals the unitarity of the discrete dynamical system. Since Ar is skew-symmetric matrix, the eigenvalues (AAr) are purely imaginary, meaning the above equation simplifies further in the polar form of A
t-1
Σ A A 1=
(-1
v(t = C'
Σ
exp
i
A.B'u(K) + B'u()..
(18)
where 12 = 1 is the complex iota and 3r.) is the function that obtains the imaginary component of a complex number. Since C" and B' are also complex due multiplication with P, we use polar forms for them too to finally obtain
A S(A(AA)) E [nn-1
1-1
)) ARRn exp(i0n)u(()) (19)
v() Rc exp(i0c)
Σ 2m
exp
Rc exp(iec) Rg exp(i0g) u(t).
To handle a d-dimensional input, this formulation is replicated d times for each input dimension. For adaptivity, we use where A(A)At); E, Xjr u, Xj and Di -Exxe (,) + x; Abias We use the above formulation in our experiments and parameterize the following for learning: Rc, 0c, RB, 8B, Xjr, Xj X Xÿr-
bias Abias
D Eigenvalue Analysis
Lemma 3 (Exponential of Skew-Symmetric Matrix is Orthogonal). Let A E R be real skew-symmn memmmsymmttiic matrix, i.e., A AA Then the matrix exponential exp( AA) is orthogonal for any
A E R, i.e.
exp(AA) exp(AA)= I
Proof. Let U exp(AA). Then,
UT (cxp(AA)) = exp(AAT) = exp(-AA),
since A = -AA Therefore,
UT- (A00-AA)xxpAAA) exp(0) I,
which shows that U is orthogonal.
Lemma 4 (Marginal Stability of Discrete-Tiiee Dynamics). Let A E guxn be a real skew-symment matrix and define t = (xp(AA) for some A > Then all eigenvalues of t lie on the complex unit circle. In particular, the discrete-ting time linear system
x xt))
Φr(( - 1)
marginally stable.
Proof. The eigenvalues of a real skew-symmetric matrix A are purely imaginary, i.e., Aú = iwsy E i.. The eigenvalues of exp( AA) then
Hy exp(AA,) = exp(iAwy),
which all lie on the complex unit circle since |exp(i0)]= for all R. Hence, the system exhibits marginal stability.
22

---

Lemma's 5 (Norm Preservation under Skew-Symment Dynamics). Let A E xR" be a real skew- symmetric matrix, and let exp(AA) Then for any x E R".
| ||
Hence, the transformation does not amplify or diminish the norm of the state vector, preventing both gradient explosion and vanishing during backpropagation through time.
Proof. Since & is orthogonal by Lemma 1, we have:
|| 2 (fb.:) (fxr) = x " or==z =| 10:1
Taking the square root yields || ||
Lemma 6 Modulated Rotation Frequencies via Skew Skew-Symmen Symmetric Generator). Let A R RaxR be smooth function such that A(u) is skew-symmetric for all u E R. Then for each u E IR, all eigenvalues of A(u) lie on the imaginary axis, and the eigenvalues of the discrete-time transition matrix ((u)= exp( A())) lie on the complex unit circle.
Furthermore, the eigenvalues of A(u) depend continuously on u, and thus the angular frequency of state-space rotation is smoothly and directly modulated by the input.
Proof. Let A(u) R"x" be symmetrii for all u E E. i... A(u) A(u). It a well-known result from linear algebra that real "kew-symmetne matrices have purely imaginary eigenvalues or zero.
Let (u) E C be an eigenvalue of A(u). Since A(u) is real and skew-symmetric, Xy(u) = iwy(u) for some wi (u) IR, and the eigenvalues come in complex-conjugatt pairs if nonzero.
Now, consider the discrete-time transition matrix:
(u)) == exp (w)))
Because the exponential of a skew-symmetric matrix is orthogonal (by Lemma 1), +(u) is an orthogonal matrix. The eigenvalues of orthogonal matrix with determinant lie on the complex unit circle, i.e..
|m;(u)]=1 for eigenvalues (y(4) of ((u).
Furthermore, the eigenvalues of (u) given by
Py (u) exp( = exp(iAw,(u))
so their arguments (i.e., angular velocitics) are precisely modulated by the real-valued frequencies ww), which in turn depend the input
To show that the rotational frequencies vary continuously with u, recall that the eigenvalues of smooth matrix function A(u) depend continuously u, provided that A(u) has distinct eigenvalues or that perturbations are small (which holds generically due to structure skew-symmetric matrices). Since A(u) is assumed to be smooth, all wy(u)) vary continuously with u, and therefore so do the corresponding angles Awj(u) of the discrete-time rotation matrix.
E Formal Language Expressivity
Our formal expressivity analysis uses the setting and proofs of [1] as a starting point. That is, we abstract away architectural details without loss of generality, and directly work with the already discretized form of the SSM. We assume floating- point arithmetic where the precision is logarithmically bounded in the sequence length, i.e., at most (log n. bits of precision on inputs of length n. Here, we briefly reiterate a somewhat abstract definition of our SSM to simplify the expressivity proofs.
Definition 1 (SSM layer). A single SSM layer is sequence-to-sequence map Rd - Rd, (ux) + (3r) for E [T] for sequence length T. It is defined recurrently by
Zt A 0 B:
(20)
23

---

where e is elementwise multiplication, To E C with m == n d, and A, B: + Cm are smooth, dependee maps with A, A(u+) and B, B(u,). Note that already subsumes the discretization variable A, which is itself a function of the input, as introduced in 19.. The output of the layer is computed as
Yt o( Uz)
(21)
where
Ø: C x IR Rd, (zt, u.) + Mixi (Tt, ue)), us)
(22)
Mix; and Mix; contain linear maps and a non-linearity (either silu or softplus) f Note that in our implementation, unlike 19), we do not apply normalization between the two Mix blocks but before the input enters the layer (see Def. F). For ease of notation, we subsume C, into Mix; without loss of generality. Mix; also usually contains convolution of the input before the SSM recurrence, which we ignore in expressivity unalyses following [TT] Remark 18J.
Definition (Mamba layer). A Mamba layer is an SSM layer where A and Bt, are input- dependde and,
and real-valued and, additionally. A E R is non-negative.
Definition 3 (AUSSM layer). An AUSSM layer is an SSM layer where Br and Ct are fixed constant functions (nol input dependent) and A, inpul dependent, complex valued, and each entry has unit magnitude, i..,
Vj [d],
|Ac/-= v ((Ass)) S(Ac,)²=1
Definition 4 (Full SSM), For full SSM, we usually stack multiple layers (I....,L) L) on top of each other, and indicate the layer we mean by a superscript, e.g., xt is the hidden state at time in layer f. The input to the first layer ut (1) is some embedding of the input of the full SSM computed by some injective embedding function : E - TR& where E is the alphabet of possible imput values at a single timestep, and the input to layer E E [or E 1 is the normalized output of the previous layer & 1:
U1 (8)
Norm(y) (
(23)
We use R.M.Norm the Norm, defined by
RMSNorm(2)
9 0 x
(24)
Σn x
where x E IR" and E Rd is a learned gain parameter. Importantly, like 19) our implementation
uses skip connections between consecutive layers, i.e..
(8) olxt, uz) + (t-1)
(25)
The final layer applies another RMSNorm and then a final output function.
We now introduce some notions from automata theory that are necessary for our expressivity results.
Definition 5. A deterministic finite-state automaton (FSA) A is a tuple (E, Q, where E is an alphabet (fimite, non-empty set), a finite set of states, and 8 Q Σ + Q is a transition function. The transition function can be lifted from symbols to symbol sequences as
8:Q x + Q2, 8(g,E)= =4
8(g, g = (5(g, oe))
where e is the empty string, Σ* is the Kleene closure over E,, and we use boldface to mark sequences of zero more symbols from Er.
The extended transition function 8 forms a transformation monoid under composition, called the transition monoid of the FSA.
Definition 6. A set-resel antomaton is FSA whose transition function maps all states to a single state for each imput symbol, that is, Vo E Σ, 3p Q s.l.
8(g () = p, E
Here, silu(z) [+exp(-x) and softplus(z) = x log (1 + exp(x))
'Note that in Mamba, B, is directly a function of the input while A. is input dependent through A, which is itself moon-lnneat) function of the input.
24

---

Note that the transition monoid of a set-reset automaton is aperiodic 67
Definition 7. A cyclic group automaton is an automaton whose transitions are permutations over states, where every input symbol acts as some power of fixed k-cycle with k = [Q] That is, for every symbol a E E, the symbol-xpecific transition map 6 Q - Q is bijection, and at least one of the symbols forms a cycle of order exactly k, i.e. for some a E Σ,, 8 id and 8% # id Vn E [1, 1].
All other symbol-transition matrices are powers of the same k-cycle, i.e. Ab 2, 6b 8% for some 72 E [0, k - 1}, where 8% id.
The transition monoid of a k-cyclic group automaton the cyclic group C2 68
We start by showing that our AUSSM architecture overcomes the limitation of most SSMs pointed out in [T] by showing that it can perform modulo counting, and therefore, can simulate cyclie group automata.
Lemma 1. For any l E Z one can construct a single-layer AUSSM that counts modulo k, which means AUSSMs can simulate arbitrary cyclic group automata.
Proof. Let A = 8. be cyclic group automaton. Now we define the input alphabet of the AUSSM to be Σ and choose its hidden dimension to be d = [ Let a E E be the symbol whose transition function 6a has order k. Then we set the parameters of the AUSSM as follows: Let B(u)= = Y'uu = e(o), a E Σ. Let A(c(a)) = exp(2ri/k). For each other symbol b E E, we know there exists E [0, k 1] such that 6; = 6% we can set A(e(a)) there is a trivial isomorphism vp between the values of x and the states of the FSA A: Just define vi: (exp(2rin//)
Now, exp(2rin/f
n [k]] + 2/kZZ exp(2rin/ + ni, which maps every hidden state the corresponding state of the automaton (arranged in the order of cycle traversal by 6,). Now there are n distinct possible hidden states which can be read out logarithmic precision.
A note on numerical precision. Floating-pont periin operations introduce rounding errors when com- puting the exponential function and repeated products thereof. A single complex multiplication introduces a relative error of at most v5u [69 where u is the unit roundoff (u = 2-25 for 32-bit single and u 2-53 for 64-bit double precision). This yields relative error bounds of v5 2 and V5 2 2-53 respectively. This means after N multiplications, the accumulated relative error is approximately v5uNN to the first order. Two adjacent kth roots of unity are separated by 2n/ /k seg- ment of the unit cirele; by the chord theorem, the distance between them is A = 2 sin(t//k) R 2 kh Approximations start overlapping if the accumulated error surpasses A/2, which occurs when:
1.26 x 10
π
N Z
(26)
2V5u
/5kk
k
For example, with 64-bit double precision and a modulo counter as large as lt 10%, it would take over 12 billion tokens (N z 1.26 x 10²0) for counts to become indistinguishable. This exceeds the sequence length of most datasets currently used in practice and is an order of magnitude larger than the human genome (= 3 x 10 base pairs). This means that whenever higher counters or longer sequence lengths are required, one can simply switch to the next higher precision. Since bit-depth is inversely proportional to the logarithm of we only require logarithmic precision in the sequence length.
The second requirement for transcending the expressivity limits of common SSMs is the ability to implement cascade products of FSAs (see (23 70) for more details on cascade products);
Definition 8. Let A₁ (Σi, Q1, Ag (22,, Q2, d2) be FSAs such that Σn = Q1 x Σ, Then
Ag
the cascade product Ai A2 is the FSA C -21, Q1 x Q2, 5c) with d, defined as
6.((gn (.(())192))0) .) (82 (y1, (12, r)), & (1)
(27)
Here, we use tuples of states taken from the state sets of the component FSAs to denote the state of the cascade. Intuitively, the state of the cascade at any given time is the combination of the states that the component FSAs are in at that point.
Assuming underflow, overflow, subnormal numbers.
25

---

Note that for transitioning to the next state, A₂ requires access to the state that A₁ was in before starting the current transition, meaning at time t. an FSA higher up in a cascade needs access the state the lower-level FSAs were in at time / 1.
We will hence the following crucial fact used for constructing FSA cascade products in Mamba SSMs:
Fact 2 (Sarrof et al. [TT], Lemma 17). For any alphabet E. there exists a single-layer Mamba SSM such that the last-but-one input symbol can be read out from the hidden state at finite precision.
We will also need the following fact about our hybrid architecture, allowing us to disregard the particular alternating ordering of layer types:
Note 3. We can always add an idempotent Mamba or AUSSM layer in the cascade without changing the model's behavior. This can be done by setting the output projection of the SSM block question to map everything to zero. Then the input to the next layer will just be the output of the last but one layer (via the skip connection). This meuns that for any Mamba or AUSSM with a specific behavior, there is hybrid LUSSM+Mamdte with the same behavior.
Now we have the necessary building blocks to show that our construction fulfills the main requirement for increased expressivity, the ability to implement cascades of the two SSM layer types:
Lemma 2. An SSM consisting of interleaved Mamba and AUSSM blocks (hybrid Mamba+AUSSM) can implement cascade products of automata simulated by Mamba SSMs and AUSSMs.
Proof. We want to show that the hybrid Manba+ architecture with alternating Mamba and AUSSM layers can implement cascade products of FSAs. In the following, we take a hybrid ManSa+AUSSM to mean stack of alternating Mamba and AUSSM layers, ignoring the initial encoding and final normalization and output map. Without loss of generality, assume that the first layer is always an AUSSM layer, and the last layer is always Mamba layer (we can achieve this by adding idempotent layers where necessary, see Note 3)
Also note that, by Note any Mamba SSM and any AUSSM can be converted equivalent hybrid Mamba+AUSSM.
It remains to be shown that a hybrid AUSSM AUSSM can simulate the cascade of two FSAs simulated by hybrid Mamba+AUSSMs. This is simply an extension of 11. Lemma 19] to our hybrid Mamba+AUSSM architecture.
Let Ar (Σ,, Q1,81), 8, A = (Σ², Q2 52) be FSAs such that Σg = -Q x E, Assume that there are hybrid AUSSA+AUSSM models Si, S2 that map input sequences (I1, T2.............................. IT) to the sequences of states under A,, A2. at logarithmic precision
Let Se be the hybrid MumSa+AUSSM we want to simulate the cascade A₁ o Ag. The lower layers of Se are just the layers of Sj. We add d dimensions that just copy the input via skip connection. We then add a Mamba layer (preceded by an idempotent AUSSM layer) that reads out the second-to-lase output of S₁ in new dimensions (by Fact 2. while again forwarding the input via the skip connection. Here, we also add a dummy dimension that is always 1, which avoids normalization, making different inputs indistinguishable. Now we have the input and the second-to-lase output of S1, corresponding to the last state of A,. Now the remaining layers of S. are just those of S2, which take this input and compute the transition and state of A2, again adding dimensions such that the state of A₂ is separate from the state of A, and the input to the overall SSM. Now, S, maps each w to the state sequence under A1 o Ag. again at logarithmic precision. This can be inductively extended to a cascade product of arbitrarily many FSAs.
Fact 3 (Consequence of Krohn-Rhodes Theorem 23 and the decomposition series of groups 27 ). Any solvable language is recognized by cascade of set-reset and cyclic group automata.
Theorem 2. Hybrid AUSSM can recognize any solvable language, that is, any regular language whose syntactic monoid does not contain non-solvable subgroups.
Proof. By Lem. D an AUSSM can simulate cyclic group automata. By M Lem. 19], a Mamba SSM can simulate set-reset automata. By Lem. 2 hybrid AUSSM+Mamba can simulate a cascade
Note here we implicitly assume a bijection exists between intervals on ge (the input at time t, æx) and the alphabet symbols of the relevant FSA.
26

---

Additionaa Memory
B
K2
X
30
K₂
X1
Rigi
Ksa
%
X2
X2
13
Kau
Kax Kiu
K4
K3
KC
K
34
D
Additional Memory
G1
Gg
Gi
0
K2:
KGa
99
G d
G
fx
Kan
Kai
Vs
G
44
Kn
Kizz
Bac
Figure 3: Space Complexity of SSM formulations: The figure illustrates an example convolution kernel for an SSM provided with four inputs different timesteps (ux). The convolution is visualized as matrix multiplication operation over the input sequence. A. In LTI SSMs, the convolution kernel (K1, K2, K3, K4) is precomputed and applied the input different timesteps obtain the output. B. In general LTV SSMs with time-varying recurrence, the convolution kernel has O(L") elements, each unique the input and output being considered each timestep. The use of convolution in this scenario leads to quadratic complexity in space (akin to the transformers). C. In the separable convolution case, the quadratic matrix of the general SSM can actually be obtained by the outer product between Se for each timestep and the cumulative sums of a function gs independent of t. D. Computing the convolution kernel can be achieved in just an additional ((2T.) space.
of automata simulated by Mamba and AUSSM SSMs. Together with Fact this means that hybrid AUSSM+Mamba can recognize any solvable language.
The importance of counting for other tasks. Note that the ability to count modulo & does not just allow SSMs model regular languages but also approximate languages higher up the Chomsky hierarchy. For example, it allows the recognition or generation of bounded Dyck languages, i.e., the correct parenthesization up to a certain depth (see 7 in the case of RNNs). Even context-sensitive language tasks can benefit from counting: For instance, sorting sequence (the bucket sort task in $5) can be done by maintaining counters for all alphabet symbols and then outputting the symbols in order, according to their count (see counting sort and direct-address tables [72] Chapters 8 and 1 1]). Note that this works as long as the number occurrences any given symbol smaller than the highest count expressible by the SSM, e.g., & when using modulo & counting.
F Complexity Analysis
SSMs leverage logarithmic complexity algorithms like FFT and parallel prefix sum to compute the convolution. Prior to this, the convolution kernel needs to be pre-computed and stored, which the main bottleneck in computing the convolution. We will show below the space complexities for computing and storing the convolutions, Further, we show how the quadratic space complexity blowup of pure LTV systems can be managed using the separable convolution framework.
27

---

F1
SSM Convolution
The convolution operation of general SSM given by the following
ytt)
C (A B3 u(k)
(28)
kst
There are two cases for the above convolution we consider:
Linear Time Invariant (LTI) : In the LTI case, the matrices in the SSM are constant over time, and the following holds
v(t) C
(29)
4:t
Now, the convolution kernel (tt CA B be precomputed, and since A is common for many settings of t and for which their difference is constant, the weights can be shared. In fact, there only (L) unique entries in the convolution kernel (see Figure A). The other entries are duplicates of these entries. Once the convolution kernel obtained, efficient algorithms like FFT or Parallel Scan can be used to compute the convolution in (log L) time for each dimension, for total of OL log( L)) time complexity. Therefore, the total time complexity for computing the kernel (L log L) with space complexity of (L)).
Linear Time Varying (LTV) In the LTV case, the matrices in the SSM can vary over time. This introduces additional complexity in representing the convolution kernel in O(L²) space, matching the quadratic complexity of computing self-attention in transformers. The reason for the quadratic complexity is that the entries in the convolution kernel (tt, k) unique for each setting of k. In the case of separable convolution kernels (e.g, the case of simultancously diagonalizable matrices), the resulting (t,,, k) matrix has further rank-1 factorization (this is discussed in detail in the main text). This factorization enables the convolution kernel to be represented with only an additional O(2T) memory, where the factor comes from each vector element in the outer product.
F.2
Parallel Scan
The reason for precomputing the convolution kernel is that we can apply one of the fast convolution algorithms- FFT or parallel scan. In our case, we perform the parallel prefix sums for computing cumulative sums. Here, we analyze the time and space complexity of the parallel prefix sum (scan) algorithm, where the goal is to compute the prefix sums of an array A = [uo, ..,...... ar.1] such that the output array S satisfies
aj for Osi L.
(30)
We assume a parallel computation model such as the PRAM (Parallel Random Access Machine) or shared-memory memory model, and are given P processors.
shared-me
The parallel prefix sum algorithm typically consists of two main phases:
1. Upsweep phase (Reduction): Build binary tree over the array and compute partial sums from leaves to the root.
2. Downsweep phase: Propagate prefix sums from the root back down the compute the final result.
Both phases traverse binary tree structure of height log, L, assuming for simplicity that L is power of two. Each level of the tree can be processed in parallel.
Work. The total number of operations (work) in both phases is: (T. - (L. (1-1) = 27. 2 O(T
W(L)
(31)
upsweep
downsweep
This is the same amount of work as the sequential prefix sum algorithm, which confirms that the parallel algorithm is work-efficient.
28

---

Time Complexity with P Processors. Using Brent's Theorem (work-spar model), the parallel time Tp on P processors bounded by:
- + W(L) +(L - o (p
(-p+ log L
Tp
(32)
This means that when the number of parallel processors grows in the sequence length according to P = (L/log L), the parallel prefix sum runs in optimal time Ollog L).
29

---

Space Complexity
The space used by the algorithm includes:
The original imput array A, of size L.
An auxiliary array store intermediate results, typically of size T..
Additional temporary variables per processor (constant per processor).
Hence, the total space complexity is:
P) = O(L) (since typically P L).
(33)
is important to note that although the algorithm requires additional O(L) space for the auxiliary variables, the CUDA kernel implementation hides these variables within the multiprocessor registers and shared memory. As a result, this complexity does not show up the plots of either Mamba or auSSM. Existing GPU hardware for the 2080ti enables parallel processing of sequences up to 2048. For longer sequences, the input chunked into batches of I. 2048.
G Implementation
The theoretical analysis of the separable kernel formulation shows that the adaptive kernel can be implemented in only additional linear space. However, the factor associated with the linear space bdn, where b the batch size, d is the input dimension, and n is the hidden state dimension. In this section, first show a PyTorch implementation of the AUSSM kernel and Mamba kernel that can be easily coded, with the higher cost of the constant factors. Next, we show how we implement the AUSSM kernel in practice so that the additional complexity is hidden within the computations of CUDA kernel.
G.1 PyTorch
One of the most useful aspects of the theory of separable convolutions that there is relatively efficient PyTorch formulation for computing SSM kernels, when the SSM is partially/fully time varying. However, an additional constant-time penalty will be incurred. Nevertheless, the existence such implementation will still be interesting as it can enable fast prototyping LTV SSMs, without dealing with the complexity of building a CUDA kemel. Here, we show two PyTorch implementations of the partial LTV Mamba kerne and the separable AUSSM kernel.
def mamba_ssm (u, dt, A, B, C, D, 2):
params:
u: input Tensor (b,
dt: Delta Tensor (b,d,1)
A: Tensor (n)
B: Tensor (b,n,l)
Tensor (b,n,1)
D: Tensor (d)
10
2: Tensor (b,d,1)
Returns:
12
y: (b, d, 1)
13
14
A einsum (A, dt, "n, bdl bd
G torch. cumsum (axis=1)
17
g einsum (-G) dt, B, u, "bdnl bdl, bnl, bål bdnl" torch. cumsum (g, axis=-1) einsum (C, exp(G), "bnl, b>bd>>
18
19
y einsum (f, g, "bdml bdnl y sill
22
23
return
30

---

The implementation of Mamba using the separable kernel formulation has fewer than 10 lines of PyTorch code. The PyTorch implementation AUSSM similar, except now have account for the time-varying A matrix, and B and C are relaxed.
def aussm (u, dt, chi, B, D, z): unn
params:
u: input Tensor (b,d,1)
dt: Delta Tensor (b,d,1)
chi: adaptivity matrix (d,n,d)
B: Tensor (n)
C: Tensor (n)
D: Tensor (d)
Tensor (b,d,1)
11
Returns:
12
y: (b,d,1)
13
A einsum (chi, "dnr, blr-br--bbl blr->bidn") einsum (dt, A, *bdl, bidd-lddndn -bldn"
16
torch. cumsum (axis=1)
17
18
einsum (exp (-G), dt, B, "bdnl, bdl bdl bdn1" torch. cumsum (g, axis=-1) einsum (c, (G) "n, bdnl---dddn1"
20
21
22
einsum (f, 8, "bdml bdnl- *u y (illu(z)
25
return y
this implementation, Mamba and PyTorch have the same space and time complexity as the hidden state is realized for both, albeit at only fraction of the cost.
G.2 CUDA Kernel
In pure PyTorch, the additional complexity of realizing the hidden state is unavoidable, even though the computation does not have quadratic memory costs. The additional complexity of realizing the hidden state can be avoided by creating CUDA kernel for the AUSSM equation. We use the following equation for the AUSSM, which we introduced in the main text:
A
Usi = R
ΣΣ C exp RSL.
Aas
(k(k))
(34)
exp i OAisS)
Each thread of the CUDA implementation computes the array inside the nested summation, which results in O(L) memory requirement for storing each the variables (A, f, g) for the forward pass. These variables not realized the same time in the GPU memory. but in registers within the streaming multiprocessors (SM), each processor holding to 16 items each array. For the 2080Ti GPU,
the CUDA kernel on, the allowable maximum sequence length that be processed by the kernel was 2048, after which the register and shared memory costs start to show up. We found that this sequence length is ideal for the hardware and tasks tested The separable convolution trick is not restricted by the hardware and can scale well for GPUs that can be released in the future with larger registers and shared memory resources.
Backpropagation: For the CUDA kernel, we implemented a custom backpropagation operation. Implementing backpropagation requires the variables computed during the storage to be stored, which creates issues because the reason we are writing the CUDA kernel is so that we do not have to realize the memory-intenste hidden state. We therefore recompute the forward pass during backpropagation. The low complexity of implementing the AUSSM in CUDA means the recomputation does not incur a heavy penalty.
31

---

Table 6: Best Hyperparameters.
Task Group
Task
Layers
d
weight decay
learning rate
repetition
64
32
0.0
0.01
bucket sort
am
64
32
0,0
0.01
majority count
ma
64
32
0.1
0.01
majority
ma
64
32
0,1
0.01
Algorithmic
solve equation
ma
64
32
0.0
0.01
mod arith
16
8
0.0
0.01
mod arith wo bra
ma
8
16
0.0
0.01
cycle nav parity
16
8
0.0
0.01
ma
16
8
0.0
0.01
Heartbeat
ma
64
64
0.0
0.0001
SCP1
16
128
0.0
0.001
Timeseries
SCP2
ma
16
128
0.0
0.0001
Classification
Ethanol
anmama
16
64
0.001
0.00001
Motor
ma
16
128
0.0
0.0001
Worms
amma
16
16
0.0
0.001
Timeseries
weather ma
16
128
0.0
0.001
Regression
H
Experiments
We conduct three sets of experiments: (1) to evaluate the tine/memory complexities of the different AUSSM implementations, (2) to evaluate the performance AUSSM algorithmic tasks enabling insights into the expressive power, and (3) evaluate real-world performance implications in range of long time series benchmarks. For each the tasks involving training models (2 and 3), we perform pipeline processes obtain the final test accuracies. The first pipeline is the training and model selection pipeline with only the training and validation sets that are preselected based on the same criteria used by prior literature. The second pipeline is the test pipeline and is entirely separate and performed starting 10 days prior to paper submission to avoid model selection based on the test results. The classification tasks evaluated using the scaled test accuracy metric, where the obtained accuracy values are scaled with respect to the baseline performance of uniform random distribution, as shown below.
scaled accuracy score accuracy score baseline accuracy score
baseline accuracy score
All the models were run in supercomputing cluster, where we used 40 2080Ti GPUs for all except the dataset Eigenworms dataset that required higher memory. This is the lowest GPU available the cluster, with at least a CUDA compute of 7.5 required to run the Mamba and AUSSM CUDA kernels. For larger memory Eigenworms workload, we used the L4 GPU, which has a VRAM of 23GB. Higher VRAM GPUs were available in the cluster, but they were in high demand and unnecessary, as our optimized CUDA kernel was able to handle even the large-scale tasks in modest hardware.
H.1 Scalability Evaluation
To evaluate scalability in a fair manner, we report only the time spent in computations, ignoring the latencies associated with moving variables between the GPU and the CPU. This provides fair evaluation of the algorithmic performance. 5 runs are used to warm the GPU before starting the evaluation remove transient start-up effects. The run-time values averaged over 50 runs, where each run computes a forward and backward pass for each of the implementations. The peak memory used during each run is also similarly recorded and averaged for each of the 50 runs.
H.2
Time Series benchmark
For time series classification and regression benchmarks, we follow the train-validation protocol for model selection, following prior works on the benchmark. For testing, modified the procedure
32

---

as the five arbitrary random seeds used to evaluate test performance in prior works may introduce unwanted biases due to the low number random samples. Also, prior works used JAX for implementations, while we used PyTorrc,, and the random seed does not create the same train- validation-test sets due to differences in the pseudorandom number generators. We thus decided to evaluate on sein-validation-test splits created with 20 different seeds. We anticipated that the higher samples would help in providing a better estimation of the test accuracy than what the five arbitrary seeds provide. For each task, we performed hyperparameter search over the following grid: {16, 64, 128}, n E {16, 64, 128], learning rate E {0.00001, 0.0001, 0.001 }. and five different seeds for model selection. The model hyperparameters with the highest mean validation accuracy chosen for evaluation in the test set.
H.3 Algorithmic Tasks
For algorithmic tasks, we used the results from [31] for comparing against baseline models. We used a grid search for hyperparameter tuning with a grid search over d E {8, 16, 32, 64), n E {8, 16, 32}, weight decay E {0.0, 0.001, 0.01 J. learning rate in {0.0001, 0.001, 0.01 } and five seeds. The batch size was fixed at 256. For pure AUSSM blocks, we tested networks with depth 2, 4, and 6. For hybrid AUSSM blocks, tested all possible 2-block configurations of Mamba (represented m) and AUSSM blocks (represented as a. (ma, am, mm, aa). For each of the evaluated algorithmic tasks, randomly sampled 10000 samples from train up length-40 sequences. The validation set is sampled independently from 40-256 sequence lengths and had 1,000 samples. The test set had 10,000 samples from sequences of up 256 sequence lengths.
The tasks use the same vocabulary size and configuration used in [31. Some samples from the tasks are shown below as timeline. Here, the mask is applied to the output to determine the output of interest for computing the loss and output.
33

---

repetition
time
0
1
2
8
9
10
imput
output
mask
0
bucketsort
11
12
time
0
8
9
10
input
15
output
ACT
PAD
16
mask
0
18
19
Task: noderithmetiewwbraess
20
23
input
0
2
0
24
oniput
2
PAD
mask
27
28
Tack: eyclenav
30
time
31
32
STAY
0
PAD
35
36
modatithmetic
38
19
time
0
8
9
10
40
input
42
output
43
mask
0
0
0
0
44
46
Task: sulveequation
47
48
time
1
8
9
50
input
ACT
3
PAD
PAD
51
output
ACT
PAD
PAD
PAD
52
mask
54
55
Task: parity
56
time
58
59
input
1
0
1
0
60
output
0
62
63
Task: mejoritycount
66
time
0
8
9
67
68
input
43
51
34
ouiput
20
mask
0
0
0
71
72
73
Task: majority
74
25
time
0
1
2
I
1
8
9
10
76
input
78
output
51
34
10
46
54
79
mask
0
0
1
0
0
82
Task: set
83
86
input
3
5
ACT
5
7
PADD
87
output
0
ACT
PAD
PAD
mask
34

---

NeurIPS Paper Checklist
1. Claims
Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Question:
Answer: [Yes]
Justification: We make the following claims in the abstract: (1) AUSSMs are maximally
expressive in the class diagonal SSMs proved in Section 3.11 (2) Separable convolution
kernel formulation enables scalability. Theoretical exposition in Section 4 and plots in
Figure 2 (3) Unitary properties analyzed Section 3.1 (4) The ability to solve general
class of regular languages experimental validation in Table I (5) Competent performance
on real-world benchmarks in Table 2
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No
NA answer to this question will not perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results
be expected to generalize other settings.
is fine to include aspirational goals as motivation as long is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: Limitations are discussed in Section
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g.. independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed not penalize honesty concerning limitations.
3. Theory assumptions and proofs
35

---

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?
Answer: [Yes]
Justification:
The sufficient conditions for applying the separable convolution formualation
in Section discusses related works and the associated assumptions used by them, which we implicitly assume.
is detailed in Section 4 The proofs
Guidelines:
The answer NA means that the paper does not include theoretical results. All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Answer: [Yes]
Short descriptions of the methodology is provided in the main text along with the discussion of the results in Section 5 where the related work that used identical experimental patterns are also discussed. More detailed descriptions are the Appendix.
Justification:
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution a specific model and empirical evaluation, may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, other means that appropriate to the research performed.
also be provided via detailed
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how to reproduce that algorithm.
(b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
36

---

(d)
The authors should provide scripts to reproduce all experimental results for the new reproducible, they We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification: The code is made available as part the supplementary information. We use data that is publicly available except for algorithmic tasks. For these tasks, we release the dataloaders along with the code. The code will be made public following the publication of the manuscript.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips.com cc/
public/guie SeetiiiggideeeCCooe wttps//oooSonnniiioooolllie....e
for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" an acceptable answer. Papers cannot be rejected simply for including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https
/pips..ccppulic//uune wttp://giieeee/eeeoee /
for more details.
The authors should provide instructions on data access and preparation, including how access the raw data, preprocessed data, intermediate data, and generated data,
proposed method and baselines. If only subset of experiments should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible supplemental material (appended the paper) recommended, but including URLs to data and code is permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (c.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: Short form descriptions experimental procedure are in the main paper. Long-form details of the precise hyperparameter tuning protocol and dain-validation-teer procedure are in the Appendix.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core the paper to a level of detail that is necessary appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment statistical significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
37

---

Justification: Standard Deviations are reported alongside the Long Time series benchmark. The algorithmic tasks do not contain standard deviations, as this is synthetic benchmark. For weather, we follow prior work and do not report standard deviation in the table; the standard deviation we obtained is 0.0173.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report I-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are range (e.g. negative eтror rates).
If error bars reported in tables plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: Resources, including the type of graphics card and the available VRAM, detailed in the experiments. Time of execution and further details of the compute resources are in the appendix.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers, CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary failed experiments that didn't make into the paper).
9. Code of ethics
Question: Does the research conducted the paper conform, in every respect, with the NeurlPS Code of Ethics https: /tctps://neurips.. http://///eaaessscc/ppb////hhhht
Answer: [Yes]
Justification: Our research introduces a new computational model, and this does not use human subjects for the experiments. We do not create any data and only publicly available datasets or standard synthetic benchmarks. There are no societal concerns we are aware of, as this is relatively small-scale study on computational research question.
Guidelines:
The answer NA means that the authors have not reviewed the NeurlPS Code of Ethics.
38

---

If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there a special consid- eration due to laws or regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: The study is performed and impacts only the academic community interested in conducting further research SSMs. The work is primarily foundational in creating a new computational algorithm for existing models.
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deeloyment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: The paper uses publicly available data that identified as risk free and typically used in conducting academic research.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
39

---

Question: Are the creators or original owners of assets (e.g.. code, data, models), used in the paper. properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification:
The code and data are publicly released as open source software. the code bases we used for compiling our code is attributed to the respective authors.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset used and, possible, include a URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset. For scraped data from a particular source (e.g.. website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperssi thcode. com datasets has curated licenses for some datasets. Their licensing guide can help determine the license dataset.
For existing datasets that are re- re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. If this information is available online, the authors are encouraged reach out to the asset's 's creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer:
[Yes]
Justification: release.
README is available how install, test and the CUDA kernel
Guidelines:
The answer NA means that the paper does release new assets.
Researchers should communicate the details the dataset/codeemoodll as part their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether how consent was obtained from people whose asset used.
At submission time, remember anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification: we do not use this experimental protocol.
Guidelines:
The answer NA means that the paper does involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According to the NeurlPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
40

---

15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based the requirements of your country institution) were obtained?
Answer: [NA]
Justification: The experimental we do does not use human subjects and do not require IRB approval.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with
human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent)
may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
recognize that the procedures for this may vary significantly between institutions and locations, and expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such the institution conducting the review.
16. Declaration of LLM usage
Question: Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods this research? Note that if the LLM used only for writing, editing, or formatting purposes and does not impact the core methodology, scientific rigorousness, or originality the research, declaration not required.
Answer: [NA]
Justification: LLM was not used in formulating the research. Only use of LLMs was in editing.
Guidelines:
The answer NA means that the method development in this research does not involve LLMs as any important, original, or non-standard components.
Please refer to our LLM policy /tttps://neuriii Cwaturiisscc/Cooffeee / /LLM) for what should or should be described.
41


Reviews:
[{"review_id": "5AKbKImrIu", "reviewer": "Reviewer_A2pm", "summary": "The paper introduces Adaptive Unitary State Space Model (AUSSM), a sub class of general SSMs to find a middle ground between theoretical expressivity limitations of the architecture and the performance / scalability on downstream tasks. The recurrent matrix $A_t$ is similar to Mamba styled models (S6), where the matrix is input dependent. Unlike Mamba, the authors here propose to make the matrix skew symmetric, as well for the overall discrete product $exp(\\Delta A)$to have Eigen values lie on a unit circle. This change ensures stability (avoiding exploding / vanishing gradients), while also allowing for separable kernel formulation. The authors also provide a CUDA kernel implementation to make these computations faster. The authors test this AUSSM block in isolation, or in a hybrid setup along with a Mamba block as well. Theoretical analysis of these reveals that the hybrid setup overcomes some of the theoretical limitations of the original Mamba block, specially for modular counting (for eg. Parity, which Mamba provably can't solve at arbitrary lengths, but AUSSM blocks can). To complement the theory, the authors conduct some experiments on the formal languages for which their claims were made as well as on some time series dataset to show that their architecture can be scaled to model long term dependencies as well.", "questions": "1. **Hybrid performance dips**\nTable 1 shows a few tasks where the AUSSM + Mamba hybrid lags behind either component on its own, even though Note 1 (lines 619–632) argues the behavior of hybrid models can be made identical to either a Mamba / AUSSM block on its own. Could you provide a bit more detail about what you think could be the reason behind slightly worse performance on these tasks vs the others. \n\n2. **Neuroscience connection**\nThe connection to the Neuroscience literature to me seemed not very clear / very limited. Is the connection to the neuroscience literature restricted to the fact some jPCA procedure where such skew symmetric ODE is used or is the derivation of the AUSSM block also inspired from it ? \n\n3. **Contextualizing results compared to recent SSM advancements** As mentioned in weaknesses, other recent work [1], [2] has also aimed to enhance SSM expressivity. How does AUSSM compare to these, could you contextualize your contributions vis a vis theirs ?\n\n*(Side note: I was able to check most appendix material, but Section F is the only section I couldn't verify properly, because of lack of background knowledge on my part there)*\n\n[1] Grazzi et al 2025, Unlocking State-Tracking in Linear RNNs Through Negative Eigenvalues, ICLR 2025. \nhttps://openreview.net/pdf?id=UvTo3tVBk2\n\n[2] Peng et al, 2025 RWKV-7 \"Goose\" with Expressive Dynamic State Evolution, Arxiv\nhttps://arxiv.org/pdf/2503.14456", "limitations": "yes", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 3, "significance": 3, "originality": 4, "strengths_and_weaknesses": "Strengths \n- The paper contributes to the space of SSMs by incorporating small changes in the parameterization of an SSM and thus achieving a theoretically more expressive architecture. All the Theorems and Lemmas seem correct. I was able to verify most proofs. (see questions for things I didn't understand). \n- The paper contributes both on the theoretical side, as well as provides efficient CUDA Kernel implementations to practical relevance. \n- Adequate empirical evidence for a theory paper. Algorithmic tasks confirm the formal-language claims, and the UEA/Weather benchmarks show the gains transfer to realistic long-range settings, all on modest hardware.\n\nWeaknesses \n- The only relevant weakness I could think of was -- inadequate comparison / discussion about similar changes proposed by some of the other recent works. For example the papers - [1], [2], achieve somewhat similar results, compared to what the authors get. I can intuitively feel that the changes the authors are proposing here, should be even better than the change proposed at least in [1], but I am not sure about [2], so it would be good from the authors point of view, to contextualize their contribution compared to these other ones. \n- Minor presentation issues (just a subset, requesting the authors to do a better passthrough for the next version) .\nTypos: Line 63 “Crdinary → Ordinary”; Line 67 “respectivelt → respectively”; Line 69 and Line 196 difficult to parse.\n- Missing citations: The appendix references papers 36, 37, … that are absent from the main list (only 1–33 appear).\nThis made understanding some of the points unclear. For example, I am unfamiliar with some of the precision results cited, Line 598 - 599 -I am unfamiliar with the paper that states that complex multiplication incurs an error $\\lt \\sqrt{5} \\epsilon$, and so couldn't verify what the authors meant there. Maybe put those additional references in the rebuttal response, so that it's easier to understand the results authors were referring to. \n\n[1] Grazzi et al 2025, Unlocking State-Tracking in Linear RNNs Through Negative Eigenvalues, ICLR 2025. \nhttps://openreview.net/pdf?id=UvTo3tVBk2\n\n[2] Peng et al, 2025 RWKV-7 \"Goose\" with Expressive Dynamic State Evolution, Arxiv\nhttps://arxiv.org/pdf/2503.14456", "was_revised": true, "final_justification": "I am satisfied with the authors responses, and if they include the details about the new experiments (that also take into account, other similar changes proposed by other recent work in this area) in the main paper, I believe this would be a good contribution to the SSM architecture landscape.", "rebuttal": "We thank the reviewer for comments and suggestions for improving the manuscript and highlighting our **core contributions to the theory and a computationally efficient CUDA kernel implementation**. We thank the reviewer for also highlighting that the benchmarks confirm our theoretical predictions, and the real-world settings also adequately show performance improvements. \n\n## Relevant Changes\n\n### 1. Experiments with general-eigenvalue baselines\n\nWe thank the reviewer for providing relevant works. We added new results in the formal language benchmarks to compare with two versions of Mamba with general eigenvalue spectrum (NOTE: Original Mamba has eigenvalue range between \\[0,1\\])\n\n|                   | Mamba with Complex Eigenvalues | Mamba in \\[-1,1\\] range |\n| ----------------- | ------------------------------ | ----------------------- |\n| repetition             | 0.09                           | 0.1                     |\n| bucket sort          | 0.21                           | 0.91                    |\n| majority count         | 0.19                           | 0.31                    |\n| majority          | 0.13                           | 0.64                    |\n| solve equation             | 0.43                           | 0.24                    |\n| modarithmetic          | 0.12                           | 0.116                   |\n| modarithmetic wo bracket | 0.23                           | 0.24                    |\n| cyclenav          | 0.42                           | 0.91                    |\n| parity            | 0.27                           | 1.0                     |\n\nAs rightly pointed out, Mamba \\[-1,1\\] performs better compared to Mamba on the parity, cyclenav, solve equation, and bucketsort. However, AUSSM and the AUSSM Hybrid model still perform better in most tasks except bucketsort and solve equation. The tested complex version of Mamba performed poorly compared to the Grazzi version, which iterates the finding in the Mamba paper that real-eigenvalued parameterizations performed better. In AUSSM, the additionally introduced adaptivity and unitarity provide sufficient requirements for enabling learnability and expressivity.\n\nFor the weather benchmark, we added S5 to our list of baselines\n\n| model | MAE  |\n| ----- | ---- |\n| S5    | 0.36 |\n\nS5 is found to have better MAE than Mamba and LinOSS, but it our proposed AUSSM Hybrid model performs the best in the benchmark.\n\n## Addressing Weaknesses and Questions\n\n### Weakness 1: The only relevant weakness I could think of was -- inadequate comparison / discussion about similar changes proposed by some of the other recent works. For example the papers - \\[1\\],\\[2\\], achieve somewhat similar results, compared to what the authors get. I can intuitively feel that the changes the authors are proposing here, should be even better than the change proposed at least in\\ [1\\], but I am not sure about \\[2\\], so it would be good from the authors point of view, to contextualize their contribution compared to these other ones.\n\nThank you for the suggestion. We have added the modified Mamba implementations (Mamba\\[-1, 1\\] Grazzi, Mamba Complex) in the formal language evaluations. Note that negative eigenvalues only give a small expressivity benefit over non-negative eigenvalues. The results we obtained (see Changes #1 above) are consistent with the prior results - Mamba\\[-1,1\\] improves on the base Mamba for most tasks but falls behind AUSSM hybrid in some. Mamba Complex, however, performs poorly in most tasks compared to both Mamba\\[-1,1\\] and AUSSM hybrid despite higher expressivity. This contradiction was found in the original Mamba paper, too. It is possible that Mamba with pure complex spectrum may have poorer learnability for many tasks - there are no free lunches.  \nRWKV-7 is an LRNN that does not have a diagonal transition matrix, which in theory could make it more expressive than our architecture in certain settings, though at the price of additional computation cost.\n\n### Weakness 3: Minor presentation issues (just a subset, requesting the authors to do a better passthrough for the next version). Typos: Line 63 “Crdinary → Ordinary”; Line 67 “respectivelt → respectively”; Line 69 and Line 196 difficult to parse.\n\nWe thank the reviewer for pointing these out. These are now corrected in the manuscript.\n\n### Weakness 4: Missing citations: The appendix references papers 36, 37, … that are absent from the main list (only 1–33 appear). This made understanding some of the points unclear. For example, I am unfamiliar with some of the precision results cited, Line 598 - 599 -I am unfamiliar with the paper that states that complex multiplication incurs an error, and so couldn't verify what the authors meant there. Maybe put those additional references in the rebuttal response, so that it's easier to understand the results authors were referring to.\n\nWe apologize for this oversight. The missing citations are:  \n\\[34\\] Babcsányi. Automata with finite congruence lattices. Acta Cybernetica, 2007.  \n\\[35\\] Salkoff et al. Propagation of traveling waves in the visual cortex requires microcircuit coordination. Neuron, 2020.  \n\\[36\\] Brent et al. Error bounds on complex floating-point multiplication. Mathematics of Computation, 2007.  \n\\[37\\] Zimmermann. On krohn-rhodes theory for semiautomata.  \n\\[38\\] Hewitt et al. RNNs can generate bounded hierarchical languages with optimal memory. EMNLP, 2020  \n\\[39\\] Cormen et al. Introduction to Algorithms, Third Edition. The MIT Press, 2009.\n\n### Question 1: Hybrid performance dips\n\nWe fix the number of blocks as two and try 4 different combinations - AM MA AA MM. A: AUSSM block and M: Mamba Block. The pure AUSSM (AA) is a good middle ground, but not the most expressive possible - see our theoretical analysis (Lemma 1 and Lemma 2 in the paper). As such, there may be tasks that require two layers of mamba for non-unitary recurrence or full adaptivity. The details of why exactly this contradiction is present are not clear at the moment and are described as one of the limitations of the current study. \n\n### Question 2: Neuroscience connection\n\nThe AUSSM ODE is derived from the jPCA procedure (detailed in appendix). The jPCA procedure, introduced by Churchland et al. (2012), was designed to study rotational dynamics in neural activity, and has been used widely in neuroscience to uncover such dynamics from neural recordings during tasks such as motor planning and execution. By fitting low-dimensional projections that maximize antisymmetric structure in the population dynamics, jPCA reveals patterns consistent with conservative, unitary evolution. Finding such conservative dynamics strongly represented in real neural data suggests that portions of cortical activity evolve with minimal dissipation or loss of energy -- exactly analogous to the AUSSM, and thereby inspiring both the method and the goal. The block structure is not biologically inspired but is used in the Mamba class of models.\n\n### Question 3: Contextualizing results compared to recent SSM advancements\n\n\\[1\\] and \\[2\\] both use non-diagonal recurrences, sacrificing some computational efficiency for enhanced expressivity, allowing them, in theory, to recognize all regular languages. Note, however, that a single Householder reflection per step is insufficient for full regular language expressivity or even to perform arbitrary rotations. As \\[1\\] points out, one needs products of multiple reflections to realize arbitrary orthogonal transformations, which can be costly. By contrast, AUSSM can perform rotations cheaply with a single diagonal update. Our contribution, therefore, is making *diagonal SSMs* maximally formally expressive while keeping computational overhead small, extending their expressivity to all solvable languages (a subset of the regular languages), which is also widely believed to be the upper expressivity bound for transformers \\[3\\]. Due to time limitations, we have added only results for running \\[1\\] on the formal language benchmarks, though we do compare to xLSTM, which is also non-diagonal like \\[2\\], but is additionally a non-linear RNN. We have added a related work section to the appendix contextualizing our expressivity results with regard to other approaches, which unfortunately does not fit within the character limit of our response.\n\n## References\n\n\\[1\\] Grazzi et al., 2025. Unlocking State-Tracking in Linear RNNs Through Negative Eigenvalues. ICLR (2025)\n\\[2\\] Peng et al., 2025. RWKV-7 \"Goose\" with Expressive Dynamic State Evolution \n\\[3\\] Merrill et al., 2022. Saturated Transformers are Constant-Depth Threshold Circuits"}, {"review_id": "Sp2H41aYb8", "reviewer": "Reviewer_MTfs", "summary": "This paper proposes the Adaptive Unitary State Space Model (AUSSM), which is built upon a skew-symmetric ordinary differential equation. While AUSSM introduces a slight increase in computational complexity, the overhead is not significant. Compared to state space models (SSMs) with real-valued eigenspectra, AUSSM offers higher expressivity with unitary eigenspectra. In addition, the paper presents a hybrid architecture, AUSSM Hybrid, which integrates the AUSSM with Mamba blocks. The method primarily targets improvements in tasks where counting is crucial, and validation is conducted across various sequence modeling benchmarks.", "questions": "1. What is the theoretical or practical benefit of an SSM having the conservation property?\n2. Aren’t complex-valued eigenvalues used not only in LinOSS but also in S4D and S5? In Mamba, complex-valued variants were also tested in ablation studies, and real-valued variants were selected due to slightly better performance. Emphasizing real eigenspectra as a prevailing trend in prior SSMs seems somewhat misaligned with the actual research landscape.\n3. In the paper, Mamba is referred to as mildly adaptive, since it fixes $A$ (actually, it eventually varies through the use of $\\Delta$), and adapts only $B$ and $C$. In some parts of the paper, however, Mamba is described as non-adaptive. Would it be more consistent to refer to it as *mildly adaptive*, and could a more detailed discussion on the role of $\\Delta$ be included?\n4. Can the baseline models (e.g., S5, S6, LinOSS, xLSTM, Mamba) and the proposed models (AUSSM, AUSSM Hybrid) be unified across all experiments? Is there a particular reason why different subsets of models are selected for different tasks?\n5. The results on the UEA dataset differ significantly from those reported in the LinOSS paper. Can this be explained?\n6. AUSSM Hybrid achieves the best performance, not the pure AUSSM model (especially in Tables 2 and 3). This raise concerns about whether AUSSM alone sufficiently supports the paper’s main claim of resolving the expressivity-efficiency tradeoff. If AUSSM indeed offers superior expressivity, shouldn't it also achieve strong results beyond counting-focused tasks?", "limitations": "yes", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 2, "significance": 2, "originality": 3, "strengths_and_weaknesses": "- Strength\n    - The paper provides a rigorous analysis of the underlying systems modeled by SSMs and explores a principled direction for improvement.\n- Weaknesses\n    1. The selection of prior works and baseline models in the experiments appears biased, especially in terms of real vs. complex eigenspectra.\n    2. Since some experiments only report results for the hybrid model, it is difficult to isolate and evaluate the impact of the proposed AUSSM.", "was_revised": true, "final_justification": "I have raised my score, taking into account both the improvements in clarity and the remaining limitations due to the codebase.", "rebuttal": "We thank the reviewer for the critical evaluation of our work and for providing constructive criticism that ultimately improved the manuscript. Thank you for highlighting the paper **provides a rigorous analysis of the underlying systems modeled by SSMs and explores a principled direction for improvement** and for the comment that the overhead posed by introducing adaptive time varying matrices was not significant. \n\n## Relevant Changes\n\n### 1. Experiments with general eigenvalue baselines\n\nWe added new results in the formal language benchmarks to compare with two versions of Mamba with general eigenvalue spectrum (NOTE: Mamba has eigenvalue range between \\[0,1\\])\n\n|                   | Mamba with Complex Eigenvalues | Mamba in \\[-1,1\\] range |\n| ----------------- | ------------------------------ | ----------------------- |\n| repetition             | 0.09                           | 0.1                     |\n| bucket sort          | 0.21                           | 0.91                    |\n| majority count         | 0.19                           | 0.31                    |\n| majority          | 0.13                           | 0.64                    |\n| solve equation             | 0.43                           | 0.24                    |\n| modarithmetic          | 0.12                           | 0.116                   |\n| modarithmetic wo bracket | 0.23                           | 0.24                    |\n| cyclenav          | 0.42                           | 0.91                    |\n| parity            | 0.27                           | 1.0                     |\n\nAs reported in prior works \\[1\\], Mamba \\[-1,1\\] performs better compared to Mamba on the parity, cyclenav, solve equation, and bucketsort. However, AUSSM and the AUSSM Hybrid model still perform better in most tasks except bucketsort and solve equation. The tested complex version of Mamba performed poorly compared to the Grazzi version, which iterates the finding in the Mamba paper that real-eigenvalue parameterizations performed better. In AUSSM, the additionally introduced adaptivity and unitarity provide sufficient requirements for enabling learnability and expressivity.\n\nFor the weather benchmark, we added S5 to our list of baselines\n\n| model | MAE  |\n| ----- | ---- |\n| S5    | 0.36 |\n\nS5 is found to have better MAE than Mamba and LinOSS, but our proposed AUSSM Hybrid model performs the best in the benchmark.\n\n## Addressing Weaknesses and Questions\n\n### Weakness 1: The selection of prior works and baseline models in the experiments appears biased, especially in terms of real vs. complex eigenspectra.\n\nTo address this concern, we have added new results after running models on our suite of benchmarks. Now, our results for formal language benchmarks include Mamba (real), Mamba Complex (complex), and Mamba \\[-1,1\\] (\\[1\\]). For UEA benchmarks, we have Mamba (real), S5, S6, Linoss (Complex), and for the weather benchmark, we compare Mamba (real), S4, S5, LinOSS (complex)\n\n### Weakness 2: Since some experiments only report results for the hybrid model, it is difficult to isolate and evaluate the impact of the proposed AUSSM.\n\nAUSSM Hybrid was chosen for real-world benchmarks since our theoretical analyses (Lemma 2) and empirical formal language results (Table 1) point to the AUSSM+S6 (mamba) hybrid model consistently performing better than a pure AUSSM block model. Note also that without the AUSSM block, the Hybrid model is a pure Mamba model for which results are already included in the tables. \n\n### Question 1: What is the theoretical or practical benefit of an SSM having the conservation property?\n\nConserved dynamics enable running SSMs on longer sequences without blowing up or decaying hidden states, and further enable stable gradient propagation during training. Prior to the introduction of SSMs in the recurrent neural network space, conserved dynamics were preferred in RNNs to avoid the diminishing/exploding gradient problems when processing long sequences (e.g. Arjovsky et al. (2016)). With the introduction of SSM initialization schemes such as HiPPO (Gu et al, 2020), such conservation schemes became less relevant since long-term stability was maintained through other means. However, newer SSMs like the LinOSS discussed in the paper still have conserved dynamics, suggesting that they help improve even newer SSMs to capture long-term dependencies.\n\n### Question 2: Aren’t complex-valued eigenvalues used not only in LinOSS but also in S4D and S5? In Mamba, complex-valued variants were also tested in ablation studies, and real-valued variants were selected due to slightly better performance. Emphasizing real eigenspectra as a prevailing trend in prior SSMs seems somewhat misaligned with the actual research landscape.\n\nWe have included new comparisons to S5, in UEA and weather, and Mamba with S4D initialization in formal language benchmarks (see Changes #1 above). Complex eigenvalue spectra have been used in many frontier SSM models - we discussed two main ones (LRU and LinOSS) in our background. However, the practice of real-valued eigenvalues is still a prevalent trend in SSMs, as evidenced by landmark papers - Mamba \\[2\\] rightly pointed out that real eigenvalues were better than complex ones, and MEGA \\[3\\] also shows similar results. The Mamba repository, which is widely adopted currently, also uses only positive real-valued recurrence (see open github issue #106 in the Mamba repo that shows that complex eigenvalue training is not supported out-of-the-box by the open source implementation). \n\n### Question 3: In the paper, Mamba is referred to as mildly adaptive, since it fixes $A$ (actually, it eventually varies through the use of $\\Delta$), and adapts only $B$ and $C$. In some parts of the paper, however, Mamba is described as non-adaptive. Would it be more consistent to refer to it as _mildly adaptive_, and could a more detailed discussion on the role of $\\Delta$ be included?\n\nWe apologize for the oversight in line 41 where we described Mamba as non-adaptive. **Mamba is partially (or mildly) adaptive** as rightly pointed out. We have corrected this. Mamba is mildly adaptive - i.e, the $A$ in the ODE formulation (eqn1) is time invariant, but the $A'$ in the discrete formulation (eqn 2) can vary based on the discretization timestep ($\\Delta$). The role of $\\Delta$ is only to scale the dynamics (fast or slow) depending on input, but if the nature of dynamical behavior needs to be changed - e.g., rotate left instead of right, or convergent instead of divergent, then the matrix $A$ itself needs to change. AUSSM is thus fully adaptive in the sense that you can control the dynamical behavior to that level of granularity.\n\n### Question 4: Can the baseline models (e.g., S5, S6, LinOSS, xLSTM, Mamba) and the proposed models (AUSSM, AUSSM Hybrid) be unified across all experiments? Is there a particular reason why different subsets of models are selected for different tasks?\n\nThe baselines were chosen depending on the currently reported state-of-the-art results in the respective tasks that were tested. To balance the choice of baselines, we have added additional results for S5, Mamba\\[-1,1\\] MambaComplex in the formal language benchmarks and S5 to the weather benchmark (see Changes #1 above). There are technical limitations that prevent rerunning certain models like LinOSS that are written in JAX because our code base is Pytorch. There is a nontrivial effort required to do the porting, and the significant risk of introducing silent bugs that may skew the results. With the newly introduced results, we have covered SSMs that use different parts of the eigenvalue spectrum. \n\n### Question 5: The results on the UEA dataset differ significantly from those reported in the LinOSS paper. Can this be explained?\n\n**The results are the same as reported in the LinOSS paper**, the values we report are the scaled accuracies instead of the raw accuracies. Scaled accuracy scales the raw accuracy between baseline performance (scaled accuracy=0) and perfect performance (scaled accuracy=1). It is computed as (raw_accuracy - baseline)/(1-baseline) where baseline is the performance of a random estimator = (1/num_classes)\n\n### Question 6: AUSSM Hybrid achieves the best performance, not the pure AUSSM model (especially in Tables 2 and 3). This raises concerns about whether AUSSM alone sufficiently supports the paper’s main claim of resolving the expressivity-efficiency tradeoff. If AUSSM indeed offers superior expressivity, shouldn't it also achieve strong results beyond counting-focused tasks?\n\nThe claim is not that AUSSM has superior expressivity by itself, but rather that it adds a different kind of expressivity to the hybrid architecture. We have adjusted the diagram in Fig 1. to make this clear. AUSSM, however, takes a step toward general LTV systems which were not possible without sacrificing computational efficiency. In formal representation theory perspective, the complex unitary eigenvalues of AUSSM give it the ability to perform cyclic group operations, whereas Mamba components can perform the operations of aperiodic transformation monoids. Together, this gives the hybrid architecture the expressivity required to model solvable monoids, or equivalently, solvable languages. The Hybrid model achieves maximal expressivity in the class of computationally efficient diagonal SSMs, any more increase in expressivity require relaxing the diagonal constraint.\n\n## References\n\n\\[1\\] Grazzi et. al. Unlocking State-Tracking in Linear RNNs Through Negative Eigenvalues. ICLR (2025)\n\\[2\\] Gu and Dao. Mamba 2023: Linear-time sequence modeling with selective state spaces. COLM (2024)\n\\[3\\] Ma et. al. Mega: Moving Average Equipped Gated Attention. ICLR (2023)"}, {"review_id": "CJ7gFuMSQD", "reviewer": "Reviewer_MtYg", "summary": "The presented model extends the eigenvalues of the state-transition matrices of Mamba 1 from the positive reals to the unit circle. The authors demonstrate that this change allows the model to learn state-tracking tasks such as parity or cyclic navigation which Mamba 1 using only real eigenvalues provable cannot learn in a length extrapolatible way.", "questions": "-", "limitations": "Yes.", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 2, "significance": 2, "originality": 2, "strengths_and_weaknesses": "## Strengths\n\n- Extending the eigenvalues of the diagonal state-transition matrix to complex values makes sense to me, however I am unsure about keeping them purely unitary, see weaknesses.\n- The experiments demonstrate that this allows them to model state-tracking tasks which are unsolvable by Mamba 1 with purely real and positive eigenvalues.\n- The architecture shows promise for long-term time-series forecasting.\n\n\n## Weaknesses\n- Models such as RWKV-7 and DeltaNet use non-diagonal state-transition matrices which were shown to be significantly better at in-context retrieval and also state-tracking when using negative eigenvalues as demonstrated by Grazzi et al. 2025. Since these are also linear RNNs it would be important to compare to them. In particular DeltaNet with negative eigenvalues can solve modular arithmetic w/o brackets (Grazzi et al. 2025) in addition to parity. Grazzi et al. 2025 also provide code to run Mamba with negative eigenvalues which you could compare to in order to show the need for unitary eigenvalues.\n\n-  It is not clear to me why the eigenvalues need to be unitary and could not also lie in the interior of the unit circle in order to allow the model to maintain signal when using the boundaries of the circle or attenuate signal when using more of the interior of the circle in order to forget better. \n\n- As I understand the critical difference between Mamba and the proposed AUSSM model is that the eigenvalue range is extended from the positive reals in Mamba to the complex unit circle. Since Mamba 1 already used complex eigenvalues in the state-transition matrix for some of their experiments on audio sequences, please compare with their implementation to isolate your contribution.\n\n- I think the key changes to mamba could be emphasized more straightforwardly. The usefulness of coining the term \"adaptive\" is unclear to me, simply calling it a linear RNN with token-dependent diagonal state-transition matrix having unitary entries would be more clear or did I misunderstand?\n\n- Nonlinear Unitary RNNs have a rich research history, please at least give a small literature overview for example on those using rotation or reflection matrices. The following come to mind: Z. Mhammedi, A. Hellicar, A. Rahman, and J. Bailey. Efficient orthogonal parametrisation of recurrent neural networks using householder reflections. Victor D. Dorobantu, Per Andre Stromhaug, and Jess Renteria. DizzyRNN: Reparameterizing Recurrent Neural Networks for Norm-Preserving Backpropagation. L. Jing, Y. Shen, T. Dubcek, J. Peurifoy, S. Skirlo, Y. LeCun, M. Tegmark, and M. Soljacic. Tunable Efficient Unitary Neural Networks (EUNN) and their application to RNNs. Rumen Dangovski, Li Jing, Preslav Nakov, Miśo Tatalović, and Marin Soljačić. Rotational unit of memory: a novel representation unit for rnns with scalable applications. C. Jose, M. Cisse, and F. Fleuret. Kronecker recurrent units.", "was_revised": true, "final_justification": "The authors addressed my questions. However, I think the paper is somewhat misleading in that it makes the case for using parallel-scans while these have been shown to be much slower on current hardware than chunk-wise parallel forms.", "rebuttal": "We thank the reviewer for the constructive comments and critical evaluation of our manuscript. Thank you for highlighting AUSSM’s **promise for long term time-series forecasting**. We want to clarify that we make **two extensions** to Mamba - (1) allowing for unitary eigenvalues **and** (2) enabling the eigenvalues to adapt based on the instantaneous inputs. We want to emphasize that the S6 SSM ODE used in Mamba has only B and C adaptive, while A is adaptive only in the discrete form indirectly through a time step factor $\\Delta$. Thus, Mamba dynamics can speed up or slow down depending on input, but not fundamentally change. Our AUSSM formulation enables the full $A$ also to be adaptive, enabling input-dependent (adaptive) dynamical behaviors, while encountering minimal penalty (only a small constant factor) to computational efficiency.\n\nWe address the concerns of comparing with modified Mamba versions (negative and complex) by adding new formal language result comparisons. We address the concern of a literature review of non-linear unitary RNNs by adding a Review section in our appendix. We summarize the changes below and provide detailed clarifications to each question in the subsequent section.\n## Relevant Changes\n\n### 1. Experiments with Negative and Complex versions of Mamba\n\nThank you for the suggestion to use the general eigenvalue Mamba versions. We implemented new experiments with both Mamba\\[-1,1\\] from the Grazzi paper and Mamba Complex from the original Mamba paper. Please note that the original Mamba paper did not detail how B and C are defined for the complex version, nor provide code to replicate Mamba results with complex eigenvalue spectrum, as evidenced by this issue#106 in the Mamba github repo. Nevertheless, we implemented our own version by modifying A to have S4D-Lin initialization and increasing the dimensionality of B to account for the imaginary part. We report the results of running the models on the Formal Language Benchmark below\n\n|                   | Mamba with Complex Eigenvalues | Mamba in \\[-1,1\\] range |\n| ----------------- | ------------------------------ | ----------------------- |\n| repetition             | 0.09                           | 0.1                     |\n| bucket sort          | 0.21                           | 0.91                    |\n| majority count         | 0.19                           | 0.31                    |\n| majority          | 0.13                           | 0.64                    |\n| solve equation             | 0.43                           | 0.24                    |\n| modarithmetic          | 0.12                           | 0.116                   |\n| modarithmetic wo bracket | 0.23                           | 0.24                    |\n| cyclenav          | 0.42                           | 0.91                    |\n| parity            | 0.27                           | 1.0                     |\n\nAs reported in prior works, Mamba[-1,1] performs better compared to Mamba on the parity, cyclenav, solve equation, and bucketsort. However, AUSSM and the AUSSM Hybrid model still perform better in most tasks except bucketsort and solve equation. The tested complex version of Mamba performed poorly compared to the Grazzi version, which iterates the finding in the Mamba paper that real-eigenvalue parameterizations performed better. In AUSSM, the additionally introduced adaptivity and unitarity provide sufficient requirements for enabling learnability and expressivity.\n\n### 2. Literature Review for Non-linear Unitary RNNs\n\nWe thank the reviewer for this suggestion. As our contribution was mainly in the development of SSMs, our current version of the paper only contained relevant SSM implementations. We add a new section in the appendix where we provide a detailed review of Non-linear Unitary RNNs. We show a brief snippet here due to character limits in the rebuttal responses.\n\n\"Unitary and orthogonal RNNs were introduced to mitigate vanishing and exploding gradients by preserving the norm of hidden states over time. Early work by Arjovsky et al. (2016) proposed Unitary RNNs with parameterizations based on structured unitary matrices, followed by Full-Capacity Unitary RNNs (Wisdom et al., 2016) which optimized directly over the class of all unitary matrices (Jing et al., 2017). Orthogonal RNNs have also been explored using soft constraints (Vorontsov et al., 2017), matrix exponentials of skew-symmetric matrices (Lezcano-Casado & Martínez-Rubio, 2019), and other structured parametrizations such as Householder reflections (Mhammedi et al., 2017)....\"\n\nArjovsky et al. (2016). Unitary Evolution RNN. ICML.\n\nWisdom et al. (2016). Full-Capacity Unitary RNNs. NeurIPS.\n\nVorontsov et al. (2017). On orthogonality and learning recurrent networks with long term dependencies. ICML.\n\nLezcano-Casado & Martínez-Rubio (2019). Cheap Orthogonal Constraints in Neural Networks: A Simple Parametrization of the Orthogonal and Unitary Group. ICML.\n\nMhammedi et al. (2017). Efficient Orthogonal Parametrization Using Householder Reflections. ICML.\n\nJing et al. (2017). Tunable Efficient Unitary Neural Networks (EUNN) and their application to RNNs . JMLR.\n\n\n\n\n## Addressing Weaknesses and Questions\n\n### Models such as RWKV-7 and DeltaNet use non-diagonal state-transition matrices...\n\nWe thank the reviewer for the constructive suggestions to improve the manuscript. To address the suggestions, we have thus added new comparisons to the Grazzi Mamba implementation and Mamba Complex implementation (see #1 in the Changes section). RWKV-7 \\[1\\] and DeltaNet \\[2,3\\] both use non-diagonal transition matrices, trading added computation cost for expressivity.  In DeltaNet, a single Householder reflection per step is insufficient for full regular language expressivity or even to perform arbitrary rotations. DeltaProduct \\[4\\] addresses this by using products of multiple reflections, which can realize arbitrary orthogonal transformations. By contrast, AUSSM can perform rotations cheaply with a single diagonal update. The added computational cost of using non-diagonal matrices is a recurring bottleneck typical of non-linear RNNs. Our focus was to make diagonal SSMs, which have efficient parallel implementations, as expressive as possible while retaining their computational efficiency advantage. From a formal expressivity standpoint, negative eigenvalues in diagonal SSMs are not enough to perform even simple tasks such as mod 3 addition (Grazzi, Thm. 2). Unitary eigenvalues, however, lift the expressivity to all solvable languages, the upper bound possible for diagonal SSMs (assuming TC0 != NC1). Further increase in formal expressivity is only possible using non-diagonal and, hence, a computationally inefficient approach, as we discuss in our Limitations section.\n\n### It is not clear to me why the eigenvalues need to be unitary and could not also lie in the interior of the unit... \n\nUsing unitary eigenvalues in the AUSSM part reduces the number of parameters because only the polar angle is parameterized, improving training efficiency, while adding enough expressivity to allow recognition of all solvable languages. The Mamba part of the hybrid architecture takes care of the memory/forgetting part since it does not have fixed eigenvalues. This is also borne out by our experiments, which show that the hybrid model performs much better than Mamba with complex eigenvalues (which has the same expressive power). \n\n### As I understand the critical difference between Mamba and the proposed AUSSM...\n\nWe make **two changes** to the S6 ODE used in Mamba (real version) - (1) extend the eigenvalue range of A to unitary and (2) enable adaptivity in the eigenvalue spectrum. We thank the reviewer for the suggestion to compare with Mamba1 with complex eigenvalues. We have now added the results from running Mamba1 with complex eigenvalues in the formal language benchmarks (see #1 in the Changes section). Grazzi's version of Mamba performs better compared to the complex Mamba, despite the complex one having better expressivity. This reinforces the result from the Mamba paper that real-valued mamba versions performed better than the complex-valued version. Existing works like Mamba with complex eigenvalues do not have full adaptivity - the recurrence is mildly adaptive in the sense that the eigenvalues can scale up/down depending on inputs, but cannot change sign or become complex. e.g., if we want to rotate left vs right depending on input. \n\n### The usefulness of coining the term \"adaptive\" is unclear to me,...\n\nAUSSM is adaptive in the sense that the recurrence structure changes (adapts) to every input presented to it. This goes beyond the selectivity introduced in the S6 SSM used in Mamba. In AUSSM,  the $\\Delta$ is not the only driver of time-varying behavior but the recurrent matrix $A$ itself. This way, AUSSM forms a middle ground between the partial LTV structure of Mamba and a full LTV dynamical system and enables the dynamics to adapt (not just scale as in partial LTV Mamba) to changes in the input.\n\n### Nonlinear Unitary RNNs have a rich research history,...\n\nWe thank the reviewer for the suggestion. We have added a new section in the appendix that aims to cover some of the main themes in the design of non-linear unitary RNNs (See Changes #2 above)\n\n## References\n\n\\[1\\] Peng et al., 2025. RWKV-7 \"Goose\" with Expressive Dynamic State Evolution \n\\[2\\] Schlag et al., 2021. Linear transformers are secretly fast weight programmers. ICML\n\\[3\\] Grazzi et al., 2025. Unlocking State-Tracking in Linear RNNs Through Negative Eigenvalues. ICLR\n\\[4\\] Siems et al., 2025. DeltaProduct: Improving State-Tracking in Linear RNNs via Householder Products"}]

Author rebuttal:

Meta-review: The paper proposes a modification to SSM they call AUSSM which uses linear time varying recurrence and a unitary eigenvalue spectrum derived from a skew-symmetric ODE. The authors show that this modification allows hybrid models Mamba + AUSSM to represent all solvable automata, here the negative eigenvalue helps implement the counter automata. The authors additionally provide a separable convolution formulation and a CUDA implementation to efficiently implement AUSSMs. Experiments support the benefits of the hybrid models compared to other variants across simple arithmetic tasks and long time-series modeling.

The paper was generally well-received by the reviewers, and the authors cleared most of the concerns of the reviewers with additional experimentation and exposition over the rebuttal. While the impact of the proposed modification is unclear for more general tasks and at larger scale, and whether the computational cost would be justified for expressivity benefits, this paper still adds a novel result in the area of improving SSMs that would be of interest to the NeurIPS audience. Therefore I lean towards accepting the paper. I encourage the authors to include the additional experiments and discussion that happened during the rebuttal period, and improve the exposition of the paper.
Decision: Accept (poster)