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

Summary

1. This paper studies the expressive capacity of state-space models from a formal language perspective. 2. It shows that for flip-flop, SSM can do well. And for parity, SSM cannot do well.

Strengths

1. State-space models have the advantage of low inference cost, therefore using it to learn language models and formal languages are important research directions. 2. The perspective of formal language and the proof is interesting

Weaknesses

1. The theorem 2 seems to be problematic: 1. Given sequence $\mathbf{x} = x_1, \dots, x_T$, consider function $f(\mathbf{x}) = (e^{i \pi \sum_{i=1}^n x_i} + 1)/2$. This continuous function can achieve the goal that if $\sum_{i=1}^n x_i$ is even, then the result is 1, otherwise the result is 0. 2. Based on a cumulative sum layer, it only requires the approximation using two-layer state-space models for a function $f(x) = (e^{i \pi x} + 1) / 2$ . There are various universal approximation theorem to prove this: 1. https://openreview.net/forum?id=i0OmcF14Kf 2. https://arxiv.org/abs/2307.11888 3. They are also related works in the expressive sense. 4. Still it is possible to have the unboundedness issue from $\sum_{i=1}^n x_i$. 5. I understand the notion of finite precision assumption is necessary in practice. However, another question is whether increasing the number of fractional bits p can be a effective method to reduce the problem proposed in this paper. As the cost of increasing p is still linear. (And increasing p can relax the unboundedness issue in $\sum_{i=1}^n x_i$ in an exponential sense. 6. I am willing to raise the score if this issue can be better discussed. 3. The experiments for 28 seem to justify that three layer mamba can work with certain accuracy. A natural question would be whether increasing the hidden dimension and layers improves the performance. 2. The role of parameterization discussed in page 3 seems a bit lacking of related work 1. HIPPO(https://arxiv.org/abs/2008.07669), Approximate Diagonalization(https://arxiv.org/abs/2310.01698), StableSSM(https://arxiv.org/abs/2311.14495). 3. Line 682, typo - have shown.

Questions

1. I understand the notion of finite precision assumption is necessary in practice. However, a natural question is whether increasing the number of fractional bits p can be a effective method to reduce the problem proposed in this paper. As the cost of increasing p is still linear. 2. In the same argument, it seems transformer and temporal convolution, in the setting of convolution, also have the same problems as Theorem 2. I tend to believe this issue comes from the setup of finite precision and infinite sequence length. Any insight about this issue?

Rating

6

Confidence

3

Soundness

2

Presentation

3

Contribution

2

Limitations

1. The result derived in this paper is limited to formal languages. It is not clear what is the generalization to natural languages.

Summary

The paper presents several results on the expressiveness of *state space models (SSMs)*, viewed as *language acceptors* or a weak form of *language predictors*, from the perspective of formal language theory. These results are parametrized to subsume currently common SSM architectures, including *non-negative SSMs*, which encode prefixes using positive reals only, *time-invariant SSMs*, where prefixes are encoded independent of the currently considered token, and *finite precision SSMs*, meaning that internal computations are limited by a fixed number of fractional bits. In more detail, the paper presents the following results: - The FlipFlop language is predicted by a two-layer SSM with finite precision (Theorem 1), and there is no non-negative, finite precision SSM that recognizes the PARITY language (Theorem 2). - The class of non-negative, finite precision SSMs predicts exactly the star-free languages (Theorem 4). - Various kinds of languages that require unbounded counting, such as Dyck languages, $a^nb^nc^n$, or boolean expressions, are predicted by an SSM. In summary, the paper establishes foundational results categorizing the expressive power of SSMs within the well-established framework of formal language theory. Additionally, this allows for a comparison of SSM with their imminent competitor, transformers, for which similar architectures exist. This comparison is empirically supported by experiments.

Strengths

The paper considers the expressiveness of recent SSMs from the perspective of formal language theory. While it is not new to analyze sequence-classifying or sequence-to-sequence models from a formal language viewpoint, this paper is one of the first to do so for SSMs on a broad scale. Additionally, this is a reasonable continuation of preceding work, particularly for understanding the differences in expressiveness between SSMs and transformers. The main result, at least from my perspective, Theorem 4, states that SSMs predict exactly the star-free languages. This is fundamental, making it well worth considering. The remaining results are more specialized, focusing on particular languages. However, these also provide insights for a more detailed understanding of the expressive power of SSMs, especially compared to transformer. Thus, these are of clear relevance as well. In general, all results are supported by technical proofs and some intuition building, allowing the reader to verify the soundness of the results in sufficient detail without unreasonably extensive effort.

Weaknesses

