Bridging Expressivity and Scalability with Adaptive Unitary SSMs

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.

We 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.

Relevant Changes

1. Experiments with Negative and Complex versions of Mamba

Thank 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

| | Mamba with Complex Eigenvalues | Mamba in $-1,1$ range | | ----------------- | ------------------------------ | ----------------------- | | repetition | 0.09 | 0.1 | | bucket sort | 0.21 | 0.91 | | majority count | 0.19 | 0.31 | | majority | 0.13 | 0.64 | | solve equation | 0.43 | 0.24 | | modarithmetic | 0.12 | 0.116 | | modarithmetic wo bracket | 0.23 | 0.24 | | cyclenav | 0.42 | 0.91 | | parity | 0.27 | 1.0 |

As 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.

2. Literature Review for Non-linear Unitary RNNs

We 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.

"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)...."

Arjovsky et al. (2016). Unitary Evolution RNN. ICML.

Wisdom et al. (2016). Full-Capacity Unitary RNNs. NeurIPS.

Vorontsov et al. (2017). On orthogonality and learning recurrent networks with long term dependencies. ICML.

Lezcano-Casado & Martínez-Rubio (2019). Cheap Orthogonal Constraints in Neural Networks: A Simple Parametrization of the Orthogonal and Unitary Group. ICML.

Mhammedi et al. (2017). Efficient Orthogonal Parametrization Using Householder Reflections. ICML.

Jing et al. (2017). Tunable Efficient Unitary Neural Networks (EUNN) and their application to RNNs . JMLR.

Addressing Weaknesses and Questions

Models such as RWKV-7 and DeltaNet use non-diagonal state-transition matrices...

We 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.

It is not clear to me why the eigenvalues need to be unitary and could not also lie in the interior of the unit...

Using 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).

As I understand the critical difference between Mamba and the proposed AUSSM...

We 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.

The usefulness of coining the term "adaptive" is unclear to me,...

AUSSM 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.

Nonlinear Unitary RNNs have a rich research history,...

We 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)

References

$1$ Peng et al., 2025. RWKV-7 "Goose" with Expressive Dynamic State Evolution

$2$ Schlag et al., 2021. Linear transformers are secretly fast weight programmers. ICML

$3$ Grazzi et al., 2025. Unlocking State-Tracking in Linear RNNs Through Negative Eigenvalues. ICLR

$4$ Siems et al., 2025. DeltaProduct: Improving State-Tracking in Linear RNNs via Householder Products

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.

Relevant Changes

1. Experiments with general eigenvalue baselines

We 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$ )

| | Mamba with Complex Eigenvalues | Mamba in $-1,1$ range | | ----------------- | ------------------------------ | ----------------------- | | repetition | 0.09 | 0.1 | | bucket sort | 0.21 | 0.91 | | majority count | 0.19 | 0.31 | | majority | 0.13 | 0.64 | | solve equation | 0.43 | 0.24 | | modarithmetic | 0.12 | 0.116 | | modarithmetic wo bracket | 0.23 | 0.24 | | cyclenav | 0.42 | 0.91 | | parity | 0.27 | 1.0 |

As 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.

For the weather benchmark, we added S5 to our list of baselines

S5 is found to have better MAE than Mamba and LinOSS, but our proposed AUSSM Hybrid model performs the best in the benchmark.

Addressing Weaknesses and Questions

Weakness 1: The selection of prior works and baseline models in the experiments appears biased, especially in terms of real vs. complex eigenspectra.

To 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)

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.

AUSSM 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.

Question 1: What is the theoretical or practical benefit of an SSM having the conservation property?

Conserved 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.

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.

We 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).

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?

We 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.

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?

The 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.

Question 5: The results on the UEA dataset differ significantly from those reported in the LinOSS paper. Can this be explained?

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)

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?

The 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.

References

$1$ Grazzi et. al. Unlocking State-Tracking in Linear RNNs Through Negative Eigenvalues. ICLR (2025)

$2$ Gu and Dao. Mamba 2023: Linear-time sequence modeling with selective state spaces. COLM (2024)

$3$ Ma et. al. Mega: Moving Average Equipped Gated Attention. ICLR (2023)

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.

Relevant Changes

1. Experiments with general-eigenvalue baselines

We 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$ )

| | Mamba with Complex Eigenvalues | Mamba in $-1,1$ range | | ----------------- | ------------------------------ | ----------------------- | | repetition | 0.09 | 0.1 | | bucket sort | 0.21 | 0.91 | | majority count | 0.19 | 0.31 | | majority | 0.13 | 0.64 | | solve equation | 0.43 | 0.24 | | modarithmetic | 0.12 | 0.116 | | modarithmetic wo bracket | 0.23 | 0.24 | | cyclenav | 0.42 | 0.91 | | parity | 0.27 | 1.0 |

As 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.

For the weather benchmark, we added S5 to our list of baselines

S5 is found to have better MAE than Mamba and LinOSS, but it our proposed AUSSM Hybrid model performs the best in the benchmark.

Addressing Weaknesses and Questions

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.

Thank 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.
RWKV-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.

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.

We thank the reviewer for pointing these out. These are now corrected in the manuscript.

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.

We apologize for this oversight. The missing citations are:

$34$ Babcsányi. Automata with finite congruence lattices. Acta Cybernetica, 2007.

$35$ Salkoff et al. Propagation of traveling waves in the visual cortex requires microcircuit coordination. Neuron, 2020.

$36$ Brent et al. Error bounds on complex floating-point multiplication. Mathematics of Computation, 2007.

$37$ Zimmermann. On krohn-rhodes theory for semiautomata.

$38$ Hewitt et al. RNNs can generate bounded hierarchical languages with optimal memory. EMNLP, 2020

$39$ Cormen et al. Introduction to Algorithms, Third Edition. The MIT Press, 2009.

Question 1: Hybrid performance dips

We 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.

Question 2: Neuroscience connection

The 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.

Question 3: Contextualizing results compared to recent SSM advancements

$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.

References

$1$ Grazzi et al., 2025. Unlocking State-Tracking in Linear RNNs Through Negative Eigenvalues. ICLR (2025)

$2$ Peng et al., 2025. RWKV-7 "Goose" with Expressive Dynamic State Evolution

$3$ Merrill et al., 2022. Saturated Transformers are Constant-Depth Threshold Circuits