Theoretical Foundations of Deep Selective State-Space Models

We thank the reviewer for their thoughtful feedback on our paper. We appreciate the positive assessment of our work's soundness, presentation, and contribution. Below, we address the raised points.

## Weaknesses

• Rates of convergence: While our current theoretical framework in general provides only density-type results, we agree that the characterization of convergence rates is a crucial question. We envisioned this paper as one establishing the language and theoretical groundwork for further analysis in such a direction.
  We want to point out that for Theorem 4.2, the bounds are explicit in $N$, as given in equation (55). Nevertheless, a study of the rates in the general case, particularly for the stacked diagonal systems, would be more technically involved and is considered beyond the scope of the current work. However, we acknowledge its importance and will highlight it as a critical direction for future research.

• Experiments: We would like to clarify a potential misunderstanding: our approach does not utilize the path signature as part of the method itself. Instead, we are simply using specific terms in the path signature as the labels for the synthetic datasets. We opted for low-dimensional, synthetic datasets to provide empirical evidence for the theoretical results in our paper in a simple setting. These experiments were not intended to comment on the effectiveness of the path signature when handling high-dimensional data. There are a number of interesting works which apply variants of the signature to real-world data, with some recent examples being the signature kernel (https://arxiv.org/abs/2006.14794, https://arxiv.org/abs/2006.05805), the path development layer (https://arxiv.org/abs/2204.00740), the randomized signature (https://arxiv.org/abs/2201.00384), and log neural controlled differential equations (https://arxiv.org/abs/2402.18512).

Questions

• See above

• Gating: Our definition of gating is aligned to that of forget and input gate in the LSTM, GRU and SSM literature. In particular, what the reviewer refers to is known as output gate: what instead $\omega$ and $\xi$ modulate is how inputs are selected and forgotten. The multiplicative interaction crucial in this setting is $Z_t^X d\omega_t^X$ - this can be thought of as a forget gate on the hidden state, with input-dependent gating function induced by $\omega$. We note that the gate terminology was introduced in the LSTM paper, but is also used in recent SSMs paper such as in Mamba and Griffin.

• NCDE: NCDE stands for Neural CDE. In fact a linear NCDE is just a linear CDE, hence there was no need to mention the “Neural” part – we used this terminology since it is linked to some results on rough path theory applied to neural networks (e.g. https://arxiv.org/abs/2005.08926) but here this connection is not necessary. We thank the reviewer for pointing this out. We will revise the manuscript to use the term “CDE” instead of “NCDE” in the relevant sections to avoid any confusion.

Limitations

We agree with the reviewer on this point. We believe that our casting of SSMs in the language of Rough Paths theory might open the way for such a study of discretizations, using the honed tools of the field.

We thank the reviewer for their thoughtful feedback on our paper. We appreciate the positive remarks regarding the timeliness and importance of the research problem we address. We would like to address the points raised concerning weaknesses and questions.

Weaknesses

• Preliminaries on Rough Path Theory: We understand the importance of making our paper accessible to readers unfamiliar with Rough Path Theory. In Appendix A (referenced at line 268), we provide an introduction to the Signature Transform and reference several in-depth studies and presentations of Rough Path Theory in the context of machine learning. Additionally, in Appendix E, we provide a self-contained theory studying the solutions to general Linear CDEs. To enhance clarity and ease of access, we will ensure to reference both Appendix A and Appendix E more prominently in the main body of the paper.

• Implications in Practice: Our results fit into an ever-growing theoretical literature on the implicit limitations deriving from architectural choices in the context of SSMs (see e.g. https://arxiv.org/pdf/2404.08819, https://arxiv.org/abs/2406.05045), which we enhance by providing explicit, analytical characterizations and by establishing a theoretical foundation, in terms of Rough Path theory, which we envision as a base for further study. The take-away from this literature is clear: devising a non-diagonal, input-dependent, and efficiently computable transition mechanism would allow you to overcome the expressive limitations of present SSMs without the need for an arbitrarily high number of layers. A promising avenue is using low-rank or even highly sparse weights, as suggested by recent works (e.g. https://arxiv.org/abs/2310.16597, https://arxiv.org/abs/2407.08459, https://arxiv.org/abs/2406.05045).

• Further experiments: As a further example of the practical implications of our theoretical results, we have performed additional experiments focused on the A5 benchmark introduced by https://arxiv.org/pdf/2404.08819. This benchmark is designed to evaluate models on their state-tracking, a crucial ability for solving problems that involve permutation composition, such as tracking chess moves. A key result of their paper is that state-space models such as Mamba require stacking in order to perform well on this benchmark. Our experiments have shown that even on the longest sequence length in the benchmark, where Mamba requires 4 stacked layers to achieve >90% test accuracy, a linear CDE with a trainable transition matrix requires only one layer to achieve >90% test accuracy. We intend to include a full discussion of these results in the final version of our paper.

Questions

• Derivation of Eq.3 and the Appearance of the Exponential Term: We appreciate the reviewer's attention to the details of our derivation. A self-contained explanation of Zero-Order Hold (ZOH) discretization is provided in Appendix F (referenced just before 164). To increase clarity, we will highlight the reference to this section in the main body of the paper. We also note that ZOH discretization is a conventional nomenclature for such a scheme in the SSM literature (see e.g. https://arxiv.org/abs/2303.06349).

We hope these clarifications address the reviewer's concerns and enhance the overall readability and impact of our paper.

We thank the reviewer for their thoughtful feedback on our paper. We appreciate the positive remarks on the theory's presentation and contribution. We would like to address the points raised regarding “closing the loop”.

Our paper shows how non-diagonal transitions lead to a substantial increase in expressivity. However, this theoretical insight is currently impeded by the reality of computation with dense layers, which is practically infeasible due to the associated computational costs. At the same time, recent literature (e.g. https://arxiv.org/pdf/2402.19427, section 4.2) has shed light on how the computation of linear diagonal RNNs is dominated by memory transfers. This indicates that there is leeway allowing for an increase in the complexity of the sequential mechanism, presenting itself as a very promising area of research.

Such a mechanism would have to be non-diagonal but efficiently computable. Recent works (e.g. https://arxiv.org/abs/2310.16597, https://arxiv.org/abs/2407.08459, https://arxiv.org/abs/2406.05045) suggest that low-rank or even highly sparse weights should lead to the same limiting behaviors found in the dense case. This is a promising avenue, which we have not studied in this work, as it would require substantial theoretical justification as well as thorough empirical evaluations at scale. As the purpose of this paper is to outline the hypothesis class and power of SSM variants, we decided to leave this avenue for future research.

However, given the importance and interest of such observations, we will augment the concluding section with them in the camera ready version, using the additional available space.

We thank the reviewer for their constructive feedback on our theoretical work. We are pleased you found our expressivity results important and our proofs accessible. We spent a considerable amount of time making our manuscript easy to parse despite the high technicality of the content.

Weaknesses

The reviewer is right in saying that our work does not explore the finite-regime setting and does not compare architectural options in terms of model parameters needed to achieve a desired level of expressivity. This would be highly desirable but is not an easy task; in addition results can be confounded by issues such as optimization, inductive biases (OOD generalization), GPU bottlenecks (e.g. memory dependency on sequence length) etc. Indeed, current research did not yet unveil such clean comparisons even for Transformers vs SSMs : the attention mechanism has more parameters compared to the S6 block in Mamba, but how exactly expressivity compares at a fixed compute budget or model dimension is unclear.

While touching on the issues above is for sure needed when proposing a new model, we like to point out that the purpose of our work is not suggesting a new mixing component (i.e. dense RNNs linear in the state), but instead exploring the source of limitations of modern sequence mixing components at generic widths. To do so, we study dense linear (in the state) recurrences, and then dive into the diagonal setting. As such, our dense linear CDE framework is not intended for direct implementation in a deep model, but provides an upper bound on the expressivity which is achievable with one RNN (SSM) linear in the state but nonlinear in the input.

Note that

• What happens at finite width. Our results prove that one linear CDE layer can approximate any nonlinear transformation, while diagonal CDEs (e.g. Mamba) cannot. This is however only a subset of our results. Leaving out chaining from the discussion here, our theorems identify the functional form for dense and diagonal settings in closed form – this holds at any width. Specifically, why the output in the dense setting can be written as $\int_0^t \Phi(\omega^{X}_{[s,t]}) \cdot d\xi^{X}_s$ (eq 7), for width-dependent function $\Phi$ determined by the model parameters. Instead, in the diagonal setting, the functional form for the output is restricted to $\int_0^t \phi(\omega^{X}_t - \omega^{X}_s) \cdot d\xi^{X}_s$. Note that width can only modify the complexity of the nonlinear functions $\Phi$ and $\phi$, but cannot modify the output structure.

• Experiments. While we believe very much that experiments can provide valuable insights in some setting, our objective here is to provide a theoretical foundation, hence the title. We remark that our results are not bound - they provide a tight description of the fundamental elements interacting when outputs are constructed, with model dimension only controlling functions $\Phi$ and $\phi$ above. We also note that our results are novel - also in the realm of Rough Paths theory as an independent mathematical subject. We refer the reviewer to the supplementary PDF for augmented and new results, as discussed in the general comment.

Questions:

• NCDE: NCDE stands for Neural CDE. In fact a linear NCDE is just a linear CDE, hence there was no need to mention the “Neural” part – we used this terminology since it is linked to some results on rough path theory applied to neural networks (e.g. https://arxiv.org/abs/2005.08926) but here this connection is not necessary. We thank the reviewer for pointing this out. We will revise the manuscript to use the term “CDE” instead of “NCDE” in the relevant sections to avoid any confusion.

• S5. As you stated, we expect that a dense CDE with a trainable transition matrix would be the best performing model on this benchmark. However, the computational burden of training such a model is very high. In practice, S5 is parameterized using a diagonal transition matrix, as discussed in Section 3.2 of the S5 paper (https://arxiv.org/abs/2208.04933). Consequently, the same theoretical results on expressivity apply to both S4 and S5. In fact, any model where $d\omega$ is $1$-dimensional (i.e. any model with only one transition matrix $A_1$) is equivalent over $\mathbb{C}$ to a diagonal model, by virtue of $A_1$ being similar to a diagonal matrix. To obtain the full expressivity results more than one $A_i$ must be present, and these matrices must not be simultaneously diagonalizable. As demonstrated in the proof of Theorem 4.3, in the diagonal case, the relevant terms of the signature are those appearing in its symmetric part. The signature of a $1$-dimensional path being fully symmetric shows consistency of our findings. On reflection, given this equivalence, we believe introducing S5 in our experiments is redundant. To streamline our presentation and strengthen the connection between our theoretical and experimental results, we have chosen to replace S5 with S4 in our experiments. We hope this revision addresses your concerns.

## Limitations:

We will augment the list of limitations in the camera ready version, using the additional available space.