The clarity and quality of the paper could be improved. This is mainly due to three reasons: - 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 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 $w_{t-1}$ can be read out from $z_t$ 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, $w_{1..t} = …00$), it is not clear how equalities (15) and (16) are derived or what *”… running these in parallel, …”* means.

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.

Rating

7

Confidence

4

Soundness

3

Presentation

2

Contribution

4

Limitations

The authors state the limitations of their work, mainly focusing on the lack of experimental evaluation, which is sufficient regarding this part. However, I would have appreciated a clearer statement on open questions focusing on the theoretical contributions. For example, they could discuss related work on other models like RNNs or Transformers and which results about formal language recognition or prediction exist in those contexts, which have not yet been considered in the context of SSMs.

Summary

This paper presents a comprehensive theoretical and empirical analysis of the expressive capacity of modern State Space Models (SSMs) within the framework of formal languages and automata theory. The authors establish important theoretical results, demonstrating that SSMs can effectively model star-free languages, various counter languages, and bounded hierarchical structures without relying on explicit stack mechanisms. By leveraging the Krohn-Rhodes theorem and concepts from algebraic theory, the paper provides rigorous proofs showing the capabilities and limitations of SSMs in capturing different language classes. By connecting theoretical computer science (TCS) with machine learning (ML), the paper highlights the importance of understanding the fundamental limitations and capabilities of computational models.

Strengths

- The paper introduces a novel theoretical framework that connects formal language theory with the expressive capabilities of State Space Models (SSMs), a relatively unexplored area in machine learning research. - By leveraging the Krohn-Rhodes theorem and algebraic theory, the authors provide proofs and characterizations of the language classes that SSMs can model, including star-free languages and bounded hierarchical structures. - The integration of TCS concepts with modern SSM architectures is important because through TCS we study the limits and what is possible or not possible - Theoretical contributions are rigorously presented, with clear and detailed proofs supporting the claims about the expressive power of SSMs. I checked all the proofs thoroughly and I am sure they are valid.

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

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? - 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?

Rating

8

Confidence

5

Soundness

3

Presentation

4

Contribution

4

Limitations

Already addressed by authors.

Summary

This study investigates the expressive power of SSMs compared to Transformers and RNNs from the perspective of formal language classes (or circuit complexities), and reports their distinct strengths. Previous studies have shown that both SSM and Transformers are in TC0, suggesting some state tracking problems are hard for both models. This study shows that both models cover overlapping but distinct parts of TC0. Theorem 1 claims that two-layer SSM can predictively model the Flip Flop (a star-free language) state tracking, and Theorem 2 claims that no SSM with non-negative gates can recognize PARITY (a non-star-free language). Since Flip-Flop and PARITY are building blocks of star-free and non-star-free languages, Theorems 1 and 2 are summarized to Theorem 4 claiming that SSMs with non-negative gates can predictively model a regular language L iff it is star-free. Further Theorems 5 and 6 claim several context-free and context-sensitive are predictively modeled by SSMs. Experimental results support theoretical claims. Due to the limitation of resources, I did not check the proofs. In particular, I am not familiar with formal language theory and I did not have enough time to learn about it. This is a __suggestion for chairs__ for future avoidance of mismatches, that the reviewers should be __examined__ if they have fundamental knowledge/background/understandings in the field. I am an expert of expressive power analysis, but not at all of formal language theory. Unfortunately, this kind of mismatch happens every year, so I am usually skeptical to any mathematical ''theorems'' published in machine learning conferences.

Strengths

- Previous studies are precisely reviewed. - Theorem 4 strictly refines the previous results that SSMs are in TC0, a subclass of regular languages. - Theorem 5 and 6 are new results in context-free and context-sensitive languages.

Weaknesses

I am really interested in reading the manuscript. Just because I am not familiar with formal language theory, some expressions were a bit hard to parse for me. For example, - in Eq.4, what does the subscript of w_{1…t} mean? In the previous sentence, w does not have any subscript. Are they the same or distinguished? - at l.215: “L is non-star-free if and only if recognizing it involves counting modulo some K.” What is K? - It would improve the readability if the authors could explain what is TC0. In my opinion, NeurIPS audiences (are familiar with sample complexity but) are less familiar with computational complexity. - What is Chomsky–Schützenberger theorem?

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)? - Some authors claim that neural networks are Turing complete. For example, [PMB19] _Jorge Pérez, Javier Marinković, Pablo Barceló, On the Turing Completeness of Modern Neural Network Architectures, ICLR2019_ If I understand correctly, TC0 is a subclass of regular language, and strictly smaller than Turing machine. What are essential differences from their claims?

Rating

7

Confidence

2

Soundness

3

Presentation

2

Contribution

3

Limitations

n/a