Understanding the differences in Foundation Models: Attention, State Space Models, and Recurrent Neural Networks

Summary

The paper concerns using the proposed "Dynamical Systems Framework" (DSF) to better understand the differences between linear SSMs, RNNs, linear attention and Softmax attention. The DSF shows a way to write each of these models as a linear time-varying recurrence. The paper aims to use this to answer questions about the differences between the three types of models, the role of state expansion, and how ideas from SSMs can be used to improve RNNs. Empirical experiments on the multi-query associative recall (MQAR) task and Wikitext are used to support the findings.

Strengths

- The goal of better understanding the differences between the different models, and in particular trying to understand the gap between Softmax attention and its more efficient alternatives is important - The paper is well-written and clear. I particularly liked Figure 4 in Appendix A which made the dimensions being referred to throughout the paper clear. - The different formulations and results in the paper appear to be technically correct - Insights are used from SSMs to propose a different normalization scheme for linear attention which appears to improve performance on the tasks considered - The empirical results on the MQAR and Wikitext experiments support the claims

Weaknesses

- The proposed DSF is nice in that it allows comparing each of the different types of models, but it's originality is weakened by the fact that most of the main parts of the framework have been discussed before in prior work - The connections between linear attention and linear RNNs/SSMs are clear and have been discussed in prior works as mentioned in related work. - In addition, the infinite dimensional representation of softmax attention as discussed in Section 3.2.1 is also discussed in Katharopoulos et al. 2020 among others - Softmax attention has also been formulated as a dynamical system in https://arxiv.org/abs/2401.06104 which should be cited - The findings that are suggested as a result of the framework also do not seem to be new or particularly surprising - The fact that larger state sizes lead to increased expressivity is not surprising and has been explored before, e.g. in https://arxiv.org/abs/2312.04927, https://arxiv.org/abs/2312.00752 - The proposed improvement to RNNs is to change the parameterization of a linear RNN (the quasi LSTM) by replacing its transition parameterization with that of S6. But this seems to just make it more like the already existing linear RNN, RG-LRU (with a slight different in parameterization from S6). This seems to be a much weaker finding than the claimed "What do selective SSMs teach us about improving RNN architectures" from line 53 of the Intro. - As a counter to this point I am making, the proposed improvement to the normalization of linear attention appears to be novel and appears to improve performance on the tasks considered. - The focus and claims around expressivity and performance in the paper appear to be very language focused. - This is obviously of interest, but perhaps narrows the ability for the DSF to provide new and interesting insights as I mention in the point above. - In addition, only 2 small tasks are considered to support the claims. This limits the ability to understand if the claims hold in more general settings. - The authors state in the limitations that to strengthen the insights, a larger and more complex language task is needed. This could be great, however I would suggest perhaps also considering more diverse data modalities could also strengthen the insights of the proposed framework.

Questions

- Perhaps considering data modalities and examples other than language could help to highlight insights that can be drawn from the DSF? E.g. perhaps modalities that more naturally arise from dynamical systems could help to further highlight differences between softmax attention and its more efficient alternatives? Or could help highlight differences between different linear attention/SSM/RNN variants? - Are the differences between softmax attention and the other methods on the MQAR task, e.g. in Figure 5, simply due to the difference in state size? One can think of the state size of softmax attention as the size of its KV cache required for the sequence length of the task. If you make the recurrent alternatives have a state size that matches the softmax attention KV cache size for the task sequence length, do they then perform as well? Or are there other things happening in Softmax Attention as well? Perhaps exploring this through the lens of the DSF framework could also provide insights? Minor: - It is stated the Matrix in Equation 12 is $L\times L$, but isn't this only the case when $B$ and $C$ are $N \times 1$ and $1 \times N$? The presentation in the paragraph before suggests these are general matrices. - Line 203 says "The DSF makes apparent how a necessary condition for separability is for the map ζ(·) to be expressed by a finite-dimensional kernel, ". But in line 154 it says "Softmax attention also satisfies the assumption of separability". I realize the point that is trying to be made regarding "practical separability", but nonetheless, it appears to be a contradiction as currently written.

Rating

6

Confidence

4

Soundness

3

Presentation

3

Contribution

2

Limitations

Adequately discussed.

Summary

This paper shows that many sequence models (including attention, SSMs, and recurrent models) can be viewed as linear time-varying dynamical systems. This is helpful for answering some questions about the differences and similarities between these architectures.

Strengths

**Importance of unifying framework**. The present state of sequence model research is quite muddy, with tons of new architectures being proposed. On the surface these architectures often look quite complicated and differentiated, when in reality they are closely related to one another. Because of this, I think unifying frameworks are useful for helping push progress in the space. **The role of normalization in S6 and Linear Attention.** The paper includes a very interesting discussion of the parameterization of the normalization in S6 and Linear attention.

Weaknesses

