The Expressive Capacity of State Space Models: A Formal Language Perspective

Given the formal nature of the paper’s title, it may be beneficial to reference the related paper available at https://arxiv.org/abs/1606.06737, which discusses formal language and the hidden Markov model. As the state-space model is a specific type of hidden Markov model, a brief discussion on the relationship between star-free state tracking, probabilistic regular grammar, and context-free grammar would be a valuable addition to the paper.

We appreciate the reviewer's detailed feedback and constructive criticism. We would like to clarify/address the issues raised as follows:

Weakness

(1.1), (1.2) We thank the reviewer for the interesting construction for PARITY, and the universal approximation references.

We would like to clarify why the suggested construction isn't a contradiction to Theorem 2. One could implement the suggested construction in an SSM-like architecture. However, computing the function $f(x) = (e^{i \pi x}+1)/2$ requires a layer-wise nonlinear function (eg an MLP) implementing sine/cosine functions at arbitrarily large inputs. Importantly, to obtain a construction working for all input lengths (as our positive results do), a single such operation would need to work at arbitrarily large input numbers.

As stated in line 56, our analysis assumes the layer-wise operations to be linear or (Swi)GLU (line 56) (as these are used in the real-world SSMs we surveyed in Appendix A). This is relevant to the proof of Theorem 2. A single GLU or SwiGLU is likely not able to represent sine/cosine for unboundedly large inputs, and we believe the same holds for more classical MLPs (ReLU/sigmoid MLPs). This is because, to the extent we are aware, universal approximation results for feedforward networks guarantee approximation in the sense of compact convergence [1], uniformly on the compactification of $\mathbb{R}$, [2], or in $L^p$ [3], none of which cover uniform approximation of sine/cosine on $\mathbb{R}$. Indeed, very recently, [4] studied universal approximation in $C_b(\mathbb{R})$, exactly what would be needed for approximating sine/cosine at any input magnitude. We believe their Proposition 5.5 entails that sine/cosine cannot be uniformly approximated with certain activation functions.

Thus it's not to be expected that $f(x) = (e^{i \pi x}+1)/2$ would be implemented by a typical MLP uniformly for arbitrarily large inputs. Hence, we believe any implementation of the proposed construction would either require custom (e.g., periodic) activation functions or require model size to vary with the input length, both resolving an apparent contradiction with Theorem 2.

Thanks further for linking the references on universal approximation, which we will cite and discuss in the Related Work section.

We would like to clarify again, there's no contradiction of our negative result in Theorem 2 to these universal approximation results. Both Wang et al and Orvieto et al provide universal approximation guarantees, but importantly, the size of the approximating networks depends on input length. One can see this from the statements of the core formal results, Propositions 3.6 and 3.9, in Wang et al. It is clearly stated by Orvieto et al in their Remark 2. In contrast, our results (both positive and negative) concern the existence of a single SSM recognizing a formal language at any input length.

Overall, we thank the reviewer again for raising these interesting points. We will address them as follows:

(1) Cite and discuss Wang et al and Orvieto et al results.

(2) Make more salient, & mention in the Theorem 2 statement that it relies on properties of the activation functions as stated in line 56.

(3) Discuss the question whether SSMs could benefit from more general, e.g., periodic activation functions in the nonlinearities.

We thank the reviewer for this insightful observation and for helping improve our paper with such constructive criticism.

(1.3) For languages where Mamba failed to predictively model perfectly with 3 layers, we did increase number of layers to 12 and hidden dimension to 256, as documented in Appendix Section D.1. Across the languages tested, we obtained no further benefits beyond 3 layers.

(2) In the "Role of Parameterization" section, we aimed to delineate the scope of our work and point out that our results aren't primarily affected by technical details regarding parameterization, as we allowed each of A, B, and $\rho$ to be arbitrary functions. However, we acknowledge the importance of citing relevant papers that focus on parameterization and will make our point about our results being unaffected by parameterization details more explicit.

(3) We thank the reviewer for pointing out the typo. We'll fix it.

Questions

