Title: Bridging Expressivity and Scalability with Adaptive Unitary SSMs
Venue: NeurIPS 2025
Accepted: yes
Keywords: State Space Models

Unitary Recurrent Neural Networks

Adaptive Recurrent Neural Network
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


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS AND WEAKNESSES:\n## 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.\n\nQUESTIONS:\n-\n\nLIMITATIONS:\nYes.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 2, "IMPACT": 2, "ORIGINALITY": 2, "DATE": "2025-06-21"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS AND WEAKNESSES:\n- 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.\n\nQUESTIONS:\n1. 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?\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "CLARITY": 2, "IMPACT": 2, "ORIGINALITY": 3, "DATE": "2025-07-02"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths \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\n\nQUESTIONS:\n1. **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\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 4, "DATE": "2025-07-16"}