**Presentation and clarity.** The introduction and abstract do not provide any details on how the Dynamical Systems Framework works and assumes a reader is familiar with dynamical systems. To improve the presentation and accessibility, I would recommend including a high-level discussion of what the DSF is and what challenges it addresses. **Theoretical Novelty.** In the first paragraph of the introduction the authors state, “Although these models show great promise in boosting efficiency, current comparisons with attention are merely empirical.” This isn’t correct: there have been many theoretical studies comparing these new recurrent architectures with Softmax attention. To name a few: - [Arora *et al.](https://arxiv.org/abs/2402.18668)* include a theoretical analysis grounded in results from communication complexity that highlights differences in models ability to perform associative recall. - Using tools from circuit complexity, [Merill *et al.*](https://arxiv.org/pdf/2404.08819) show that Attention and SSMs are in the same complexity class ($\text{TC}^0$). - Using tools from communication complexity, [Bhattamishra *et al.*](https://arxiv.org/pdf/2406.09347) provide separation results between attention and recurrent architectures on numerous tasks. **Empirical Novelty.** The main empirical result in this paper is that increased state size (across architectures) leads to improved performance on MQAR and that separable attention matches Softmax attention with sufficiently large recurrent size. This result has already been shown theoretically and empirically in [Arora *et al.*](https://arxiv.org/pdf/2402.18668) (ICML ‘24) - see Figure 2, Theorem 3.1, and Section 3. The work under review builds upon Arora *et al.* by including a few additional architectures to the analysis. But this prior work should be mentioned in the related work and/or the introduction given the significant overlap in results . **Claims around Lemma 2.** Some of the language around Lemma 2 is a bit incautious. The authors state that because of Lemma 2, “the larger the state expansion n, the more expressivity the architecture has”. In other words, they are claiming that increasing state expansion strictly increases expressivity. To support this claim, the theorem would need to additionally show that a dynamical system of state dimension $N$ ***cannot*** recover all dynamical systems with state dimension $\hat{N}$. Additionally, I think its important to add the qualification that this holds assuming everything else about the architecture is held constant. It’s of course possible for an architecture with smaller state size to admit solutions that a larger state size model cannot represent if the parameterizations are different. **Minor typos and clarifications.** - Line 32: “framework that allows to evaluate” → “allows us to evaluate” - Line 61: “We use $\sigma(\cdot)$ to denote is the sigmoid function. - In the list of questions answered by the paper (Line 40), I would include hyperlinks to the relevant sections.

Questions

It would help to add additional exposition on the results in Fig. 2. What are the takeaways? It looks like exponentiating improves performance, but this isn’t stated explicitly.

Rating

6

Confidence

4

Soundness

3

Presentation

2

Contribution

2

Limitations

Yes the authors discuss limitations.

Summary

This paper provides a dynamical system based framework for principled comparisons between various existing recurrent architectures (linear attention, SSMs, etc.). There are focuses on the formulation and the role of the hidden state dimension. Some experimental results on the multi-query associative recall (MQAR) benchmark and the WikiText-103 dataset are provided to demonstrate the differences between the considered models.

Strengths

- The paper is overall well written and organized - Systematic comparisons between the formulations of various sequence models are provided, which are neat and useful for researchers without much background in the area

Weaknesses

- The two proposed experiments (the MQAR benchmark and the WikiText-103 dataset) seem random and limited. There are no experiments on the more standard benchmark tasks such as the LRA task (focusing on smaller tasks such as the sequential CIFAR task would be valuable) and time series forecasting tasks - The theoretical insights are actually quite limited. There are unaddressed natural questions such as: (1) one might wonder about finite-dimensional approximation of the softmax attention, (2) how Lemma 1 could be related to the kernel method and RKHS - There are also other aspects that are left out for the comparative studies: e.g., how do the training dynamics differ between the various models (which architecture will converge faster), as well as the stability of the models (how prone are they to the infamous vanishing/exploding gradient problem). Also, can we quantify and compare the long range dependency learning capability of these models? - No practical guidance in terms of which architecture is suitable for what kind of task is provided. After all, this should be the main goal of the comparisons

Questions

See above

Rating

5

Confidence

3

Soundness

2

Presentation

3

Contribution

3

Limitations

Yes

Summary

This paper introduces the Dynamical Systems Framework (DSF), a theoretical approach for analyzing and comparing various foundation models in AI. The DSF reformulates attention-based models, State Space Models (SSMs), and Recurrent Neural Networks (RNNs) into a common dynamical systems representation. This allows for principled comparisons and generates new insights into their similarities and differences. The authors leverage the DSF to compare linear attention and selective SSMs, provide conditions for approximating softmax attention, analyze state expansion in RNNs and SSMs, investigate differences between linear attention and S6 (Mamba) models, and apply SSM insights to RNN architectures. The paper combines theoretical results with empirical validations on the MQAR benchmark and WikiText-103 dataset.

Strengths

- Novel unified perspective on different foundation model architectures - Enables new theoretical insights and comparisons between previously isolated model classes - Sound theoretical analysis with supporting lemmas and proofs - Empirical validation on MQAR and WikiText-103 provides practical grounding - Clear potential for guiding future development of efficient, scalable models - Innovative application of SSM insights to improve RNN architectures

Weaknesses

- Limited empirical evaluation, focusing primarily on MQAR and a single language modeling task - Lack of statistical significance reporting for experimental results - Some theoretical results (e.g., Lemma 2) are relatively straightforward - Dense theoretical sections may be challenging for broader audience - Immediate practical impact on model performance is somewhat limited - Doesn't fully explore implications for training or inference efficiency

Questions

1. How well might the DSF insights scale to much larger models (billions of parameters) used in state-of-the-art applications? 2. Are there theoretical limits to the convergence of linear attention to softmax attention performance with increased state expansion? 3. Have you explored alternative normalization schemes that might bridge the gap between S6 and linear attention while maintaining computational efficiency? 4. How does the S6-inspired state transition in qLSTM affect its long-term memory capabilities? 5. Can the theoretical insights from DSF reformulations be leveraged to develop new, efficient implementations?

Rating

7

Confidence

3

Soundness

3

Presentation

3

Contribution

3

Limitations

The authors acknowledge that the DSF parametrization, while theoretically insightful, doesn't necessarily lead to efficient implementations. Empirical validation is limited to synthetic tasks and a relatively small language modeling task, constraining generalizability. The paper also lacks comprehensive analysis of computational efficiency implications for different architectures under the DSF lens.

Thank you for the response

I am satisfied with the response and have raised my score.