(1) We agree that the question of precision is interesting. Increasing $p$ to any finite bound might improve expressiveness on short inputs, but not mitigate limitations on unboundedly long inputs. An open question is to what extent increasing $p$ with input length ($N$), e.g., $\log N$ could mitigate expressiveness limitations in theory. We will add discussion of the same.

(2) Indeed, both transformers and convolutional networks face similar challenges as SSMs regarding PARITY. For transformers, [5] shows that such limitations provably persist even with infinite precision (PARITY creates very sharp minima on long inputs). We conjecture that this reflects a general phenomenon in highly parallellizable architectures, and that similar phenomena may apply with SSMs even with infinite precision.

Limitations

We acknowledge that our results are derived for formal languages. However, Dyck, Flip-Flop, $a^nb^n$ etc. have linguistic motivations and provide insights into the structure/processing of natural language, which we'll expand on in the camera ready version.

References

[1] Cybenko, G. (1989). Approximation by superpositions of a sigmoidal function. Mathematics of control, signals and systems

[2] Ito, Y. (1992). Approximation of continuous functions on Rd by linear combinations of shifted rotations of a sigmoid function with and without scaling. Neural Networks

[3] Arora, et al (2018) "Understanding Deep Neural Networks with Rectified Linear Units." ICLR

[4] van Nuland, T. D. (2024). Noncompact uniform universal approximation. Neural Networks

[5] Hahn and Rofin (2024). Why are Sensitive Functions Hard for Transformers? ACL

We thank the reviewer for their positive score and encouraging review. Regarding the issues pointed out:

Response to Weaknesses:

This first reason is minor compared to the others. There are many results in this paper, which is desirable. However, it is a bit difficult to get an overview of the results. For example, it is unclear which should be considered the main result. Of course, it could be the case that there is not a single main result, which I assume applies here. However, the structure of the results seems odd, as we first show what is (not) possible for some languages (Theorem 1, Theorem 2), then consider a general result (Theorem 4), and then, again, focus on specific languages (Theorem 5 and Theorem 6). I think the main problem is that all results are squashed into a single “Theoretical Results” section, which is also mixed with basic definitions like finite-state automata and (star-free) regular languages. Giving the reader a better hint on the central theme of the results and a clear separation of preliminaries and results would improve the clarity of the paper.

• The paper indeed does not have a single main result; rather, we have four major results that we aimed to make explicit and contextualize in the Take-Aways section on Page 9. We acknowledge that combining some of the background information with our main results in the Theoretical Results section was not ideal. We will separate the background information from the results into distinct sections to improve clarity.

The authors tend to state lemmas and theorems informally, making it hard to grasp the concrete setting. For example, Theorem 2 states, ”… can recognize PARITY at arbitrary input lengths …”. It is not precisely clear what this means, especially the “arbitrary length” part. Or take Lemma 13 in Appendix B.3. This lemma states that ”… the character can be read out from at finite precision.”. It is hard to understand what is exactly meant by this, too. The same goes for the overall statement of Lemma 22. While I appreciate that the main part of the paper focuses on a low-technical presentation, such a theoretical work cannot be separated from its technical appendix. Therefore, statements in the appendix should be precise. This goes beyond these three examples and is a general problem. This also applies to the exact SSM architecture considered in respective results. In most cases, one can guess what is meant, but I don’t think such things should be left to the reader’s imagination to this degree.

Some proofs could be more detailed. For example, the proof of Lemma 17 includes previously unused notation (line 824, ), it is not clear how equalities (15) and (16) are derived or what ”… running these in parallel, …” means.

• We agree that understanding the concrete settings of our results might require the reader to refer to the appendix. We appreciate that the reviewer acknowledges our focus on intuition in the main paper, leaving technical details in the appendix. We also thank the reviewer for pointing out that some of our proofs, although correct, might require more detail. For each theorem or lemma, we will provide a fully formally precise statement in the Appendix. We will revisit all proofs and provide additional steps to better elucidate them.

Response to Questions:

The notion of finite precision you are using seems to have been considered before. Could you elaborate on why it is a reasonable one? At first sight, it seems artificial, as you either consider practical settings where you are limited by some amount of bits for representing a number, or you do not care for practical realizations and thus everything can be unbounded. The notion you used feels a bit like a convenience choice, as fixing fractional parts helps keeping fuzziness under control but still allows for unbounded counting using unbounded integer parts. Especially, in the context of your positive, predictive results regarding finite precision. I assume these do not hold for an overall bounded number of bits.

We understand that our choice of finite precision for representing fractional components versus allowing unbounded counting, and thus unbounded integer values, might seem like a convenience choice. However, we would like to clarify that we do not consider unbounded integer values in any of our theorems except for the unbounded counting case: All our results except for the unbounded counting case in fact have overall bounded numbers of bits. For unbounded counting, the task definition requires unbounded integer values, and any length-generalizing recurrent solution will require an unbounded overall number of bits. While, theoretically, a solution for unbounded counting can be constructed with SSMs, as shown in our paper, it might be unlikely that such a solution would be practical owing to such requirements of unbounded overall number of bits. Indeed, in our experiments across all counter languages in Figure 3, we were unable to get the SSM to learn a length-generalizable solution. We will expand our discussion of the precision assumptions and their implications in Appendix E, and add appropriate notes in the main paper.

Thank you for your clarifications. I appreciate your willingness to improve the clarity of some of your proofs.

However, I have a follow-up question as I am not fully satisfied with your comment on your definition of "finite precision":

If I understand your comment correctly, you imply that results like Theorem 4, which are not the "unbounded counting case" are also valid if we assume "full" finite precision, meaning an overall bounded number of bits for representing numerical values?

We thank the reviewer for their high score and positive comments about our paper. We appreciate the reviewer's thorough examination of our proofs and their constructive insights and suggestions. We address the points raised as follows:

Response to Weaknesses:

While the paper does an admirable job of explaining complex theoretical concepts, some sections may still be challenging for readers without a strong background in theorem proving. Additional intuitive explanations or visual aids could help make the content more accessible.

• We acknowledge that some sections of the paper may be challenging for readers without a strong background in theorem proving. To make the content more accessible, we will include additional intuitive explanations and visual aids. Specifically, we will add visual representations to help elucidate the SSM constructions described in our theorems.

In the related work section add more works that study this, in particular: Limits of Deep Learning: Sequence Modeling through the Lens of Complexity Theory. Nikola Zubić, Federico Soldá, Aurelio Sulser, Davide Scaramuzza - https://arxiv.org/abs/2405.16674 . Mention also how does your work overlap with it and is there any interesting connection?

• The paper by Zubić et al. indeed represents significant work in the field of analyzing expressivity of SSMs. Their approach differs from ours in that they focus on a range of algorithmic tasks, such as function composition, whereas we focus on fine-grained study of length-generalizable expressions for core formal languages. We do not see any overlap in results, but will make sure to include the comparison in the related work section of our main paper.

There are typos, use for example Grammarly to fix them. Also there are some typos in math, for example Lemma 22 should have w_j inside the sum, not w_i etc. Check and make sure indexes are well defined. Theorems are proven well, but there are typos.

We apologize for the typographical errors and mistakes in the mathematical notation, such as the incorrect index in Lemma 22. We will thoroughly proofread the paper again and carefully correct all the typos we missed.

Response to Questions:

While the experiments on synthetic datasets and formal languages are comprehensive, additional experiments on more complex and diverse real-world tasks could strengthen the empirical validation. Are there any plans to extend the experimental validation to other types of datasets or applications that are not necessarily big, but challenging and prove the points?

• We agree that additional experiments on more complex and diverse real-world tasks would further strengthen our empirical validation, and we do plan to do the same in follow up work.

The paper briefly mentions the practical implications of the theoretical insights for designing new ML architectures. Could the authors provide concrete examples or guidelines on how to implement these insights in real-world systems? What are the key considerations for practitioners looking to apply these findings?

• One of the main practical implications that can be taken from the theoretical insights of our paper in designing new ML architectures is for future architectures to combine the strengths of attention with that of the SSM recurrence update, and indeed some work [1][2][3] that has been published since our submission has found empirical evidence of the same. We do agree that more concrete examples and guidelines could be given and would include some examples in the appendix of our paper due to space considerations.
  • [1] Lieber, Opher, et al. "Jamba: A hybrid transformer-mamba language model." arXiv preprint arXiv:2403.19887 (2024).
  • [2] Waleffe, Roger, et al. "An Empirical Study of Mamba-based Language Models." arXiv preprint arXiv:2406.07887 (2024).
  • [3] Ren, Liliang, et al. "Samba: Simple Hybrid State Space Models for Efficient Unlimited Context Language Modeling." arXiv preprint arXiv:2406.07522 (2024).

We once again appreciate the reviewer’s positive comments and valuable feedback and will incorporate the suggestions to enhance the clarity, completeness, and practical relevance of our paper.

1. Authors said that they "will include additional intuitive explanations and visual aids" in the paper.
2. Zubić et al. work will be included in the Related Work section.
3. Authors said that they "will thoroughly proofread the paper again and carefully correct all the typos we missed".
4. In the follow-up work they can also do more experiments for various tasks.
5. Authors discussed that probably Attention + SSMs will be the future of Neural Net Architectures.

Therefore, everything I had as problem is addressed, and I will finalize my rating as:

8: Strong Accept

We thank the reviewer for their thoughtful comments and questions. We address the raised points as follows:

Response to Weaknesses:

• Equation 4 Notation: We acknowledge that the subscript $w_{1...t}$ in Equation 4 was intended to indicate that the prefix $w$ consists of a sequence of $t$ inputs, but it is superfluous. We will revise the notation for consistency and clarity, matching the representation in the previous sentence.
• Modulo Some K: The reference to "modulo some K" at line 215 indicates that $K$ is a natural number ($K \in \mathbf{N}$). For instance, in the case of Parity, $K$ is 2 because we use modulo 2 to determine the sequence output.
• Explanation of TC0: While TC0 is a common term in computational complexity literature, we agree that it might not be familiar to everyone in the diverse NeurIPS audience. Although space constraints prevented us from including a detailed explanation in the main paper, we will add a brief explanation in the appendix to improve readability.
• Chomsky–Schützenberger Theorem: The Chomsky–Schützenberger theorem is a foundational result in formal language theory. It states that any context-free language can be represented using a combination of a Dyck language and a regular language. More formally, any context-free language can be expressed as a homomorphic image of the intersection of a Dyck language with a regular language.

Response to Questions:

Is the assumption of non-negative gates critical in Theorem 4? If so, can Mamba be improved by replacing non-negative gates (such as ReLU or Swish) with signed gates (such as Leaky ReLU)?

• Our proof for Theorem 4 does make the assumption of non-negative gates. With more general gates, we show in Theorem 21 (page 22 in the Appendix) that expressiveness would increase substantially. In Mamba, the non-negativity is created by parameterizing gates as exponentials. We have attempted to change the implementation to allow signed gates, but found training to stop working, suggesting the non-negativity may be useful for training. We believe that studying SSMs without the nonnegativity constraint presents an interesting oppportunity for follow-up work.

Some authors claim that neural networks are Turing complete. For example, [...] If I understand correctly, TC0 is a subclass of regular language, and strictly smaller than Turing machine. What are essential differences from their claims?

• Regarding the Pérez et al. paper, it is indeed a seminal piece of work, and we do cite it in our paper. Its results show the Turing completeness of transformers when they are allowed to generate an unbounded number of intermediate tokens before making their decision -- akin to chain-of-thought and scratchpad techniques. This is a key difference to our setup: We are interested in the expressiveness of neural models when they have to make their decision immediately upon processing the input, without additional intermediate steps. This is foundational because any additional steps would result in higher computational cost. Indeed, there is work [1][2] showing that transformers become more expressive than TC0 when allowed intermediate steps (while increasing computational cost due to intermediate steps) -- and we believe that the same arguments will hold for SSMs.

[1] Li et al, Chain of Thought Empowers Transformers to Solve Inherently Serial Problems, ICLR 2024

[2] Feng et al, Towards Revealing the Mystery behind Chain of Thought: A Theoretical Perspective, NeurIPS 2023